Wrap emitters with ExceptT monad

app/Assembler/Emitter.hs

@@ -1,6 +1,10 @@
 module Assembler.Emitter where
 
-import Control.Monad.State
+import Control.Monad (when)
+import Control.Monad.Trans (lift)
+import Control.Monad.State (State, execState, get, put)
+import Control.Monad.Except (throwError)
+import Control.Monad.Trans.Except (ExceptT, runExceptT)
 import Data.Word (Word8)
 
 import qualified Data.Map as M
@@ -15,51 +19,52 @@ data Bean = Byte Word8
 data Context = Context { _beans   :: [Bean]
                        , _labels  :: M.Map String Int 
                        } deriving (Show, Eq)
-type Emitter = AST -> State Context ()
+type Emitter = AST -> ExceptT String (State Context) ()
 
 empty :: Context
 empty = Context { _beans = [], _labels = M.fromList [] }
 
-emitLabelDef :: Emitter
+emitBean :: Bean -> ExceptT String (State Context) ()
-emitLabelDef (LabelDef label) = do
+emitBean bean = lift $ do
   ctx <- get
-  let labels = _labels ctx
+  put ctx { _beans = _beans ctx ++ [bean] }
-  let current = length (_beans ctx)
-  put ctx { _labels = M.insert label current labels }
   return ()
-emitLabelDef _ = return()
+
+emitByte :: Word8 -> ExceptT String (State Context) ()
+emitByte byte = emitBean $ Byte $ byte
 
 emitParam :: Emitter
 emitParam (Param (Integer x))  = emitByte $ fromIntegral $ x
-emitParam (Param (LabelRef label)) = emitBean $ Reference $ label
+emitParam (Param (LabelRef l)) = emitBean $ Reference $ l
-emitParam _ = return()
+emitParam _ = throwError "Number or label reference expected"
+
+emitLabelDef :: Emitter
+emitLabelDef (LabelDef label) = do
+  ctx <- lift get
+  let labels = _labels ctx
+  let current = length (_beans ctx)
+  when (label `M.member` labels) (throwError $ "Label '" ++ (label) ++ "' is already defined")
+  put ctx { _labels = M.insert label current labels }
+  return ()
+emitLabelDef _ = throwError "Label definition expected"
 
 emitInstr :: Emitter
 emitInstr (Instruction (Operator op) Empty) = emitByte $ fromIntegral . fromEnum $ op
 emitInstr (Instruction (Operator op) (Params params)) = do
-  emitByte $ fromIntegral . fromEnum $ op
+  emitByte $ fromIntegral $ fromEnum op
   mapM_ emitParam params
   return ()
-emitInstr _ = return()
+emitInstr _ = throwError "Instruction expected"
 
 emitLine :: Emitter
 emitLine (Line labelDef Empty) = emitLabelDef labelDef
 emitLine (Line Empty instr) = emitInstr instr
 emitLine (Line labelDef instr) = emitLabelDef labelDef >> emitInstr instr >> return ()
-emitLine _ = return()
+emitLine _ = throwError "Line of code expected"
 
 emitProgram :: Emitter
 emitProgram (Program progLines) = mapM emitLine progLines >> return ()
-emitProgram _ = return()
+emitProgram _ = throwError "Program code expected"
-
-emitByte :: Word8 -> State Context ()
-emitByte byte = emitBean $ Byte $ byte
-
-emitBean :: Bean -> State Context ()
-emitBean bean = do
-  ctx <- get
-  put ctx { _beans = _beans ctx ++ [bean] }
-  return ()
 
 resolveLabels :: M.Map String Int -> [Bean] -> Either String [Bean]
 resolveLabels labels beans = sequence $ foldr folder [] beans
@@ -74,7 +79,7 @@ resolveLabel labels (Reference label) = case M.lookup label labels of
 
 emit :: AST -> Either String [Word8]
 emit root = do  
-  let ctx = execState (emitProgram root) empty
+  ctx <- flip evalState empty $ runExceptT $ emitProgram root >> lift get 
   let labels = _labels ctx
   let beans = _beans ctx
   resolved <- resolveLabels labels beans

