packages feed

lua-bc (empty) → 0.1

raw patch · 10 files changed

+1477/−0 lines, 10 filesdep +basedep +binarydep +bytestringsetup-changed

Dependencies added: base, binary, bytestring, containers, data-binary-ieee754, pretty, text, vector

Files

+ CHANGELOG.md view
@@ -0,0 +1,4 @@+0.1+---++* Initial release
+ LICENSE view
@@ -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.
+ README.md view
@@ -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.
+ Setup.hs view
@@ -0,0 +1,2 @@+import Distribution.Simple+main = defaultMain
+ lua-bc.cabal view
@@ -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
+ src/Language/Lua/Bytecode.hs view
@@ -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) <?= R(A+1) then { pc+=sBx; R(A+3)=R(A) }-}+  | OP_FORPREP  Reg Int         {- ^    A sBx   R(A)-=R(A+2); pc+=sBx                           -}++  | OP_TFORCALL Reg Int         {- ^    A C     R(A+3), ... ,R(A+2+C) := R(A)(R(A+1), R(A+2));  -}+  | OP_TFORLOOP Reg Int         {- ^    A sBx   if R(A+1) ~= nil then { R(A)=R(A+1); pc += sBx }-}++  | OP_SETLIST  Reg Int Int   {- ^    A B C   R(A)[(C-1)*FPF+i] := R(A+i), 1 <= i <= B        -}++  | OP_CLOSURE  Reg ProtoIx     {- ^    A Bx    R(A) := closure(KPROTO[Bx])                     -}++  | OP_VARARG   Reg Count {- ^    A B     R(A), R(A+1), ..., R(A+B-2) = vararg            -}++  | OP_EXTRAARG Int             {- ^    Ax      extra (larger) argument for previous opcode     -}+  deriving (Read, Show, Eq, Ord)++data Count = CountInt Int | CountTop+  deriving (Read, Show, Eq, Ord)++data Function = Function+  { funcSource          :: !(Maybe ByteString)+  , funcLineDefined     :: !Int+  , funcLastLineDefined :: !Int+  , funcNumParams       :: !Int+  , funcIsVararg        :: !Bool+  , funcMaxStackSize    :: !Int+  , funcCode            :: !(Vector OpCode)+  , funcConstants       :: !(Vector Constant)+  , funcUpvalues        :: !(Vector Upvalue)+  , funcProtos          :: !(Vector Function)+  , funcDebug           :: !DebugInfo+  }+  deriving (Read, Show, Eq)++data Constant = KNil | KBool Bool | KNum Double | KInt Int | KString ByteString | KLongString ByteString+  deriving (Read, Show, Eq, Ord)++data Upvalue = UpReg Reg | UpUp UpIx+  deriving (Read, Show, Eq, Ord)++type LineNumber = Int++data DebugInfo = DebugInfo+  { debugInfoLines      :: !(Vector LineNumber)+  , debugInfoVars       :: !(Vector VarInfo)+  , debugInfoUpvalues   :: !(Vector ByteString)+  }+  deriving (Read, Show, Eq)++data VarInfo = VarInfo+  { varInfoName :: !ByteString+  , varInfoStart, varInfoEnd :: !Int+  }+  deriving (Read, Show, Eq)++propagateSources :: Function -> 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 ..]+
+ src/Language/Lua/Bytecode/Debug.hs view
@@ -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)+++
+ src/Language/Lua/Bytecode/FunId.hs view
@@ -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+
+ src/Language/Lua/Bytecode/Parser.hs view
@@ -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
+ src/Language/Lua/Bytecode/Pretty.hs view
@@ -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)+++