Creating a Stack Machine with Haskell

Posted on May 28, 2016

About the Project

This is a small demonstration project showing how a simple byte code interpreting stack machine (virtual machine) can be built with Haskell. It is not a production VM nor of any particular practical use but is rather a simple demonstration of how a stack machine can be built. All code is available in the github repo https://github.com/andrevdm/SimpleHaskellStackVM

I built this for mainly as a project for learning Haskell, i.e. something a little bigger to work on. So NB this is probably not idiomatic Haskell, and may have some newbie mistakes. Hopefully it is interesting enough despite this…

Stack Machines

A stack machine is a real or emulated computer that uses a pushdown stack rather than individual machine registers to evaluate each sub-expression in the program. A stack computer is programmed with a reverse Polish notation instruction set. - wikipedia [1]

Stack machines are simpler to implement than register machines but are still practical even for production VMs.

Writing the VM in Haskell

Virtual machines are typically written in a low level language like C for maximum efficiency. They also are typically written using mutable data structures. Here I’m using Haskell and pure functional data structures. I have absolutely no performance goals, so there are no constraints I need to worry about.

A few design decisions

  1. I’m using a Data.Sequence rather than a list. While I’m not concerned about performance using a linked list to do random access is still a bad idea

  2. Try to avoid bottom (⊥), so use “safe” versions whenever ⊥ could otherwise be returned (toEnum, index, head etc)

  3. I’m using Stephen Diehl’s Protolude [2] and removing the default prelude

Using the immutable data structures like Data.Sequence had worked out nicely and means that as the VM runs you can keep a full history of the VM’s state. So you have everything (apart from time) you need to build a nice visualiser / ‘time travelling’ debugger.

The VM

Opcodes & byte code

The opcodes are the operations that the virtual CPU can perform. These are simple operations like push, pop, add and jump. The opcodes are encoded as bytes together with the opcode parameters (e.g. the value to push) this forms the byte code.

Here is the current set of opcodes understood by the CPU

data Operation = Nop       -- No operation
               | Halt      -- Stop CPU execution
               | Push      -- Push a value onto the stack
               | Pop       -- Pop the most recent value from the stack
               | PopPrev   -- Pop n values before the most recent value from the stack
               | Add       -- Add the top two items on the stack
               | Inc       -- Increment the top item on the stack
               | Dup       -- Duplicate the most recent item on the stack
               | Jmp       -- Jump unconditionally to a location
               | Bne       -- Pop the top two items, compare and branch if the values are not equal
               | Beq       -- Pop the top two items, compare and branch if the values are equal
               | Bgt       -- Pop the top two items, compare and branch if value 1 is greater than value 2
               | Bgte      -- Pop the top two items, compare and branch if value 1 is greater or equal to value 2
               | Blt       -- Pop the top two items, compare and branch if value 1 is less than value 2
               | Blte      -- Pop the top two items, compare and branch if value 1 is less than or equal to value 2
               | Call      -- Call a function
               | Ret       -- Return from a function
               | LdArg     -- Push local value n onto the stack
               deriving (Show, Eq, Ord, Enum, Bounded)

The virtual CPU

The data structure representing the CPU is defined below

data Cpu = Cpu { ip :: Int               -- Instruction pointer
               , fp :: Int               -- Frame pointer
               , cpuStack :: S.Seq Int   -- The stack
               , cpuGlobals :: S.Seq Int -- Gloal variables
               , ranOp :: Int            -- The last opcode that was executed
               , state :: Text           -- Debugging message
               , debug :: Bool           -- Enable/disable debugging
               , panic :: Bool           -- Is the CPU in a faulted state
               }
         deriving (Show, Eq)

-- | Default empty/initial CPU state
emptyCpu :: Cpu
emptyCpu = Cpu { ip = -1
               , fp = -1
               , cpuStack = S.empty
               , cpuGlobals = S.empty
               , state = ""
               , debug = True
               , ranOp = 0
               , panic = False 
               }

This is fairly minimal but it’s more than enough for a simple VM like this.

Some things to note

The byte code assembler

Rather than writing the byte code in hex a very simple assembler is used. Later on additional assemblers and compliers can be layer on top of this low level assembler which does little more than convert opcode mnemonics to byte code.

An operation can take parameters from the byte code stream. For instance a push instruction takes a parameter that is the value to push onto the stack. The type that represents this is the inventively named OpAndParam

