{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}

{- | Generate a stack map table for a method.
This process MUST run last in the high level stage,
modifications to the code after this point will invalidate the stack map table and cause invalid class files to be generated.
-}
module H2JVM.Analyse.StackMap (
    calculateStackMapFrames,
    BasicBlock (..),
    Frame (..),
    LocalVariable (..),
    analyseBlockDiff,
    diffFrames,
    splitIntoBasicBlocks,
    topFrame,
    StackMapError (..),
)
where

import Control.Monad (foldM, (>=>))
import Data.Foldable (foldlM)
import Data.List
import Data.List.NonEmpty (NonEmpty (..))
import Data.Map.Strict (Map)
import Data.Maybe (isNothing, mapMaybe, maybeToList)
import Data.Set (Set)
import Data.Text (Text)
import Data.Traversable (for)
import Effectful (Eff, runPureEff)
import Effectful.Error.Static
import Effectful.Reader.Static (Reader, ask, runReader)
import Effectful.State.Static.Local (State, execState, get, gets, modify, put)
import Witch

import Data.List.NonEmpty qualified as NE
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set

import H2JVM.Builder.Label
import H2JVM.ClassFile.AccessFlags (MethodAccessFlag (..))
import H2JVM.ClassFile.Method
import H2JVM.Descriptor (MethodDescriptor (..), returnDescriptorType)
import H2JVM.Instruction
import H2JVM.Internal.Pretty (Pretty (pretty))
import H2JVM.Internal.Raw.Types
import H2JVM.Internal.Util (bug)
import H2JVM.Name (QualifiedClassName)
import H2JVM.Type (FieldType (..), PrimitiveType (..), classInfoTypeToFieldType, fieldTypeToClassInfoType)

{- | A basic block is a sequence of instructions with a single entry and exit point.
Control flow only enters at the beginning and exits at the end.
-}
data BasicBlock = BasicBlock
    { BasicBlock -> Int
index :: Int
    -- ^ The index of this block in the method's block list
    , BasicBlock -> [Instruction]
instructions :: [Instruction]
    -- ^ The instructions in this block (excluding the starting label)
    , BasicBlock -> Maybe Label
start :: Maybe Label
    -- ^ The label at the start of this block, if any
    , BasicBlock -> Maybe Label
end :: Maybe Label
    -- ^ The label that the next block starts with, if any
    }
    deriving (BasicBlock -> BasicBlock -> Bool
(BasicBlock -> BasicBlock -> Bool)
-> (BasicBlock -> BasicBlock -> Bool) -> Eq BasicBlock
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: BasicBlock -> BasicBlock -> Bool
== :: BasicBlock -> BasicBlock -> Bool
$c/= :: BasicBlock -> BasicBlock -> Bool
/= :: BasicBlock -> BasicBlock -> Bool
Eq, Int -> BasicBlock -> ShowS
[BasicBlock] -> ShowS
BasicBlock -> String
(Int -> BasicBlock -> ShowS)
-> (BasicBlock -> String)
-> ([BasicBlock] -> ShowS)
-> Show BasicBlock
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> BasicBlock -> ShowS
showsPrec :: Int -> BasicBlock -> ShowS
$cshow :: BasicBlock -> String
show :: BasicBlock -> String
$cshowList :: [BasicBlock] -> ShowS
showList :: [BasicBlock] -> ShowS
Show)

-- | Represents the JVM frame state at a particular program point.
data Frame = Frame
    { Frame -> [LocalVariable]
locals :: [LocalVariable]
    -- ^ The local variable slots
    , Frame -> [StackEntry]
stack :: [StackEntry]
    -- ^ The operand stack, with the top of stack at the head
    }
    deriving (Frame -> Frame -> Bool
(Frame -> Frame -> Bool) -> (Frame -> Frame -> Bool) -> Eq Frame
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Frame -> Frame -> Bool
== :: Frame -> Frame -> Bool
$c/= :: Frame -> Frame -> Bool
/= :: Frame -> Frame -> Bool
Eq, Int -> Frame -> ShowS
[Frame] -> ShowS
Frame -> String
(Int -> Frame -> ShowS)
-> (Frame -> String) -> ([Frame] -> ShowS) -> Show Frame
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Frame -> ShowS
showsPrec :: Int -> Frame -> ShowS
$cshow :: Frame -> String
show :: Frame -> String
$cshowList :: [Frame] -> ShowS
showList :: [Frame] -> ShowS
Show)

-- | Represents the type of a local variable slot in the stack frame
data LocalVariable
    = -- | The slot has not been initialised or contains an unusable value
      Uninitialised
    | -- | The slot contains a value of the given type
      LocalVariable FieldType
    deriving (LocalVariable -> LocalVariable -> Bool
(LocalVariable -> LocalVariable -> Bool)
-> (LocalVariable -> LocalVariable -> Bool) -> Eq LocalVariable
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: LocalVariable -> LocalVariable -> Bool
== :: LocalVariable -> LocalVariable -> Bool
$c/= :: LocalVariable -> LocalVariable -> Bool
/= :: LocalVariable -> LocalVariable -> Bool
Eq, Int -> LocalVariable -> ShowS
[LocalVariable] -> ShowS
LocalVariable -> String
(Int -> LocalVariable -> ShowS)
-> (LocalVariable -> String)
-> ([LocalVariable] -> ShowS)
-> Show LocalVariable
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> LocalVariable -> ShowS
showsPrec :: Int -> LocalVariable -> ShowS
$cshow :: LocalVariable -> String
show :: LocalVariable -> String
$cshowList :: [LocalVariable] -> ShowS
showList :: [LocalVariable] -> ShowS
Show)

instance Pretty LocalVariable where
    pretty :: forall ann. LocalVariable -> Doc ann
pretty LocalVariable
Uninitialised = Doc ann
"uninitialised"
    pretty (LocalVariable FieldType
ft) = FieldType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. FieldType -> Doc ann
pretty FieldType
ft

-- | Represents the type of a stack slot.
data StackEntry
    = -- | A typed value on the stack
      StackEntry FieldType
    | -- | An unusable or uninitialised stack slot
      StackEntryTop
    | -- | A null reference on the stack
      StackEntryNull
    deriving (StackEntry -> StackEntry -> Bool
(StackEntry -> StackEntry -> Bool)
-> (StackEntry -> StackEntry -> Bool) -> Eq StackEntry
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: StackEntry -> StackEntry -> Bool
== :: StackEntry -> StackEntry -> Bool
$c/= :: StackEntry -> StackEntry -> Bool
/= :: StackEntry -> StackEntry -> Bool
Eq, Int -> StackEntry -> ShowS
[StackEntry] -> ShowS
StackEntry -> String
(Int -> StackEntry -> ShowS)
-> (StackEntry -> String)
-> ([StackEntry] -> ShowS)
-> Show StackEntry
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StackEntry -> ShowS
showsPrec :: Int -> StackEntry -> ShowS
$cshow :: StackEntry -> String
show :: StackEntry -> String
$cshowList :: [StackEntry] -> ShowS
showList :: [StackEntry] -> ShowS
Show)

instance Pretty StackEntry where
    pretty :: forall ann. StackEntry -> Doc ann
pretty StackEntry
StackEntryTop = Doc ann
"top"
    pretty StackEntry
StackEntryNull = Doc ann
"null"
    pretty (StackEntry FieldType
ft) = FieldType -> Doc ann
forall a ann. Pretty a => a -> Doc ann
forall ann. FieldType -> Doc ann
pretty FieldType
ft

