Skip to content

Commit

Permalink
partial rewrite of SMT algorithm, definitely doesnt work well
Browse files Browse the repository at this point in the history
  • Loading branch information
@bristermitten
bristermitten committed Dec 5, 2023
1 parent 51e6162 commit 58c14eb
Show file tree
Hide file tree
Showing 3 changed files with 188 additions and 19 deletions.
111 changes: 108 additions & 3 deletions src/JVM/Data/Analyse/StackMap.hs
View file
Original file line number Diff line number Diff line change
@@ -1,24 +1,129 @@
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MultiWayIf #-}

{- | 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 ((^.), (^?), _Just)
import Control.Lens.Extras (is)
import Control.Lens.Fold
import Data.Generics.Sum (AsAny (_As))
import Data.List.Split (split, splitOn, splitWhen)
import Data.List
import Data.Maybe (fromJust)
import JVM.Data.Abstract.Builder.Label
import JVM.Data.Abstract.ClassFile.Method
import JVM.Data.Abstract.Descriptor (MethodDescriptor (MethodDescriptor))
import JVM.Data.Abstract.Instruction
import JVM.Data.Abstract.Type (ClassInfoType (..), FieldType (..), PrimitiveType (..), fieldTypeToClassInfoType)

data BasicBlock = BasicBlock
{ index :: Int
, instructions :: [Instruction]
, end :: Maybe Label
}
deriving (Show, Eq)

data Frame = Frame
{ locals :: [LocalVariable]
, stack :: [StackEntry]
}
deriving (Show, Eq)

data LocalVariable = Uninitialised | LocalVariable FieldType
deriving (Show, Eq)

data StackEntry = StackEntry FieldType | StackEntryTop | StackEntryNull
deriving (Show, Eq)

lvToStackEntry :: LocalVariable -> StackEntry
lvToStackEntry Uninitialised = StackEntryTop
lvToStackEntry (LocalVariable ft) = StackEntry ft

stackEntryToLV :: StackEntry -> LocalVariable
stackEntryToLV StackEntryTop = Uninitialised
stackEntryToLV StackEntryNull = Uninitialised
stackEntryToLV (StackEntry ft) = LocalVariable ft

splitIntoBasicBlocks :: [Instruction] -> [BasicBlock]
splitIntoBasicBlocks [] = []
splitIntoBasicBlocks l =
let blocks = splitWhen (is (_As @"Label")) l
in zipWith BasicBlock [0 ..] blocks
let blockToInstAndLabel = splitOnLabels l
in zipWith (\i (l, b) -> BasicBlock i b l) [0 ..] blockToInstAndLabel

splitOnLabels :: [Instruction] -> [(Maybe Label, [Instruction])]
splitOnLabels xs = go xs []
where
go :: [Instruction] -> ([Instruction]) -> [(Maybe Label, [Instruction])]
go [] acc = [(Nothing, acc)]
go (x : xs) acc = case x ^? _As @"Label" of
Just l' -> (Just l', acc) : go xs []
Nothing -> go xs (acc <> [x])

topFrame :: MethodDescriptor -> Frame
topFrame (MethodDescriptor args _) = Frame (map LocalVariable args) []

analyseBlockDiff :: Frame -> BasicBlock -> Frame
analyseBlockDiff current block = do
foldl (flip analyseInstruction) current (takeWhileInclusive (not . isConditionalJump) block.instructions)
where
isConditionalJump :: Instruction -> Bool
isConditionalJump (IfEq _) = True
isConditionalJump (IfNe _) = True
isConditionalJump (IfLt _) = True
isConditionalJump (IfGe _) = True
isConditionalJump (IfGt _) = True
isConditionalJump (IfLe _) = True
isConditionalJump _ = False

