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