-- | Convert a local variable to a stack entry
lvToStackEntry :: LocalVariable -> StackEntry
lvToStackEntry :: LocalVariable -> StackEntry
lvToStackEntry LocalVariable
Uninitialised = StackEntry
StackEntryTop
lvToStackEntry (LocalVariable FieldType
ft) = FieldType -> StackEntry
StackEntry FieldType
ft

-- | Convert a stack entry to a local variable
stackEntryToLV :: StackEntry -> LocalVariable
stackEntryToLV :: StackEntry -> LocalVariable
stackEntryToLV StackEntry
StackEntryTop = LocalVariable
Uninitialised
stackEntryToLV StackEntry
StackEntryNull = LocalVariable
Uninitialised
stackEntryToLV (StackEntry FieldType
ft) = FieldType -> LocalVariable
LocalVariable FieldType
ft

-- | Split a list of instructions into basic blocks, by splitting at labels and branch instructions.
splitIntoBasicBlocks :: HasCallStack => NonEmpty Instruction -> NonEmpty BasicBlock
splitIntoBasicBlocks :: HasCallStack => NonEmpty Instruction -> NonEmpty BasicBlock
splitIntoBasicBlocks NonEmpty Instruction
l =
    let blockData :: [(Maybe Label, [Instruction])]
blockData = [Instruction] -> [(Maybe Label, [Instruction])]
splitOnLabels (NonEmpty Instruction -> [Instruction]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty Instruction
l)
        startLabels :: [Maybe Label]
startLabels = ((Maybe Label, [Instruction]) -> Maybe Label)
-> [(Maybe Label, [Instruction])] -> [Maybe Label]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe Label, [Instruction]) -> Maybe Label
forall a b. (a, b) -> a
fst [(Maybe Label, [Instruction])]
blockData
        instructions :: [[Instruction]]
instructions = ((Maybe Label, [Instruction]) -> [Instruction])
-> [(Maybe Label, [Instruction])] -> [[Instruction]]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Maybe Label, [Instruction]) -> [Instruction]
forall a b. (a, b) -> b
snd [(Maybe Label, [Instruction])]
blockData
        endLabels :: [Maybe Label]
endLabels = Int -> [Maybe Label] -> [Maybe Label]
forall a. Int -> [a] -> [a]
drop Int
1 [Maybe Label]
startLabels [Maybe Label] -> [Maybe Label] -> [Maybe Label]
forall a. [a] -> [a] -> [a]
++ [Maybe Label
forall a. Maybe a
Nothing]
        blocksList :: [BasicBlock]
blocksList = (Int -> [Instruction] -> Maybe Label -> Maybe Label -> BasicBlock)
-> [Int]
-> [[Instruction]]
-> [Maybe Label]
-> [Maybe Label]
-> [BasicBlock]
forall a b c d e.
(a -> b -> c -> d -> e) -> [a] -> [b] -> [c] -> [d] -> [e]
zipWith4 Int -> [Instruction] -> Maybe Label -> Maybe Label -> BasicBlock
BasicBlock [Int
0 ..] [[Instruction]]
instructions [Maybe Label]
startLabels [Maybe Label]
endLabels
     in case [BasicBlock]
blocksList of
            (BasicBlock
b : [BasicBlock]
bs) -> BasicBlock
b BasicBlock -> [BasicBlock] -> NonEmpty BasicBlock
forall a. a -> [a] -> NonEmpty a
:| [BasicBlock]
bs
            [] -> String -> NonEmpty BasicBlock
forall a. HasCallStack => String -> a
bug String
"splitIntoBasicBlocks: splitOnLabels produced 0 blocks from NonEmpty input"

-- | Build a map from labels to the index of the block that starts with that label.
buildLabelToBlockMap :: NonEmpty BasicBlock -> Map Label Int
buildLabelToBlockMap :: NonEmpty BasicBlock -> Map Label Int
buildLabelToBlockMap NonEmpty BasicBlock
blocks =
    [(Label, Int)] -> Map Label Int
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
        [(Label
label, BasicBlock
block.index) | BasicBlock
block <- NonEmpty BasicBlock -> [BasicBlock]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty BasicBlock
blocks, Just Label
label <- [BasicBlock
block.start]]

{- | Split instructions into groups, each starting with an optional label.
Also splits after branch instructions to ensure proper block boundaries.
-}
splitOnLabels :: [Instruction] -> [(Maybe Label, [Instruction])]
splitOnLabels :: [Instruction] -> [(Maybe Label, [Instruction])]
splitOnLabels [Instruction]
xs = [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [Instruction]
xs [] Maybe Label
forall a. Maybe a
Nothing
  where
    go :: [Instruction] -> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
    go :: [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [] [Instruction]
acc Maybe Label
label
        | [Instruction] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Instruction]
acc Bool -> Bool -> Bool
&& Maybe Label -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Label
label = []
        | Bool
otherwise = [(Maybe Label
label, [Instruction] -> [Instruction]
forall a. [a] -> [a]
reverse [Instruction]
acc)]
    go (Instruction
x : [Instruction]
xs) [Instruction]
acc Maybe Label
label = case Instruction
x of
        Label Label
l' ->
            if [Instruction] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Instruction]
