From b9eb9506ef65add252a8224e666489db5df71656 Mon Sep 17 00:00:00 2001
From: Olle Fredriksson <fredriksson.olle@gmail.com>
Date: Sat, 25 May 2024 19:37:02 +0200
Subject: [PATCH] wip

---
 src/LowToLLVM.hs | 67 +++++++++++++++++++++++++++++-------------------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/src/LowToLLVM.hs b/src/LowToLLVM.hs
index c0bb425..463cdd1 100644
--- a/src/LowToLLVM.hs
+++ b/src/LowToLLVM.hs
@@ -25,9 +25,12 @@ import Low.PassBy (PassBy)
 import qualified Low.PassBy as PassBy
 import qualified Low.Representation as Representation
 import qualified Low.Syntax as Syntax
+import Monad hiding (freshVar)
 import Name (Name)
 import qualified Name
 import Protolude hiding (IntMap, cast, local, moduleName, repr)
+import qualified Query
+import Rock.Core
 
 newtype Var = Var Text
   deriving (Eq, Ord, Show, Hashable)
@@ -78,7 +81,7 @@ wordBits = 64
 wordSizedInt :: Builder
 wordSizedInt = "i" <> Builder.intDec wordBits
 
-type Assembler = State AssemblerState
+type Assembler = StateT AssemblerState M
 
 data AssemblerState = AssemblerState
   { usedGlobals :: HashSet Name.Lifted
@@ -89,11 +92,11 @@ data AssemblerState = AssemblerState
   , basicBlocks :: Builder
   }
 
-runAssembler :: Assembler a -> (a, (HashSet Name.Lifted, HashMap Text Builder))
+runAssembler :: Assembler a -> M (a, (HashSet Name.Lifted, HashMap Text Builder))
 runAssembler =
-  second (\s -> (s.usedGlobals, s.usedLLVMGlobals))
+  fmap (second (\s -> (s.usedGlobals, s.usedLLVMGlobals)))
     . flip
-      runState
+      runStateT
       AssemblerState
         { usedLocals = mempty
         , usedGlobals = mempty
@@ -165,15 +168,30 @@ declareLLVMGlobal name decl =
 
 -------------------------------------------------------------------------------
 
-assembleModule :: [(Name.Lifted, Syntax.Definition)] -> Lazy.ByteString
+saveStack :: Assembler Var
+saveStack = do
+  declareLLVMGlobal "llvm.stackrestore" "declare ccc void @llvm.stackrestore(ptr)"
+  var <- freshVar "stack"
+  emitInstruction $ varName var <> " = call ccc ptr @llvm.stacksave()"
+  pure var
+
+restoreStack :: Var -> Assembler ()
+restoreStack var = do
+  declareLLVMGlobal "llvm.stacksave" "declare ccc ptr @llvm.stackesave()"
+  emitInstruction $ "call ccc void @llvm.stackrestore" <> parens ["ptr " <> varName var]
+
+-------------------------------------------------------------------------------
+
+assembleModule :: [(Name.Lifted, Syntax.Definition)] -> M Lazy.ByteString
 assembleModule definitions = do
-  let (assembledDefinitions, allUsedGlobals) =
-        unzip $ foreach definitions $ runAssembler . uncurry assembleDefinition
-      (usedGlobals, usedLLVMGlobals) = bimap mconcat mconcat $ unzip allUsedGlobals
+  (assembledDefinitions, allUsedGlobals) <-
+    unzip <$> forM definitions (runAssembler . uncurry assembleDefinition)
+  let (usedGlobals, usedLLVMGlobals) = bimap mconcat mconcat $ unzip allUsedGlobals
       assembledDefinitions' = concat assembledDefinitions
-  Builder.toLazyByteString $
-    separate "\n\n" $
-      HashMap.elems usedLLVMGlobals <> map snd assembledDefinitions'
+  pure $
+    Builder.toLazyByteString $
+      separate "\n\n" $
+        HashMap.elems usedLLVMGlobals <> map snd assembledDefinitions'
 
 type Environment v = Index.Seq v (PassBy, Operand)
 
@@ -195,7 +213,8 @@ assembleFunction functionName env = \case
     assembleFunction functionName (env Index.Seq.:> (passBy, Local var)) function
   Syntax.Body passReturnBy term -> do
     let parameters = second fromLocal <$> Index.Seq.toSeq env
-    (result, restore) <- assembleTerm env term
+    (result, stack) <- assembleTerm env Nothing term
+    mapM_ restoreStack stack
     emitInstruction $ "ret " <> llvmType passReturnBy <> " " <> operand result
     basicBlocks <- gets (.basicBlocks)
     pure
@@ -228,15 +247,15 @@ assembleFunction functionName env = \case
     alignment :: (Num a) => a
     alignment = 8
 
-assembleTerm :: Environment v -> Syntax.Term v -> Maybe Name -> Assembler (Operand, Maybe Var)
+assembleTerm :: Environment v -> Maybe Name -> Syntax.Term v -> Assembler (Operand, Maybe Var)
 assembleTerm env nameSuggestion = \case
-  Syntax.Operand operand -> (,Nothing) <$> assembleOperand env operand
+  Syntax.Operand o -> (,Nothing) <$> assembleOperand env o
   Syntax.Let passBy name term body -> do
     (termResult, termStack) <- assembleTerm env (Just name) term
-    (bodyResult, bodyStack) <- assembleTerm (env Index.Seq.:> (passBy, termResult)) body
+    (bodyResult, bodyStack) <- assembleTerm (env Index.Seq.:> (passBy, termResult)) nameSuggestion body
     mapM_ restoreStack termStack
     mapM_ restoreStack bodyStack
-    pure (result, Nothing)
+    pure (bodyResult, Nothing)
   Syntax.Seq term1 term2 -> do
     (_, stack1) <- assembleTerm env Nothing term1
     (result, stack2) <- assembleTerm env nameSuggestion term2
@@ -256,15 +275,11 @@ assembleTerm env nameSuggestion = \case
 
 assembleOperand :: Environment v -> Syntax.Operand v -> Assembler Operand
 assembleOperand env = \case
-  Syntax.Var index -> _
-  Syntax.Global global -> _
-  Syntax.Literal literal -> _
+  Syntax.Var index -> pure $ snd $ Index.Seq.index env index
+  Syntax.Global global -> pure $ Global global
+  Syntax.Literal (Literal.Integer i) -> pure $ Constant $ Builder.integerDec int
   Syntax.Representation repr -> _
-  Syntax.Tag tag -> _
+  Syntax.Tag constr -> do
+    (_, maybeTag) <- fetch $ Query.ConstructorRepresentation constr
+    pure $ Constant $ Builder.intDec $ fromMaybe 0 maybeTag
   Syntax.Undefined repr -> _
-
-saveStack :: Assembler Var
-saveStack = undefined
-
-restoreStack :: Var -> Assembler ()
-restoreStack = undefined