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