blip-0.2.1: src/StackDepth.hs
{-# LANGUAGE RecordWildCards #-}
-----------------------------------------------------------------------------
-- |
-- Module : StackDepth
-- Copyright : (c) 2012, 2013 Bernie Pope
-- License : BSD-style
-- Maintainer : florbitous@gmail.com
-- Stability : experimental
-- Portability : ghc
--
-- Compute an upper bound on the stack usage of a block of bytecode.
-- It is safe to make the stack too big (but
-- it would waste memory), but if it is too small then the interpreter will
-- probably crash (or worse, keep running in an undefined state).
--
-----------------------------------------------------------------------------
module StackDepth (maxStackDepth) where
import Types (AnnotatedCode (..))
import Utils (isJumpBytecode, isRelativeJump, isConditionalJump)
import Blip.Bytecode (Bytecode (..), BytecodeArg (..), Opcode (..), bytecodeSize)
import Data.Word (Word32, Word16)
import Control.Monad.RWS.Strict (RWS, runRWS, ask, local, gets, modify, when)
import qualified Data.Map as Map
import qualified Data.Set as Set (insert, member, Set, empty)
import Data.Bits ((.&.), shiftR)
type StackDepth = Word32
type InstructionIndex = Word16
type InstructionSeen = Set.Set InstructionIndex
-- Mapping from byte address (jump target) to sequence of bytecode
-- from that address onwards.
type BytecodeMap = Map.Map InstructionIndex [AnnotatedCode]
type StackDepthCache = Map.Map InstructionIndex StackDepth
type CalcStackDepth = RWS InstructionSeen () StackDepthState
maxStackDepth :: [AnnotatedCode] -> StackDepth
maxStackDepth code =
stackDepth_maxDepth finalState
where
(_, finalState, _) = runRWS (maxStackDepthM 0 code)
Set.empty $ initStackDepthState $
makeBytecodeMap code
makeBytecodeMap :: [AnnotatedCode] -> BytecodeMap
makeBytecodeMap = makeBytecodeMapAcc Map.empty
where
makeBytecodeMapAcc :: BytecodeMap -> [AnnotatedCode] -> BytecodeMap
makeBytecodeMapAcc map [] = map
makeBytecodeMapAcc map code@(instruction@(AnnotatedCode {..}) : rest)
| isLabelled instruction = do
let newMap = Map.insert annotatedCode_index code map
makeBytecodeMapAcc newMap rest
| otherwise = makeBytecodeMapAcc map rest
data StackDepthState =
StackDepthState
{ stackDepth_bytecodeMap :: BytecodeMap
, stackDepth_maxDepth :: !StackDepth
, stackDepth_cache :: StackDepthCache
}
initStackDepthState :: BytecodeMap -> StackDepthState
initStackDepthState bytecodeMap =
StackDepthState
{ stackDepth_bytecodeMap = bytecodeMap
, stackDepth_maxDepth = 0
, stackDepth_cache = Map.empty }
isLabelled :: AnnotatedCode -> Bool
isLabelled (AnnotatedCode {..}) = not $ null annotatedCode_labels
isLoopBack :: InstructionIndex -> CalcStackDepth Bool
isLoopBack index = do
seen <- ask
return (index `Set.member` seen)
-- record that we've visited this jump target at this depth
-- in case we visit it again in any path. There is no point
-- traversing further if the previous visit was at an equal
-- or greater depth.
visitedDeeper :: InstructionIndex -> StackDepth -> CalcStackDepth Bool
visitedDeeper index newDepth = do
stackDepthCache <- gets stackDepth_cache
case Map.lookup index stackDepthCache of
-- not been here before at any depth
Nothing -> return False
Just oldDepth -> return (oldDepth >= newDepth)
recordDepth :: InstructionIndex -> StackDepth -> CalcStackDepth ()
recordDepth index depth = do
stackDepthCache <- gets stackDepth_cache
let newCache = Map.insert index depth stackDepthCache
modify $ \s -> s { stackDepth_cache = newCache }
maxStackDepthM :: StackDepth -> [AnnotatedCode] -> CalcStackDepth ()
maxStackDepthM _depth [] = return ()
maxStackDepthM depth code@(instruction@(AnnotatedCode {..}) : _rest) = do
-- check if this instruction is a jump target
if isLabelled instruction
then do
seenBeforeOnPath <- isLoopBack annotatedCode_index
seenDeeper <- visitedDeeper annotatedCode_index depth
if seenBeforeOnPath || seenDeeper
-- we've seen this instruction before on this path, or
-- we've visisted it on any path at this depth or
-- deeper, no point in going further down this path.
then return ()
else local (Set.insert annotatedCode_index) $ do
recordDepth annotatedCode_index depth
maxStackDepthFurther depth code
else
maxStackDepthFurther depth code
where
maxStackDepthFurther :: StackDepth -> [AnnotatedCode] -> CalcStackDepth ()
maxStackDepthFurther depth (instruction@(AnnotatedCode {..}) : rest) = do
let newDepth = depth + codeStackEffect annotatedCode_bytecode
updateMaxDepth newDepth
when (isJumpBytecode annotatedCode_bytecode) $
-- follow the path of the jump
maxStackDepthJump newDepth instruction
-- follow the remaining instructions
-- unless the current instruction is an unconditional jump
when (isConditionalBytecode annotatedCode_bytecode) $
maxStackDepthM newDepth rest
maxStackDepthFurther _depth [] =
error $ "maxStackDepthFurther called on empty sequence of code"
isConditionalBytecode :: Bytecode -> Bool
isConditionalBytecode (Bytecode {..}) = isConditionalJump opcode
{-
from CPython, compile.c:
if (instr->i_opcode == FOR_ITER) {
target_depth = depth-2;
} else if (instr->i_opcode == SETUP_FINALLY ||
instr->i_opcode == SETUP_EXCEPT) {
target_depth = depth+3;
if (target_depth > maxdepth)
maxdepth = target_depth;
}
-}
maxStackDepthJump :: StackDepth -> AnnotatedCode -> CalcStackDepth ()
maxStackDepthJump depth instruction@(AnnotatedCode {..}) = do
let targetDepth =
case opcode annotatedCode_bytecode of
FOR_ITER -> depth - 2
SETUP_FINALLY -> depth + 3
SETUP_EXCEPT -> depth + 3
_other -> depth
updateMaxDepth targetDepth
code <- getJumpToCode instruction
maxStackDepthM targetDepth code
getJumpToCode :: AnnotatedCode -> CalcStackDepth [AnnotatedCode]
getJumpToCode instruction@(AnnotatedCode {..}) = do
let jumpTarget =
if isRelativeJump $ opcode annotatedCode_bytecode
then relativeTarget instruction
else absoluteTarget instruction
bytecodeMap <- gets stackDepth_bytecodeMap
case Map.lookup jumpTarget bytecodeMap of
Nothing -> error $ "Jump to uknown target: " ++ show instruction
Just code -> return code
relativeTarget :: AnnotatedCode -> InstructionIndex
relativeTarget instruction@(AnnotatedCode {..}) =
target + (annotatedCode_index + instructionSize)
where
instructionSize = fromIntegral $ bytecodeSize annotatedCode_bytecode
target = getJumpTarget instruction
absoluteTarget :: AnnotatedCode -> InstructionIndex
absoluteTarget instruction@(AnnotatedCode {..})
= getJumpTarget instruction
getJumpTarget :: AnnotatedCode -> InstructionIndex
getJumpTarget instruction@(AnnotatedCode {..}) =
case args annotatedCode_bytecode of
Nothing -> error $ "Jump instruction without argument: " ++ show instruction
Just (Arg16 label) -> label
updateMaxDepth :: StackDepth -> CalcStackDepth ()
updateMaxDepth depth = do
currentMaxDepth <- gets stackDepth_maxDepth
when (depth > currentMaxDepth) $
modify $ \s -> s { stackDepth_maxDepth = depth }
-- Compute the effect of each opcode on the depth of the stack.
-- This is used to compute an upper bound on the depth of the stack
-- for each code object. It is safe to over-estimate the depth of the
-- effect, but it is unsafe to underestimate it. Over-estimation will
-- potentially result in the stack being bigger than needed, which would
-- waste memory but otherwise be safe. Under-estimation will likely result
-- in the stack being too small and a serious fatal error in the interpreter, such
-- as segmentation fault (or reading/writing some other part of memory).
-- Some opcodes have different effect on depth depending on other factors, this function
-- convservatively takes the largest possible value.
-- This function is supposed to be identical in behaviour to opcode_stack_effect
-- in Python/compile.c.
codeStackEffect :: Bytecode -> StackDepth
codeStackEffect bytecode@(Bytecode {..}) =
case opcode of
POP_TOP -> -1
ROT_TWO -> 0
ROT_THREE -> 0
DUP_TOP -> 1
DUP_TOP_TWO -> 2
UNARY_POSITIVE -> 0
UNARY_NEGATIVE -> 0
UNARY_NOT -> 0
UNARY_INVERT -> 0
SET_ADD -> -1
LIST_APPEND -> -1
MAP_ADD -> -2
BINARY_POWER -> -1
BINARY_MULTIPLY -> -1
BINARY_MODULO -> -1
BINARY_ADD -> -1
BINARY_SUBTRACT -> -1
BINARY_SUBSCR -> -1
BINARY_FLOOR_DIVIDE -> -1
BINARY_TRUE_DIVIDE -> -1
INPLACE_FLOOR_DIVIDE -> -1
INPLACE_TRUE_DIVIDE -> -1
INPLACE_ADD -> -1
INPLACE_SUBTRACT -> -1
INPLACE_MULTIPLY -> -1
INPLACE_MODULO -> -1
STORE_SUBSCR -> -3
STORE_MAP -> -2
DELETE_SUBSCR -> -2
BINARY_LSHIFT -> -1
BINARY_RSHIFT -> -1
BINARY_AND -> -1
BINARY_XOR -> -1
BINARY_OR -> -1
INPLACE_POWER -> -1
GET_ITER -> 0
PRINT_EXPR -> -1
LOAD_BUILD_CLASS -> 1
INPLACE_LSHIFT -> -1
INPLACE_RSHIFT -> -1
INPLACE_AND -> -1
INPLACE_XOR -> -1
INPLACE_OR -> -1
BREAK_LOOP -> 0
SETUP_WITH -> 7
WITH_CLEANUP -> -1 -- Sometimes more
STORE_LOCALS -> -1
RETURN_VALUE -> -1
IMPORT_STAR -> -1
YIELD_VALUE -> 0
YIELD_FROM -> -1
POP_BLOCK -> 0
POP_EXCEPT -> 0 -- -3 except if bad bytecode
END_FINALLY -> -1 -- or -2 or -3 if exception occurred
STORE_NAME -> -1
DELETE_NAME -> 0
UNPACK_SEQUENCE -> withArg $ \oparg -> oparg - 1
UNPACK_EX -> withArg $ \oparg -> (oparg .&. 0xFF) + (oparg `shiftR` 8)
FOR_ITER -> 1 -- or -1, at end of iterator
STORE_ATTR -> -2
DELETE_ATTR -> -1
STORE_GLOBAL -> -1
DELETE_GLOBAL -> 0
LOAD_CONST -> 1
LOAD_NAME -> 1
BUILD_TUPLE -> withArg $ \oparg -> 1 - oparg
BUILD_LIST -> withArg $ \oparg -> 1 - oparg
BUILD_SET -> withArg $ \oparg -> 1 - oparg
BUILD_MAP -> 1
LOAD_ATTR -> 0
COMPARE_OP -> -1
IMPORT_NAME -> -1
IMPORT_FROM -> 1
JUMP_FORWARD -> 0
JUMP_IF_TRUE_OR_POP -> 0 -- -1 if jump not taken
JUMP_IF_FALSE_OR_POP -> 0 -- ditto
JUMP_ABSOLUTE -> 0
POP_JUMP_IF_FALSE -> -1
POP_JUMP_IF_TRUE -> -1
LOAD_GLOBAL -> 1
CONTINUE_LOOP -> 0
SETUP_LOOP -> 0
SETUP_EXCEPT -> 6
SETUP_FINALLY -> 6 -- can push 3 values for the new exception
-- plus 3 others for the previous exception state
LOAD_FAST -> 1
STORE_FAST -> -1
DELETE_FAST -> 0
RAISE_VARARGS -> withArg $ \oparg -> -1 * oparg
CALL_FUNCTION -> withArg $ \oparg -> -1 * nargs oparg
CALL_FUNCTION_VAR -> withArg $ \oparg -> (-1 * nargs oparg) - 1
CALL_FUNCTION_KW -> withArg $ \oparg -> (-1 * nargs oparg) - 1
CALL_FUNCTION_VAR_KW -> withArg $ \oparg -> (-1 * nargs oparg) - 2
MAKE_FUNCTION -> withArg $ \oparg -> -1 - (nargs oparg) - ((oparg `shiftR` 16) .&. 0xffff)
MAKE_CLOSURE -> withArg $ \oparg -> -2 - (nargs oparg) - ((oparg `shiftR` 16) .&. 0xffff)
BUILD_SLICE -> withArg $ \oparg -> if oparg == 3 then -2 else -1
LOAD_CLOSURE -> 1
LOAD_DEREF -> 1
STORE_DEREF -> -1
DELETE_DEREF -> 0
_other -> error $ "unexpected opcode in codeStackEffect: " ++ show bytecode
where
-- #define NARGS(o) (((o) % 256) + 2*(((o) / 256) % 256))
nargs :: Word32 -> Word32
nargs o = (o `mod` 256) + (2 * ((o `div` 256) `mod` 256))
withArg :: (Word32 -> Word32) -> Word32
withArg f
= case args of
Nothing -> error $ "codeStackEffect: " ++ (show opcode) ++ " missing argument"
Just (Arg16 word16) -> f $ fromIntegral word16
-- other -> error $ "codeStackEffect unexpected opcode argument: " ++ show other