diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,4 @@ +0.1 +--- + +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +The MIT License (MIT) +Copyright (c) 2016 Galois, Inc. + +Permission is hereby granted, free of charge, to any person obtaining a +copy of this software and associated documentation files (the +"Software"), to deal in the Software without restriction, including +without limitation the rights to use, copy, modify, merge, publish, +distribute, sublicense, and/or sell copies of the Software, and to +permit persons to whom the Software is furnished to do so, subject to +the following conditions: + +The above copyright notice and this permission notice shall be included +in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. +IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, +TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE +SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,9 @@ +lua-bc +====== + +Parser and generator for the Lua bytecode format. + +Currently versions 5.2 and 5.3 can be parsed and 5.3 can be generated. + +This package provides a Haskell datastructure for the format as well +as a few helpers for extracting debug information. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/lua-bc.cabal b/lua-bc.cabal new file mode 100644 --- /dev/null +++ b/lua-bc.cabal @@ -0,0 +1,41 @@ +name: lua-bc +version: 0.1 +synopsis: Lua bytecode parser +description: Lua bytecode parser +license: MIT +license-file: LICENSE +author: Eric Mertens +maintainer: emertens@galois.com +copyright: 2015 Galois, Inc. +category: Language +build-type: Simple +cabal-version: >=1.10 +homepage: https://github.com/GaloisInc/lua-bc +bug-reports: https://github.com/GaloisInc/lua-bc/issues + +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/GaloisInc/lua-bc + +library + exposed-modules: Language.Lua.Bytecode + Language.Lua.Bytecode.Parser + Language.Lua.Bytecode.Pretty + Language.Lua.Bytecode.Debug + Language.Lua.Bytecode.FunId + + build-depends: base >=4.8 && <4.10, + binary >=0.7 && <0.9, + data-binary-ieee754 >=0.4 && <0.5, + bytestring >=0.10 && <0.11, + text >=1.2 && <1.3, + vector >=0.10 && <0.12, + containers >=0.5 && <0.6, + pretty >=1.1 && <1.2 + + hs-source-dirs: src + default-language: Haskell2010 diff --git a/src/Language/Lua/Bytecode.hs b/src/Language/Lua/Bytecode.hs new file mode 100644 --- /dev/null +++ b/src/Language/Lua/Bytecode.hs @@ -0,0 +1,152 @@ +module Language.Lua.Bytecode where + +import Data.Vector(Vector) +import Data.ByteString (ByteString) +import qualified Data.ByteString as B + +newtype Reg = Reg Int + deriving (Read,Show,Eq,Ord) + +instance Enum Reg where + toEnum = Reg + fromEnum (Reg r) = r + +newtype UpIx = UpIx Int + deriving (Read,Show,Eq,Ord) + +newtype ProtoIx = ProtoIx Int + deriving (Read,Show,Eq,Ord) + +newtype Kst = Kst Int + deriving (Read,Show,Eq,Ord) + +data RK = RK_Reg Reg | RK_Kst Kst + deriving (Read,Show,Eq,Ord) + +data Chunk = Chunk Int Function -- ^ number of upvalues and function body + deriving (Read,Show,Eq) + +data OpCode + = OP_MOVE Reg Reg {- ^ A B R(A) := R(B) -} + | OP_LOADK Reg Kst {- ^ A Bx R(A) := Kst(Bx) -} + | OP_LOADKX Reg {- ^ A R(A) := Kst(extra arg) -} + | OP_LOADBOOL Reg Bool Bool {- ^ A B C R(A) := (Bool)B; if (C) pc++ -} + | OP_LOADNIL Reg Int {- ^ A B R(A), R(A+1), ..., R(A+B) := nil -} + | OP_GETUPVAL Reg UpIx {- ^ A B R(A) := UpValue[B] -} + + | OP_GETTABUP Reg UpIx RK {- ^ A B C R(A) := UpValue[B][RK(C)] -} + | OP_GETTABLE Reg Reg RK {- ^ A B C R(A) := R(B)[RK(C)] -} + + | OP_SETTABUP UpIx RK RK {- ^ A B C UpValue[A][RK(B)] := RK(C) -} + | OP_SETUPVAL Reg UpIx {- ^ A B UpValue[B] := R(A) -} + | OP_SETTABLE Reg RK RK {- ^ A B C R(A)[RK(B)] := RK(C) -} + + | OP_NEWTABLE Reg Int Int {- ^ A B C R(A) := {} (size = B,C) -} + + | OP_SELF Reg Reg RK {- ^ A B C R(A+1) := R(B); R(A) := R(B)[RK(C)] -} + + | OP_ADD Reg RK RK {- ^ A B C R(A) := RK(B) + RK(C) -} + | OP_SUB Reg RK RK {- ^ A B C R(A) := RK(B) - RK(C) -} + | OP_MUL Reg RK RK {- ^ A B C R(A) := RK(B) * RK(C) -} + | OP_MOD Reg RK RK {- ^ A B C R(A) := RK(B) % RK(C) -} + | OP_POW Reg RK RK {- ^ A B C R(A) := RK(B) ^ RK(C) -} + | OP_DIV Reg RK RK {- ^ A B C R(A) := RK(B) / RK(C) -} + | OP_IDIV Reg RK RK {- ^ A B C R(A) := RK(B) // RK(C) -} + | OP_BAND Reg RK RK {- ^ A B C R(A) := RK(B) & RK(C) -} + | OP_BOR Reg RK RK {- ^ A B C R(A) := RK(B) | RK(C) -} + | OP_BXOR Reg RK RK {- ^ A B C R(A) := RK(B) ~ RK(C) -} + | OP_SHL Reg RK RK {- ^ A B C R(A) := RK(B) << RK(C) -} + | OP_SHR Reg RK RK {- ^ A B C R(A) := RK(B) >> RK(C) -} + | OP_UNM Reg Reg {- ^ A B R(A) := -R(B) -} + | OP_BNOT Reg Reg {- ^ A B R(A) := ~R(B) -} + | OP_NOT Reg Reg {- ^ A B R(A) := not R(B) -} + | OP_LEN Reg Reg {- ^ A B R(A) := length of R(B) -} + + | OP_CONCAT Reg Reg Reg {- ^ A B C R(A) := R(B).. ... ..R(C) -} + + | OP_JMP (Maybe Reg) Int {- ^ A sBx pc+=sBx; if (A) close all upvalues >= R(A - 1) -} + | OP_EQ Bool RK RK {- ^ A B C if ((RK(B) == RK(C)) ~= A) then pc++ -} + | OP_LT Bool RK RK {- ^ A B C if ((RK(B) < RK(C)) ~= A) then pc++ -} + | OP_LE Bool RK RK {- ^ A B C if ((RK(B) <= RK(C)) ~= A) then pc++ -} + + | OP_TEST Reg Bool {- ^ A C if not (R(A) <=> C) then pc++ -} + | OP_TESTSET Reg Reg Bool {- ^ A B C if (R(B) <=> C) then R(A) := R(B) else pc++ -} + + | OP_CALL Reg Count Count {- ^ A B C R(A), ... ,R(A+C-2) := R(A)(R(A+1), ... ,R(A+B-1)) -} + | OP_TAILCALL Reg Count Count {- ^ A B C return R(A)(R(A+1), ... ,R(A+B-1)) -} + | OP_RETURN Reg Count {- ^ A B return R(A), ... ,R(A+B-2) (see note) -} + + | OP_FORLOOP Reg Int {- ^ A sBx R(A)+=R(A+2); if R(A) Function +propagateSources = go B.empty + where + go name func = + case funcSource func of + Nothing -> + func { funcSource = (Just name) + , funcProtos = fmap (go name) (funcProtos func) } + Just src -> + func { funcProtos = fmap (go src) (funcProtos func) } + + +-- | Compute a register relative to another. +plusReg :: Reg -> Int {- ^ offset -} -> Reg +plusReg (Reg i) j = Reg (i+j) + +-- | Compute a list of registers given a startin register and length. +regRange :: Reg {- ^ start -} -> Int {- ^ length -} -> [Reg] +regRange start n = take n [start ..] + diff --git a/src/Language/Lua/Bytecode/Debug.hs b/src/Language/Lua/Bytecode/Debug.hs new file mode 100644 --- /dev/null +++ b/src/Language/Lua/Bytecode/Debug.hs @@ -0,0 +1,170 @@ +{-# LANGUAGE RecordWildCards #-} +module Language.Lua.Bytecode.Debug where + +import Data.ByteString (ByteString, append) +import qualified Data.Vector as Vector +import Data.String(fromString) +import Data.Map ( Map ) +import qualified Data.Map as Map +import Data.Maybe (mapMaybe) + +import Language.Lua.Bytecode +import Language.Lua.Bytecode.FunId + + + +------------------------------------------------------------------------ +-- Debugging functions +------------------------------------------------------------------------ + +lookupLineNumber :: + Function -> + Int {- ^ program counter -} -> + Maybe Int +lookupLineNumber Function{..} pc = debugInfoLines Vector.!? pc + where + DebugInfo{..} = funcDebug + + +{- | Given a function, compute a map from line numbers to +op-codes in the function. This is useful for adding break-points +identified by line number. Does not consider nested functions. -} +shallowLineNumberMap :: Function -> Map Int [Int] +shallowLineNumberMap f = Map.fromListWith (++) + (mapMaybe lookupPc [ 0 .. ops - 1 ]) + where + ops = Vector.length (funcCode f) + lookupPc pc = do ln <- lookupLineNumber f pc + return (ln,[pc]) + +{- | Given a function, compute a map from line numbers to op-codes +in this function or a nested function. For each line number we +return a list of pairs (typically just 1). The first element in the +pair is the path to the nested function---empty if not nested---and +the second one are the PC locations associated with that function. -} +deepLineNumberMap :: Function -> Map Int [ (FunId, [Int]) ] +deepLineNumberMap f0 = Map.unionsWith (++) + $ me : zipWith child [ 0 .. ] + (Vector.toList (funcProtos f0)) + where + me = fmap (\pc -> [ (noFun, pc) ]) (shallowLineNumberMap f0) + ext n (path,pc) = (subFun path n,pc) + child n f = fmap (map (ext n)) (deepLineNumberMap f) + + + + + +-- | Compute the locals at a specific program counter. +-- This function is memoized, so it is fast to lookup things many times. +lookupLocalName :: Function -> Int -> Reg -> Maybe ByteString +lookupLocalName func = \pc (Reg x) -> do vs <- memo Vector.!? pc + vs Vector.!? x + where + memo = Vector.generate (Vector.length (funcCode func)) locals + locals pc = Vector.map varInfoName + $ Vector.filter (\x -> pc < varInfoEnd x) + $ Vector.takeWhile (\x -> varInfoStart x <= pc) + $ debugInfoVars (funcDebug func) + + + + +-- | Compute the names for the functions defined withing the given function. +-- The 'Int' is the index of the sub-function's prototype. +inferSubFunctionNames :: Function -> [ (Int, ByteString) ] +inferSubFunctionNames fun = foldr checkOp [] (Vector.indexed (funcCode fun)) + where + checkOp (pc,op) rest = + case op of + OP_CLOSURE _ (ProtoIx k) + | Just nm <- inferFunctionName fun pc -> (k,nm) : rest + _ -> rest + + + + +-- | Figure out a name for the function defined at the given program +-- counter in a function. Note that this operation could be fairly expensive, +-- so probably a good idea to cache the results. +inferFunctionName :: Function -> Int -> Maybe ByteString +inferFunctionName fun pc0 = + do op <- getOp pc0 + case op of + OP_CLOSURE r (ProtoIx _) -> findForward r (pc0 + 1) + _ -> Nothing + where + getOp pc = funcCode fun Vector.!? pc + getLoc = lookupLocalName fun + getUp (UpIx n) = debugInfoUpvalues (funcDebug fun) Vector.!? n + getLab k = case k of + RK_Kst (Kst n) -> + do v <- funcConstants fun Vector.!? n + case v of + KString lab -> return lab + _ -> Nothing + RK_Reg _ -> Nothing + + jn x y = append x (append (fromString ".") y) + + + -- Look forward for an instraction that puts a given register in a table. + findForward r pc = + case getLoc pc r of + Just x -> return x + Nothing -> + do op <- getOp pc + case op of + OP_MOVE r1 r2 + | r == r2 -> findForward r1 (pc + 1) + + OP_SETTABLE r1 k (RK_Reg r2) + | r2 == r -> do nm <- getNameFor r1 pc + lab <- getLab k + return (jn nm lab) + + OP_SETUPVAL r2 u + | r2 == r -> getUp u + + OP_SETTABUP u k (RK_Reg r2) + | r2 == r -> do nm <- getUp u + lab <- getLab k + return (jn nm lab) + + OP_CALL r2 _ _ + | r2 <= r -> return (fromString "_") + + OP_TAILCALL r2 _ _ + | r2 <= r -> return (fromString "_") + + _ -> findForward r (pc + 1) + + findBack r pc = + do op <- getOp pc + case op of + + OP_GETTABLE r1 r2 k + | r1 == r -> do nm <- getNameFor r2 pc + lab <- getLab k + return (jn nm lab) + + OP_GETUPVAL r1 u + | r1 == r -> getUp u + + OP_GETTABUP r1 u k + | r1 == r -> do nm <- getUp u + lab <- getLab k + return (jn nm lab) + + OP_NEWTABLE r1 _ _ + | r1 == r -> findForward r (pc + 1) + + _ -> findBack r (pc - 1) + + getNameFor r pc = + case getLoc pc r of + Just x -> return x + Nothing -> findBack r (pc - 1) + + + diff --git a/src/Language/Lua/Bytecode/FunId.hs b/src/Language/Lua/Bytecode/FunId.hs new file mode 100644 --- /dev/null +++ b/src/Language/Lua/Bytecode/FunId.hs @@ -0,0 +1,45 @@ +module Language.Lua.Bytecode.FunId where + +import Data.List(intercalate) +import Text.PrettyPrint(text) +import Text.Read(readMaybe) + + +-- | Path from root to function. +newtype FunId = FunId [Int] + deriving (Eq,Ord,Show) + +isRootFun :: FunId -> Bool +isRootFun (FunId x) = case x of + [_] -> True + _ -> False + +getRoot :: FunId -> Maybe Int +getRoot (FunId xs) = if null xs then Nothing else Just (last xs) + +noFun :: FunId +noFun = FunId [] + +isNoFun :: FunId -> Bool +isNoFun (FunId xs) = null xs + +rootFun :: Int -> FunId +rootFun n = FunId [n] + +subFun :: FunId -> Int -> FunId +subFun (FunId xs) y = FunId (y : xs) + +funIdString :: FunId -> String +funIdString (FunId xs) = intercalate "_" (map show (reverse xs)) + +funIdFromString :: String -> Maybe FunId +funIdFromString = fmap (FunId . reverse) . mapM readMaybe . words . map cvt + where + cvt x = if x == '_' then ' ' else x + +funIdList :: FunId -> [Int] +funIdList (FunId xs) = reverse xs + +funNestDepth :: FunId -> Int +funNestDepth (FunId xs) = length xs + diff --git a/src/Language/Lua/Bytecode/Parser.hs b/src/Language/Lua/Bytecode/Parser.hs new file mode 100644 --- /dev/null +++ b/src/Language/Lua/Bytecode/Parser.hs @@ -0,0 +1,763 @@ +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RecordWildCards #-} + +module Language.Lua.Bytecode.Parser + ( -- * Parsers + parseLuaBytecode + , parseLuaBytecodeFile + + -- * Dump + , dumpLuaBytecode + , dumpLuaBytecodeFile + ) where + +import Control.Applicative +import Control.Monad (when, guard) +import Control.Exception +import Data.Binary.Get +import Data.Binary.Put +import Data.Binary.IEEE754 (putFloat64le, getFloat64le) +import Data.Bits +import Data.ByteString (ByteString) +import Data.Foldable (traverse_) +import Data.Maybe (fromMaybe) +import qualified Data.Text as Text +import Data.Text.Encoding (decodeUtf8, encodeUtf8) +import Data.Vector (Vector) +import qualified Data.Vector as Vector +import Data.Word (Word8, Word32) +import Numeric (showHex) +import System.Environment (lookupEnv) +import System.Exit +import System.IO +import System.IO.Error +import qualified Data.ByteString.Lazy as L +import qualified Data.ByteString as B + +import Language.Lua.Bytecode + +type Puts a = a -> Put + +data LuacVersion = Luac52 | Luac53 + deriving (Eq, Ord, Show, Read) + +data Sizet = Sizet32 | Sizet64 + +sizetSize :: Sizet -> Int +sizetSize Sizet32 = 4 +sizetSize Sizet64 = 8 + +data BytecodeMode = BytecodeMode + { bytecodeVersion :: !LuacVersion + , bytecodeSizet :: !Sizet + } + +parseLuaBytecode :: Maybe String {- ^ Optional source name -} -> + L.ByteString {- ^ Bytecodes -} -> + Either String Chunk +parseLuaBytecode mbName bc = + case runGetOrFail loadChunk bc of + Left (_,_,e) -> Left e + Right (_,_,chunk) -> Right (maybe chunk (setName chunk) mbName) + where + setName (Chunk n f) nm = Chunk n (setSourceName nm f) + +dumpLuaBytecode :: BytecodeMode -> Chunk -> L.ByteString +dumpLuaBytecode mode = runPut . saveChunk mode + +------------------------------------------------------------------------ + +parseLuaBytecodeFile :: FilePath -> IO (Either String Chunk) +parseLuaBytecodeFile fp = parseLuaBytecode (Just fp) <$> L.readFile fp + +dumpLuaBytecodeFile :: FilePath -> Chunk -> IO () +dumpLuaBytecodeFile fp luafile = L.writeFile fp (dumpLuaBytecode mode53 luafile) + +mode53 = BytecodeMode Luac53 Sizet64 + +------------------------------------------------------------------------ + +setSourceName :: String -> Function -> Function +setSourceName name func = func { funcSource = Just (packUtf8 name) } + +------------------------------------------------------------------------ + +unpackUtf8 :: ByteString -> String +unpackUtf8 = Text.unpack . decodeUtf8 + +packUtf8 :: String -> ByteString +packUtf8 = encodeUtf8 . Text.pack + +-- magic numbers ------------------------------------------------------- + +luaSignature :: ByteString +luaSignature = "\x1bLua" + +mkLuaVersion :: Word8 -> Word8 -> Word8 +mkLuaVersion major minor + | major < 16, minor < 16 = major * 16 + minor + | otherwise = error "mkLuaVersion: Bad arguments" + +luacFormat :: Word8 +luacFormat = 0 + +luacData :: ByteString +luacData = "\x19\x93\r\n\x1a\n" + +luacInt :: Int +luacInt = 0x5678 + +luacNum :: Double +luacNum = 370.5 + +intSize :: Word8 +intSize = 4 + +instructionSize :: Word8 +instructionSize = intSize + +luaIntegerSize :: Word8 +luaIntegerSize = 8 + +luaNumberSize :: Word8 +luaNumberSize = 8 + +-- top-level lua chunk file codec -------------------------------------- + +loadChunk :: Get Chunk +loadChunk = + do mode <- loadHeader + case bytecodeVersion mode of + Luac52 -> do f <- loadFunction mode + return (Chunk (Vector.length (funcUpvalues f)) f) + Luac53 -> do n <- fromIntegral <$> loadByte + f <- loadFunction mode + return (Chunk n f) + +saveChunk :: BytecodeMode -> Puts Chunk +saveChunk mode (Chunk n f) = + do saveHeader mode + saveByte (fromIntegral n) + saveFunction mode f + +-- size_t codec -------------------------------------------------------- + +loadSizeT :: BytecodeMode -> Get Int +loadSizeT mode = + case bytecodeSizet mode of + Sizet64 -> fromIntegral <$> getWord64host + Sizet32 -> fromIntegral <$> getWord32host + +saveSizeT :: BytecodeMode -> Puts Int +saveSizeT mode = + case bytecodeSizet mode of + Sizet64 -> putWord64host . fromIntegral + Sizet32 -> putWord32host . fromIntegral + +-- byte codec ---------------------------------------------------------- + +loadString :: BytecodeMode -> Get (Maybe ByteString) +loadString mode = + case bytecodeVersion mode of + Luac52 -> + do n <- loadSizeT mode + if n == 0 + then pure Nothing -- NULL? + else Just . B.init <$> getByteString n + + Luac53 -> + do n <- loadByte + mb <- case n of + 0 -> return Nothing -- NULL? + 0xff -> Just <$> loadSizeT mode + _ -> return (Just (fromIntegral n)) + traverse (\n' -> getByteString (n'-1)) mb + +loadString' :: String -> BytecodeMode -> Get ByteString +loadString' name mode = + do mb <- loadString mode + case mb of + Nothing -> fail name + Just x -> return x + +saveString :: BytecodeMode -> Puts (Maybe ByteString) +saveString _ Nothing = saveByte 0 +saveString mode (Just x) = + do let n = B.length x + if n < 0xfe + then saveByte (fromIntegral n+1) + else saveByte 0xff >> saveSizeT mode (n+1) + putByteString x + +saveString' :: BytecodeMode -> Puts ByteString +saveString' mode = saveString mode . Just + +-- byte codec ---------------------------------------------------------- + +loadByte :: Get Word8 +loadByte = getWord8 + +saveByte :: Puts Word8 +saveByte = putWord8 + +-- lua integer value codec --------------------------------------------- + +loadInteger :: Get Int +loadInteger = fromIntegral <$> getWord64host + +saveInteger :: Puts Int +saveInteger = putWord64host . fromIntegral + +-- lua int codec ------------------------------------------------------- + +loadInt :: Get Int +loadInt = fromIntegral <$> getWord32host + +saveInt :: Puts Int +saveInt = putWord32host . fromIntegral + +-- lua number codec ---------------------------------------------------- + +loadNumber :: Get Double +loadNumber = getFloat64le -- I'd use the "host" version if there was one + +saveNumber :: Puts Double +saveNumber = putFloat64le -- I'd use the "host" version if there was one + +-- boolean codec ------------------------------------------------------- + +loadBool :: Get Bool +loadBool = fmap (/= 0) loadByte + +saveBool :: Puts Bool +saveBool x = saveByte (if x then 1 else 0) + +-- array codec --------------------------------------------------------- + +loadVectorOf :: Get a -> Get (Vector a) +loadVectorOf p = + do n <- loadInt + Vector.replicateM n p + +saveVectorOf :: Puts a -> Puts (Vector a) +saveVectorOf put xs = + do saveInt (Vector.length xs) + traverse_ put xs + +-- chunk header codec -------------------------------------------------- + + +() :: Get a -> String -> Get a +m str = m <|> fail str + +string :: ByteString -> Get () +string xs = + do ys <- getByteString (B.length xs) + guard (xs == ys) + +word8 :: Word8 -> Get () +word8 x = + do y <- getWord8 + guard (x == y) + +loadHeader :: Get BytecodeMode +loadHeader = + do string luaSignature "signature" + version <- getWord8 + if version == mkLuaVersion 5 3 then loadHeader53 + else if version == mkLuaVersion 5 2 then loadHeader52 + else fail "Unknown version" + +loadHeader52 :: Get BytecodeMode +loadHeader52 = + do word8 luacFormat "luac format" + word8 1 "endianess" + word8 intSize "int size" + sizet <- getSizetSize + word8 instructionSize "Instruction size" + word8 luaNumberSize "Number size" + word8 0 "Integral Number" + string luacData "corruption check" + return BytecodeMode + { bytecodeVersion = Luac52 + , bytecodeSizet = sizet + } + +loadHeader53 :: Get BytecodeMode +loadHeader53 = + do word8 luacFormat "luac format" + string luacData "corruption check" + word8 intSize "int size" + sizet <- getSizetSize + word8 instructionSize "Instruction size" + word8 luaIntegerSize "Integer size" + word8 luaNumberSize "Number size" + luaInteger luacInt "Integer encoding" + luaNumber luacNum "Number encoding" + return BytecodeMode + { bytecodeVersion = Luac53 + , bytecodeSizet = sizet + } + where + + luaInteger x = + do y <- loadInteger + guard (x == y) + + luaNumber x = + do y <- loadNumber + guard (x == y) + +getSizetSize :: Get Sizet +getSizetSize = + do n <- getWord8 + case n of + 8 -> return Sizet64 + 4 -> return Sizet32 + _ -> fail ("Bad size_t size: " ++ show n) + +saveHeader :: BytecodeMode -> Put +saveHeader mode = + case bytecodeVersion mode of + Luac52 -> error "saveHeader not implemented for version 5.2" + Luac53 -> + do putByteString luaSignature + saveByte (mkLuaVersion 5 3) + saveByte luacFormat + putByteString luacData + saveByte intSize + saveByte (fromIntegral (sizetSize (bytecodeSizet mode))) + saveByte instructionSize + saveByte luaIntegerSize + saveByte luaNumberSize + saveInteger luacInt + saveNumber luacNum + +-- function codec ------------------------------------------------------ + +loadFunction :: BytecodeMode -> Get Function +loadFunction mode = + case bytecodeVersion mode of + Luac52 -> + do funcLineDefined <- loadInt + funcLastLineDefined <- loadInt + funcNumParams <- fromIntegral <$> loadByte + funcIsVararg <- loadBool + funcMaxStackSize <- fromIntegral <$> loadByte + funcCode <- loadVectorOf (loadInstruction mode) + funcConstants <- loadVectorOf (loadConstant mode) + funcProtos <- loadVectorOf (loadFunction mode) + funcUpvalues <- loadVectorOf loadUpvalue + funcSource <- loadString mode + funcDebug <- loadDebugInfo mode + return Function{..} + + Luac53 -> Function + <$> loadString mode + <*> loadInt + <*> loadInt + <*> (fromIntegral <$> loadByte) + <*> loadBool + <*> (fromIntegral <$> loadByte) + <*> loadVectorOf (loadInstruction mode) + <*> loadVectorOf (loadConstant mode) + <*> loadVectorOf loadUpvalue + <*> loadVectorOf (loadFunction mode) + <*> loadDebugInfo mode + +saveFunction :: BytecodeMode -> Puts Function +saveFunction mode Function{..} = + do saveString mode funcSource + saveInt funcLineDefined + saveInt funcLastLineDefined + saveByte (fromIntegral funcNumParams) + saveBool funcIsVararg + saveByte (fromIntegral funcMaxStackSize) + saveVectorOf saveInstruction funcCode + saveVectorOf (saveConstant mode) funcConstants + saveVectorOf saveUpvalue funcUpvalues + saveVectorOf (saveFunction mode) funcProtos + saveDebugInfo mode funcDebug + +-- function debug information codec ------------------------------------ + +loadDebugInfo :: BytecodeMode -> Get DebugInfo +loadDebugInfo mode = DebugInfo + <$> loadVectorOf loadInt + <*> loadVectorOf (loadVarInfo mode) + <*> loadVectorOf (fromMaybe "UNKNOWN" <$> loadString mode) + +saveDebugInfo :: BytecodeMode -> Puts DebugInfo +saveDebugInfo mode DebugInfo{..} = + do saveVectorOf saveInt debugInfoLines + saveVectorOf (saveVarInfo mode) debugInfoVars + saveVectorOf (saveString' mode) debugInfoUpvalues + +-- function constant codec --------------------------------------------- + +loadConstant :: BytecodeMode -> Get Constant +loadConstant mode = + do tag <- loadByte + case tag of + 0 -> return KNil + 1 -> KBool <$> loadBool + 3 -> KNum <$> loadNumber + 4 -> KString <$> loadString' "short constant" mode -- short + 19 -> KInt <$> loadInteger + 20 -> KLongString <$> loadString' "long constant" mode -- long + _ -> fail ("unknown constant type: " ++ show tag) + +saveConstant :: BytecodeMode -> Puts Constant +saveConstant mode x = + case x of + KNil -> saveByte 0 + KBool b -> saveByte 1 >> saveBool b + KNum n -> saveByte 3 >> saveNumber n + KString s -> saveByte 4 >> saveString' mode s + KInt i -> saveByte 19 >> saveInteger i + KLongString s -> saveByte 20 >> saveString' mode s + +-- function upvalue codec ---------------------------------------------- + +loadUpvalue :: Get Upvalue +loadUpvalue = + do instack <- loadBool + idx <- fromIntegral <$> loadByte + return (if instack then UpReg (Reg idx) else UpUp (UpIx idx)) + +saveUpvalue :: Puts Upvalue +saveUpvalue x = + case x of + UpReg (Reg idx) -> saveBool True >> saveByte (fromIntegral idx) + UpUp (UpIx idx) -> saveBool False >> saveByte (fromIntegral idx) + +-- debug var info codec ------------------------------------------------ + +loadVarInfo :: BytecodeMode -> Get VarInfo +loadVarInfo mode = VarInfo + <$> loadString' "local variable" mode + <*> loadInt + <*> loadInt + +saveVarInfo :: BytecodeMode -> Puts VarInfo +saveVarInfo mode VarInfo{..} = + do saveString' mode varInfoName + saveInt varInfoStart + saveInt varInfoEnd + +-- instruction codec --------------------------------------------------- + +loadInstruction :: BytecodeMode -> Get OpCode +loadInstruction mode = + case bytecodeVersion mode of + Luac52 -> loadInstruction52 + Luac53 -> loadInstruction53 + +loadInstruction52 :: Get OpCode +loadInstruction52 = + do raw <- getWord32host + let op = extractPart size_OP pos_OP raw + a = extractPart size_A pos_A raw + ax = extractPart size_Ax pos_Ax raw + b = extractPart size_B pos_B raw + bx = extractPart size_Bx pos_Bx raw + c = extractPart size_B pos_C raw + rk x | testBit x (size_B-1) = RK_Kst (Kst (clearBit x (size_B-1))) + | otherwise = RK_Reg (Reg x) + sBx = bx - ((1 `shiftL` (size_Bx-1)) - 1) + case op of + 0 -> return $ OP_MOVE (Reg a) (Reg b) + 1 -> return $ OP_LOADK (Reg a) (Kst bx) + 2 -> return $ OP_LOADKX (Reg a) + 3 -> return $ OP_LOADBOOL (Reg a) (b /= 0) (c /= 0) + 4 -> return $ OP_LOADNIL (Reg a) b + 5 -> return $ OP_GETUPVAL (Reg a) (UpIx b) + + 6 -> return $ OP_GETTABUP (Reg a) (UpIx b) (rk c) + 7 -> return $ OP_GETTABLE (Reg a) (Reg b) (rk c) + + 8 -> return $ OP_SETTABUP (UpIx a) (rk b) (rk c) + 9 -> return $ OP_SETUPVAL (Reg a) (UpIx b) + 10 -> return $ OP_SETTABLE (Reg a) (rk b) (rk c) + + 11 -> return $ OP_NEWTABLE (Reg a) (fb2int b) (fb2int c) + + 12 -> return $ OP_SELF (Reg a) (Reg b) (rk c) + + 13 -> return $ OP_ADD (Reg a) (rk b) (rk c) + 14 -> return $ OP_SUB (Reg a) (rk b) (rk c) + 15 -> return $ OP_MUL (Reg a) (rk b) (rk c) + 16 -> return $ OP_DIV (Reg a) (rk b) (rk c) + 17 -> return $ OP_MOD (Reg a) (rk b) (rk c) + 18 -> return $ OP_POW (Reg a) (rk b) (rk c) + 19 -> return $ OP_UNM (Reg a) (Reg b) + 20 -> return $ OP_NOT (Reg a) (Reg b) + 21 -> return $ OP_LEN (Reg a) (Reg b) + + 22 -> return $ OP_CONCAT (Reg a) (Reg b) (Reg c) + + 23 -> let r | a == 0 = Nothing + | otherwise = Just (Reg (a-1)) + in return $ OP_JMP r sBx + 24 -> return $ OP_EQ (a /= 0) (rk b) (rk c) + 25 -> return $ OP_LT (a /= 0) (rk b) (rk c) + 26 -> return $ OP_LE (a /= 0) (rk b) (rk c) + + 27 -> return $ OP_TEST (Reg a) (c /= 0) + 28 -> return $ OP_TESTSET (Reg a) (Reg b) (c /= 0) + + 29 -> return $ OP_CALL (Reg a) (mkCount b) (mkCount c) + 30 -> return $ OP_TAILCALL (Reg a) (mkCount b) (mkCount c) + 31 -> return $ OP_RETURN (Reg a) (mkCount b) + + 32 -> return $ OP_FORLOOP (Reg a) sBx + 33 -> return $ OP_FORPREP (Reg a) sBx + + 34 -> return $ OP_TFORCALL (Reg a) c + 35 -> return $ OP_TFORLOOP (Reg a) sBx + + 36 -> return $ OP_SETLIST (Reg a) b c + + 37 -> return $ OP_CLOSURE (Reg a) (ProtoIx bx) + + 38 -> return $ OP_VARARG (Reg a) (mkCount b) + + 39 -> return $ OP_EXTRAARG ax + + _ -> fail $ "Bad instruction: 0x" ++ showHex raw "" + +loadInstruction53 :: Get OpCode +loadInstruction53 = + do raw <- getWord32host + let op = extractPart size_OP pos_OP raw + a = extractPart size_A pos_A raw + ax = extractPart size_Ax pos_Ax raw + b = extractPart size_B pos_B raw + bx = extractPart size_Bx pos_Bx raw + c = extractPart size_B pos_C raw + rk x | testBit x (size_B-1) = RK_Kst (Kst (clearBit x (size_B-1))) + | otherwise = RK_Reg (Reg x) + sBx = bx - ((1 `shiftL` (size_Bx-1)) - 1) + case op of + 0 -> return $ OP_MOVE (Reg a) (Reg b) + 1 -> return $ OP_LOADK (Reg a) (Kst bx) + 2 -> return $ OP_LOADKX (Reg a) + 3 -> return $ OP_LOADBOOL (Reg a) (b /= 0) (c /= 0) + 4 -> return $ OP_LOADNIL (Reg a) b + 5 -> return $ OP_GETUPVAL (Reg a) (UpIx b) + + 6 -> return $ OP_GETTABUP (Reg a) (UpIx b) (rk c) + 7 -> return $ OP_GETTABLE (Reg a) (Reg b) (rk c) + + 8 -> return $ OP_SETTABUP (UpIx a) (rk b) (rk c) + 9 -> return $ OP_SETUPVAL (Reg a) (UpIx b) + 10 -> return $ OP_SETTABLE (Reg a) (rk b) (rk c) + + 11 -> return $ OP_NEWTABLE (Reg a) (fb2int b) (fb2int c) + + 12 -> return $ OP_SELF (Reg a) (Reg b) (rk c) + + 13 -> return $ OP_ADD (Reg a) (rk b) (rk c) + 14 -> return $ OP_SUB (Reg a) (rk b) (rk c) + 15 -> return $ OP_MUL (Reg a) (rk b) (rk c) + 16 -> return $ OP_MOD (Reg a) (rk b) (rk c) + 17 -> return $ OP_POW (Reg a) (rk b) (rk c) + 18 -> return $ OP_DIV (Reg a) (rk b) (rk c) + 19 -> return $ OP_IDIV (Reg a) (rk b) (rk c) + 20 -> return $ OP_BAND (Reg a) (rk b) (rk c) + 21 -> return $ OP_BOR (Reg a) (rk b) (rk c) + 22 -> return $ OP_BXOR (Reg a) (rk b) (rk c) + 23 -> return $ OP_SHL (Reg a) (rk b) (rk c) + 24 -> return $ OP_SHR (Reg a) (rk b) (rk c) + 25 -> return $ OP_UNM (Reg a) (Reg b) + 26 -> return $ OP_BNOT (Reg a) (Reg b) + 27 -> return $ OP_NOT (Reg a) (Reg b) + 28 -> return $ OP_LEN (Reg a) (Reg b) + + 29 -> return $ OP_CONCAT (Reg a) (Reg b) (Reg c) + + 30 -> let r | a == 0 = Nothing + | otherwise = Just (Reg (a-1)) + in return $ OP_JMP r sBx + 31 -> return $ OP_EQ (a /= 0) (rk b) (rk c) + 32 -> return $ OP_LT (a /= 0) (rk b) (rk c) + 33 -> return $ OP_LE (a /= 0) (rk b) (rk c) + + 34 -> return $ OP_TEST (Reg a) (c /= 0) + 35 -> return $ OP_TESTSET (Reg a) (Reg b) (c /= 0) + + 36 -> return $ OP_CALL (Reg a) (mkCount b) (mkCount c) + 37 -> return $ OP_TAILCALL (Reg a) (mkCount b) (mkCount c) + 38 -> return $ OP_RETURN (Reg a) (mkCount b) + + 39 -> return $ OP_FORLOOP (Reg a) sBx + 40 -> return $ OP_FORPREP (Reg a) sBx + + 41 -> return $ OP_TFORCALL (Reg a) c + 42 -> return $ OP_TFORLOOP (Reg a) sBx + + 43 -> return $ OP_SETLIST (Reg a) b c + + 44 -> return $ OP_CLOSURE (Reg a) (ProtoIx bx) + + 45 -> return $ OP_VARARG (Reg a) (mkCount b) + + 46 -> return $ OP_EXTRAARG ax + + _ -> fail $ "Bad instruction: 0x" ++ showHex raw "" + +mkCount :: Int -> Count +mkCount 0 = CountTop +mkCount x = CountInt (x-1) + +saveInstruction :: Puts OpCode +saveInstruction x = putWord32host $ sum $ + case x of + OP_MOVE a b -> [asOp 0, asA a, asB b] + OP_LOADK a bx -> [asOp 1, asA a, asBX bx] + OP_LOADKX a -> [asOp 2, asA a] + OP_LOADBOOL a b c -> [asOp 3, asA a, asB b, asC c] + OP_LOADNIL a b -> [asOp 4, asA a, asB b] + OP_GETUPVAL a b -> [asOp 5, asA a, asB b] + + OP_GETTABUP a b c -> [asOp 6, asA a, asB b, asC c] + OP_GETTABLE a b c -> [asOp 7, asA a, asB b, asC c] + + OP_SETTABUP a b c -> [asOp 8, asA a, asB b, asC c] + OP_SETUPVAL a b -> [asOp 9, asA a, asB b] + OP_SETTABLE a b c -> [asOp 10, asA a, asB b, asC c] + + OP_NEWTABLE a b c -> [asOp 11, asA a, asB (int2fb b), asC (int2fb c)] + + OP_SELF a b c -> [asOp 12, asA a, asB b, asC c] + + OP_ADD a b c -> [asOp 13, asA a, asB b, asC c] + OP_SUB a b c -> [asOp 14, asA a, asB b, asC c] + OP_MUL a b c -> [asOp 15, asA a, asB b, asC c] + OP_MOD a b c -> [asOp 16, asA a, asB b, asC c] + OP_POW a b c -> [asOp 17, asA a, asB b, asC c] + OP_DIV a b c -> [asOp 18, asA a, asB b, asC c] + OP_IDIV a b c -> [asOp 19, asA a, asB b, asC c] + OP_BAND a b c -> [asOp 20, asA a, asB b, asC c] + OP_BOR a b c -> [asOp 21, asA a, asB b, asC c] + OP_BXOR a b c -> [asOp 22, asA a, asB b, asC c] + OP_SHL a b c -> [asOp 23, asA a, asB b, asC c] + OP_SHR a b c -> [asOp 24, asA a, asB b, asC c] + OP_UNM a b -> [asOp 25, asA a, asB b] + OP_BNOT a b -> [asOp 26, asA a, asB b] + OP_NOT a b -> [asOp 27, asA a, asB b] + OP_LEN a b -> [asOp 28, asA a, asB b] + + OP_CONCAT a b c -> [asOp 29, asA a, asB b, asC c] + + OP_JMP Nothing sbx -> [asOp 30, asSBX sbx] + OP_JMP (Just (Reg a)) sbx -> [asOp 30, asA (a+1), asSBX sbx] + + OP_EQ a b c -> [asOp 31, asA a, asB b, asC c] + OP_LT a b c -> [asOp 32, asA a, asB b, asC c] + OP_LE a b c -> [asOp 33, asA a, asB b, asC c] + + OP_TEST a c -> [asOp 34, asA a, asC c] + OP_TESTSET a b c -> [asOp 35, asA a, asB b, asC c] + + OP_CALL a b c -> [asOp 36, asA a, asB b, asC c] + OP_TAILCALL a b c -> [asOp 37, asA a, asB b, asC c] + OP_RETURN a b -> [asOp 38, asA a, asB b] + + OP_FORLOOP a sbx -> [asOp 39, asA a, asSBX sbx] + OP_FORPREP a sbx -> [asOp 40, asA a, asSBX sbx] + + OP_TFORCALL a c -> [asOp 41, asA a, asC c] + OP_TFORLOOP a sbx -> [asOp 42, asA a, asSBX sbx] + + OP_SETLIST a b c -> [asOp 43, asA a, asB b, asC c] + + OP_CLOSURE a bx -> [asOp 44, asA a, asBX bx] + + OP_VARARG a b -> [asOp 45, asA a, asB b] + + OP_EXTRAARG ax -> [asOp 46, asAX ax] + +class OpArg a where opArg :: a -> Word32 +instance OpArg Reg where opArg (Reg x) = fromIntegral x +instance OpArg UpIx where opArg (UpIx x) = fromIntegral x +instance OpArg Kst where opArg (Kst x) = fromIntegral x +instance OpArg ProtoIx where opArg (ProtoIx x) = fromIntegral x +instance OpArg Int where opArg = fromIntegral +instance OpArg Word32 where opArg = id +instance OpArg Bool where opArg True = 1 + opArg False = 0 +instance OpArg RK where opArg (RK_Reg r) = opArg r + opArg (RK_Kst k) = setBit (opArg k) (size_B-1) +instance OpArg Count where opArg CountTop = 0 + opArg (CountInt x) = fromIntegral x + 1 + +asOp :: Word32 -> Word32 +asOp = createPart size_OP pos_OP + +asA :: OpArg a => a -> Word32 +asA = createPart size_A pos_A . opArg + +asB :: OpArg a => a -> Word32 +asB = createPart size_B pos_B . opArg + +asC :: OpArg a => a -> Word32 +asC = createPart size_C pos_C . opArg + +asSBX :: Int -> Word32 +asSBX = createPart size_Bx pos_Bx . fromIntegral . correct + where + correct x = x + ((1 `shiftL` (size_Bx-1)) - 1) + +asBX :: OpArg a => a -> Word32 +asBX = createPart size_Bx pos_Bx . opArg + +asAX :: OpArg a => a -> Word32 +asAX = createPart size_Ax pos_Ax . opArg + +extractPart :: Int -> Int -> Word32 -> Int +extractPart size pos x = fromIntegral ((x `shiftL` (32-size-pos)) `shiftR` (32-size)) + +createPart :: Int -> Int -> Word32 -> Word32 +createPart size pos x = (mask .&. x) `shiftL` pos + where + mask = 1 `shiftL` size - 1 + +{- | + converts an integer to a "floating point byte", represented as + (eeeeexxx), where the real value is (1xxx) * 2^(eeeee - 1) if + eeeee != 0 and (xxx) otherwise. +-} +fb2int :: Int -> Int +fb2int x + | x < 8 = x + | otherwise = (8 + (x .&. 7)) * 2^(shiftR x 3 - 1) + +int2fb :: Int -> Word32 +int2fb x0 + | x0 < 0 = error "int2fb: negative argument" + | x0 < 8 = fromIntegral x0 + | otherwise = coarse 0 (fromIntegral x0) + where + coarse e x + | x >= 8 `shiftL` 4 = coarse (e+4) ((x + 0xf) `shiftR` 4) + | otherwise = fine e x + fine e x + | x >= 8 `shiftL` 1 = fine (e+1) ((x + 1) `shiftR` 1) + | otherwise = finish e x + finish e x = (e+1) `shiftL` 3 .|. (x - 8) + +size_A, size_Ax, size_B, size_Bx, size_C, size_OP :: Int +pos_A, pos_Ax, pos_B, pos_Bx, pos_C, pos_OP :: Int +pos_A = (pos_OP + size_OP) +pos_Ax = pos_A +pos_B = (pos_C + size_C) +pos_Bx = pos_C +pos_C = (pos_A + size_A) +pos_OP = 0 +size_A = 8 +size_Ax = (size_C + size_B + size_A) +size_B = 9 +size_C = 9 +size_Bx = (size_C + size_B) +size_OP = 6 diff --git a/src/Language/Lua/Bytecode/Pretty.hs b/src/Language/Lua/Bytecode/Pretty.hs new file mode 100644 --- /dev/null +++ b/src/Language/Lua/Bytecode/Pretty.hs @@ -0,0 +1,270 @@ +module Language.Lua.Bytecode.Pretty where + +import Language.Lua.Bytecode +import Language.Lua.Bytecode.Debug(lookupLocalName) +import Language.Lua.Bytecode.FunId + +import Text.PrettyPrint +import Data.Maybe(fromMaybe) +import Data.Text.Lazy() +import Data.Text.Encoding(decodeUtf8With) +import Data.Text.Encoding.Error(lenientDecode) +import qualified Data.ByteString.Char8 as BS +import qualified Data.Vector as Vector + +data PPInfo = PPInfo + { ppVarAt :: Int -> Reg -> Maybe Doc + , ppConstant :: Kst -> Maybe Doc + , ppUpIx :: UpIx -> Maybe Doc + , ppFun :: ProtoIx -> Maybe Doc + , ppExtraArg :: Maybe Int + , ppPC :: !Int + } + +ppVar :: PPInfo -> Reg -> Maybe Doc +ppVar p = ppVarAt p (ppPC p) + +blankPPInfo :: PPInfo +blankPPInfo = PPInfo + { ppVarAt = \_ _ -> Nothing + , ppConstant = \_ -> Nothing + , ppUpIx = \_ -> Nothing + , ppFun = \_ -> Nothing + , ppExtraArg = Nothing + , ppPC = 0 + } + +ppNextPC :: PPInfo -> PPInfo +ppNextPC p = p { ppPC = 1 + ppPC p } + +ppCode :: Function -> Doc +ppCode f = vcat $ map (ppOpCode f) [ 0 .. Vector.length (funcCode f) - 1 ] + +ppOpCode :: Function -> Int -> Doc +ppOpCode f pc0 = case funcCode f Vector.!? pc0 of + Just op -> pp us op + Nothing -> parens (text ("op code at " ++ show pc0)) + where + us = PPInfo { ppVarAt = lkpVar + , ppConstant = lkpConstant + , ppUpIx = lkpUpIx + , ppFun = lkpFun + , ppExtraArg = lkpExtra + , ppPC = pc0 + } + + -- This is written like this so that `lkp` memoizes. + lkpVar = \pc r -> do bs <- lkp pc r + return (text (BS.unpack bs)) + where lkp = lookupLocalName f + + lkpUpIx (UpIx u) = + do bs <- debugInfoUpvalues (funcDebug f) Vector.!? u + return (text (BS.unpack bs)) + + lkpConstant (Kst k) = fmap (pp us) (funcConstants f Vector.!? k) + + lkpExtra = do OP_EXTRAARG a <- funcCode f Vector.!? (pc0 + 1) + return a + + lkpFun (ProtoIx x) = + do f1 <- funcProtos f Vector.!? x + let nm = case funcSource f1 of + Just bs -> "function from " ++ BS.unpack bs + Nothing -> "function at" + return (text nm <+> parens (int (funcLineDefined f1) <> text "-" <> + int (funcLastLineDefined f1))) + + + +instance PP Constant where + pp _ cnst = + case cnst of + KNil -> text "nil" + KBool b -> text (if b then "true" else "false") + KNum d -> double d + KInt n -> int n + KString bs -> str bs + KLongString b -> str b + + where str = text . show . decodeUtf8With lenientDecode + + +class PP a where + pp :: PPInfo -> a -> Doc + +instance PP FunId where + pp n = text . funIdString + +instance PP Reg where + pp i r@(Reg x) = fromMaybe d (ppVar i r) + where d = text "R" <> brackets (int (x+1)) + +instance PP Kst where + pp i k@(Kst x) = fromMaybe d (ppConstant i k) + where d = text "K" <> brackets (int (x+1)) + +instance PP UpIx where + pp i u@(UpIx x) = fromMaybe d (ppUpIx i u) + where d = text "U" <> brackets (int (x+1)) + +instance PP ProtoIx where + pp i p@(ProtoIx x) = fromMaybe d (ppFun i p) + where d = text "KPROTO" <> brackets (int (x+1)) + +instance PP RK where + pp i rk = + case rk of + RK_Reg r -> pp i r + RK_Kst k -> pp i k + +instance PP Doc where + pp _ d = d + +instance PP Bool where + pp _ b = text (if b then "true" else "false") + +ppRegRange :: PPInfo -> Reg -> Int -> Doc +ppRegRange _ _ 0 = empty +ppRegRange i r 1 = pp i r +ppRegRange i r@(Reg x) n = + case ppVar i r of + Just v -> v <> comma <+> ppRegRange i (succ r) (n-1) + Nothing -> pp i r <+> text ".." <+> pp i (Reg (x + n - 1)) + + +ppRegRangeInf :: PPInfo -> Reg -> Doc +ppRegRangeInf i r0 = hsep (punctuate comma (ppE r0)) <> text ".." + where + ppE r@(Reg x) = case ppVar i r of + Just v -> v : ppE (Reg (x+1)) + Nothing -> [ pp i r ] + +ppRegRangeCount :: PPInfo -> Reg -> Count -> Doc +ppRegRangeCount i r ct = + case ct of + CountTop -> ppRegRangeInf i r + CountInt n -> ppRegRange i r n + + + + + +instance PP OpCode where + pp i opCode = + let (=:) :: (PP a, PP b) => a -> b -> Doc + x =: y = let d1 = pp (ppNextPC i) x + d2 = pp i y + in if isEmpty d1 then d2 else d1 <+> text "=" <+> d2 + + lkp :: (PP a, PP b) => a -> b -> Doc + lkp x y = pp i x <> brackets (pp i y) + + op1 :: (PP a) => String -> a -> Doc + op1 o x = text o <> pp i x + + op2 :: (PP a, PP b) => String -> a -> b -> Doc + op2 o x y = pp i x <+> text o <+> pp i y + + cond a x mb = text "if" <+> expr <+> text "then" <+> text "pc++" + <+> mbElse + where expr = if a then text "not" <+> parens x else x + mbElse = maybe empty (\y -> text "else" <+> y) mb + + + in + case opCode of + OP_MOVE a b -> a =: b + + OP_LOADK a b -> a =: b + + OP_LOADKX a -> a =: case ppExtraArg i of + Nothing -> text "extra" + Just x -> pp i (Kst x) + + OP_LOADBOOL a b c -> (a =: b) <> more + where more = if c then semi <+> text "pc++" else empty + + OP_LOADNIL a b -> ppRegRange (ppNextPC i) a (b+1) =: text "nil" + + OP_GETUPVAL a b -> a =: b + + OP_GETTABUP a b c -> a =: lkp b c + OP_GETTABLE a b c -> a =: lkp b c + + OP_SETTABUP a b c -> lkp a b =: c + OP_SETUPVAL a b -> a =: b + OP_SETTABLE a b c -> lkp a b =: c + + OP_NEWTABLE a b c -> a =: (text "{}" <+> parens (text "array size =" <+> + int b <> comma <+> + text "table size =" <+> int c)) + + OP_SELF a b c -> (succ a =: b) <> semi <+> (a =: lkp b c) + + OP_ADD a b c -> a =: op2 "+" b c + OP_SUB a b c -> a =: op2 "-" b c + OP_MUL a b c -> a =: op2 "*" b c + OP_MOD a b c -> a =: op2 "%" b c + OP_POW a b c -> a =: op2 "^" b c + OP_DIV a b c -> a =: op2 "/" b c + OP_IDIV a b c -> a =: op2 "//" b c + OP_BAND a b c -> a =: op2 "&" b c + OP_BOR a b c -> a =: op2 "|" b c + OP_BXOR a b c -> a =: op2 "~" b c + OP_SHL a b c -> a =: op2 "<<" b c + OP_SHR a b c -> a =: op2 ">>" b c + OP_UNM a b -> a =: op1 "-" b + OP_BNOT a b -> a =: op1 "~" b + OP_NOT a b -> a =: (text "not" <+> pp i b) + OP_LEN a b -> a =: op1 "#" b + + OP_CONCAT a b@(Reg x) (Reg y) -> + a =: text "concat" <+> ppRegRange i b (y - x + 1) + + OP_JMP mb b -> text "pc +=" <+> int b <+> + case mb of + Nothing -> empty + Just a -> semi <+> text "close" <+> ppRegRangeInf i a + + OP_EQ a b c -> cond a (op2 "==" b c) Nothing + OP_LT a b c -> cond a (op2 "<" b c) Nothing + OP_LE a b c -> cond a (op2 "<=" b c) Nothing + OP_TEST a c -> cond c (pp i a) Nothing + OP_TESTSET a b c -> cond c (pp i b) (Just (a =: b)) + OP_CALL a b c -> ppRegRangeCount (ppNextPC i) a c =: + pp i a <> parens (ppRegRangeCount i (succ a) b) + + OP_TAILCALL a b _c -> + text "return" <+> pp i a <> parens (ppRegRangeCount i (succ a) b) + + OP_RETURN a b -> text "return" <+> ppRegRangeCount i a b + + OP_FORPREP a b -> text "FORPREP" <+> pp i a <+> int b + OP_FORLOOP a b -> text "FORLOOP" <+> pp i a <+> int b + OP_TFORCALL a b -> text "FORCALL" <+> pp i a <+> int b + OP_TFORLOOP a b -> text "FORLOOP" <+> pp i a <+> int b + + OP_SETLIST a b c -> + pp (ppNextPC i) a <> brackets rng =: ppRegRange i (succ a) b + where + rng | c == 0 = conc (ppExtraArg i) + | otherwise = conc (Just c) + + conc x | b == 1 = num x 1 + conc x = num x 1 <> text ".." <> num x b + num x j = case x of + Just y -> int ((y-1) * 50 + j) + Nothing -> text "50*(extra-1) +" <+> int j + + + OP_CLOSURE a b -> a =: (text "closure" <+> pp i b) + OP_VARARG a b -> ppRegRangeCount (ppNextPC i) a b =: text "..." + + OP_EXTRAARG n -> + case ppExtraArg i of + Just _ -> empty + Nothing -> nest 2 (text "where extra =" <+> int n) + + +