From f600c3b57cafb0c3450a7aa947f23fc57f8b283b Mon Sep 17 00:00:00 2001
From: bunny <bunnyanonn@proton.me>
Date: Sun, 7 Sep 2025 01:20:16 +0100
Subject: [PATCH] added conditions and if else branching

---
 app/Compiler.hs | 71 ++++++++++++++++++++++++++++++++++++++++++-------
 app/Parser.hs   | 11 +++++++-
 2 files changed, 72 insertions(+), 10 deletions(-)

diff --git a/app/Compiler.hs b/app/Compiler.hs
index f328d03..7855c4e 100644
--- a/app/Compiler.hs
+++ b/app/Compiler.hs
@@ -46,7 +46,9 @@ forthPrelude h = do
              \    ret\n\
              \_start:\n"
 
-type Environment = M.Map String ZorthAST
+data Environment = Environment { label :: Int
+                               , environment :: M.Map String ZorthAST
+                               }
 
 handleSymbol :: Handle -> ZorthExpr -> StateT Environment IO ()
 
@@ -99,26 +101,77 @@ handleSymbol h (ZorthASTWord "drop") = do
 
 handleSymbol h (ZorthASTWord ".") = do
     liftIO $ hPutStr h 
-              "    pop rbx\n\
-              \    call print_number\n"
+             "    pop rbx\n\
+             \    call print_number\n"
+    return ()
+
+handleSymbol h (ZorthASTWord "=") = do
+    liftIO $ truthOperator h "sete"
+    return ()
+
+handleSymbol h (ZorthASTWord ">") = do
+    liftIO $ truthOperator h "setg"
+    return ()
+
+handleSymbol h (ZorthASTWord "<") = do
+    liftIO $ truthOperator h "setl"
+    return ()
+
+handleSymbol h (ZorthASTWord ">=") = do
+    liftIO $ truthOperator h "setge"
+    return ()
+
+handleSymbol h (ZorthASTWord "<=") = do
+    liftIO $ truthOperator h "setle"
     return ()
 
 handleSymbol h (ZorthASTWord w) = do
     state <- get
-    liftIO $ compileZorthAST h (state M.! w) state -- parent env to child and discard
+    liftIO $ compileZorthAST h ((environment state) M.! w) state -- parent env to child and discard
     return ()
 
 handleSymbol _ (ZorthASTWordDecl (name,ast)) = do
-    state <- get
-    put $ M.insert name ast state
+    (Environment labels state) <- get
+    put $ Environment labels (M.insert name ast state) -- maybe use `lens`?..
     return ()
 
+-- honestly, this feels kinda janky, the way I manipulate state
+handleSymbol h (ZorthASTIfElse (ifBranch,elseBranch)) = do
+    (Environment l state) <- get
+    let l1 = show $ l+1
+    let l2 = show $ l+2
+    put $ Environment (l+2) state
+    liftIO $ hPutStr h $
+             "    pop rax\n\
+             \    cmp rax, 0\n\
+             \    je .L"<>l1<>"\n"
+    compileZorthASTState h ifBranch
+    (Environment l' _) <- get
+    put $ Environment l' state
+    liftIO $ hPutStr h $ "    jmp .L"<>l2<>"\n.L"<>l1<>":\n"
+    compileZorthASTState h elseBranch
+    (Environment l'' _) <- get
+    put $ Environment l'' state
+    liftIO $ hPutStr h $ ".L"<>l2<>":\n"
+
+truthOperator :: Handle -> String -> IO ()
+truthOperator h s =
+    hPutStr h $
+    "    pop rbx\n\
+    \    pop rax\n\
+    \    cmp rax, rbx\n\
+    \    "<>s<>" al\n\
+    \    movsx rax, al\n\
+    \    push rax\n"
+
+compileZorthASTState :: Handle -> ZorthAST -> StateT Environment IO ()
+compileZorthASTState h ast = (foldl (\m x -> m >> handleSymbol h x) (return ()) ast)
+
 compileZorthAST :: Handle -> ZorthAST -> Environment -> IO ()
-compileZorthAST h ast state = runStateT (foldl (\m x -> m >> handleSymbol h x) (return ()) ast) state >> return ()
+compileZorthAST h ast state = runStateT (compileZorthASTState h ast) state >> return ()
 
 compileZorth :: Handle -> ZorthAST -> IO ()
 compileZorth _ [] = return ()
 compileZorth h xs = do
     forthPrelude h
-    compileZorthAST h xs M.empty
-    
\ No newline at end of file
+    compileZorthAST h xs $ Environment 0 M.empty
\ No newline at end of file
diff --git a/app/Parser.hs b/app/Parser.hs
index 59a4ecd..6cedafb 100644
--- a/app/Parser.hs
+++ b/app/Parser.hs
@@ -10,6 +10,7 @@ type ZorthAST = [ZorthExpr]
 data ZorthExpr = ZorthASTInteger Int
                | ZorthASTWord String
                | ZorthASTWordDecl (String,ZorthAST)
+               | ZorthASTIfElse (ZorthAST,ZorthAST)
                deriving Show
 
 word1 :: Parser String
@@ -55,8 +56,16 @@ pZorthWordDecl = do
     xs <- manyTill pZorthExpr (do { ZorthASTWord ";" <- pZorthWord; return () })
     return $ ZorthASTWordDecl (name,xs)
 
+
+pZorthIfElse :: Parser ZorthExpr
+pZorthIfElse = do
+    ZorthASTWord "if" <- pZorthWord
+    ifBranch <- manyTill pZorthExpr (do { ZorthASTWord "else" <- pZorthWord; return () })
+    elseBranch <- manyTill pZorthExpr (do { ZorthASTWord "fi" <- pZorthWord; return () })
+    return $ ZorthASTIfElse (ifBranch,elseBranch)
+
 pZorthExpr :: Parser ZorthExpr
-pZorthExpr = pZorthWordDecl <++ pZorthInteger <++ pZorthWord
+pZorthExpr = pZorthIfElse <|> pZorthWordDecl <++ pZorthInteger <++ pZorthWord
 
 pZorth :: Parser ZorthAST 
 pZorth = some pZorthExpr