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