vmcompiler: make functional command compile

app/VMCompiler.hs

@@ -2,15 +2,17 @@
 
 module Main (main) where
 
-import           Control.Monad         (forM_)
+import           Control.Monad         (forM, forM_)
 import qualified Data.Text.IO          as TIO
 import           Nand2Tetris.Hack      (AsmLine (..), HackInstruction (..),
                                         canonicalRepresentation)
 import           Nand2Tetris.VM        (AnnotatedAsmLine (..),
-                                        compileInstructions, parseInstructions)
+                                        bootstrapPreamble, compileInstructions,
-import           Options.Applicative   (execParser, fullDesc, header, help,
+                                        parseInstructions)
-                                        helper, info, long, metavar, progDesc,
+import           Options.Applicative   (argument, execParser, fullDesc, header,
-                                        short, strOption, (<**>))
+                                        help, helper, info, long, metavar,
+                                        progDesc, short, some, str, strOption,
+                                        (<**>))
 import           System.IO             (IOMode (WriteMode), hPutStrLn, withFile)
 import           Text.Megaparsec       (runParser)
 import           Text.Megaparsec.Error (ParseErrorBundle,
@@ -20,20 +22,22 @@ import           Text.Megaparsec.Error (ParseErrorBundle,
 data Options =
   Options
   {
-    inputFile  :: FilePath,
+    inputFiles :: [FilePath],
     outputFile :: FilePath
   }
 
 main :: IO ()
 main = do
   options <- execParser opts
-  programText <- TIO.readFile (inputFile options)
+  parsed <- mconcat <$> forM (inputFiles options) (\fname -> do
-  let result = runParser parseInstructions (inputFile options) programText
+    txt <- TIO.readFile fname
+    let result = runParser parseInstructions fname txt
     case result of
-    Left err -> putStrLn $ errorBundlePretty err
+      Left err -> putStrLn (errorBundlePretty err) >> return []
-    Right parsed -> do
+      Right p  -> return p)
-      let compiled = compileInstructions parsed
+  let compiled = bootstrapPreamble <> compileInstructions parsed
-      withFile (outputFile options) WriteMode $ \h -> forM_ compiled (\x -> TIO.hPutStrLn h (printLine x))
+  withFile (outputFile options) WriteMode $ \h -> forM_ (compiled) (\x -> TIO.hPutStrLn h (printLine x))
+
   where
     printLine (Comment t) = "// " <> t
     printLine (Code i)    = canonicalRepresentation (LineInstruction i)
@@ -44,10 +48,7 @@ main = do
             <> progDesc "Compiles Hack VM code to Hack assembly"
             <> header "vmcompiler - Hack VM compiler (nand2tetris)" )
     hackAsmOptParser = Options <$>
-      strOption ( long "input"
+      some (argument str (help "Input files" )) <*>
-               <> short 'i'
-               <> metavar "FILENAME"
-               <> help "Input file" ) <*>
       strOption ( long "output"
                <> short 'o'
                <> metavar "FILENAME"

src/Nand2Tetris/VM.hs

@@ -15,7 +15,8 @@ module Nand2Tetris.VM
     compileInstructions,
     defaultCompilerEnv,
     runCompiler,
-    runCompilerDef
+    runCompilerDef,
+    bootstrapPreamble
   ) where
 
 import           Control.Monad.State.Strict (State, evalState, forM, gets,
@@ -27,7 +28,7 @@ import           Nand2Tetris.Hack           (HackInstruction (A, C),
                                              Imm (Imm, Label),
                                              Jump (JEq, JGt, JLt, JNe, JNone, Jmp),
                                              Reg (RegA, RegD, RegM),
-                                             Source (S0, SDAndM, SDOrM, SDPlusA, SDPlusM, SMMinusD, SNeg, SNeg1, SNot, SReg, SRegMinus1, SRegPlus1),
+                                             Source (S0, SDAndM, SDMinusA, SDOrM, SDPlusA, SDPlusM, SMMinusD, SNeg, SNeg1, SNot, SReg, SRegMinus1, SRegPlus1),
                                              canonicalRepresentation)
 import           Text.Megaparsec            (MonadParsec (try), Parsec, choice,
                                              empty, eof, many, (<|>))
@@ -40,6 +41,13 @@ data VMInstruction =
     VMArithmetic ArithmeticCmd
   | VMMemAccess AccessType Segment
   | VMBranching BranchingCmd
+  | VMFunction FunctionCmd
+  deriving (Show, Eq)
+
+data FunctionCmd =
+    FFunction T.Text Int
+  | FCall T.Text Int
+  | FReturn
   deriving (Show, Eq)
 
 data BranchingCmd =
@@ -108,7 +116,9 @@ instrAnnotation (VMMemAccess accessType segment) =
 instrAnnotation (VMBranching (BLabel label)) = "label " <> label
 instrAnnotation (VMBranching (BGoto label)) = "goto " <> label
 instrAnnotation (VMBranching (BIfGoto label)) = "if-goto " <> label
-
+instrAnnotation (VMFunction (FFunction name nvars)) = "function " <> name <> " " <> (T.pack . show $ nvars)
+instrAnnotation (VMFunction (FCall name nargs)) = "call " <> name <> " " <> (T.pack . show $ nargs)
+instrAnnotation (VMFunction FReturn) = "return"
 
 type Parser = Parsec Error T.Text
 
@@ -122,9 +132,17 @@ parseInstruction :: Parser VMInstruction
 parseInstruction =
   try parseMemAccess <|>
   try parseArithmetic <|>
-      parseBranching
+  try parseBranching <|>
+      parseFunction
 
   where
+    parseFunction :: Parser VMInstruction
+    parseFunction = VMFunction <$>
+      choice [ string "function" *> space1 *> (FFunction <$> (parseId <* space1) <*> L.decimal),
+               string "call" *> space1 *> (FCall <$> (parseId <* space1) <*> L.decimal),
+               string "return" $> FReturn
+             ]
+
     parseBranching :: Parser VMInstruction
     parseBranching = VMBranching <$>
       choice [ string "label" *> space1 *> (BLabel <$> parseId),
@@ -192,7 +210,8 @@ data AnnotatedAsmLine =
 
 data CompilerEnv = CompilerEnv
   {
-    ceLabelCounter :: Int
+    ceLabelCounter    :: Int,
+    ceCurrentFunction :: Maybe T.Text
   }
 
 type Compiler = State CompilerEnv
@@ -205,7 +224,7 @@ nextLabelId = do
 
 
 defaultCompilerEnv :: CompilerEnv
-defaultCompilerEnv = CompilerEnv 0
+defaultCompilerEnv = CompilerEnv 0 Nothing
 
 runCompiler :: CompilerEnv -> Compiler a -> a
 runCompiler = flip evalState
@@ -279,7 +298,114 @@ compileInstruction instr = do
     compileInstruction' (VMBranching (BIfGoto label)) =
       return [ Code (A (Label "SP")),
                Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
-               Code (C [] (SReg RegM) JNe)
+               Code (C [RegD] (SReg RegM) JNone),
+               Code (A (Label label)),
+               Code (C [] (SReg RegD) JNe)
+             ]
+    compileInstruction' (VMFunction (FFunction name nvars)) = do
+      modify' (\s -> s { ceCurrentFunction = Just name })
+      return ([ ALabel name,
+                Code (A (Label "SP")) ]<>
+               (concat . replicate nvars $
+                 [ Code (C [RegM] S0 JNone),
+                   Code (C [RegA] (SRegPlus1 RegA) JNone)
+                 ]))
+
+    compileInstruction' (VMFunction (FCall name nargs)) = do
+      lab <- nextLabelId
+      mbCurrentFunction <- gets ceCurrentFunction
+      let labelId =
+            case mbCurrentFunction of
+              Just fname -> fname <> ".ret." <> lab
+              Nothing    -> lab
+      return $ pushLabel labelId <>
+               pushLabelValue "LCL" <>
+               pushLabelValue "ARG" <>
+               pushLabelValue "THIS" <>
+               pushLabelValue "THAT" <>
+               -- ARG = SP - 5 - nargs
+               [ Code (A (Label "SP")),
+                 Code (C [RegD] (SReg RegM) JNone),
+                 Code (A (Imm 5)),
+                 Code (C [RegD] SDMinusA JNone),
+                 Code (A (Imm nargs)),
+                 Code (C [RegD] SDMinusA JNone),
+                 Code (A (Label "ARG")),
+                 Code (C [RegM] (SReg RegD) JNone)
+               ] <>
+               -- LCL = SP
+               [ Code (A (Label "SP")),
+                 Code (C [RegD] (SReg RegM) JNone),
+                 Code (A (Label "LCL")),
+                 Code (C [RegM] (SReg RegD) JNone)
+               ] <>
+               -- goto function & and generate return label
+               [
+                 Code (A (Label name)),
+                 Code (C [] S0 Jmp),
+                 ALabel labelId
+               ]
+
+    compileInstruction' (VMFunction FReturn) = do
+      return $
+        -- endFrame (R15) = LCL
+        [ Code (A (Label "LCL")),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Imm 15)),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- retAddr (R14) = *(endFrame - 5)
+        [ Code (A (Imm 5)),
+          Code (C [RegA] SDMinusA JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Imm 14)),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        --  *ARG = pop()
+        [ Code (A (Label "SP")),
+          Code (C [RegA] (SRegMinus1 RegM) JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Label "ARG")),
+          Code (C [RegA] (SReg RegM) JNone),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        --  SP = ARG+1
+        [ Code (C [RegD] (SRegPlus1 RegA) JNone),
+          Code (A (Label "SP")),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- THAT = *(endFrame - 1), endFrame--
+        [ Code (A (Imm 15)),
+          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Label "THAT")),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- THIS = *(endFrame - 1), endFrame--
+        [ Code (A (Imm 15)),
+          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Label "THIS")),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- ARG = *(endFrame - 1), endFrame--
+        [ Code (A (Imm 15)),
+          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Label "ARG")),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- ARG = *(endFrame - 1), endFrame--
+        [ Code (A (Imm 15)),
+          Code (C [RegA, RegM] (SRegMinus1 RegM) JNone),
+          Code (C [RegD] (SReg RegM) JNone),
+          Code (A (Label "LCL")),
+          Code (C [RegM] (SReg RegD) JNone)
+        ] <>
+        -- goto retAddr
+        [ Code (A (Imm 14)),
+          Code (C [RegA] (SReg RegM) JNone),
+          Code (C [] S0 Jmp)
         ]
 
     compileConditional jumpType = do
