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