{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE TemplateHaskell #-}

-- | Provides a monadic interface for building class files in a high-level format.
module H2JVM.Builder (
    ClassBuilder,
    addAccessFlag,
    setName,
    getName,
    setVersion,
    setSuperClass,
    addInterface,
    addField,
    addMethod,
    addMethodWithCode,
    addMethodWithCode_,
    addAttribute,
    addBootstrapMethod,
    runClassBuilder,
)
where

import Data.Text (Text)
import Data.Tuple (swap)
import Effectful
import Effectful.Dispatch.Dynamic
import Effectful.Error.Static
import Effectful.State.Static.Local
import Effectful.TH

import H2JVM (CodeBuilder)
import H2JVM.Analyse.StackMap (StackMapError, calculateStackMapFrames)
import H2JVM.Builder.Code (runCodeBuilder)
import H2JVM.ClassFile (ClassFile (..), ClassFileAttribute (BootstrapMethods))
import H2JVM.ClassFile.AccessFlags (ClassAccessFlag, MethodAccessFlag)
import H2JVM.ClassFile.Field
import H2JVM.ClassFile.Method
import H2JVM.ConstantPool
import H2JVM.Descriptor (MethodDescriptor)
import H2JVM.JVMVersion
import H2JVM.Name

import H2JVM.Data.TypeMergingList qualified as TML

-- | The class builder effect for constructing a 'ClassFile' statefully.
data ClassBuilder m a where
    ModifyClass :: (ClassFile -> ClassFile) -> ClassBuilder m ()
    GetClass :: ClassBuilder m ClassFile

makeEffect ''ClassBuilder

-- | Add an access flag modifier to the class being built.
addAccessFlag :: ClassBuilder :> r => ClassAccessFlag -> Eff r ()
addAccessFlag :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
ClassAccessFlag -> Eff r ()
addAccessFlag ClassAccessFlag
flag = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{accessFlags = flag : c.accessFlags})

-- | Set the fully qualified name of the class being built.
setName :: ClassBuilder :> r => QualifiedClassName -> Eff r ()
setName :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
QualifiedClassName -> Eff r ()
setName QualifiedClassName
n = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{name = n})

-- | Retrieve the current fully qualified name of the class being built.
getName :: ClassBuilder :> r => Eff r QualifiedClassName
getName :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
Eff r QualifiedClassName
getName = (.name) (ClassFile -> QualifiedClassName)
-> Eff r ClassFile -> Eff r QualifiedClassName
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Eff r ClassFile
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
Eff es ClassFile
getClass

-- | Set the target JVM version of the class being built.
setVersion :: ClassBuilder :> r => JVMVersion -> Eff r ()
setVersion :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
JVMVersion -> Eff r ()
setVersion JVMVersion
v = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{version = v})

-- | Set the superclass name of the class being built.
setSuperClass :: ClassBuilder :> r => QualifiedClassName -> Eff r ()
setSuperClass :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
QualifiedClassName -> Eff r ()
setSuperClass QualifiedClassName
s = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{superClass = Just s})

-- | Add an implemented interface to the class being built.
addInterface :: ClassBuilder :> r => QualifiedClassName -> Eff r ()
addInterface :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
QualifiedClassName -> Eff r ()
addInterface QualifiedClassName
i = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{interfaces = i : c.interfaces})

-- | Add a field definition to the class being built.
addField :: ClassBuilder :> r => ClassFileField -> Eff r ()
addField :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
ClassFileField -> Eff r ()
addField ClassFileField
f = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{fields = f : c.fields})

-- | Add a method definition to the class being built.
addMethod :: ClassBuilder :> r => ClassFileMethod -> Eff r ()
addMethod :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
ClassFileMethod -> Eff r ()
addMethod ClassFileMethod
m = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{methods = m : c.methods})

