1 files changed, 142 insertions, 22 deletions

diff --git a/2-kyu/assembler-interpreter.hs b/2-kyu/assembler-interpreter.hs
index cf18b68..a7a6bb7 100644
--- a/2-kyu/assembler-interpreter.hs
+++ b/2-kyu/assembler-interpreter.hs
@@ -1,8 +1,16 @@
 module AssemblerInterpreter where
 
+import Debug.Trace
+
+import Data.Function((&))
 import Data.Maybe (fromMaybe)
-import Text.Parsec
+import Text.Parsec hiding (State)
 import Text.Parsec.Char
+import qualified Data.Map as M
+
+import Control.Monad.Trans.Maybe
+import Control.Monad.Trans.State
+import Control.Monad.Trans.Class(lift)
 
 interpret :: String -> Maybe String
 interpret prog = Nothing
@@ -14,26 +22,6 @@ data Op =
   Jge | Jg  | Jle | Jl  | Call | Ret | Msg | End 
   deriving (Show)
 
-argsNum Mov = 2
-argsNum Inc = 1
-argsNum Dec = 1
-argsNum Add = 2
-argsNum Sub = 2
-argsNum Mul = 2
-argsNum Div = 2
-argsNum Jmp = 1
-argsNum Cmp = 2
-argsNum Jne = 1
-argsNum Je  = 1
-argsNum Jge = 1
-argsNum Jg  = 1
-argsNum Jle = 1
-argsNum Jl  = 1
-argsNum Call= 1
-argsNum Ret = 0
-argsNum Msg = -1
-argsNum End = 0
-
 opStrP :: Parser String
 opStrP = choice $ map string' [
   "mov", "inc", "dec", "add", "sub", "mul", "div", "jmp", "cmp", "jne",
@@ -153,13 +141,145 @@ progP = do
         return (++)
   stmts <- chainl stmtLineP newlineP []
   spaces
+  eof
   return stmts
 
+processLabel :: [Stmt] -> ([Stmt], M.Map String Int)
+processLabel stmts = go stmts [] M.empty 0 where
+  go [] processed labelMap i = (reverse processed, labelMap)
+  go (instr@(InstrStmt op args):xs) processed labelMap i =
+    go xs (instr:processed) labelMap (i+1)
+  go ((LabelStmt (Label ident)):xs) processed labelMap i =
+    let newLabelMap = M.insert ident i labelMap 
+    in
+      go xs processed newLabelMap i
+
+data MachineState = MachineState
+  { machineRegisters :: M.Map String Int
+  , machineStack :: [Int]
+  , machineProgCnt :: Int
+  , machineOutput :: Maybe String
+  , machineProg :: [Stmt]
+  , machineLabels :: M.Map String Int
+  , machineEnd :: Bool
+  }
+  deriving (Show)
+
+type MachineM = MaybeT (State MachineState)
+
+nextInstr :: MachineM ()
+nextInstr = do
+  ms <- lift get
+  lift $ put (ms {machineProgCnt = machineProgCnt ms + 1})
+
+fetchInstr :: MachineM (Maybe Stmt)
+fetchInstr = do
+  ms <- lift get
+  let pc = machineProgCnt ms
+      prog = machineProg ms
+  if pc >= length prog || pc < 0 || machineEnd ms then return Nothing
+  else return $ Just (prog !! pc)
+
+setReg :: String -> Int -> MachineM ()
+setReg r n = do
+  ms <- lift get
+  let regs = machineRegisters ms
+  lift $ put ms {machineRegisters = M.insert r n regs}
+
+getReg :: String -> MachineM Int
+getReg r = do
+  ms <- lift get
+  let regs = machineRegisters ms
+  case M.lookup r regs of
+    Just n -> return n
+    Nothing -> fail []
+
+
+execInstr :: Stmt -> MachineM ()
+execInstr (InstrStmt op args) = decodeOp op args
+execInstr _ = fail []
+
+getResult :: MachineM String
+getResult = do
+  ms <- lift get
+  if not (machineEnd ms) then fail []
+  else case machineOutput ms of
+    Nothing -> fail []
+    Just msg -> return msg
+
+runMachine :: MachineM String
+runMachine = do
+  ms <- lift get
+  instr <- fetchInstr
+  case instr of
+    Nothing -> getResult
+    Just stmt -> do execInstr stmt ; runMachine
+
+execMov :: [Arg] -> MachineM ()
+execMov [IdentifierArg reg, NumberArg n] = do setReg reg n
+execMov [IdentifierArg r1, IdentifierArg r2] = do
+  n <- getReg r1
+  setReg r2 n
+  nextInstr
+execMov _ = fail []
+
+execMsg :: [Arg] -> MachineM ()
+execMsg [] = fail []
+execMsg args = go args "" where
+  go [] msg = do
+    ms <- lift get
+    lift $ put ms { machineOutput = Just msg }
+    nextInstr
+  go ((NumberArg n):xs) msg = go  xs (msg ++ show n)
+  go ((StringArg s):xs) msg = go xs (msg ++ s)
+  go ((IdentifierArg reg):xs) msg = do
+    n <- getReg reg
+    go xs (msg ++ show n)
+
+execEnd :: [Arg] -> MachineM ()
+execEnd [] = do
+  ms <- lift get
+  lift $ put ms { machineEnd = True }
+execEdn _ = fail []
+
+decodeOp :: Op -> [Arg] -> MachineM ()
+decodeOp Mov = execMov
+decodeOp Msg = execMsg
+decodeOp End = execEnd
+
+buildMachine :: String -> Maybe MachineState
+buildMachine code = 
+  let
+    parseResult = runParser progP () "" code
+  in
+    case parseResult of
+      Right lines ->
+        let (stmts, labels) = processLabel lines
+        in
+          Just MachineState
+            { machineRegisters = M.empty
+            , machineEnd = False
+            , machineLabels = labels
+            , machineOutput = Nothing
+            , machineProg = stmts
+            , machineProgCnt = 0
+            , machineStack = []
+            }
+      _ -> Nothing
+
+
+interp :: String -> Maybe String
+interp code = do
+  ms <- buildMachine code
+  let (res, _) = runState (runMaybeT runMachine) ms
+  res
+
 main = do
   print $ runParser opP () "" "msg" 
   print $ runParser stmtLineP () "" "ret" 
   print $ runParser progP () "" "ret \n  mov a , 12\n uwu: \n ret" 
   print $ runParser progP () "" "ret ; 123 ; 123\n  mov a , 12\n uwu: \n ret" 
-  print $ runParser progP () "" "\nmov   a, 2            ; value1\nmov   b, 10           ; value2\nmov   c, a            ; temp1\nmov   d, b            ; temp2\ncall  proc_func\ncall  print\nend\n\nproc_func:\n    cmp   d, 1\n    je    continue\n    mul   c, a\n    dec   d\n    call  proc_func\n\ncontinue:\n    ret\n\nprint:\n    msg a, '^', b, ' = ', c\n    ret\n"
+  print $ interp "msg \'hello\'\n end"
+