JVM/Builder/Monad.hs

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