analyseInstruction :: Instruction -> Frame -> Frame
analyseInstruction (Label _) ba = error "Label should not be encountered in analyseInstruction"
analyseInstruction (ALoad i) ba = ba{stack = lvToStackEntry (ba.locals !! fromIntegral (i - 1)) : ba.stack}
analyseInstruction (ILoad i) ba = ba{stack = lvToStackEntry (ba.locals !! fromIntegral (i - 1)) : ba.stack}
analyseInstruction (AStore i) ba = ba{locals = replaceAtOrGrow (fromIntegral (i - 1)) (stackEntryToLV $ head ba.stack) ba.locals, stack = tail ba.stack}
analyseInstruction (IStore i) ba = ba{locals = replaceAtOrGrow (fromIntegral (i - 1)) (stackEntryToLV $ head ba.stack) ba.locals, stack = tail ba.stack}
analyseInstruction AReturn ba = ba{stack = tail ba.stack}
analyseInstruction Return ba = ba
analyseInstruction (LDC (LDCInt _)) ba = ba{stack = StackEntry (PrimitiveFieldType Int) : ba.stack}
analyseInstruction AConstNull ba = ba{stack = StackEntryNull : ba.stack}
analyseInstruction Dup ba = ba{stack = head ba.stack : ba.stack}
analyseInstruction (IfEq _) ba = ba{stack = tail ba.stack}
analyseInstruction (IfNe _) ba = ba{stack = tail ba.stack}
analyseInstruction (IfLt _) ba = ba{stack = tail ba.stack}
analyseInstruction (IfGe _) ba = ba{stack = tail ba.stack}
analyseInstruction (IfGt _) ba = ba{stack = tail ba.stack}
analyseInstruction (IfLe _) ba = ba{stack = tail ba.stack}
analyseInstruction other ba = error $ "Instruction not supported: " <> show other

frameDiffToSMF :: Frame -> BasicBlock -> StackMapFrame
frameDiffToSMF f1@(Frame locals1 stack1) bb = do
let f2@(Frame locals2 stack2) = analyseBlockDiff f1 bb
if
| locals1 == locals2 && stack1 == stack2 -> SameFrame (fromJust bb.end)
| stack1 == stack2 && locals1 `isPrefixOf` locals2 -> AppendFrame (map lvToVerificationTypeInfo (drop (length locals1) locals2)) (fromJust bb.end)
| otherwise -> error (show f1 <> show f2)

lvToVerificationTypeInfo :: LocalVariable -> VerificationTypeInfo
lvToVerificationTypeInfo Uninitialised = TopVariableInfo
lvToVerificationTypeInfo (LocalVariable ft) = case ft of
PrimitiveFieldType Int -> IntegerVariableInfo
PrimitiveFieldType Float -> FloatVariableInfo
PrimitiveFieldType Long -> LongVariableInfo
PrimitiveFieldType Double -> DoubleVariableInfo
_ -> ObjectVariableInfo (fieldTypeToClassInfoType ft)

replaceAtOrGrow :: Int -> LocalVariable -> [LocalVariable] -> [LocalVariable]
replaceAtOrGrow i x xs
| i < length xs = replaceAt i x xs
| otherwise = xs <> replicate (i - length xs) Uninitialised <> [x]

replaceAt :: Int -> a -> [a] -> [a]
replaceAt i x xs = take i xs <> [x] <> drop (i + 1) xs

takeWhileInclusive :: (a -> Bool) -> [a] -> [a]
takeWhileInclusive _ [] = []
takeWhileInclusive p (x : xs)
| p x = x : takeWhileInclusive p xs
| otherwise = [x]
94 changes: 78 additions & 16 deletions test/Analyse.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,10 @@ import JVM.Data.Abstract.Name
import JVM.Data.Abstract.Type

import JVM.Data.Abstract.Builder.Code
import JVM.Data.Abstract.ClassFile.Method ( StackMapFrame (..), VerificationTypeInfo (..))
import JVM.Data.Abstract.Descriptor
import JVM.Data.Abstract.Instruction
import JVM.Data.Analyse.StackMap (BasicBlock (BasicBlock), splitIntoBasicBlocks)
import JVM.Data.Analyse.StackMap (BasicBlock (BasicBlock), Frame (..), LocalVariable (..), analyseBlockDiff, frameDiffToSMF, splitIntoBasicBlocks, topFrame)
import Test.Hspec
import Test.Hspec.Hedgehog

