JVM/Builder/Monad.hs

   1 {-# LANGUAGE FlexibleContexts, TypeFamilies, OverloadedStrings #-}
   2 -- | This module defines Generate monad, which helps generating JVM code and
   3 -- creating Java class constants pool.
   4 module JVM.Builder.Monad
   5   (GState (..),
   6    emptyGState,
   7    Generate,
   8    addToPool,
   9    i0, i1, i8,
  10    newMethod,
  11    setStackSize, setMaxLocals,
  12    withClassPath,
  13    getClassField, getClassMethod,
  14    generate
  15   ) where
  16
  17 import Control.Monad.State as St
  18 import Data.Word
  19 import Data.List
  20 import Data.Binary
  21 import qualified Data.Map as M
  22 import qualified Data.Set as S
  23 import qualified Data.ByteString.Lazy as B
  24
  25 import JVM.Common ()  -- import instances only
  26 import JVM.ClassFile
  27 import JVM.Assembler
  28 import Java.ClassPath
  29
  30 -- | Generator state
  31 data GState = GState {
  32   generated :: [Instruction],               -- ^ Already generated code (in current method)
  33   currentPool :: Pool Direct,             -- ^ Already generated constants pool
  34   doneMethods :: [Method Direct],         -- ^ Already generated class methods
  35   currentMethod :: Maybe (Method Direct), -- ^ Current method
  36   stackSize :: Word16,                      -- ^ Maximum stack size for current method
  37   locals :: Word16,                         -- ^ Maximum number of local variables for current method
  38   classPath :: [Tree CPEntry]
  39   }
  40   deriving (Eq,Show)
  41
  42 -- | Empty generator state
  43 emptyGState ::  GState
  44 emptyGState = GState {
  45   generated = [],
  46   currentPool = M.empty,
  47   doneMethods = [],
  48   currentMethod = Nothing,
  49   stackSize = 496,
  50   locals = 0,
  51   classPath = []}
  52
  53 -- | Generate monad
  54 type Generate a = StateT GState IO a
  55
  56 -- | Update ClassPath
  57 withClassPath :: ClassPath () -> Generate ()
  58 withClassPath cp = do
  59   res <- liftIO $ execClassPath cp
  60   st <- St.get
  61   St.put $ st {classPath = res}
  62
  63 -- | Append a constant to pool
  64 appendPool :: Constant Direct -> Pool Direct -> (Pool Direct, Word16)
  65 appendPool c pool =
  66   let size = fromIntegral (M.size pool)
  67       pool' = M.insert size c pool
  68   in  (pool', size)
  69
  70 -- | Add a constant to pool
  71 addItem :: Constant Direct -> Generate Word16
  72 addItem c = do
  73   pool <- St.gets currentPool
  74   case lookupPool c pool of
  75     Just i -> return (i+1)
  76     Nothing -> do
  77       let (pool', i) = appendPool c pool
  78       st <- St.get
  79       St.put $ st {currentPool = pool'}
  80       return (i+1)
  81
  82 -- | Lookup in a pool
  83 lookupPool :: Constant Direct -> Pool Direct -> Maybe Word16
  84 lookupPool c pool =
  85   fromIntegral `fmap` findIndex (== c) (M.elems pool)
  86
  87 addNT :: HasSignature a => NameType a -> Generate Word16
  88 addNT (NameType name sig) = do
  89   let bsig = encode sig
  90   x <- addItem (CNameType name bsig)
  91   addItem (CUTF8 name)
  92   addItem (CUTF8 bsig)
  93   return x
  94
  95 addSig :: MethodSignature -> Generate Word16
  96 addSig c@(MethodSignature args ret) = do
  97   let bsig = encode c
  98   addItem (CUTF8 bsig)
  99
 100 -- | Add a constant into pool
 101 addToPool :: Constant Direct -> Generate Word16
 102 addToPool c@(CClass str) = do
 103   addItem (CUTF8 str)
 104   addItem c
 105 addToPool c@(CField cls name) = do
 106   addToPool (CClass cls)
 107   addNT name
 108   addItem c
 109 addToPool c@(CMethod cls name) = do
 110   addToPool (CClass cls)
 111   addNT name
 112   addItem c
 113 addToPool c@(CIfaceMethod cls name) = do
 114   addToPool (CClass cls)
 115   addNT name
 116   addItem c
 117 addToPool c@(CString str) = do
 118   addToPool (CUTF8 str)
 119   addItem c
 120 addToPool c@(CNameType name sig) = do
 121   addItem (CUTF8 name)
 122   addItem (CUTF8 sig)
 123   addItem c
 124 addToPool c = addItem c
 125
 126 putInstruction :: Instruction -> Generate ()
 127 putInstruction instr = do
 128   st <- St.get
 129   let code = generated st
 130   St.put $ st {generated = code ++ [instr]}
 131
 132 -- | Generate one (zero-arguments) instruction
 133 i0 :: Instruction -> Generate ()
 134 i0 = putInstruction
 135
 136 -- | Generate one one-argument instruction
 137 i1 :: (Word16 -> Instruction) -> Constant Direct -> Generate ()
 138 i1 fn c = do
 139   ix <- addToPool c
 140   i0 (fn ix)
 141
 142 -- | Generate one one-argument instruction
 143 i8 :: (Word8 -> Instruction) -> Constant Direct -> Generate ()
 144 i8 fn c = do
 145   ix <- addToPool c
 146   i0 (fn $ fromIntegral ix)
 147
 148 -- | Set maximum stack size for current method
 149 setStackSize :: Word16 -> Generate ()
 150 setStackSize n = do
 151   st <- St.get
 152   St.put $ st {stackSize = n}
 153
 154 -- | Set maximum number of local variables for current method
 155 setMaxLocals :: Word16 -> Generate ()
 156 setMaxLocals n = do
 157   st <- St.get
 158   St.put $ st {locals = n}
 159
 160 -- | Start generating new method
 161 startMethod :: [AccessFlag] -> B.ByteString -> MethodSignature -> Generate ()
 162 startMethod flags name sig = do
 163   addToPool (CString name)
 164   addSig sig
 165   setStackSize 4096
 166   setMaxLocals 100
 167   st <- St.get
 168   let method = Method {
 169     methodAccessFlags = S.fromList flags,
 170     methodName = name,
 171     methodSignature = sig,
 172     methodAttributesCount = 0,
 173     methodAttributes = AR M.empty }
 174   St.put $ st {generated = [],
 175                currentMethod = Just method }
 176
 177 -- | End of method generation
 178 endMethod :: Generate ()
 179 endMethod = do
 180   m <- St.gets currentMethod
 181   code <- St.gets genCode
 182   case m of
 183     Nothing -> fail "endMethod without startMethod!"
 184     Just method -> do
 185       let method' = method {methodAttributes = AR $ M.fromList [("Code", encodeMethod code)],
 186                             methodAttributesCount = 1}
 187       st <- St.get
 188       St.put $ st {generated = [],
 189                    currentMethod = Nothing,
 190                    doneMethods = doneMethods st ++ [method']}
 191
 192 -- | Generate new method
 193 newMethod :: [AccessFlag]               -- ^ Access flags for method (public, static etc)
 194           -> B.ByteString               -- ^ Method name
 195           -> [ArgumentSignature]        -- ^ Signatures of method arguments
 196           -> ReturnSignature            -- ^ Method return signature
 197           -> Generate ()                -- ^ Generator for method code
 198           -> Generate (NameType Method)
 199 newMethod flags name args ret gen = do
 200   let sig = MethodSignature args ret
 201   startMethod flags name sig
 202   gen
 203   endMethod
 204   return (NameType name sig)
 205
 206 -- | Get a class from current ClassPath
 207 getClass :: String -> Generate (Class Direct)
 208 getClass name = do
 209   cp <- St.gets classPath
 210   res <- liftIO $ getEntry cp name
 211   case res of
 212     Just (NotLoaded p) -> fail $ "Class file was not loaded: " ++ p
 213     Just (Loaded _ c) -> return c
 214     Just (NotLoadedJAR p c) -> fail $ "Class was not loaded from JAR " ++ p ++ ": " ++ c
 215     Just (LoadedJAR _ c) -> return c
 216     Nothing -> fail $ "No such class in ClassPath: " ++ name
 217
 218 -- | Get class field signature from current ClassPath
 219 getClassField :: String -> B.ByteString -> Generate (NameType Field)
 220 getClassField clsName fldName = do
 221   cls <- getClass clsName
 222   case lookupField fldName cls of
 223     Just fld -> return (fieldNameType fld)
 224     Nothing  -> fail $ "No such field in class " ++ clsName ++ ": " ++ toString fldName
 225
 226 -- | Get class method signature from current ClassPath
 227 getClassMethod :: String -> B.ByteString -> Generate (NameType Method)
 228 getClassMethod clsName mName = do
 229   cls <- getClass clsName
 230   case lookupMethod mName cls of
 231     Just m -> return (methodNameType m)
 232     Nothing  -> fail $ "No such method in class " ++ clsName ++ ": " ++ toString mName
 233
 234 -- | Convert Generator state to method Code.
 235 genCode :: GState -> Code
 236 genCode st = Code {
 237     codeStackSize = stackSize st,
 238     codeMaxLocals = locals st,
 239     codeLength = len,
 240     codeInstructions = generated st,
 241     codeExceptionsN = 0,
 242     codeExceptions = [],
 243     codeAttrsN = 0,
 244     codeAttributes = AP [] }
 245   where
 246     len = fromIntegral $ B.length $ encodeInstructions (generated st)
 247
 248 -- | Start class generation.
 249 initClass :: B.ByteString -> Generate Word16
 250 initClass name = do
 251   addToPool (CClass "java/lang/Object")
 252   addToPool (CClass name)
 253   addToPool (CString "Code")
 254
 255 -- | Generate a class
 256 generate :: [Tree CPEntry] -> B.ByteString -> Generate () -> IO (Class Direct)
 257 generate cp name gen = do
 258   let generator = do
 259         initClass name
 260         gen
 261   res <- execStateT generator (emptyGState {classPath = cp})
 262   let code = genCode res
 263       d = defaultClass :: Class Direct
 264   return $ d {
 265         constsPoolSize = fromIntegral $ M.size (currentPool res),
 266         constsPool = currentPool res,
 267         accessFlags = S.fromList [ACC_PUBLIC, ACC_STATIC],
 268         thisClass = name,
 269         superClass = "java/lang/Object",
 270         classMethodsCount = fromIntegral $ length (doneMethods res),
 271         classMethods = doneMethods res }
 272