-- | A single CPU operator and its parameters
type OpAndParam = (Operation, [Int])

Having a type that defines the number of parameters an op takes and how many values it pops off the stack allows the assembler to perform some basic checks. The interpreter can also use this definition when running the code.

-- | Configuration for an operation
-- | opParamCount = number of paramaters taken from the code stream
-- | opPopsCount = number of values this op pops from the stack
-- | opSimple = determines if the op needs full access to cpu state to change things like the fp and ip
-- |            note that 'complex' instructions do not need to honour opParamCount and opPopsCount
-- |            e.g. a 'ret' instruction pops a variable number of parameters
data Instruction = Instruction { opCode :: Operation
                               , opPopsCount :: Int
                               , opParamCount :: Int
                               , opSimple :: Bool
                               }
                 deriving (Show, Eq)

The instructions can then be setup and a map created from opcode to Instruction

-- | Config for the op codes
instructions :: [Instruction]
instructions = [ Instruction { opCode = Nop, opParamCount = 0, opPopsCount = 0, opSimple = True }
                -- ...
               , Instruction { opCode = Call, opParamCount = 1, opPopsCount = 0, opSimple = False }
               ]

-- | Instructions indexed by opcode
instrByOp :: Map.Map Operation Instruction
instrByOp = Map.fromList $ map (\i -> (opCode i, i)) instructions

The assembler then does nothing more than converting the opcode enum to a byte (an Int in the code but it would be serialised as a byte) checking the number of parameters for each opcode. It is small enough to be pasted in full here

-- | A single assembler error
data AssemblerError = AssemblerError Integer Operation Text deriving (Show, Eq)

-- | Compiles the list to byte code
-- | Returns as many errors as possible rather than just first error
assembleByteCode :: [(Operation, [Int])] -> Either [AssemblerError] [Int]	
assembleByteCode code =
  let res = foldl assemble [] code in
  case lefts res of
    [] -> Right $ concat $ rights res
    errors -> Left errors
  
  where
    assemble :: [Either AssemblerError [Int]] -> OpAndParam -> [Either AssemblerError [Int]]
    assemble res (op, prms) =
      res ++ case Map.lookup op instrByOp of
               Nothing -> [Left $ AssemblerError (toInteger $ length res) op "unknown op code"]
               Just i ->
                 if opParamCount i == length prms
                 then [Right $ fromEnum (opCode i) : prms]
                 else [Left $ AssemblerError (toInteger $ length res) op "incorrect number of parameters"]

The interpreter

With all of that in place the interpreter can finally be written.

-- | Interpreter for the byte code
-- | Given a byte code stream will 'run' the code
-- | If debug is enabled then the full history (all states) will be returned. 
interpretByteCode :: S.Seq Int -> [Cpu]

The interpreter takes the output of the assembler or bytes loaded from a file and runs it producing CPU state along the way. In debug mode the interpreter stores all the states as it interprets, in “non-debug” mode only the last state is kept. Note that currently the interpreter is always in debug mode

Before looking at the implementation of the interpreter its worth going over a few examples of how it should operate

Push & Pop

In the first example the value 123 is pushed onto an empty stack. Then the value 456 is pushed. The head of the stack is at the “bottom”

A pop is the opposite of a push. The pop operation get the most recent value from the stack (FIFO) in this case 456 leaving 123 on the stack.

Add

An add operation pops the top two items from the stack, adds them and pushes the result back onto the stack

Look at the definition of the instructions for these three operators

Instruction { opCode = Push, opParamCount = 1, opPopsCount = 0, opSimple = True }
Instruction { opCode = Pop, opParamCount = 0, opPopsCount = 1, opSimple = True }
Instruction { opCode = Add, opParamCount = 0, opPopsCount = 2, opSimple = True }

From this you can see that the Instruction shows that a push takes one parameter, a pop takes no parameters but pops a single value off the stack and an Add pops two values off the stack. As noted in the code comments opSimple indicates that these are simple operators with fixed stack effects.

Jmp

The Jmp operator performs an unconditional jump to a fixed location relative to the start of the code stream. I.e. Jmp 100 sets the instruction pointer to 100 and execution continues from there.

Consider the following simple list of ops in Haskel

 [ (Jmp, [3])
 , (Nop, [])
 , (Halt, [])
 ]

