JVM/Builder/Monad.hs

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