From 5c3b7c2c3016f9fdf36d89e668944d2e108c0edf Mon Sep 17 00:00:00 2001
From: crumbtoo <crumb@disroot.org>
Date: Mon, 4 Dec 2023 13:50:48 -0700
Subject: [PATCH] strict arith

---
 docs/src/references/gm-state-transitions.rst | 25 +++++++-
 src/Core/Syntax.hs                           |  8 +--
 src/GM.hs                                    | 60 +++++++++++++++++++-
 3 files changed, 87 insertions(+), 6 deletions(-)

diff --git a/docs/src/references/gm-state-transitions.rst b/docs/src/references/gm-state-transitions.rst
index ac39602..e44c909 100644
--- a/docs/src/references/gm-state-transitions.rst
+++ b/docs/src/references/gm-state-transitions.rst
@@ -147,6 +147,29 @@ Core Transition Rules
       & m
       }
 
+#. Again, building on the previous rules, this rule makes the machine consider
+   unapplied supercombinators to be in WHNF
+
+   .. math::
+      \gmrule
+      { \mathtt{Unwind} : \nillist
+      & a_0 : \ldots : a_n : \nillist
+      & \langle i, s \rangle : d
+      & h
+      \begin{bmatrix}
+            a_0 : \mathtt{NGlobal} \; k \; c
+      \end{bmatrix}
+      & m
+      }
+      { i
+      & a_n : s
+      & d
+      & h
+      & m \\
+      \SetCell[c=2]{c}
+      \text{when $n < k$}
+      }
+
 #. If an application is on top of the stack, :code:`Unwind` continues unwinding
 
    .. math::
@@ -339,7 +362,7 @@ Core Transition Rules
       \end{bmatrix}
       & m
       }
-   
+
 ***************
 Extension Rules
 ***************
diff --git a/src/Core/Syntax.hs b/src/Core/Syntax.hs
index 16d9567..03d2f6f 100644
--- a/src/Core/Syntax.hs
+++ b/src/Core/Syntax.hs
@@ -37,14 +37,14 @@ data Expr = Var Name
           | Lam [Name] Expr
           | App Expr Expr
           | IntE Int
-          deriving (Show, Lift)
+          deriving (Show, Lift, Eq)
 
 infixl 2 :$
 pattern (:$) :: Expr -> Expr -> Expr
 pattern f :$ x = App f x
 
 data Binding = Binding Name Expr
-    deriving (Show, Lift)
+    deriving (Show, Lift, Eq)
 
 infixl 1 :=
 pattern (:=) :: Name -> Expr -> Binding
@@ -55,12 +55,12 @@ data Rec = Rec
          deriving (Show, Eq, Lift)
 
 data Alter = Alter Int [Name] Expr
-    deriving (Show, Lift)
+    deriving (Show, Lift, Eq)
 
 type Name = String
 
 data ScDef = ScDef Name [Name] Expr
-    deriving (Show, Lift)
+    deriving (Show, Lift, Eq)
 
 data Module = Module (Maybe (Name, [Name])) Program
     deriving (Show, Lift)
diff --git a/src/GM.hs b/src/GM.hs
index c9655be..c581284 100644
--- a/src/GM.hs
+++ b/src/GM.hs
@@ -310,6 +310,15 @@ step st = case head (st ^. gmCode) of
             NAp f x     -> st
                          -- leave the Unwind instr; continue unwinding
                          & gmStack %~ (f:)
+
+            NGlobal k c
+                | n < k -> st
+                         & gmCode  .~ i
+                         & gmStack .~ s
+                         & gmDump  .~ d
+                where
+                    ((i,s) : d) = st ^. gmDump
+                    n = st ^. gmStack & length
             -- assumes length s < d (i.e. enough args have been supplied)
             NGlobal n c -> st
                          -- 'jump' to global's code by replacing our current
@@ -348,6 +357,7 @@ primitive1 box unbox f st
         & f
         & box (st & gmStack .~ s)
         & advanceCode
+        & gmStats . stsPrimReductions %~ succ
     where
         putNewStack = gmStack .~ s
         (a:s) = st ^. gmStack
@@ -361,6 +371,7 @@ primitive2 :: (GmState -> b -> GmState) -- boxing function
 primitive2 box unbox f st
         = st'
         & advanceCode
+        & gmStats . stsPrimReductions %~ succ
     where
         (ax:ay:s) = st ^. gmStack
         putNewStack = gmStack .~ s
@@ -434,10 +445,11 @@ buildInitialHeap (Program ss) = mapAccumL allocateSc mempty compiledScs
         -- << [ref/compileSc]
 
         compileR :: Env -> Expr -> Code
-        compileR g e = compileC g e <> [Update d, Pop d, Unwind]
+        compileR g e = compileE g e <> [Update d, Pop d, Unwind]
             where
                 d = length g
 
+        -- compile an expression in a lazy context
         compileC :: Env -> Expr -> Code
         compileC g (Var k)
             | k `elem` domain  = [Push n]
@@ -483,6 +495,52 @@ buildInitialHeap (Program ss) = mapAccumL allocateSc mempty compiledScs
                 compileBinder :: (Binding, Int) -> Code
                 compileBinder (k := v, a) = compileC g' v <> [Update a]
 
+        -- compile an expression in a strict context such that a pointer to the
+        -- expression is left on top of the stack in WHNF
+        compileE :: Env -> Expr -> Code
+        compileE g (IntE n) = [PushInt n]
+        compileE g (Let NonRec bs e) =
+                -- we use compileE instead of compileC
+                mconcat binders <> compileE g' e <> [Slide d]
+            where
+                d = length bs
+                (g',binders) = mapAccumL compileBinder (argOffset d g) addressed
+                -- kinda gross. revisit this
+                addressed = bs `zip` reverse [0 .. d-1]
+
+                compileBinder :: Env -> (Binding, Int) -> (Env, Code)
+                compileBinder m (k := v, a) = (m',c)
+                    where
+                        m' = (k,a) : m
+                        -- make note that we use m rather than m'!
+                        c = compileC m v
+
+        compileE g (Let Rec bs e) =
+                Alloc d : initialisers <> body <> [Slide d]
+            where
+                d = length bs
+                g' = fmap toEnv addressed ++ argOffset d g
+                toEnv (k := _, a) = (k,a)
+                -- kinda gross. revisit this
+                addressed = bs `zip` reverse [0 .. d-1]
+                initialisers = mconcat $ compileBinder <$> addressed
+
+                -- we use compileE instead of compileC
+                body = compileE g' e
+
+                -- we use compileE instead of compileC
+                compileBinder :: (Binding, Int) -> Code
+                compileBinder (k := v, a) = compileC g' v <> [Update a]
+
+        -- special cases for prim functions
+        compileE g ("negate#" :$ a) = compileE g a <>                 [Neg]
+        compileE g ("+#" :$ a :$ b) = compileE g a <> compileE g b <> [Add]
+        compileE g ("-#" :$ a :$ b) = compileE g a <> compileE g b <> [Sub]
+        compileE g ("*#" :$ a :$ b) = compileE g a <> compileE g b <> [Mul]
+        compileE g ("/#" :$ a :$ b) = compileE g a <> compileE g b <> [Div]
+
+        compileE g e = compileC g e ++ [Eval]
+
         -- | offset each address in the environment by n
         argOffset :: Int -> Env -> Env
         argOffset n = each . _2 %~ (+n)