This gets assembled into the following byte code

The Jmp instruction causes the CPU to set the instruction pointer (ip) to 3. In this example that means that the Nop at offset 2 is skipped and execution continues with the Halt operation at offset 3

Branching (Beq, Bne, Blr, Blte, Bgt, Bgte)

Branching is a conditional jump. The top two values are popped off the stack compared based on the type of conditional operator.

In the following example the values 1 and to are pushed. Bgt is executed and if 2 is greater than 1 then the CPU jumps to location 7. If no then it continues executing at the location after the branch (6)

 [ (Push, [1])
 , (Push, [2])
 , (Bgt, [7])
 , (Nop, [])
 , (Halt, [])
 ]

This is what the history of the CPU would look like for the above example

Cpu {ip = -1, cpuStack =    [], ranOp = 0,  state = "",     panic = False}
Cpu {ip =  1, cpuStack =   [1], ranOp = 3,  state = "Push", panic = False}
Cpu {ip =  3, cpuStack = [2,1], ranOp = 3,  state = "Push", panic = False}
Cpu {ip =  5, cpuStack =    [], ranOp = 12, state = "Bgt",  panic = False}
Cpu {ip =  6, cpuStack =    [], ranOp = 0,  state = "Nop",  panic = False}
Cpu {ip =  7, cpuStack =    [], ranOp = 2,  state = "Halt", panic = True}
Call & Ret

A function call is much like a jmp except that you have to store an address to return to. You could have two stacks, one for values and one for return address. It’s more common however to have a single stack with “stack frames”.

As a trivial example consider a the case when there are no parameters

The CPU does the following

However this simple scheme does not work when you have variable numbers of parameters, locals etc. This is where the frame pointer (fp) and stack frames come in.

A stack frame is the set of data stored on the stack for each method call. In this virtual machine that is 1. The return address 2. Parameters for the function 3. The previous frame pointer value

As an example consider a function that adds two numbers and returns the sum.

The important thing to notice is that since the old frame pointer is stored on the stack you are able to call multiple functions and always be able to return to the previous function. Also having the fp lets you unwind the stack to the frame no matter how many items the current function may have pushed onto the stack, i.e. you don’t need to try and track that

Here is the output for the code above

The interpreter code
-- | Interpreter for the byte code
-- | Given a byte code stream will run the code
interpretByteCode :: S.Seq Int -> [Cpu]
  interpret [emptyCpu] byteCode
  where
    interpret :: [Cpu] -> S.Seq Int -> [Cpu]

    -- Ensure that this function is not called with an empty CPU list
    interpret [] _ = [emptyCpu { state = "INVALID: no start CPU" }]

    -- Start interpreting 
    interpret cpus@(cpu:_) code =
      -- Move to next op code
      let atIp = ip cpu + 1 in

      -- Try get the code byte at index atIp 
      case indexMay code atIp of
        Nothing ->
          -- No byte at expected index, return error
          cpu { ip = atIp, ranOp = 0, state = "INVALID: invalid ip index, reading past end of byte stream", panic = True } : cpus
          
        Just opByte ->
          -- Ensure this is a valid opcode

          case toEnumMay opByte :: Maybe Operation of
            Nothing ->
              -- This is not a valid opcode, error
              cpu { ip = atIp, ranOp = opByte, state = "INVALID: Unknown op", panic = True } : cpus

            Just op ->
              -- Get the instruction for the op code
              case (op, Map.lookup op instrByOp) of
                (_, Nothing) ->
                  -- The byte was an opcode enum but was not configured as an instruction
                  cpu { ip = atIp, ranOp = opByte, state = "INVALID: Op not found", panic = True } : cpus
  1. Check that there is a start CPU
  2. Check that there is an opcode in the sequence at the ip index
  3. Check that the opcode is valid, i.e. belongs to the Operation enum
  4. Check that the opcode is in the instruction map instrByOp
                (_, Just instr) ->
                  -- 'params' are the bytes from the code stream that are used as parameters for the op, e.g. to be pushed onto stack
                  -- get the params from the byte stream into a list
                  let paramsCount = opParamCount instr in 
                  let params = S.take paramsCount $ S.drop (atIp + 1) code in

                  -- 'Pops' are the bytes popped from the stack and used by the current instruction
                  -- get the values from the stack into the list
                  let popsCount = opPopsCount instr in 
                  let (pops, remainingStack) = S.splitAt popsCount $ cpuStack cpu in

                  if S.length params > paramsCount 
                  then
                    cpu { ip = atIp, ranOp = opByte, state = "Code underflow", panic = True } : cpus
                  else
                    if length pops < popsCount
                    then cpu { ip = atIp, ranOp = opByte, state = "Stack underflow", panic = True } : cpus
                    else
                      -- Interpret the opcode using the simple/complex interpreter as indicated by the instruction
                      let next = if opSimple instr then
                                   let res = interpretSimple emptyCpu { ip = atIp + paramsCount, state = show op } op (toList pops) params in
                                   cpu { ranOp = opByte
                                       , cpuStack = cpuStack res <> remainingStack
                                       , ip = ip res
                                       , panic = panic res
                                       , state = state res
                                       }
                                 else
                                   interpretComplex cpu { ip = atIp + paramsCount, state = show op, ranOp = opByte } op (toList pops) params remainingStack
                                  
                      in if panic next
                      then
                        -- In panic state, add current CPU and stop executing
                        next : cpus
                      else
                        -- Everything fine, add CPU and interpret nex byte
                        interpret (next : cpus) code

