JVM/Builder/Monad.hs

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