Expand Down Expand Up @@ -47,26 +49,86 @@ genFieldType =
spec :: Spec
spec = describe "Analysis checks" $ do
describe "Does StackDiff concatenation correctly" $ do
it "Can identify basic blocks properly" $ do
it "Can identify sameframe blocks properly" $ do
let (l, _, code) = runCodeBuilder' $ do
label <- newLabel
emit $ LDC (LDCInt 0) -- [0]
emit $ IfEq label
emit $ LDC (LDCInt 0)
emit $ IStore 1
emit $ LDC (LDCInt 0)
emit $ IStore 2
emit $ ILoad 1
emit $ IfLe label
emit $ LDC (LDCInt 0)
emit $ IStore 3
emit AReturn
emit $ Label label
emit $ LDC (LDCInt 0)
emit $ IStore 3
emit Return
emit $ LDC (LDCInt 1)
emit AReturn

pure label
hedgehog $ do
let x = splitIntoBasicBlocks code
let blocks = splitIntoBasicBlocks code

x
=== [ BasicBlock 0 [LDC (LDCInt 0), IStore 1, LDC (LDCInt 0), IStore 2, ILoad 1, IfLe l, LDC (LDCInt 0), IStore 3]
, BasicBlock 1 [LDC (LDCInt 0), IStore 3, Return]
blocks
=== [ BasicBlock 0 [LDC (LDCInt 0), IfEq l, LDC (LDCInt 0), AReturn] (Just l)
, BasicBlock 1 [LDC (LDCInt 1), AReturn] Nothing
]

let top = topFrame (MethodDescriptor [] (TypeReturn (PrimitiveFieldType Int)))
let nextFrame = analyseBlockDiff top (head blocks)

nextFrame
=== Frame
{ locals = []
, stack = []
}

let nextFrame' = analyseBlockDiff nextFrame (blocks !! 1)

nextFrame'
=== Frame
{ locals = []
, stack = []
}

frameDiffToSMF top (head blocks)
=== SameFrame l

it "Can identify append frame blocks properly" $ do
let (l, _, code) = runCodeBuilder' $ do
label <- newLabel
emit $ LDC (LDCInt 0) -- [0]
emit $ IStore 1 -- []
emit $ LDC (LDCInt 0) -- [0]
emit $ IStore 2 -- []
emit $ ILoad 1 -- [0]
emit $ IfLe label -- []
emit $ LDC (LDCInt 0) -- [0]
emit $ IStore 3 -- []
emit $ Label label -- []
emit $ LDC (LDCInt 0) -- [0]
emit $ IStore 3 -- []
emit Return -- []
pure label
hedgehog $ do
let blocks = splitIntoBasicBlocks code

blocks
=== [ BasicBlock 0 [LDC (LDCInt 0), IStore 1, LDC (LDCInt 0), IStore 2, ILoad 1, IfLe l, LDC (LDCInt 0), IStore 3] (Just l)
, BasicBlock 1 [LDC (LDCInt 0), IStore 3, Return] Nothing
]

let top = topFrame (MethodDescriptor [] VoidReturn)
let nextFrame = analyseBlockDiff top (head blocks)

nextFrame
=== Frame
{ locals = [LocalVariable (PrimitiveFieldType Int), LocalVariable (PrimitiveFieldType Int)]
, stack = []
}

let nextFrame' = analyseBlockDiff nextFrame (blocks !! 1)

nextFrame'
=== Frame
{ locals = [LocalVariable (PrimitiveFieldType Int), LocalVariable (PrimitiveFieldType Int), LocalVariable (PrimitiveFieldType Int)]
, stack = []
}

frameDiffToSMF top (head blocks)
=== AppendFrame [IntegerVariableInfo, IntegerVariableInfo] l
2 changes: 2 additions & 0 deletions test/Builder.hs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,3 +32,5 @@ spec = describe "test code building" $ do
, Raw.IfEq 3
, Raw.Return
]


0 comments on commit 58c14eb

Please sign in to comment.