{-# 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 JVM.Data.Analyse.StackMap where

import Control.Lens.Fold
import Control.Monad ((>=>))
import Data.Generics.Sum (AsAny (_As))
import Data.List
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (isNothing, mapMaybe, maybeToList)
import Data.Set (Set)
import Data.Set qualified as Set
import Effectful (Eff, runPureEff)
import Effectful.Reader.Static (Reader, ask, runReader)
import Effectful.State.Static.Local (State, execState, get, gets, modify, put)
import GHC.Stack (HasCallStack)
import JVM.Data.Abstract.Builder.Label
import JVM.Data.Abstract.ClassFile.AccessFlags (MethodAccessFlag (..))
import JVM.Data.Abstract.ClassFile.Method
import JVM.Data.Abstract.Descriptor (MethodDescriptor (..), returnDescriptorType)
import JVM.Data.Abstract.Instruction
import JVM.Data.Abstract.Name (QualifiedClassName)
import JVM.Data.Abstract.Type (FieldType (..), PrimitiveType (..), classInfoTypeToFieldType, fieldTypeToClassInfoType)
import JVM.Data.Pretty (Pretty (pretty), showPretty)
import JVM.Data.Raw.Types
import Prettyprinter (vsep)

{- | 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 (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, 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)

{- | Represents the JVM frame state at a particular program point.
Used for computing stack map frames required by the JVM verifier.
-}
data Frame = Frame
    { Frame -> [LocalVariable]
locals :: [LocalVariable]
    -- ^ The local variable slots, indexed from 0
    , Frame -> [StackEntry]
stack :: [StackEntry]
    -- ^ The operand stack, with the top of stack at the head
    }
    deriving (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, 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)

-- | 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 (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, 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)

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 (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, 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)

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.
Blocks are split at labels and after branch/return instructions.
-}
splitIntoBasicBlocks :: [Instruction] -> [BasicBlock]
splitIntoBasicBlocks :: [Instruction] -> [BasicBlock]
splitIntoBasicBlocks [] = []
splitIntoBasicBlocks [Instruction]
l =
    let blockData :: [(Maybe Label, [Instruction])]
blockData = [Instruction] -> [(Maybe Label, [Instruction])]
splitOnLabels [Instruction]
l
        startLabels :: [Maybe Label]
startLabels = ((Maybe Label, [Instruction]) -> Maybe Label)
-> [(Maybe Label, [Instruction])] -> [Maybe Label]
forall a b. (a -> b) -> [a] -> [b]
map (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]
map (Maybe Label, [Instruction]) -> [Instruction]
forall a b. (a, b) -> b
snd [(Maybe Label, [Instruction])]
blockData
        endLabels :: [Maybe Label]
endLabels = [Maybe Label] -> [Maybe Label]
forall a. HasCallStack => [a] -> [a]
tail [Maybe Label]
startLabels [Maybe Label] -> [Maybe Label] -> [Maybe Label]
forall a. [a] -> [a] -> [a]
++ [Maybe Label
forall a. Maybe a
Nothing]
     in (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

-- | Build a map from labels to the index of the block that starts with that label.
buildLabelToBlockMap :: [BasicBlock] -> Map Label Int
buildLabelToBlockMap :: [BasicBlock] -> Map Label Int
buildLabelToBlockMap [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 <- [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 (IfEq Label
_) = Bool
True
    isBranchInstruction (IfNe Label
_) = Bool
True
    isBranchInstruction (IfLt Label
_) = Bool
True
    isBranchInstruction (IfGe Label
_) = Bool
True
    isBranchInstruction (IfGt Label
_) = Bool
True
    isBranchInstruction (IfLe 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.
Includes @this@ reference for instance methods, followed by method parameters.
-}
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) -> [FieldType] -> [LocalVariable]
forall a b. (a -> b) -> [a] -> [b]
map FieldType -> LocalVariable
LocalVariable ([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

{- | Compute the frame state after executing all instructions in a basic block.
Takes the frame state at block entry and returns the state at block exit.
-}
analyseBlockDiff :: (HasCallStack) => Frame -> BasicBlock -> Frame
analyseBlockDiff :: HasCallStack => Frame -> BasicBlock -> Frame
analyseBlockDiff Frame
current BasicBlock
block = (Frame -> Instruction -> Frame) -> Frame -> [Instruction] -> Frame
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((Instruction -> Frame -> Frame) -> Frame -> Instruction -> Frame
forall a b c. (a -> b -> c) -> b -> a -> c
flip HasCallStack => Instruction -> Frame -> Frame
Instruction -> Frame -> Frame
analyseInstruction) Frame
current BasicBlock
block.instructions
  where
    analyseInstruction :: (HasCallStack) => Instruction -> Frame -> Frame
    analyseInstruction :: HasCallStack => Instruction -> Frame -> Frame
analyseInstruction Instruction
inst Frame
frame =
        Eff '[] Frame -> Frame
forall a. HasCallStack => Eff '[] a -> a
runPureEff (Eff '[] Frame -> Frame) -> Eff '[] Frame -> Frame
forall a b. (a -> b) -> a -> b
$
            BasicBlock -> Eff '[Reader BasicBlock] Frame -> Eff '[] Frame
forall r (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
r -> Eff (Reader r : es) a -> Eff es a
runReader BasicBlock
block (Eff '[Reader BasicBlock] Frame -> Eff '[] Frame)
-> Eff '[Reader BasicBlock] Frame -> Eff '[] Frame
forall a b. (a -> b) -> a -> b
$
                Frame
-> Eff '[State Frame, Reader BasicBlock] ()
-> Eff '[Reader BasicBlock] Frame
forall s (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es s
execState Frame
frame (Eff '[State Frame, Reader BasicBlock] ()
 -> Eff '[Reader BasicBlock] Frame)
-> Eff '[State Frame, Reader BasicBlock] ()
-> Eff '[Reader BasicBlock] Frame
forall a b. (a -> b) -> a -> b
$
                    Instruction -> Eff '[State Frame, Reader BasicBlock] ()
analyse Instruction
inst
      where
        analyse :: Instruction -> Analyser
        analyse :: Instruction -> Eff '[State Frame, Reader BasicBlock] ()
analyse = \case
            (Label Label
_) -> String -> Eff '[State Frame, Reader BasicBlock] ()
forall a. HasCallStack => String -> a
error String
"Label should not be encountered in analyseInstruction"
            (ALoad U2
i) -> String -> U2 -> Eff '[State Frame, Reader BasicBlock] ()
loads String
"ALoad" U2
i
            (ILoad U2
i) -> String -> U2 -> Eff '[State Frame, Reader BasicBlock] ()
loads String
"ILoad" U2
i
            (AStore U2
i) -> Int -> Eff '[State Frame, Reader BasicBlock] ()
stores (U2 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral U2
i)
            (IStore U2
i) -> Int -> Eff '[State Frame, Reader BasicBlock] ()
stores (U2 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral U2
i)
            Instruction
AReturn -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            Instruction
IReturn -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            Instruction
AConstNull -> StackEntry -> Eff '[State Frame, Reader BasicBlock] ()
pushesEntry StackEntry
StackEntryNull
            Instruction
Return -> () -> Eff '[State Frame, Reader BasicBlock] ()
forall a. a -> Eff '[State Frame, Reader BasicBlock] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            LDC LDCEntry
t -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (LDCEntry -> FieldType
ldcEntryToFieldType LDCEntry
t)
            Instruction
IConst0 -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
Int)
            Instruction
IConst1 -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
Int)
            Instruction
Dup -> do
                s <- (Frame -> [StackEntry])
-> Eff '[State Frame, Reader BasicBlock] [StackEntry]
forall s (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, State s :> es) =>
(s -> a) -> Eff es a
gets (.stack)
                case s of
                    [] -> String -> Eff '[State Frame, Reader BasicBlock] ()
forall a. HasCallStack => String -> a
error String
"Stack underflow during dup"
                    StackEntry
head : [StackEntry]
_ -> StackEntry -> Eff '[State Frame, Reader BasicBlock] ()
pushesEntry StackEntry
head
            Instruction
IAnd -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
2
                FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
Int)
            Instruction
IOr -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
2
                FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
Int)
            IfEq Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            IfNe Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            IfLt Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            IfGe Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            IfGt Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            IfLe Label
_ -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            CheckCast ClassInfoType
ft -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
replaceTop (ClassInfoType -> FieldType
classInfoTypeToFieldType ClassInfoType
ft)
            Instanceof ClassInfoType
_ -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
replaceTop (PrimitiveType -> FieldType
PrimitiveFieldType PrimitiveType
Int)
            InvokeStatic ClassInfoType
_ Text
_ MethodDescriptor
md -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
            InvokeVirtual ClassInfoType
_ Text
_ MethodDescriptor
md -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
                Int -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
            InvokeInterface ClassInfoType
_ Text
_ MethodDescriptor
md -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
                Int -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
            InvokeDynamic BootstrapMethod
_ Text
_ MethodDescriptor
md -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
            InvokeSpecial ClassInfoType
_ Text
_ MethodDescriptor
md -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
                Int -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe (ReturnDescriptor -> Maybe FieldType
returnDescriptorType MethodDescriptor
md.returnDesc)
            PutStatic{} -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
            GetField ClassInfoType
_ Text
_ FieldType
ft -> do
                Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
1
                FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes FieldType
ft
            GetStatic ClassInfoType
_ Text
_ FieldType
ft -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes FieldType
ft
            PutField{} -> Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
2
            Goto Label
_ -> () -> Eff '[State Frame, Reader BasicBlock] ()
forall a. a -> Eff '[State Frame, Reader BasicBlock] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
            New ClassInfoType
t -> FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes (ClassInfoType -> FieldType
classInfoTypeToFieldType ClassInfoType
t)

-- | 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 (previous stack must be empty)
    | [StackEntry
x] <- [StackEntry]
stack2, [LocalVariable]
locals1 [LocalVariable] -> [LocalVariable] -> Bool
forall a. Eq a => a -> a -> Bool
== [LocalVariable]
locals2, [StackEntry] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [StackEntry]
stack1 = 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]
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]
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]
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
&& [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 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
3 =
        U1 -> Label -> StackMapFrame
ChopFrame (Int -> U1
forall a b. (Integral a, Num b) => a -> b
fromIntegral (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]
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
Int -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Byte -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Char -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Short -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Boolean -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Float -> VerificationTypeInfo
FloatVariableInfo
    PrimitiveFieldType PrimitiveType
Long -> VerificationTypeInfo
LongVariableInfo
    PrimitiveFieldType PrimitiveType
Double -> 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
Int -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Byte -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Char -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Short -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Boolean -> VerificationTypeInfo
IntegerVariableInfo
    PrimitiveFieldType PrimitiveType
Float -> VerificationTypeInfo
FloatVariableInfo
    PrimitiveFieldType PrimitiveType
Long -> VerificationTypeInfo
LongVariableInfo
    PrimitiveFieldType PrimitiveType
Double -> 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: differing types become 'Uninitialised'.
-}
mergeFrames :: Frame -> Frame -> Maybe Frame
mergeFrames :: Frame -> Frame -> Maybe Frame
mergeFrames (Frame [LocalVariable]
locals1 [StackEntry]
stack1) (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
forall a. Maybe a
Nothing
    | Bool
otherwise =
        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 = if [StackEntry]
stack1 [StackEntry] -> [StackEntry] -> Bool
forall a. Eq a => a -> a -> Bool
== [StackEntry]
stack2 then [StackEntry]
stack1 else []
                }
  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

    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 fall-through to the next block (if not terminated).
-}
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. [a] -> [a] -> [a]
++ Maybe Int -> [Int]
forall a. Maybe a -> [a]
Data.Maybe.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.
-}
computeBlockFrames :: (HasCallStack) => Frame -> [BasicBlock] -> Map Int Frame
computeBlockFrames :: HasCallStack => Frame -> [BasicBlock] -> Map Int Frame
computeBlockFrames Frame
initialFrame [BasicBlock]
blocks =
    let labelToBlock :: Map Label Int
labelToBlock = [BasicBlock] -> Map Label Int
buildLabelToBlockMap [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 <- [BasicBlock]
blocks]
        numBlocks :: Int
numBlocks = [BasicBlock] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [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
-> Map Int Frame
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> 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 -> Map Int Frame
    worklistLoop :: HasCallStack =>
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> 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
frames
        | Bool
otherwise =
            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 :: Frame
outputFrame = HasCallStack => Frame -> BasicBlock -> Frame
Frame -> BasicBlock -> Frame
analyseBlockDiff Frame
inputFrame BasicBlock
block
                successorIndices :: [Int]
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
                (Map Int Frame
frames', Set Int
worklist'') = ((Map Int Frame, Set Int) -> Int -> (Map Int Frame, Set Int))
-> (Map Int Frame, Set Int) -> [Int] -> (Map Int Frame, Set Int)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (Frame
-> (Map Int Frame, Set Int) -> Int -> (Map Int Frame, Set Int)
propagateFrame Frame
outputFrame) (Map Int Frame
frames, Set Int
worklist') [Int]
successorIndices
             in HasCallStack =>
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> Map Int Frame
Set Int
-> Map Int Frame
-> Map Label Int
-> Map Int BasicBlock
-> Int
-> Map Int Frame
worklistLoop Set Int
worklist'' Map Int Frame
frames' Map Label Int
labelToBlock Map Int BasicBlock
blockArray Int
numBlocks

    propagateFrame :: Frame -> (Map Int Frame, Set Int) -> Int -> (Map Int Frame, Set Int)
    propagateFrame :: Frame
-> (Map Int Frame, Set Int) -> Int -> (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 ->
                (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 ->
                case Frame -> Frame -> Maybe Frame
mergeFrames Frame
existingFrame Frame
outFrame 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
    [Instruction] ->
    [StackMapFrame]
calculateStackMapFrames :: HasCallStack =>
QualifiedClassName
-> [MethodAccessFlag]
-> MethodDescriptor
-> [Instruction]
-> [StackMapFrame]
calculateStackMapFrames QualifiedClassName
enclosingClassName [MethodAccessFlag]
flags MethodDescriptor
md [Instruction]
code = do
    let blocks :: [BasicBlock]
blocks = [Instruction] -> [BasicBlock]
splitIntoBasicBlocks [Instruction]
code
    let top :: Frame
top = QualifiedClassName
-> [MethodAccessFlag] -> MethodDescriptor -> Frame
topFrame QualifiedClassName
enclosingClassName [MethodAccessFlag]
flags MethodDescriptor
md
    let blockFrames :: Map Int Frame
blockFrames = HasCallStack => Frame -> [BasicBlock] -> Map Int Frame
Frame -> [BasicBlock] -> Map Int Frame
computeBlockFrames Frame
top [BasicBlock]
blocks
    let labelToBlockIdx :: Map Label Int
labelToBlockIdx = [BasicBlock] -> Map Label Int
buildLabelToBlockMap [BasicBlock]
blocks
    let orderedPairs :: [(Label, Frame)]
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 <- [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 [(Label, Frame)]
orderedPairs of
        [] -> []
        ((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
firstSMF StackMapFrame -> [StackMapFrame] -> [StackMapFrame]
forall a. a -> [a] -> [a]
: [StackMapFrame]
restSMFs

-- | 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 :: 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.
-}
type Analyser = Eff '[State Frame, Reader BasicBlock] ()

-- | Pops n items off the stack
pops :: Int -> Analyser
pops :: Int -> Eff '[State Frame, Reader BasicBlock] ()
pops Int
n = (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ())
-> (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushes FieldType
ft = (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ())
-> (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesEntry StackEntry
se = (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
(s -> s) -> Eff es ()
modify ((Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ())
-> (Frame -> Frame) -> Eff '[State Frame, Reader BasicBlock] ()
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] ()
pushesMaybe Maybe FieldType
Nothing = () -> Eff '[State Frame, Reader BasicBlock] ()
forall a. a -> Eff '[State Frame, Reader BasicBlock] a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
pushesMaybe (Just FieldType
ft) = FieldType -> Eff '[State Frame, Reader BasicBlock] ()
pushes FieldType
ft

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

-- | Loads a local variable onto the stack
loads :: String -> U2 -> Analyser
loads :: String -> U2 -> Eff '[State Frame, Reader BasicBlock] ()
loads String
instName U2
i = do
    f <- Eff '[State Frame, Reader BasicBlock] Frame
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
    block <- ask
    if i >= genericLength f.locals
        then
            indexOOBError instName i f block
        else do
            let lv = Frame
f.locals [LocalVariable] -> Int -> LocalVariable
forall a. HasCallStack => [a] -> Int -> a
!! U2 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral U2
i
            pushesEntry (lvToStackEntry lv)

-- | Stores the top of the stack into a local variable at the given index
stores :: Int -> Analyser
stores :: Int -> Eff '[State Frame, Reader BasicBlock] ()
stores Int
i = do
    f <- Eff '[State Frame, Reader BasicBlock] Frame
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
    case f.stack of
        [] -> String -> Eff '[State Frame, Reader BasicBlock] ()
forall a. HasCallStack => String -> a
error String
"Stack underflow during store"
        (StackEntry
top : [StackEntry]
rest) ->
            Frame -> Eff '[State Frame, Reader BasicBlock] ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put
                Frame
f
                    { locals = replaceAtOrGrow i (stackEntryToLV top) f.locals
                    , stack = rest
                    }

-- | Report an error when a local variable index is out of bounds.
indexOOBError :: (HasCallStack) => String -> U2 -> Frame -> BasicBlock -> a
indexOOBError :: forall a. HasCallStack => String -> U2 -> Frame -> BasicBlock -> a
indexOOBError String
instName U2
i Frame
ba BasicBlock
block =
    String -> a
forall a. HasCallStack => String -> a
error (String -> a) -> String -> a
forall a b. (a -> b) -> a -> b
$
        Doc (ZonkAny 0) -> String
forall s a. (IsString s, Pretty a) => a -> s
showPretty
            ( [Doc (ZonkAny 0)] -> Doc (ZonkAny 0)
forall ann. [Doc ann] -> Doc ann
vsep
                [ String -> Doc (ZonkAny 0)
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty String
instName Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 0)
" at index " Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> U2 -> Doc (ZonkAny 0)
forall ann. U2 -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty U2
i Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> Doc (ZonkAny 0)
" is out of bounds."
                , Doc (ZonkAny 0)
"Locals: " Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> [LocalVariable] -> Doc (ZonkAny 0)
forall ann. [LocalVariable] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty Frame
ba.locals
                , Doc (ZonkAny 0)
"Instructions: " Doc (ZonkAny 0) -> Doc (ZonkAny 0) -> Doc (ZonkAny 0)
forall a. Semigroup a => a -> a -> a
<> [Instruction] -> Doc (ZonkAny 0)
forall ann. [Instruction] -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty BasicBlock
block.instructions
                ]
            )