{- | Add a method definition to the class being built, whose body is the result of some 'CodeBuilder' monad.
This is essentially a convenience function to:
1. Run the 'CodeBuilder'
2. Calculate the 'StackMapTable' etc.
3. Setup the 'ClassFileMethod' and 'CodeAttributeData'
4. Add the method to this 'ClassBuilder'.
-}
addMethodWithCode :: (ClassBuilder :> r, Error StackMapError :> r, HasCallStack) => Text -> [MethodAccessFlag] -> MethodDescriptor -> Eff (CodeBuilder ': r) a -> Eff r a
addMethodWithCode :: forall (r :: [Effect]) a.
(ClassBuilder :> r, Error StackMapError :> r, HasCallStack) =>
Text
-> [MethodAccessFlag]
-> MethodDescriptor
-> Eff (CodeBuilder : r) a
-> Eff r a
addMethodWithCode Text
methodName [MethodAccessFlag]
methodAccessFlags MethodDescriptor
methodDescriptor Eff (CodeBuilder : r) a
codeBlock = do
    className <- Eff r QualifiedClassName
forall (r :: [Effect]).
(ClassBuilder :> r) =>
Eff r QualifiedClassName
getName
    (res, userCodeAttrs, instructions) <- runCodeBuilder codeBlock

    (frames, maxStack, maxLocals) <-
        case calculateStackMapFrames className methodAccessFlags methodDescriptor instructions of
            Left StackMapError
err -> StackMapError -> Eff r ([StackMapFrame], Int, Int)
forall e (es :: [Effect]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError StackMapError
err
            Right ([StackMapFrame], Int, Int)
res -> ([StackMapFrame], Int, Int) -> Eff r ([StackMapFrame], Int, Int)
forall a. a -> Eff r a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([StackMapFrame], Int, Int)
res

    let finalCodeAttrs = [StackMapFrame] -> CodeAttribute
StackMapTable [StackMapFrame]
frames CodeAttribute -> [CodeAttribute] -> [CodeAttribute]
forall a. a -> [a] -> [a]
: [CodeAttribute]
userCodeAttrs
        codeData =
            CodeAttributeData
                { maxStack :: U2
maxStack = Int -> U2
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
maxStack
                , maxLocals :: U2
maxLocals = Int -> U2
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
maxLocals
                , code :: NonEmpty Instruction
code = NonEmpty Instruction
instructions
                , exceptionTable :: [ExceptionTableEntry]
exceptionTable = [] -- TODO: this will be supported in more depth soon
                , codeAttributes :: [CodeAttribute]
codeAttributes = [CodeAttribute]
finalCodeAttrs
                }
        methodInfo =
            ClassFileMethod
                { methodAccessFlags :: [MethodAccessFlag]
methodAccessFlags = [MethodAccessFlag]
methodAccessFlags
                , methodName :: Text
methodName = Text
methodName
                , methodDescriptor :: MethodDescriptor
methodDescriptor = MethodDescriptor
methodDescriptor
                , methodAttributes :: TypeMergingList MethodAttribute
methodAttributes = [MethodAttribute] -> TypeMergingList MethodAttribute
forall a. (DataMergeable a, Data a) => [a] -> TypeMergingList a
TML.fromList [CodeAttributeData -> MethodAttribute
Code CodeAttributeData
codeData]
                }

    addMethod methodInfo
    pure res

-- | Like 'addMethodWithCode', but ignoring the result of the function.
addMethodWithCode_ :: (ClassBuilder :> r, Error StackMapError :> r, HasCallStack) => Text -> [MethodAccessFlag] -> MethodDescriptor -> Eff (CodeBuilder ': r) a -> Eff r ()
addMethodWithCode_ :: forall (r :: [Effect]) a.
(ClassBuilder :> r, Error StackMapError :> r, HasCallStack) =>
Text
-> [MethodAccessFlag]
-> MethodDescriptor
-> Eff (CodeBuilder : r) a
-> Eff r ()
addMethodWithCode_ Text
methodName [MethodAccessFlag]
methodAccessFlags MethodDescriptor
methodDescriptor Eff (CodeBuilder : r) a
codeBlock = do
    _ <- Text
-> [MethodAccessFlag]
-> MethodDescriptor
-> Eff (CodeBuilder : r) a
-> Eff r a
forall (r :: [Effect]) a.
(ClassBuilder :> r, Error StackMapError :> r, HasCallStack) =>
Text
-> [MethodAccessFlag]
-> MethodDescriptor
-> Eff (CodeBuilder : r) a
-> Eff r a
addMethodWithCode Text
methodName [MethodAccessFlag]
methodAccessFlags MethodDescriptor
methodDescriptor Eff (CodeBuilder : r) a
codeBlock
    pure ()

-- | Add a class file attribute to the class being built.
addAttribute :: ClassBuilder :> r => ClassFileAttribute -> Eff r ()
addAttribute :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
ClassFileAttribute -> Eff r ()
addAttribute ClassFileAttribute
a = (ClassFile -> ClassFile) -> Eff r ()
forall {k} (es :: [Effect]).
(HasCallStack, ClassBuilder :> es) =>
(ClassFile -> ClassFile) -> Eff es ()
modifyClass (\ClassFile
c -> ClassFile
c{attributes = c.attributes `TML.snoc` a})

-- | Add a bootstrap method for @invokedynamic@ calls.
addBootstrapMethod :: ClassBuilder :> r => BootstrapMethod -> Eff r ()
addBootstrapMethod :: forall (r :: [Effect]).
(ClassBuilder :> r) =>
BootstrapMethod -> Eff r ()
addBootstrapMethod BootstrapMethod
b = ClassFileAttribute -> Eff r ()
forall (r :: [Effect]).
(ClassBuilder :> r) =>
ClassFileAttribute -> Eff r ()
addAttribute ([BootstrapMethod] -> ClassFileAttribute
BootstrapMethods [BootstrapMethod
b])

dummyClass :: QualifiedClassName -> JVMVersion -> ClassFile
dummyClass :: QualifiedClassName -> JVMVersion -> ClassFile
dummyClass QualifiedClassName
name JVMVersion
version =
    ClassFile
        { name :: QualifiedClassName
name = QualifiedClassName
name
        , version :: JVMVersion
version = JVMVersion
version
        , accessFlags :: [ClassAccessFlag]
accessFlags = []
        , superClass :: Maybe QualifiedClassName
superClass = Maybe QualifiedClassName
forall a. Maybe a
Nothing
        , interfaces :: [QualifiedClassName]
interfaces = []
        , fields :: [ClassFileField]
fields = []
        , methods :: [ClassFileMethod]
methods = []
        , attributes :: TypeMergingList ClassFileAttribute
attributes = TypeMergingList ClassFileAttribute
forall a. Monoid a => a
mempty
        }

classBuilderToState :: State ClassFile :> r => Eff (ClassBuilder ': r) a -> Eff r a
classBuilderToState :: forall (r :: [Effect]) a.
(State ClassFile :> r) =>
Eff (ClassBuilder : r) a -> Eff r a
classBuilderToState = EffectHandler ClassBuilder r -> Eff (ClassBuilder : r) a -> Eff r a
forall (e :: Effect) (es :: [Effect]) a.
(HasCallStack, DispatchOf e ~ 'Dynamic) =>
EffectHandler e es -> Eff (e : es) a -> Eff es a
interpret (EffectHandler ClassBuilder r
 -> Eff (ClassBuilder : r) a -> Eff r a)
-> EffectHandler ClassBuilder r
-> Eff (ClassBuilder : r) a
-> Eff r a
forall a b. (a -> b) -> a -> b
$ \LocalEnv localEs r
_ -> \case
    ModifyClass ClassFile -> ClassFile
f -> (ClassFile -> ClassFile) -> Eff r ()
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ClassFile -> ClassFile
f
    ClassBuilder (Eff localEs) a
GetClass -> Eff r a
forall s (es :: [Effect]).
(HasCallStack, State s :> es) =>
Eff es s
get

{- | Run a class builder effect, returning the completed 'ClassFile' along with the result of the action
Any unset attributes in the class will be 'mempty' or equivalent, except 'name' and 'version' which must be present.
-}
runClassBuilder :: QualifiedClassName -> JVMVersion -> Eff (ClassBuilder : r) a -> Eff r (ClassFile, a)
runClassBuilder :: forall (r :: [Effect]) a.
QualifiedClassName
-> JVMVersion -> Eff (ClassBuilder : r) a -> Eff r (ClassFile, a)
runClassBuilder QualifiedClassName
n JVMVersion
v =
    (Eff r (a, ClassFile) -> Eff r (ClassFile, a))
-> (Eff (ClassBuilder : r) a -> Eff r (a, ClassFile))
-> Eff (ClassBuilder : r) a
-> Eff r (ClassFile, a)
forall a b.
(a -> b)
-> (Eff (ClassBuilder : r) a -> a) -> Eff (ClassBuilder : r) a -> b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        (((a, ClassFile) -> (ClassFile, a))
-> Eff r (a, ClassFile) -> Eff r (ClassFile, a)
forall a b. (a -> b) -> Eff r a -> Eff r b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a, ClassFile) -> (ClassFile, a)
forall a b. (a, b) -> (b, a)
swap)
        ( ClassFile -> Eff (State ClassFile : r) a -> Eff r (a, ClassFile)
forall s (es :: [Effect]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es (a, s)
runState (QualifiedClassName -> JVMVersion -> ClassFile
dummyClass QualifiedClassName
n JVMVersion
v)
            (Eff (State ClassFile : r) a -> Eff r (a, ClassFile))
-> (Eff (ClassBuilder : r) a -> Eff (State ClassFile : r) a)
-> Eff (ClassBuilder : r) a
-> Eff r (a, ClassFile)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Eff (ClassBuilder : State ClassFile : r) a
-> Eff (State ClassFile : r) a
forall (r :: [Effect]) a.
(State ClassFile :> r) =>
Eff (ClassBuilder : r) a -> Eff r a
classBuilderToState
            (Eff (ClassBuilder : State ClassFile : r) a
 -> Eff (State ClassFile : r) a)
-> (Eff (ClassBuilder : r) a
    -> Eff (ClassBuilder : State ClassFile : r) a)
-> Eff (ClassBuilder : r) a
-> Eff (State ClassFile : r) a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Eff (ClassBuilder : r) a
-> Eff (ClassBuilder : State ClassFile : r) a
forall (subEs :: [Effect]) (es :: [Effect]) a.
Subset subEs es =>
Eff subEs a -> Eff es a
inject
        )