acc Bool -> Bool -> Bool
&& Maybe Label -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Label
label
                then [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [Instruction]
xs [] (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
l')
                else (Maybe Label
label, [Instruction] -> [Instruction]
forall a. [a] -> [a]
reverse [Instruction]
acc) (Maybe Label, [Instruction])
-> [(Maybe Label, [Instruction])] -> [(Maybe Label, [Instruction])]
forall a. a -> [a] -> [a]
: [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [Instruction]
xs [] (Label -> Maybe Label
forall a. a -> Maybe a
Just Label
l')
        Instruction
_other ->
            let acc' :: [Instruction]
acc' = Instruction
x Instruction -> [Instruction] -> [Instruction]
forall a. a -> [a] -> [a]
: [Instruction]
acc
             in if Instruction -> Bool
isBranchInstruction Instruction
x
                    then (Maybe Label
label, [Instruction] -> [Instruction]
forall a. [a] -> [a]
reverse [Instruction]
acc') (Maybe Label, [Instruction])
-> [(Maybe Label, [Instruction])] -> [(Maybe Label, [Instruction])]
forall a. a -> [a] -> [a]
: [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [Instruction]
xs [] Maybe Label
forall a. Maybe a
Nothing
                    else [Instruction]
-> [Instruction] -> Maybe Label -> [(Maybe Label, [Instruction])]
go [Instruction]
xs [Instruction]
acc' Maybe Label
label

    isBranchInstruction :: Instruction -> Bool
    isBranchInstruction :: Instruction -> Bool
isBranchInstruction (If IfCond Label
_) = Bool
True
    isBranchInstruction (IfICmp IfCond Label
_) = Bool
True
    isBranchInstruction (Goto Label
_) = Bool
True
    isBranchInstruction Instruction
AReturn = Bool
True
    isBranchInstruction Instruction
Return = Bool
True
    isBranchInstruction Instruction
_ = Bool
False

{- | Compute the initial frame state at method entry (the "top frame"),
based on the method descriptor and access flags (static vs instance).
-}
topFrame :: QualifiedClassName -> [MethodAccessFlag] -> MethodDescriptor -> Frame
topFrame :: QualifiedClassName
-> [MethodAccessFlag] -> MethodDescriptor -> Frame
topFrame QualifiedClassName
thisType [MethodAccessFlag]
flags (MethodDescriptor [FieldType]
args ReturnDescriptor
_) =
    [LocalVariable] -> [StackEntry] -> Frame
Frame ([FieldType] -> [LocalVariable]
expandLocals ([FieldType] -> [LocalVariable]) -> [FieldType] -> [LocalVariable]
forall a b. (a -> b) -> a -> b
$ [MethodAccessFlag] -> [FieldType] -> [FieldType]
adjustForThis [MethodAccessFlag]
flags [FieldType]
args) []
  where
    adjustForThis :: [MethodAccessFlag] -> [FieldType] -> [FieldType]
    adjustForThis :: [MethodAccessFlag] -> [FieldType] -> [FieldType]
adjustForThis [MethodAccessFlag]
flags [FieldType]
params =
        if MethodAccessFlag
MStatic MethodAccessFlag -> [MethodAccessFlag] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [MethodAccessFlag]
flags
            then [FieldType]
params
            else QualifiedClassName -> FieldType
ObjectFieldType QualifiedClassName
thisType FieldType -> [FieldType] -> [FieldType]
forall a. a -> [a] -> [a]
: [FieldType]
params -- instance methods have @this@ as a local
    expandLocals :: [FieldType] -> [LocalVariable]
    expandLocals :: [FieldType] -> [LocalVariable]
expandLocals [] = []
    expandLocals (PrimitiveFieldType PrimitiveType
JDouble : [FieldType]
ts) =
        FieldType -> LocalVariable
LocalVariable (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JDouble) LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: LocalVariable
Uninitialised LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [FieldType] -> [LocalVariable]
expandLocals [FieldType]
ts
    expandLocals (PrimitiveFieldType PrimitiveType
JLong : [FieldType]
ts) =
        FieldType -> LocalVariable
LocalVariable (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JLong) LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: LocalVariable
Uninitialised LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [FieldType] -> [LocalVariable]
expandLocals [FieldType]
ts
    expandLocals (FieldType
t : [FieldType]
ts) = FieldType -> LocalVariable
LocalVariable FieldType
t LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [FieldType] -> [LocalVariable]
expandLocals [FieldType]
ts

-- | Remove any implicit 'top' entries that are used for padding wide types in the locals list
filterImplicitTops :: [LocalVariable] -> [LocalVariable]
filterImplicitTops :: [LocalVariable] -> [LocalVariable]
filterImplicitTops [] = []
filterImplicitTops (LocalVariable (PrimitiveFieldType PrimitiveType
JDouble) : LocalVariable
Uninitialised : [LocalVariable]
rest) =
    FieldType -> LocalVariable
LocalVariable (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JDouble) LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
rest
filterImplicitTops (LocalVariable (PrimitiveFieldType PrimitiveType
JLong) : LocalVariable
Uninitialised : [LocalVariable]
rest) =
    FieldType -> LocalVariable
LocalVariable (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JLong) LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
rest
filterImplicitTops (LocalVariable
x : [LocalVariable]
xs) = LocalVariable
x LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. a -> [a] -> [a]
: [LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
xs

-- | Compute the frame state after executing all instructions in a basic block.
analyseBlockDiff :: HasCallStack => Frame -> BasicBlock -> Either StackMapError Frame
analyseBlockDiff :: HasCallStack => Frame -> BasicBlock -> Either StackMapError Frame
analyseBlockDiff Frame
current BasicBlock
block = (Frame -> Instruction -> Either StackMapError Frame)
-> Frame -> [Instruction] -> Either StackMapError Frame
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ((Instruction -> Frame -> Either StackMapError Frame)
-> Frame -> Instruction -> Either StackMapError Frame
forall a b c. (a -> b -> c) -> b -> a -> c
flip HasCallStack => Instruction -> Frame -> Either StackMapError Frame
Instruction -> Frame -> Either StackMapError Frame
analyseInstruction) Frame
current BasicBlock
block.instructions
  where
    analyseInstruction :: HasCallStack => Instruction -> Frame -> Either StackMapError Frame
    analyseInstruction :: HasCallStack => Instruction -> Frame -> Either StackMapError Frame
analyseInstruction Instruction
inst Frame
frame =
        Eff '[] (Either StackMapError Frame) -> Either StackMapError Frame
forall a. HasCallStack => Eff '[] a -> a
runPureEff (Eff '[] (Either StackMapError Frame)
 -> Either StackMapError Frame)
-> Eff '[] (Either StackMapError Frame)
-> Either StackMapError Frame
forall a b. (a -> b) -> a -> b
$
            forall e (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
Eff (Error e : es) a -> Eff es (Either e a)
runErrorNoCallStack @StackMapError (Eff '[Error StackMapError] Frame
 -> Eff '[] (Either StackMapError Frame))
-> Eff '[Error StackMapError] Frame
-> Eff '[] (Either StackMapError Frame)
forall a b. (a -> b) -> a -> b
$
                BasicBlock
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
-> Eff '[Error StackMapError] Frame
forall r (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
r -> Eff (Reader r : es) a -> Eff es a
runReader BasicBlock
block (Eff '[Reader BasicBlock, Error StackMapError] Frame
 -> Eff '[Error StackMapError] Frame)
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
-> Eff '[Error StackMapError] Frame
forall a b. (a -> b) -> a -> b
$
                    Frame
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
forall s (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es s
execState Frame
frame (Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
 -> Eff '[Reader BasicBlock, Error StackMapError] Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
forall a b. (a -> b) -> a -> b
$
                        HasCallStack =>
Instruction
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
Instruction
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
analyse Instruction
inst

-- | analyse a single instruction's effect on the stack and locals
analyse :: HasCallStack => Instruction -> Analyser
analyse :: HasCallStack =>
Instruction
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
analyse = \case
    (Label Label
_) -> String
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a. HasCallStack => String -> a
bug String
"Label should not be encountered in analyseInstruction"
    inst :: Instruction
inst@Instruction
AALoad -> do
        s <- (Frame -> [StackEntry])
-> Eff
     '[State Frame, Reader BasicBlock, Error StackMapError] [StackEntry]
forall s (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, State s :> es) =>
(s -> a) -> Eff es a
gets (.stack)
        case s of
            (StackEntry
_index : StackEntry (ArrayFieldType FieldType
innerType) : [StackEntry]
_) -> do
                Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2 -- index, arrayref
                FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
innerType
            (StackEntry
_index : StackEntry
StackEntryTop : [StackEntry]
_) ->
                StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (StackMapError
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ Instruction -> [StackEntry] -> Maybe Text -> StackMapError
InvalidStackState Instruction
inst [StackEntry]
s (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"AAload: arrayref is Uninitialised (Top)")
            (StackEntry
_index : StackEntry
StackEntryNull : [StackEntry]
_) -> do
                -- this will npe at runtime but isn't technically invalid
                Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
                StackEntry
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesEntry StackEntry
StackEntryTop
            [StackEntry]
_ -> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (StackMapError
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ Instruction -> StackMapError
StackUnderflow Instruction
inst
    Instruction
ArrayLength -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1 -- arrayref
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    inst :: Instruction
inst@(ALoad U2
i) -> Instruction
-> U2
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
loads Instruction
inst U2
i
    inst :: Instruction
inst@(ILoad U2
i) -> Instruction
-> U2
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
loads Instruction
inst U2
i
    inst :: Instruction
inst@(AStore U2
i) -> Instruction
-> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
stores Instruction
inst (U2 -> Int
forall target source. From source target => source -> target
into U2
i)
    inst :: Instruction
inst@(IStore U2
i) -> Instruction
-> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
stores Instruction
inst (U2 -> Int
forall target source. From source target => source -> target
into U2
i)
    Instruction
AReturn -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
    Instruction
IReturn -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
    Instruction
AConstNull -> StackEntry
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesEntry StackEntry
StackEntryNull
    Instruction
Return -> () -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a.
a -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    LDC LDCEntry
t -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (LDCEntry -> FieldType
ldcEntryToFieldType LDCEntry
t)
    Instruction
IConst0 -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    Instruction
IConst1 -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    inst :: Instruction
inst@Instruction
Dup -> do
        s <- (Frame -> [StackEntry])
-> Eff
     '[State Frame, Reader BasicBlock, Error StackMapError] [StackEntry]
forall s (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, State s :> es) =>
(s -> a) -> Eff es a
gets (.stack)
        case s of
            [] -> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (StackMapError
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ Instruction -> StackMapError
StackUnderflow Instruction
inst
            StackEntry
head : [StackEntry]
_ -> StackEntry
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesEntry StackEntry
head
    Instruction
IAnd -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    Instruction
IOr -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    If IfCond Label
_ -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
    CheckCast ClassInfoType
ft -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
replaceTop (ClassInfoType -> FieldType
classInfoTypeToFieldType ClassInfoType
ft)
    Instanceof ClassInfoType
_ -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
replaceTop (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    InvokeStatic ClassInfoType
_ Text
_ MethodDescriptor
md -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops ([FieldType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MethodDescriptor
md.params)
        Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
    InvokeVirtual ClassInfoType
_ Text
_ MethodDescriptor
md -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [FieldType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MethodDescriptor
md.params)
        Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
    InvokeInterface ClassInfoType
_ Text
_ MethodDescriptor
md -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [FieldType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MethodDescriptor
md.params)
        Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
    InvokeDynamic BootstrapMethod
_ Text
_ MethodDescriptor
md -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops ([FieldType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MethodDescriptor
md.params)
        Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
    InvokeSpecial ClassInfoType
_ Text
_ MethodDescriptor
md -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ [FieldType] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length MethodDescriptor
md.params)
        Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
    PutStatic{} -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
    GetField ClassInfoType
_ Text
_ FieldType
ft -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
ft
    GetStatic ClassInfoType
_ Text
_ FieldType
ft -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
ft
    PutField{} -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
    Goto Label
_ -> () -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a.
a -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    New ClassInfoType
t -> FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (ClassInfoType -> FieldType
classInfoTypeToFieldType ClassInfoType
t)
    IfICmp IfCond Label
_cmp -> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
    Instruction
IAdd -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    Instruction
ISub -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    Instruction
IMul -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)
    Instruction
IDiv -> do
        Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
2
        FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
JInt)

-- | Compute the delta between two frames to produce a StackMapFrame
diffFrames :: Frame -> Frame -> Label -> StackMapFrame
diffFrames :: Frame -> Frame -> Label -> StackMapFrame
diffFrames (Frame [LocalVariable]
locals1 [StackEntry]
_stack1) (Frame [LocalVariable]
locals2 [StackEntry]
stack2) Label
label
    | [LocalVariable]
locals1 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => a -> a -> Bool
== [LocalVariable]
locals2 Bool -> Bool -> Bool
&& [StackEntry] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StackEntry]
stack2 = Label -> StackMapFrame
SameFrame Label
label
    -- same locals, one stack item
    | [StackEntry
x] <- [StackEntry]
stack2, [LocalVariable]
locals1 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => a -> a -> Bool
== [LocalVariable]
locals2 = VerificationTypeInfo -> Label -> StackMapFrame
SameLocals1StackItemFrame (StackEntry -> VerificationTypeInfo
seToVerificationTypeInfo StackEntry
x) Label
label
    -- stack empty, locals appended
    | [StackEntry] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StackEntry]
stack2
        Bool -> Bool -> Bool
&& [LocalVariable]
locals1 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [LocalVariable]
locals2
        Bool -> Bool -> Bool
&& let diff :: Int
diff = [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
locals2) Int -> Int -> Int
forall a. Num a => a -> a -> a
- [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
locals1) in Int
diff Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
diff Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3 =
        let difference :: [LocalVariable]
difference = Int -> [LocalVariable] -> [LocalVariable]
forall a. Int -> [a] -> [a]
drop ([LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LocalVariable]
locals1) [LocalVariable]
locals2
         in [VerificationTypeInfo] -> Label -> StackMapFrame
AppendFrame ((LocalVariable -> VerificationTypeInfo)
-> [LocalVariable] -> [VerificationTypeInfo]
forall a b. (a -> b) -> [a] -> [b]
map LocalVariable -> VerificationTypeInfo
lvToVerificationTypeInfo ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
difference)) Label
label
    | [StackEntry] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StackEntry]
stack2
        Bool -> Bool -> Bool
&& [LocalVariable]
locals2 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` [LocalVariable]
locals1
        Bool -> Bool -> Bool
&& let diff :: Int
diff = [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
locals1) Int -> Int -> Int
forall a. Num a => a -> a -> a
- [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
locals2)
            in Int
diff Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0 Bool -> Bool -> Bool
&& Int
diff Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3 =
        U1 -> Label -> StackMapFrame
ChopFrame (Int -> U1
forall target source.
(HasCallStack, TryFrom source target, Show source, Typeable source,
 Typeable target) =>
source -> target
unsafeInto (Int -> U1) -> Int -> U1
forall a b. (a -> b) -> a -> b
$ [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LocalVariable]
locals1 Int -> Int -> Int
forall a. Num a => a -> a -> a
- [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LocalVariable]
locals2) Label
label
    | Bool
otherwise =
        [VerificationTypeInfo]
-> [VerificationTypeInfo] -> Label -> StackMapFrame
FullFrame
            ((LocalVariable -> VerificationTypeInfo)
-> [LocalVariable] -> [VerificationTypeInfo]
forall a b. (a -> b) -> [a] -> [b]
map LocalVariable -> VerificationTypeInfo
lvToVerificationTypeInfo ([LocalVariable] -> [LocalVariable]
filterImplicitTops [LocalVariable]
locals2))
            ((StackEntry -> VerificationTypeInfo)
-> [StackEntry] -> [VerificationTypeInfo]
forall a b. (a -> b) -> [a] -> [b]
map StackEntry -> VerificationTypeInfo
seToVerificationTypeInfo [StackEntry]
stack2)
            Label
label

-- | Convert a local variable to its JVM verification type info representation.
lvToVerificationTypeInfo :: LocalVariable -> VerificationTypeInfo
lvToVerificationTypeInfo :: LocalVariable -> VerificationTypeInfo
lvToVerificationTypeInfo LocalVariable
Uninitialised = VerificationTypeInfo
TopVariableInfo
lvToVerificationTypeInfo (LocalVariable FieldType
ft) = case FieldType
ft of
    PrimitiveFieldType PrimitiveType
JInt -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JByte -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JChar -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JShort -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JBoolean -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JFloat -> VerificationTypeInfo
FloatVariableInfo
    PrimitiveFieldType PrimitiveType
JLong -> VerificationTypeInfo
LongVariableInfo
    PrimitiveFieldType PrimitiveType
JDouble -> VerificationTypeInfo
DoubleVariableInfo
    ObjectFieldType{} -> ClassInfoType -> VerificationTypeInfo
ObjectVariableInfo (FieldType -> ClassInfoType
fieldTypeToClassInfoType FieldType
ft)
    ArrayFieldType{} -> ClassInfoType -> VerificationTypeInfo
ObjectVariableInfo (FieldType -> ClassInfoType
fieldTypeToClassInfoType FieldType
ft)

-- | Convert a stack entry to its JVM verification type info representation.
seToVerificationTypeInfo :: StackEntry -> VerificationTypeInfo
seToVerificationTypeInfo :: StackEntry -> VerificationTypeInfo
seToVerificationTypeInfo StackEntry
StackEntryTop = VerificationTypeInfo
TopVariableInfo
seToVerificationTypeInfo StackEntry
StackEntryNull = VerificationTypeInfo
NullVariableInfo
seToVerificationTypeInfo (StackEntry FieldType
ft) = case FieldType
ft of
    PrimitiveFieldType PrimitiveType
JInt -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JByte -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JChar -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JShort -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JBoolean -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
JFloat -> VerificationTypeInfo
FloatVariableInfo
    PrimitiveFieldType PrimitiveType
JLong -> VerificationTypeInfo
LongVariableInfo
    PrimitiveFieldType PrimitiveType
JDouble -> VerificationTypeInfo
DoubleVariableInfo
    FieldType
_ -> ClassInfoType -> VerificationTypeInfo
ObjectVariableInfo (FieldType -> ClassInfoType
fieldTypeToClassInfoType FieldType
ft)

{- | Merge two frames that could both reach the same program point.
Returns 'Nothing' if frames are identical, 'Just' the merged frame otherwise.
Uses a conservative merge where differing types become 'Uninitialised'.
-}
mergeFrames :: HasCallStack => Frame -> Frame -> Either StackMapError (Maybe Frame)
mergeFrames :: HasCallStack =>
Frame -> Frame -> Either StackMapError (Maybe Frame)
mergeFrames frame1 :: Frame
frame1@(Frame [LocalVariable]
locals1 [StackEntry]
stack1) frame2 :: Frame
frame2@(Frame [LocalVariable]
locals2 [StackEntry]
stack2)
    | [LocalVariable]
locals1 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => a -> a -> Bool
== [LocalVariable]
locals2 Bool -> Bool -> Bool
&& [StackEntry]
stack1 [StackEntry] -> [StackEntry] -> Bool
forall a. Eq a => a -> a -> Bool
== [StackEntry]
stack2 = Maybe Frame -> Either StackMapError (Maybe Frame)
forall a. a -> Either StackMapError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe Frame
forall a. Maybe a
Nothing
    | [StackEntry] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [StackEntry]
stack1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= [StackEntry] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [StackEntry]
stack2 = StackMapError -> Either StackMapError (Maybe Frame)
forall a b. a -> Either a b
Left (StackMapError -> Either StackMapError (Maybe Frame))
-> StackMapError -> Either StackMapError (Maybe Frame)
forall a b. (a -> b) -> a -> b
$ Frame -> Frame -> StackMapError
IncompatibleFrameMerge Frame
frame1 Frame
frame2
    | Bool
otherwise =
        Maybe Frame -> Either StackMapError (Maybe Frame)
forall a. a -> Either StackMapError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe Frame -> Either StackMapError (Maybe Frame))
-> Maybe Frame -> Either StackMapError (Maybe Frame)
forall a b. (a -> b) -> a -> b
$
            Frame -> Maybe Frame
forall a. a -> Maybe a
Just (Frame -> Maybe Frame) -> Frame -> Maybe Frame
forall a b. (a -> b) -> a -> b
$
                Frame
                    { locals :: [LocalVariable]
locals = LocalVariable
-> (LocalVariable -> LocalVariable -> LocalVariable)
-> [LocalVariable]
-> [LocalVariable]
-> [LocalVariable]
forall a. a -> (a -> a -> a) -> [a] -> [a] -> [a]
zipWithDefault LocalVariable
Uninitialised LocalVariable -> LocalVariable -> LocalVariable
mergeLocal [LocalVariable]
locals1 [LocalVariable]
locals2
                    , stack :: [StackEntry]
stack = (StackEntry -> StackEntry -> StackEntry)
-> [StackEntry] -> [StackEntry] -> [StackEntry]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith StackEntry -> StackEntry -> StackEntry
mergeStack [StackEntry]
stack1 [StackEntry]
stack2
                    }
  where
    mergeLocal :: LocalVariable -> LocalVariable -> LocalVariable
    mergeLocal :: LocalVariable -> LocalVariable -> LocalVariable
mergeLocal LocalVariable
Uninitialised LocalVariable
_ = LocalVariable
Uninitialised
    mergeLocal LocalVariable
_ LocalVariable
Uninitialised = LocalVariable
Uninitialised
    mergeLocal LocalVariable
x LocalVariable
y = if LocalVariable
x LocalVariable -> LocalVariable -> Bool
forall a. Eq a => a -> a -> Bool
== LocalVariable
y then LocalVariable
x else LocalVariable
Uninitialised

    mergeStack :: StackEntry -> StackEntry -> StackEntry
    mergeStack :: StackEntry -> StackEntry -> StackEntry
mergeStack StackEntry
x StackEntry
y = if StackEntry
x StackEntry -> StackEntry -> Bool
forall a. Eq a => a -> a -> Bool
== StackEntry
y then StackEntry
x else StackEntry
StackEntryTop

    zipWithDefault :: a -> (a -> a -> a) -> [a] -> [a] -> [a]
    zipWithDefault :: forall a. a -> (a -> a -> a) -> [a] -> [a] -> [a]
zipWithDefault a
def a -> a -> a
f = [a] -> [a] -> [a]
go
      where
        go :: [a] -> [a] -> [a]
go [] [] = []
        go (a
a : [a]
as) [] = a -> a -> a
f a
a a
def a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
go [a]
as []
        go [] (a
b : [a]
bs) = a -> a -> a
f a
def a
b a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
go [] [a]
bs
        go (a
a : [a]
as) (a
b : [a]
bs) = a -> a -> a
f a
a a
b a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a] -> [a] -> [a]
go [a]
as [a]
bs

{- | Get the indices of successor blocks (blocks reachable from this block).
Includes both jump targets and fallthrough to the next block, if the last instruction is not a terminator.
-}
getSuccessors :: Map Label Int -> Int -> BasicBlock -> [Int]
getSuccessors :: Map Label Int -> Int -> BasicBlock -> [Int]
getSuccessors Map Label Int
labelToBlock Int
blockIdx BasicBlock
block =
    let jumpTargetIndices :: [Int]
jumpTargetIndices = (Instruction -> Maybe Int) -> [Instruction] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Instruction -> Maybe Label
forall label. Instruction' label -> Maybe label
jumpTarget (Instruction -> Maybe Label)
-> (Label -> Maybe Int) -> Instruction -> Maybe Int
forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> (Label -> Map Label Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
`Map.lookup` Map Label Int
labelToBlock)) BasicBlock
block.instructions
        lastInst :: Maybe Instruction
lastInst = if [Instruction] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null BasicBlock
block.instructions then Maybe Instruction
forall a. Maybe a
Nothing else Instruction -> Maybe Instruction
forall a. a -> Maybe a
Just ([Instruction] -> Instruction
forall a. HasCallStack => [a] -> a
last BasicBlock
block.instructions)
        fallThroughIdx :: Maybe Int
fallThroughIdx = if Maybe Instruction -> Bool
forall {label}. Maybe (Instruction' label) -> Bool
isTerminator Maybe Instruction
lastInst then Maybe Int
forall a. Maybe a
Nothing else Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
blockIdx Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1)
     in [Int]
jumpTargetIndices [Int] -> [Int] -> [Int]
forall a. Semigroup a => a -> a -> a
<> Maybe Int -> [Int]
forall a. Maybe a -> [a]
maybeToList Maybe Int
fallThroughIdx
  where
    isTerminator :: Maybe (Instruction' label) -> Bool
isTerminator (Just (Goto label
_)) = Bool
True
    isTerminator (Just Instruction' label
AReturn) = Bool
True
    isTerminator (Just Instruction' label
Return) = Bool
True
    isTerminator Maybe (Instruction' label)
_ = Bool
False

{- | Compute the frame state at the entry of each basic block using a worklist algorithm.
Handles control flow merges by conservatively merging frame states with 'mergeFrames'.
-}
computeBlockFrames :: HasCallStack => Frame -> NonEmpty BasicBlock -> Either StackMapError (Map Int Frame)
computeBlockFrames :: HasCallStack =>
Frame
-> NonEmpty BasicBlock -> Either StackMapError (Map Int Frame)
computeBlockFrames Frame
initialFrame NonEmpty BasicBlock
blocks = do
    let labelToBlock :: Map Label Int
labelToBlock = NonEmpty BasicBlock -> Map Label Int
buildLabelToBlockMap NonEmpty BasicBlock
blocks
        blockArray :: Map Int BasicBlock
blockArray = [(Int, BasicBlock)] -> Map Int BasicBlock
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(BasicBlock
b.index, BasicBlock
b) | BasicBlock
b <- NonEmpty BasicBlock -> [BasicBlock]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty BasicBlock
blocks]
        numBlocks :: Int
numBlocks = [BasicBlock] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (NonEmpty BasicBlock -> [BasicBlock]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty BasicBlock
blocks)
        initialFrames :: Map Int Frame
initialFrames = Int -> Frame -> Map Int Frame
forall k a. k -> a -> Map k a
Map.singleton Int
0 Frame
initialFrame
        initialWorklist :: Set Int
initialWorklist = Int -> Set Int
forall a. a -> Set a
Set.singleton Int
0
     in HasCallStack =>
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> Either StackMapError (Map Int Frame)
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> Either StackMapError (Map Int Frame)
worklistLoop Set Int
initialWorklist Map Int Frame
initialFrames Map Label Int
labelToBlock Map Int BasicBlock
blockArray Int
numBlocks
  where
    worklistLoop :: HasCallStack => Set Int -> Map Int Frame -> Map Label Int -> Map Int BasicBlock -> Int -> Either StackMapError (Map Int Frame)
    worklistLoop :: HasCallStack =>
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> Either StackMapError (Map Int Frame)
worklistLoop Set Int
worklist Map Int Frame
frames Map Label Int
labelToBlock Map Int BasicBlock
blockArray Int
numBlocks
        | Set Int -> Bool
forall a. Set a -> Bool
Set.null Set Int
worklist = Map Int Frame -> Either StackMapError (Map Int Frame)
forall a. a -> Either StackMapError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Map Int Frame
frames
        | Bool
otherwise = do
            let (Int
blockIdx, Set Int
worklist') = Set Int -> (Int, Set Int)
forall a. Set a -> (a, Set a)
Set.deleteFindMin Set Int
worklist
                block :: BasicBlock
block = Map Int BasicBlock
blockArray Map Int BasicBlock -> Int -> BasicBlock
forall k a. Ord k => Map k a -> k -> a
Map.! Int
blockIdx
                inputFrame :: Frame
inputFrame = Map Int Frame
frames Map Int Frame -> Int -> Frame
forall k a. Ord k => Map k a -> k -> a
Map.! Int
blockIdx
            outputFrame <- HasCallStack => Frame -> BasicBlock -> Either StackMapError Frame
Frame -> BasicBlock -> Either StackMapError Frame
analyseBlockDiff Frame
inputFrame BasicBlock
block
            let successorIndices = (Int -> Bool) -> [Int] -> [Int]
forall a. (a -> Bool) -> [a] -> [a]
filter (Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
numBlocks) ([Int] -> [Int]) -> [Int] -> [Int]
forall a b. (a -> b) -> a -> b
$ Map Label Int -> Int -> BasicBlock -> [Int]
getSuccessors Map Label Int
labelToBlock Int
blockIdx BasicBlock
block
            (frames', worklist'') <- foldlM (propagateFrame outputFrame) (frames, worklist') successorIndices
            worklistLoop worklist'' frames' labelToBlock blockArray numBlocks

    propagateFrame :: Frame -> (Map Int Frame, Set Int) -> Int -> Either StackMapError (Map Int Frame, Set Int)
    propagateFrame :: Frame
-> (Map Int Frame, Set Int)
-> Int
-> Either StackMapError (Map Int Frame, Set Int)
propagateFrame Frame
outFrame (Map Int Frame
frames, Set Int
worklist) Int
succIdx =
        case Int -> Map Int Frame -> Maybe Frame
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Int
succIdx Map Int Frame
frames of
            Maybe Frame
Nothing ->
                (Map Int Frame, Set Int)
-> Either StackMapError (Map Int Frame, Set Int)
forall a b. b -> Either a b
Right (Int -> Frame -> Map Int Frame -> Map Int Frame
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Int
succIdx Frame
outFrame Map Int Frame
frames, Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
succIdx Set Int
worklist)
            Just Frame
existingFrame -> do
                merged <- HasCallStack =>
Frame -> Frame -> Either StackMapError (Maybe Frame)
Frame -> Frame -> Either StackMapError (Maybe Frame)
mergeFrames Frame
existingFrame Frame
outFrame
                pure $ case merged of
                    Maybe Frame
Nothing -> (Map Int Frame
frames, Set Int
worklist)
                    Just Frame
mergedFrame ->
                        (Int -> Frame -> Map Int Frame -> Map Int Frame
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Int
succIdx Frame
mergedFrame Map Int Frame
frames, Int -> Set Int -> Set Int
forall a. Ord a => a -> Set a -> Set a
Set.insert Int
succIdx Set Int
worklist)

{- | Calculate the stack map frames for a method's bytecode.
This function works by first splitting the code into basic blocks,
then computing the frame state at the start of each block using a dataflow analysis,
and finally generating the stack map frames by comparing the frame states at block entries.
-}
calculateStackMapFrames ::
    HasCallStack =>
    -- | The class containing this method
    QualifiedClassName ->
    -- | Method access flags (to determine if static)
    [MethodAccessFlag] ->
    -- | The method descriptor
    MethodDescriptor ->
    -- | The method's instruction list
    NonEmpty Instruction ->
    -- | (stack map frames, max stack, max locals)
    Either StackMapError ([StackMapFrame], Int, Int)
calculateStackMapFrames :: HasCallStack =>
QualifiedClassName
-> [MethodAccessFlag]
-> MethodDescriptor
-> NonEmpty Instruction
-> Either StackMapError ([StackMapFrame], Int, Int)
calculateStackMapFrames QualifiedClassName
enclosingClassName [MethodAccessFlag]
flags MethodDescriptor
md NonEmpty Instruction
code = do
    let blocks :: NonEmpty BasicBlock
blocks = HasCallStack => NonEmpty Instruction -> NonEmpty BasicBlock
NonEmpty Instruction -> NonEmpty BasicBlock
splitIntoBasicBlocks NonEmpty Instruction
code
    let top :: Frame
top = QualifiedClassName
-> [MethodAccessFlag] -> MethodDescriptor -> Frame
topFrame QualifiedClassName
enclosingClassName [MethodAccessFlag]
flags MethodDescriptor
md

    blockFrames <- HasCallStack =>
Frame
-> NonEmpty BasicBlock -> Either StackMapError (Map Int Frame)
Frame
-> NonEmpty BasicBlock -> Either StackMapError (Map Int Frame)
computeBlockFrames Frame
top NonEmpty BasicBlock
blocks

    (maxStack, maxLocals) <- calculateMethodMaxes blockFrames blocks

    let labelToBlockIdx = NonEmpty BasicBlock -> Map Label Int
buildLabelToBlockMap NonEmpty BasicBlock
blocks
    let orderedPairs =
            [ (Label
label, Map Int Frame
blockFrames Map Int Frame -> Int -> Frame
forall k a. Ord k => Map k a -> k -> a
Map.! Int
blockIdx)
            | BasicBlock
block <- NonEmpty BasicBlock -> [BasicBlock]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty BasicBlock
blocks
            , Just Label
label <- [BasicBlock
block.start]
            , Just Int
blockIdx <- [Label -> Map Label Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Label
label Map Label Int
labelToBlockIdx]
            , Int -> Map Int Frame -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.member Int
blockIdx Map Int Frame
blockFrames
            , Int
blockIdx Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
0
            ]

    case orderedPairs of
        [] -> ([StackMapFrame], Int, Int)
-> Either StackMapError ([StackMapFrame], Int, Int)
forall a. a -> Either StackMapError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([], Int
maxStack, Int
maxLocals)
        ((Label
firstLabel, Frame
firstFrame) : [(Label, Frame)]
rest) ->
            let firstSMF :: StackMapFrame
firstSMF = Frame -> Frame -> Label -> StackMapFrame
diffFrames Frame
top Frame
firstFrame Label
firstLabel
                restSMFs :: [StackMapFrame]
restSMFs =
                    ((Label, Frame) -> (Label, Frame) -> StackMapFrame)
-> [(Label, Frame)] -> [(Label, Frame)] -> [StackMapFrame]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith
                        (\(Label
_, Frame
prevFrame) (Label
currLabel, Frame
currFrame) -> Frame -> Frame -> Label -> StackMapFrame
diffFrames Frame
prevFrame Frame
currFrame Label
currLabel)
                        [(Label, Frame)]
orderedPairs
                        [(Label, Frame)]
rest
             in ([StackMapFrame], Int, Int)
-> Either StackMapError ([StackMapFrame], Int, Int)
forall a. a -> Either StackMapError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StackMapFrame
firstSMF StackMapFrame -> [StackMapFrame] -> [StackMapFrame]
forall a. a -> [a] -> [a]
: [StackMapFrame]
restSMFs, Int
maxStack, Int
maxLocals)

-- | Replace element at index in list, growing with 'Uninitialised' if needed.
replaceAtOrGrow :: Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow :: Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow Int
i LocalVariable
x [LocalVariable]
xs
    | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LocalVariable]
xs = Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
forall a. Int -> a -> [a] -> [a]
replaceAt Int
i LocalVariable
x [LocalVariable]
xs
    | Bool
otherwise = [LocalVariable]
xs [LocalVariable] -> [LocalVariable] -> [LocalVariable]
forall a. Semigroup a => a -> a -> a
<> Int -> LocalVariable -> [LocalVariable]
forall a. Int -> a -> [a]
replicate (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [LocalVariable]
xs) LocalVariable
Uninitialised [LocalVariable] -> [LocalVariable] -> [LocalVariable]
forall a. Semigroup a => a -> a -> a
<> [LocalVariable
x]

{- | Replace element at index i in list. Appends to end if i is out of bounds.
>>> replaceAt 1 'x' "abc"
"axc"
>>> replaceAt 5 'x' "abc"
"abcx"
-}
replaceAt :: Int -> a -> [a] -> [a]
replaceAt :: forall a. Int -> a -> [a] -> [a]
replaceAt Int
i a
x [a]
xs = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
i [a]
xs [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a
x] [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [a]
xs

{- | Effect stack for analysing instruction effects on the frame state.
We can modify the current frame and read the current basic block.
This monad acts a mini embedded DSL for describing the effect of each instruction on the frame.
-}
type Analyser = Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()

data StackMapError
    = StackUnderflow Instruction
    | InvalidStackState Instruction [StackEntry] (Maybe Text)
    | LocalIndexOutOfBounds Instruction U2 Frame BasicBlock
    | IncompatibleFrameMerge Frame Frame
    | MissingBlockFrame Int
    deriving (Int -> StackMapError -> ShowS
[StackMapError] -> ShowS
StackMapError -> String
(Int -> StackMapError -> ShowS)
-> (StackMapError -> String)
-> ([StackMapError] -> ShowS)
-> Show StackMapError
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> StackMapError -> ShowS
showsPrec :: Int -> StackMapError -> ShowS
$cshow :: StackMapError -> String
show :: StackMapError -> String
$cshowList :: [StackMapError] -> ShowS
showList :: [StackMapError] -> ShowS
Show)

-- | Pops n items off the stack
pops :: Int -> Analyser
pops :: Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
n = (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame)
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ \Frame
f -> Frame
f{stack = drop n f.stack}

-- | Pushes a single type onto the stack
pushes :: FieldType -> Analyser
pushes :: FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
ft = (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame)
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ \Frame
f -> Frame
f{stack = StackEntry ft : f.stack}

-- | Pushes a raw StackEntry
pushesEntry :: StackEntry -> Analyser
pushesEntry :: StackEntry
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesEntry StackEntry
se = (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame)
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> (Frame -> Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ \Frame
f -> Frame
f{stack = se : f.stack}

-- | Pushes an item only if it exists (for void returns)
pushesMaybe :: Maybe FieldType -> Analyser
pushesMaybe :: Maybe FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushesMaybe Maybe FieldType
Nothing = () -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a.
a -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
pushesMaybe (Just FieldType
ft) = FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
ft

-- | Replaces the top stack entry with the given type
replaceTop :: FieldType -> Analyser
replaceTop :: FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
replaceTop FieldType
ft = do
    Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pops Int
1
    FieldType
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
pushes FieldType
ft

-- | Loads a local variable onto the stack
loads :: Instruction -> U2 -> Analyser
loads :: Instruction
-> U2
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
loads Instruction
inst U2
i = do
    frame <- Eff '[State Frame, Reader BasicBlock, Error StackMapError] Frame
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
    block <- ask
    if i >= genericLength frame.locals
        then
            throwError $ LocalIndexOutOfBounds inst i frame block
        else do
            let lv = Frame
frame.locals [LocalVariable] -> Int -> LocalVariable
forall a. HasCallStack => [a] -> Int -> a
!! U2 -> Int
forall target source. From source target => source -> target
into U2
i
            pushesEntry (lvToStackEntry lv)

-- | Stores the top of the stack into a local variable at the given index
stores :: Instruction -> Int -> Analyser
stores :: Instruction
-> Int
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
stores Instruction
inst Int
i = do
    f <- Eff '[State Frame, Reader BasicBlock, Error StackMapError] Frame
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
    case f.stack of
        [] -> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (StackMapError
 -> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ())
-> StackMapError
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall a b. (a -> b) -> a -> b
$ Instruction -> StackMapError
StackUnderflow Instruction
inst
        (StackEntry
top : [StackEntry]
rest) -> do
            let lv :: LocalVariable
lv = StackEntry -> LocalVariable
stackEntryToLV StackEntry
top
            let locals' :: [LocalVariable]
locals' = Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow Int
i LocalVariable
lv Frame
f.locals
            let finalLocals :: [LocalVariable]
finalLocals = case LocalVariable
lv of
                    LocalVariable (PrimitiveFieldType PrimitiveType
JDouble) -> Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) LocalVariable
Uninitialised [LocalVariable]
locals'
                    LocalVariable (PrimitiveFieldType PrimitiveType
JLong) -> Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) LocalVariable
Uninitialised [LocalVariable]
locals'
                    LocalVariable
_ -> [LocalVariable]
locals'
            Frame
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put
                Frame
f
                    { locals = finalLocals
                    , stack = rest
                    }

-- | How many stack slots does a given 'FieldType' take up?
fieldTypeSlotSize :: FieldType -> Int
fieldTypeSlotSize :: FieldType -> Int
fieldTypeSlotSize (PrimitiveFieldType PrimitiveType
JDouble) = Int
2
fieldTypeSlotSize (PrimitiveFieldType PrimitiveType
JLong) = Int
2
fieldTypeSlotSize FieldType
_ = Int
1

-- | How many stack slots does a given 'StackEntry' take up?
seSlotSize :: StackEntry -> Int
seSlotSize :: StackEntry -> Int
seSlotSize (StackEntry FieldType
ft) = FieldType -> Int
fieldTypeSlotSize FieldType
ft
seSlotSize StackEntry
_ = Int
1

-- | How many stack slots does a given 'LocalVariable' take up?
lvSlotSize :: LocalVariable -> Int
lvSlotSize :: LocalVariable -> Int
lvSlotSize (LocalVariable FieldType
ft) = FieldType -> Int
fieldTypeSlotSize FieldType
ft
lvSlotSize LocalVariable
_ = Int
1

-- | Physical size of the operand stack of a 'Frame'
frameStackSize :: Frame -> Int
frameStackSize :: Frame -> Int
frameStackSize Frame
f = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((StackEntry -> Int) -> [StackEntry] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map StackEntry -> Int
seSlotSize Frame
f.stack)

-- | Physical size of the local variables array of a 'Frame'
frameLocalsSize :: Frame -> Int
frameLocalsSize :: Frame -> Int
frameLocalsSize Frame
f =
    case [LocalVariable] -> [LocalVariable]
forall a. [a] -> [a]
reverse Frame
f.locals of
        [] -> Int
0
        (LocalVariable
lastLv : [LocalVariable]
_) -> [LocalVariable] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Frame
f.locals Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ LocalVariable -> Int
lvSlotSize LocalVariable
lastLv

-- | Calculate the max stack and max locals across all intermediate instruction frames.
calculateMethodMaxes ::
    -- | Map from block index to frame at block entry
    Map Int Frame ->
    -- | The method's basic blocks
    NonEmpty BasicBlock ->
    -- | (max stack, max locals)
    Either StackMapError (Int, Int)
calculateMethodMaxes :: Map Int Frame
-> NonEmpty BasicBlock -> Either StackMapError (Int, Int)
calculateMethodMaxes Map Int Frame
blockFrames NonEmpty BasicBlock
blocks = do
    allFrames <- ([[Frame]] -> [Frame])
-> Either StackMapError [[Frame]] -> Either StackMapError [Frame]
forall a b.
(a -> b) -> Either StackMapError a -> Either StackMapError b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap [[Frame]] -> [Frame]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat (Either StackMapError [[Frame]] -> Either StackMapError [Frame])
-> Either StackMapError [[Frame]] -> Either StackMapError [Frame]
forall a b. (a -> b) -> a -> b
$ [BasicBlock]
-> (BasicBlock -> Either StackMapError [Frame])
-> Either StackMapError [[Frame]]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for (NonEmpty BasicBlock -> [BasicBlock]
forall a. NonEmpty a -> [a]
NE.toList NonEmpty BasicBlock
blocks) ((BasicBlock -> Either StackMapError [Frame])
 -> Either StackMapError [[Frame]])
-> (BasicBlock -> Either StackMapError [Frame])
-> Either StackMapError [[Frame]]
forall a b. (a -> b) -> a -> b
$ \BasicBlock
block -> do
        startFrame <-
            Either StackMapError Frame
-> (Frame -> Either StackMapError Frame)
-> Maybe Frame
-> Either StackMapError Frame
forall b a. b -> (a -> b) -> Maybe a -> b
maybe
                (StackMapError -> Either StackMapError Frame
forall a b. a -> Either a b
Left (StackMapError -> Either StackMapError Frame)
-> StackMapError -> Either StackMapError Frame
forall a b. (a -> b) -> a -> b
$ Int -> StackMapError
MissingBlockFrame BasicBlock
block.index)
                Frame -> Either StackMapError Frame
forall a b. b -> Either a b
Right
                (Int -> Map Int Frame -> Maybe Frame
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup BasicBlock
block.index Map Int Frame
blockFrames)
        scanlM
            (\Frame
f Instruction
inst -> Eff '[] (Either StackMapError Frame) -> Either StackMapError Frame
forall a. HasCallStack => Eff '[] a -> a
runPureEff (Eff '[] (Either StackMapError Frame)
 -> Either StackMapError Frame)
-> Eff '[] (Either StackMapError Frame)
-> Either StackMapError Frame
forall a b. (a -> b) -> a -> b
$ Eff '[Error StackMapError] Frame
-> Eff '[] (Either StackMapError Frame)
forall e (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
Eff (Error e : es) a -> Eff es (Either e a)
runErrorNoCallStack (Eff '[Error StackMapError] Frame
 -> Eff '[] (Either StackMapError Frame))
-> Eff '[Error StackMapError] Frame
-> Eff '[] (Either StackMapError Frame)
forall a b. (a -> b) -> a -> b
$ BasicBlock
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
-> Eff '[Error StackMapError] Frame
forall r (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
r -> Eff (Reader r : es) a -> Eff es a
runReader BasicBlock
block (Eff '[Reader BasicBlock, Error StackMapError] Frame
 -> Eff '[Error StackMapError] Frame)
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
-> Eff '[Error StackMapError] Frame
forall a b. (a -> b) -> a -> b
$ Frame
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
forall s (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es s
execState Frame
f (Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
 -> Eff '[Reader BasicBlock, Error StackMapError] Frame)
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
-> Eff '[Reader BasicBlock, Error StackMapError] Frame
forall a b. (a -> b) -> a -> b
$ HasCallStack =>
Instruction
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
Instruction
-> Eff '[State Frame, Reader BasicBlock, Error StackMapError] ()
analyse Instruction
inst)
            startFrame
            block.instructions
    pure
        ( maximum (0 : map frameStackSize allFrames)
        , maximum (0 : map frameLocalsSize allFrames)
        )

scanlM :: Monad m => (b -> a -> m b) -> b -> [a] -> m [b]
scanlM :: forall (m :: * -> *) b a.
Monad m =>
(b -> a -> m b) -> b -> [a] -> m [b]
scanlM b -> a -> m b
_ b
z [] = [b] -> m [b]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [b
z]
scanlM b -> a -> m b
f b
z (a
x : [a]
xs) = do
    z' <- b -> a -> m b
f b
z a
x
    (z :) <$> scanlM f z' xs