some updates.
[hs-java.git] / JVM / Generator.hs
1 {-# LANGUAGE FlexibleContexts, TypeFamilies, OverloadedStrings #-}
2 module JVM.Generator where
3
4 import Control.Monad.State as St
5 import Data.Word
6 import Data.List
7 import Data.Binary
8 import qualified Data.Map as M
9 import qualified Data.Set as S
10 import qualified Data.ByteString.Lazy as B
11
12 import JVM.Types
13 import JVM.ClassFile
14 import JVM.Assembler
15
16 data GState = GState {
17   generated :: [Instruction],
18   currentPool :: Pool Resolved,
19   doneMethods :: [Method Resolved],
20   currentMethod :: Maybe (Method Resolved)}
21   deriving (Eq,Show)
22
23 emptyGState = GState {
24   generated = [],
25   currentPool = M.empty,
26   doneMethods = [],
27   currentMethod = Nothing }
28
29 type Generate a = State GState a
30
31 appendPool :: Constant Resolved -> Pool Resolved -> (Pool Resolved, Word16)
32 appendPool c pool =
33   let size = fromIntegral (M.size pool)
34       pool' = M.insert size c pool
35   in  (pool', size)
36
37 addItem :: Constant Resolved -> Generate Word16
38 addItem c = do
39   pool <- St.gets currentPool
40   case lookupPool c pool of
41     Just i -> return (i+1)
42     Nothing -> do
43       let (pool', i) = appendPool c pool
44       st <- St.get
45       St.put $ st {currentPool = pool'}
46       return (i+1)
47
48 lookupPool :: Constant Resolved -> Pool Resolved -> Maybe Word16
49 lookupPool c pool =
50   fromIntegral `fmap` findIndex (== c) (M.elems pool)
51
52 addNT :: Binary (Signature a) => NameType a -> Generate Word16
53 addNT (NameType name sig) = do
54   let bsig = encode sig
55   x <- addItem (CNameType name bsig)
56   addItem (CUTF8 name)
57   addItem (CUTF8 bsig)
58   return x
59
60 addSig :: MethodSignature -> Generate Word16
61 addSig c@(MethodSignature args ret) = do
62   let bsig = encode c
63   addItem (CUTF8 bsig)
64
65 addToPool :: Constant Resolved -> Generate Word16
66 addToPool c@(CClass str) = do
67   addItem (CUTF8 str)
68   addItem c
69 addToPool c@(CField cls name) = do
70   addToPool (CClass cls)
71   addNT name
72   addItem c
73 addToPool c@(CMethod cls name) = do
74   addToPool (CClass cls)
75   addNT name
76   addItem c
77 addToPool c@(CIfaceMethod cls name) = do
78   addToPool (CClass cls)
79   addNT name
80   addItem c
81 addToPool c@(CString str) = do
82   addToPool (CUTF8 str)
83   addItem c
84 addToPool c@(CNameType name sig) = do
85   addItem (CUTF8 name)
86   addItem (CUTF8 sig)
87   addItem c
88 addToPool c = addItem c
89
90 putInstruction :: Instruction -> Generate ()
91 putInstruction instr = do
92   st <- St.get
93   let code = generated st
94   St.put $ st {generated = code ++ [instr]}
95
96 i0 :: Instruction -> Generate ()
97 i0 = putInstruction
98
99 i1 :: (Word16 -> Instruction) -> Constant Resolved -> Generate ()
100 i1 fn c = do
101   ix <- addToPool c
102   i0 (fn ix)
103
104 i8 :: (Word8 -> Instruction) -> Constant Resolved -> Generate ()
105 i8 fn c = do
106   ix <- addToPool c
107   i0 (fn $ fromIntegral ix)
108
109 startMethod :: [AccessFlag] -> B.ByteString -> MethodSignature -> Generate ()
110 startMethod flags name sig = do
111   addToPool (CString name)
112   addSig sig
113   st <- St.get
114   let method = Method {
115     methodAccessFlags = S.fromList flags,
116     methodName = name,
117     methodSignature = sig,
118     methodAttributesCount = 0,
119     methodAttributes = AR M.empty }
120   St.put $ st {generated = [],
121                currentMethod = Just method }
122
123 endMethod :: Generate ()
124 endMethod = do
125   m <- St.gets currentMethod
126   code <- St.gets genCode
127   case m of
128     Nothing -> fail "endMethod without startMethod!"
129     Just method -> do
130       let method' = method {methodAttributes = AR $ M.fromList [("Code", encodeMethod code)],
131                             methodAttributesCount = 1}
132       st <- St.get
133       St.put $ st {generated = [],
134                    currentMethod = Nothing,
135                    doneMethods = doneMethods st ++ [method']}
136
137 newMethod :: [AccessFlag] -> B.ByteString -> [ArgumentSignature] -> ReturnSignature -> Generate () -> Generate (NameType Method)
138 newMethod flags name args ret gen = do
139   let sig = MethodSignature args ret
140   startMethod flags name sig
141   gen
142   endMethod
143   return (NameType name sig)
144
145 genCode :: GState -> Code
146 genCode st = Code {
147     codeStackSize = 4096,
148     codeMaxLocals = 100,
149     codeLength = len,
150     codeInstructions = generated st,
151     codeExceptionsN = 0,
152     codeExceptions = [],
153     codeAttrsN = 0,
154     codeAttributes = AP [] }
155   where
156     len = fromIntegral $ B.length $ encodeInstructions (generated st)
157
158 initClass :: B.ByteString -> Generate Word16
159 initClass name = do
160   addToPool (CClass "java/lang/Object")
161   addToPool (CClass name)
162   addToPool (CString "Code")
163
164 generate :: B.ByteString -> Generate () -> Class Resolved
165 generate name gen =
166   let generator = do
167         initClass name
168         gen
169       res = execState generator emptyGState
170       code = genCode res
171   in  Class {
172         magic = 0xCAFEBABE,
173         minorVersion = 0,
174         majorVersion = 50,
175         constsPoolSize = fromIntegral $ M.size (currentPool res),
176         constsPool = currentPool res,
177         accessFlags = S.fromList [ACC_PUBLIC, ACC_STATIC],
178         thisClass = name,
179         superClass = "java/lang/Object",
180         interfacesCount = 0,
181         interfaces = [],
182         classFieldsCount = 0,
183         classFields = [],
184         classMethodsCount = fromIntegral $ length (doneMethods res),
185         classMethods = doneMethods res,
186         classAttributesCount = 0,
187         classAttributes = AR M.empty }
188