@@ -327,6 +453,16 @@ compileInstruction instr = do
         Code (C [RegM] (SReg RegD) JNone)
       ]
 
+    pushLabel label = [
+          Code (A . Label $ label),
+          Code (C [RegD] (SReg RegA) JNone),
+          Code (A (Label "SP")),
+          Code (C [RegA] (SReg RegM) JNone),
+          Code (C [RegM] (SReg RegD) JNone),
+          Code (C [RegD] (SRegPlus1 RegA) JNone),
+          Code (A (Label "SP")),
+          Code (C [RegM] (SReg RegD) JNone) ]
+
     pushLabelValue label = [
           Code (A . Label $ label),
           Code (C [RegD] (SReg RegM) JNone),
@@ -386,3 +522,15 @@ compileInstruction instr = do
 
                            Code (A . Imm $ offset),
                            Code (C [RegM] (SReg RegD) JNone) ]
+
+bootstrapPreamble :: [AnnotatedAsmLine]
+bootstrapPreamble = [ Comment "Bootstrap preamble",
+                      Code (A (Imm 261)),
+                      Code (C [RegD] (SReg RegA) JNone),
+                      Code (A (Label "SP")),
+                      Code (C [RegM] (SReg RegD) JNone),
+                      Code (A (Label "LCL")),
+                      Code (C [RegM] (SReg RegD) JNone),
+                      Code (A (Label "Sys.init")),
+                      Code (C [] S0 Jmp)
+                    ]