Finally the core interpreter code can be called. Since the params, pops are now stored as lists and all checks performed this code is quite simple.

Simple instructions
    -- | Simple instructions, can not directly change CPU state, e.g. cant set ip/fp and they just return data to be added to stack
    interpretSimple :: Cpu -> Operation -> [Int] -> S.Seq Int -> Cpu
    interpretSimple cpu op popped params =
      case op of
        Nop -> cpu 
        Push -> cpu { cpuStack = params }
        Pop -> cpu
        Bne -> branchIf cpu popped params (/=)
        Beq -> branchIf cpu popped params (==)
        Bgt -> branchIf cpu popped params (>)
        Bgte -> branchIf cpu popped params (>=)
        Blt -> branchIf cpu popped params (<)
        Blte -> branchIf cpu popped params (<=)
    -- | Perform a branch instruction if the predicate for the op (beq, bne, bgte etc) is true
    branchIf :: Cpu -> [Int] -> S.Seq Int -> (Int -> Int -> Bool) -> Cpu
    branchIf cpu popped params prd =
      case head params of
        Nothing -> cpu { panic = True, state = "missing param" }
        Just jmp ->
          case popped of
            [b, a] -> cpu { ip = if prd a b then ip cpu + jmp else ip cpu } 
            _ -> cpu { panic = True, state = "invalid stack" }
Complex Instructions
    -- | Complex instructions have full access to the CPU and must ensure the CPU is correctly setup for the next instruction
    interpretComplex :: Cpu -> Operation -> [Int] -> S.Seq Int -> S.Seq Int -> Cpu
    interpretComplex cpu op popped params remainingStack =
      case op of
        Call ->
          -- Set up a stack frame and jump to the call location. See the comments above for details of the stack frame
          singleVal cpu (toList params) (\v -> cpu { ip = v - 1 -- 1 byte before method, next loop increments to function start byte
                                                   , fp = S.length remainingStack -- frame pointer points to start of new stack frame
                                                   , cpuStack = S.fromList [ip cpu + 1, fp cpu] <> remainingStack -- add return ip and current fp to stack
                                                   }) ""
        Ret ->
          -- Return to the calling code, get the previous frame pointer (fp) and return address from the current stack frame.
          -- A ret always 'returns' a single value, the last item on the stack 
          let stackAtFp = S.drop (S.length remainingStack - fp cpu - 2) remainingStack in
          let (retParams, retStack) = S.splitAt 2 stackAtFp in
          if S.length retParams == 2 then
            let [retIp, retFp] = toList retParams in
            cpu { cpuStack = S.fromList popped <> retStack
                , ip = retIp - 1
                , fp = retFp
                }
          else
            cpu { panic = True, state = "Stack underflow getting return frame" }

What’s next?

This explanation of the stack machine is significantly longer than the code for it :). Hopefully you’ll be able to see how easy creating a simple stack machine is. There are many things than can be built on top of this. E.g.

  1. https://en.wikipedia.org/wiki/Stack_machine
  2. https://github.com/sdiehl/protolude

See also