app/VirtualMachine/Interpreter.hs

@@ -3,6 +3,7 @@ module VirtualMachine.Interpreter where
 import Data.Word (Word8)
 import Data.List (intercalate)
 
+import Control.Monad (when, unless)
 import Control.Monad.Trans.State (get, evalStateT)
 import Control.Monad.Trans.Except (ExceptT, except, runExceptT)
 import Control.Monad.Trans (lift)
@@ -23,9 +24,8 @@ parseInstr (opCode:rest) = do
     Nothing -> Left "Unknown instruction"
   let noParams = _noParams instr
   let params = map fromIntegral $ take noParams rest :: [Word8]
-  if length params == noParams 
+  unless (length params == noParams) (Left $ "Expected " ++ (show noParams) ++ " parameter(s), got " ++ (show $ length params) ++ " for operator '" ++ (show op) ++ "'")
-  then return (instr, params)
+  return (instr, params)
-  else Left $ "Expected " ++ (show noParams) ++ " parameter(s), got " ++ (show $ length params) ++ " for operator '" ++ (show op) ++ "'"
 parseInstr [] = Left "Unexpected end of the file"
 
 parse :: [Word8] -> Either String [Unit]
@@ -51,26 +51,22 @@ interpretUnit [] = throwError "Nothing to interpret"
 interpretUnit units = do
   pc <- lift getPc
   let progSize = length units
-  if pc < progSize
+  unless (pc < progSize) (throwError $ "PC (=" ++ (show pc) ++ ") exceeds program size (=" ++ (show progSize) ++ ")")
-  then case units !! pc of
+  case units !! pc of
     (Instr instr) -> dispatchInstr units instr
     (Byte _)      -> throwError $ "PC (=" ++ (show pc) ++ ") currently points to the data byte rather than instruction" 
-  else throwError $ "PC (=" ++ (show pc) ++ ") exceeds program size (=" ++ (show progSize) ++ ")"
 
 dispatchInstr :: [Unit] -> Instruction -> ExceptT String Machine ()
 dispatchInstr units instr = do
   debug <- lift isDebug
   
-  if debug
+  when debug $ lift $ do
-  then lift $ do
     vm <- get
     pc <- getPc
     let noParams = _noParams instr
     let params = intercalate "" $ map (show . _byte) $ take noParams $ drop (pc + 1) $ units
     liftIO $ putStrLn $ show vm
     liftIO $ putStrLn $ (show pc) ++ ": " ++ (show $ _op instr) ++ " " ++ params    
-    return ()
-  else return ()
 
   case instr of
     Simple {}  -> interpretSimple units instr
@@ -85,13 +81,11 @@ interpretSimple units instr = do
   let params = map (fromIntegral . _byte) paramBytes   
   let action = _sAction instr
   pops <- lift $ pop noPops  
-  if length pops == noPops
+  unless (length pops == noPops) (throwError $ "Attempt to pop from empty stack: tried to pop " ++ (show noPops) ++ " elements, got " ++ (show $ length pops))  
-  then lift $ do
   let pushes = action params pops
-    pushS pushes
+  lift $ pushS pushes
-    forward $ noParams + 1
+  lift $ forward $ noParams + 1
   return ()
-  else throwError $ "Attempt to pop from empty stack: tried to pop " ++ (show noPops) ++ " elements, got " ++ (show $ length pops)  
 
 interpretComplex :: [Unit] -> Instruction -> ExceptT String Machine ()
 interpretComplex units instr = do
@@ -102,9 +96,8 @@ interpretComplex units instr = do
   let params = map (fromIntegral . _byte) paramBytes   
   let action = _cAction instr
   pops <- lift $ pop noPops  
-  if length pops == noPops
+  unless (length pops == noPops) (throwError $ "Attempt to pop from empty stack: tried to pop " ++ (show noPops) ++ " elements, got " ++ (show $ length pops))
-  then do action params pops
+  action params pops
-  else throwError $ "Attempt to pop from empty stack: tried to pop " ++ (show noPops) ++ " elements, got " ++ (show $ length pops)  
 
 run :: VM -> B.ByteString -> IO (Either String VM)
 run vm input = evalStateT (runExceptT machine) vm

test/Assembler/EmitterSpec.hs

@@ -3,20 +3,31 @@ module Assembler.EmitterSpec where
 import Test.Hspec
 
 import qualified Data.Map as M
-import Control.Monad.State (execState)
+import Control.Monad.Trans (lift)
+import Control.Monad.State (evalState, get)
+import Control.Monad.Trans.Except (runExceptT)
 
 import Assembler.Tokenizer (tokenize)
 import Assembler.Parser (AST(..), parse)
 import Assembler.Emitter as E
 import VirtualMachine.VM (Op(..))
 
+evalContext :: Context -> AST -> Emitter -> Either String Context
+evalContext ctx ast emitter = flip evalState ctx $ runExceptT $ emitter ast >> lift get 
+
 spec :: Spec
 spec = do
   describe "emitLabelDef" $ do
     it "inserts label definition to the context" $ do
       let ctx = E.empty
       let input = LabelDef "main"
-      execState (emitLabelDef input) ctx `shouldBe` ctx { _labels = M.fromList[("main", 0)] }
+      let expected = Right (ctx { _labels = M.fromList[("main", 0)] })
+      evalContext ctx input emitLabelDef `shouldBe` expected
+    it "does not allow to redefine label" $ do
+      let ctx = E.empty { _labels = M.fromList [("main", 0)] }
+      let input = LabelDef "main"
+      let expected = Left "Label 'main' is already defined"
+      evalContext ctx input emitLabelDef `shouldBe` expected
       
   describe "resolveLabels" $ do
     it "replaces reference with actual byte number" $ do
@@ -26,27 +37,32 @@ spec = do
     it "throws error if label does not exist" $ do
       let beans = [ Byte 1, Byte 2, Reference "not_existing_label", Byte 4 ]
       let labels = M.fromList [("main", 3)]
-      resolveLabels labels beans `shouldBe` Left "Label 'not_existing_label' is not defined"
+      let expected = Left "Label 'not_existing_label' is not defined"
+      resolveLabels labels beans `shouldBe` expected
 
   describe "emitParam" $ do
     it "emits byte for integer literal" $ do
       let ctx = E.empty
       let input = Param (Integer 4)
-      execState (emitParam input) ctx `shouldBe` ctx { _beans = [Byte 0x04] }
+      let expected = Right (ctx { _beans = [Byte 0x04] })
+      evalContext ctx input emitParam `shouldBe` expected
     it "emits reference mark for label reference" $ do
       let ctx = E.empty
       let input = Param (LabelRef "main")
-      execState (emitParam input) ctx `shouldBe` ctx { _beans = [Reference "main"] }
+      let expected = Right (ctx { _beans = [Reference "main"] })
+      evalContext ctx input emitParam `shouldBe` expected
 
   describe "emitInstr" $ do
     it "emits byte for no-param instruction" $ do
       let ctx = E.empty
       let input = Instruction (Operator Halt) Empty
-      execState (emitInstr input) ctx `shouldBe` ctx { _beans = [Byte 0x01] }
+      let expected = Right (ctx { _beans = [Byte 0x01] })
+      evalContext ctx input emitInstr `shouldBe` expected
     it "emits bytes for 2-param instruction" $ do
       let ctx = E.empty
       let input = Instruction (Operator Push) (Params [(Param (Integer 11)), (Param (LabelRef "main"))])
-      execState (emitInstr input) ctx `shouldBe` ctx { _beans = [Byte 0x02, Byte 0x0B, Reference "main"] }
+      let expected = Right (ctx { _beans = [Byte 0x02, Byte 0x0B, Reference "main"] })
+      evalContext ctx input emitInstr `shouldBe` expected
 
   describe "emit" $ do
     it "label resolution works" $ do