works albeit comically inefficiently

gyehoek.cabal

@@ -63,6 +63,7 @@ library
     Gyehoek.CPS.Stackify
     Gyehoek.CPS.Syntax
     Gyehoek.Driver
+    Gyehoek.Language
     Gyehoek.GenSym
     Gyehoek.Jalmot
     Gyehoek.Lift1

src/Gyehoek/CPS/Close.hs

@@ -12,7 +12,7 @@ import Gyehoek.Prelude
 close :: GenSym :> es => Exp -> Eff es Exp
 close = transformM \case
   ExpLetRec [(f, AbsLambda lam@(MkLambda bs kb m))] e -> do
-    f_code <- gensym' @Name $ f ^. _Wrapped'. to (<> "-code")
+    f_code <- gensym' @Name $ f ^. _Wrapped' . to (<> "-code")
     -- it would probably be most sane to generate a symbol for `env`,
     -- but we're reusing the lambda binding so we don't have to
     -- explicitly substitute recursive calls.

src/Gyehoek/CPS/Stackify.hs

@@ -27,7 +27,7 @@ runStackify = runWriter
 live :: Free a => Env -> a -> List Name
 -- TODO: free' should return an OSet lol
 live g e = nub (free' e) & filter \x ->
-  x `H.member` g.bound
+  x `elem` g.bound
   -- && not (x `elem` g.contStack)
 
 data BlockBuilder
@@ -48,14 +48,14 @@ stackify
   => Env -> Exp -> Eff es BlockBuilder
 
 stackify g (ExpLetRec [(f, AbsKappa kap)] e) = do
-  stackifyKappa g (MkLabel f) kap \g' kap' -> do
-    emitRoutine kap'
-    stackify g' e
+  kap' <- stackifyKappa g kap
+  emitRoutine (Stk.MkRoutine (MkLabel f) . buildBlock $ kap')
+  stackify g e
 
 stackify g (ExpLetRec [(f, AbsLambda lam)] e) = do
-  stackifyLambda g (MkLabel f) lam \g' lam' -> do
-    emitRoutine lam'
-    stackify g' e
+  lam' <- stackifyLambda g (MkLabel f) lam
+  emitRoutine lam'
+  stackify g e
 
 stackify g (ExpIf c t f) = do
   let c' = stackifyVal g c
@@ -63,96 +63,99 @@ stackify g (ExpIf c t f) = do
   f' <- buildBlock <$> stackify g f
   pure . Tail $ Stk.If c' t' f'
 
--- stackify g (ExpApply f xs ktail) = pure $
---     Code [ Stk.Push (Stk.ValReg l) | l <- ls ] $
---     Tail (Stk.TailCall (stackifyVal g f) (k : (stackifyVal g <$> xs)))
---   where
---     k = var g ktail
---     ls = fold $ (k ^? #ValImm . #ImmLabel)
---       >>= \klbl -> g ^. #liveness . at klbl
-
--- stackify g e@(ExpContinue k xs) = do
---   pure $
---     Code [ Stk.Push (Stk.ValReg l) | l <- ls ] $
---     Tail (Stk.TailCall k' (stackifyVal g <$> xs))
---   where
---     k' = stackifyVal g k
---     ls = fold $ (k' ^? #ValImm . #ImmLabel)
---       >>= \klbl -> g ^. #liveness . at klbl
-
-stackify g (ExpPrim p (MkKappa [x] e)) = do
-  e' <- stackify (g & #bound . at x ?~ Stk.ValReg (MkReg x)) e
+stackify g (ExpApply f xs ktail) = do
   pure $
-    Code [ Stk.Prim (MkReg x) (stackifyVal g <$> p) ] e'
+    Code [ Stk.Push $ stackifyVal g (ValVar ktail)
+         , Stk.Push $ stackifyVal g f
+         ] $
+    Code (pushArgs g xs) $
+    Tail (Stk.Call (length xs))
+
+-- assume that `k` is the continuation on top of the stack lol.
+stackify g e@(ExpContinue k xs)
+  | isn't (#_ValVar . only g.tail) k = pure $
+    Code [ Stk.Push (stackifyVal g k) ] $
+    Code (pushArgs g xs) $
+    Tail $ Stk.TailCall (length xs)
+  | otherwise = pure $
+    Code (pushArgs g xs) $
+    Tail (Stk.Return (length xs))
+
+stackify g (ExpPrim p kap) = do
+  kap' <- stackifyKappa g kap
+  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] kap'
 
 stackify _ e = error [i|unimplemented exp: #{e}|]
 
+loadArgs :: List Name -> List Stk.Instr
+loadArgs = imapOf itraversed \n x -> Stk.Load (MkReg x) n
+
+pushArgs :: Env -> List Val -> List Stk.Instr
+pushArgs g args = [ Stk.Push $ stackifyVal g x | x <- reverse args ]
+
 -- affine
 _ValName :: Traversal' Val Name
-_ValName = failing #ValVar (#ValImm . #ImmLabel . #MkLabel)
+_ValName = failing #_ValVar (#_ValImm . #_ImmLabel . #_MkLabel)
 
 stackifyKappa
   :: (Stackify :> es, GenSym :> es)
-  => Env -> Label -> Kappa
-  -> (Env -> Stk.Routine -> Eff es r)
-  -> Eff es r
-stackifyKappa g name kap@(MkKappa xs m) w = _
--- stackifyKappa g name kap@(MkKappa xs m) w = do
---   let vs = (name, Stk.ValLabel name) : (bindReg <$> xs)
---   let ls = live g kap
---   m' <- stackify (g & #bound <>~ H.fromList (vs ++ (bindReg <$> ls))) m
---   let g' = g & #bound . at name ?~ Stk.ValLabel name
---              & #liveness . at name ?~ live g kap
---   let rt = Stk.MkRoutine name xs . buildBlock $
---        -- pop in the opposite order we push
---        Code [Stk.Pop x | x <- reverse ls] m'
---   w g' rt
+  => Env -> Kappa
+  -> Eff es BlockBuilder
+stackifyKappa g (MkKappa xs m) = do
+  let g' = g & #bound <>:~ xs
+  Code (loadArgs g'.bound)
+    <$> stackify g' m
 
 stackifyLambda
   :: (Stackify :> es, GenSym :> es)
   => Env -> Label -> Lambda
-  -> (Env -> Stk.Routine -> Eff es r)
-  -> Eff es r
-stackifyLambda g name (MkLambda xs k m) w = do
-  let vs = [ (x, Stk.ValReg (MkReg x)) | x <- k:xs ]
-  m' <- stackify (g & #bound <>~ H.fromList vs) m
-  let g' = g & #bound . at (name ^. wrappedIso) ?~ Stk.ValLabel name
-  w g' $ Stk.MkRoutine name (buildBlock m')
+  -> Eff es Stk.Routine
+stackifyLambda g name (MkLambda xs k m) = do
+  m' <- stackify (g & #bound .~ xs & #tail .~ k) m
+  pure $
+    Stk.MkRoutine name . buildBlock $
+    Code (loadArgs xs) $
+    Code [Stk.Load (MkReg k) (length xs + 1)] m'
 
 stackifyVal :: Env -> Val -> Stk.Val
 stackifyVal g = \case
   ValImm imm -> Stk.ValImm imm
-  ValVar v -> var g v
+  ValVar v -> case regOf g v of
+    Just r -> Stk.ValReg r
+    Nothing -> Stk.ValLabel (MkLabel v)
   v -> error [i|unimplemented val: #{v}|]
 
-var :: Env -> Name -> Stk.Val
-var g v = case g ^. #bound . at v of
-  Just x -> x
-  Nothing -> Stk.ValLabel (MkLabel v)
-
-bindReg :: Name -> (Name, Stk.Val)
-bindReg x = (x, Stk.ValReg (MkReg x))
+regOf :: Env -> Name -> Maybe Reg
+regOf g x
+  | x `elem` g.bound || x == g.tail = Just . MkReg $ x
+  | otherwise                       = Nothing
 
 
 
 data Env = MkEnv
-  { bound :: HashMap Name Stk.Val
+  -- | `bound` tracks the stack lifetime of bound variables.
+  { bound :: List Name
   -- | for each locally-bound continuation @k@, @liveness@ has an
   -- entry @(k,ls)@ where @ls@ is the sequence of registers @k@
   -- expects to find saved on the stack.
-  , liveness :: HashMap Name (List Name)
+  , liveness :: HashMap Label (List Name)
+  , tail :: Name
   }
   deriving (Show, Generic)
 
 emptyEnv :: Env
-emptyEnv = MkEnv mempty mempty
+emptyEnv = MkEnv
+  { bound = mempty
+  , liveness = mempty
+  , tail = "halt"
+  }
 
 
 
 stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
 stackifyProgram (MkProgram lam) = do
   let g = emptyEnv
-  (_,p) <- runStackify $ stackifyLambda g "start" lam (const emitRoutine)
+  (_,p) <- runStackify $ emitRoutine =<< stackifyLambda g "start" lam
   pure p
 
 letfn :: Program
@@ -168,3 +171,21 @@ letfn = [cps|
       (continue let-body6 lambda-body1))))
 |]
 
+blah :: Program
+blah = [cps|
+(λ (ktail0)
+          (letrec ((fac (λ (n ktail)
+                             (prim (zero? n)
+                               (κ (x0)
+                                 (if x0
+                                   (continue ktail 1)
+                                   (prim (- n 1)
+                                     (κ (x1)
+                                       (letrec ((fac-k0
+                                                  (κ (x2)
+                                                    (prim (* n x2)
+                                                      (κ (x3)
+                                                        (continue ktail x3))))))
+                                         (fac x1 fac-k0))))))))))
+               (fac 6 halt)))
+|]

src/Gyehoek/Language.hs

@@ -0,0 +1,4 @@
+module Gyehoek.Language
+  (
+  ) where
+

src/Gyehoek/Stack/Syntax.hs

@@ -75,7 +75,7 @@ data Instr
   = Pop Reg
   | Push Val
   | Load Reg Int
-  | Prim Reg (Prim Val)
+  | Prim (Prim Val)
   deriving stock (Show, Generic, Data)
   deriving anyclass (NFData)
 
@@ -98,7 +98,7 @@ instance S.DatumIso Instr where
     $ S.With (S.headTagged1 "pop!" S.datumIso >>>)
     $ S.With (S.headTagged1 "push!" S.datumIso >>>)
     $ S.With (S.headTagged2 "load" S.datumIso S.datumIso >>>)
-    $ S.With (S.headTagged2 "prim" S.datumIso S.datumIso >>>)
+    $ S.With (S.headTagged1 "prim" S.datumIso >>>)
     $ S.End
     where
 

src/Gyehoek/Stack/VM.hs

@@ -122,39 +122,7 @@ stepI e vm (Load r j) = do
 stepI e vm (Push v) = traverseOf activeFrame push vm
   where push xs = cons <$> evalVal e vm v <*> pure xs
 
-stepI e vm (Prim r p) = traverse (evalVal e vm) p >>= \case
-  PrimZeroP x -> case x of
-    ObjImm (ImmInt n) -> ret . ObjImm . ImmBool $ n == 0
-    _ -> vmerror [i|bad arg to zero?: #{x}|]
-  PrimAdd x y -> arith_binop (+) x y
-  PrimMul x y -> arith_binop (*) x y
-  PrimSub x y -> arith_binop (-) x y
-  PrimDiv x y -> arith_binop div x y
-  PrimMakeClosure f env ->
-    case f of
-      ObjImm (ImmLabel l) -> ret . ObjHob $ HobClosure l env
-      _ -> vmerror [i|expected label, got #{f}|]
-  PrimEnvCode env ->
-    case env of
-      ObjHob (HobClosure l _) -> ret . ObjImm . ImmLabel $ l
-      _ -> vmerror [i|expected closure, got #{env}|]
-  PrimEnvRef env n ->
-    case env of
-      ObjHob (HobClosure _ xs) -> ret $ xs ^?! ix n
-      _ -> vmerror [i|expected closure, got #{env}|]
-  PrimCons x y -> ret $ ObjHob $ HobPair x y
-  PrimCar x -> case x of
-      ObjHob (HobPair car _) -> ret car
-      _ -> vmerror [i|expected pair, got ${x}|]
-  PrimCdr x -> case x of
-      ObjHob (HobPair _ cdr) -> ret cdr
-      _ -> vmerror [i|expected pair, got ${x}|]
-  x -> vmerror [i|unimplemented prim: #{p}|]
-  where
-    ret v = pure $ vm & #registers . at r ?~ v
-    arith_binop op (ObjImm (ImmInt x)) (ObjImm (ImmInt y)) =
-      ret $ ObjImm (ImmInt (op x y))
-    arith_binop _ x y = vmerror [i|bad arith: #{x}, #{y}|]
+stepI g vm (Prim p) = stepP g vm p
 
 stepI e vm (Pop r) = case vm ^? activeFrame . _Cons of
   Nothing -> vmerror "empty stack"
@@ -188,6 +156,9 @@ stepT g vm tc@(Return nret) = do
       vm & traverseOf #stack (fmap snd . popFrame)
          & mapped . activeFrame %~ pushes xs
          & mapped %~ jumpToRoutine rt
+         -- it is not essential we clear the registers, but it'll make
+         -- bugs more obvious.
+         & mapped . #registers .~ mempty
 
 stepT g vm tc@(TailCall nargs) = do
   (args,f,ra) <- parseTailCall nargs (vm ^. activeFrame)
@@ -203,8 +174,8 @@ stepT g vm tc@(TailCall nargs) = do
         & jumpToRoutine rt
         -- replace the active frame; don't push a new one.
         & activeFrame .~ newFrame
-        -- it is not essential we clear the registers, but it'll
-        -- make bugs more obvious.
+        -- it is not essential we clear the registers, but it'll make
+        -- bugs more obvious.
         & #registers .~ mempty
 
 stepT g vm (If c t f) = do
@@ -213,6 +184,42 @@ stepT g vm (If c t f) = do
       _ -> t
   pure $ jumpToBlock branch vm
 
+stepP :: Jalmot :> es => Env -> VM -> Prim Val -> Eff es VM
+stepP g vm p = traverse (evalVal g vm) p >>= \case
+  PrimZeroP x -> case x of
+    ObjImm (ImmInt n) -> ret1 . ObjImm . ImmBool $ n == 0
+    _ -> vmerror [i|bad arg to zero?: #{x}|]
+  PrimAdd x y -> arith_binop (+) x y
+  PrimMul x y -> arith_binop (*) x y
+  PrimSub x y -> arith_binop (-) x y
+  PrimDiv x y -> arith_binop div x y
+  PrimMakeClosure f env ->
+    case f of
+      ObjImm (ImmLabel l) -> ret1 . ObjHob $ HobClosure l env
+      _ -> vmerror [i|expected label, got #{f}|]
+  PrimEnvCode env ->
+    case env of
+      ObjHob (HobClosure l _) -> ret1 . ObjImm . ImmLabel $ l
+      _ -> vmerror [i|expected closure, got #{env}|]
+  PrimEnvRef env n ->
+    case env of
+      ObjHob (HobClosure _ xs) -> ret1 $ xs ^?! ix n
+      _ -> vmerror [i|expected closure, got #{env}|]
+  PrimCons x y -> ret1 $ ObjHob $ HobPair x y
+  PrimCar x -> case x of
+      ObjHob (HobPair car _) -> ret1 car
+      _ -> vmerror [i|expected pair, got ${x}|]
+  PrimCdr x -> case x of
+      ObjHob (HobPair _ cdr) -> ret1 cdr
+      _ -> vmerror [i|expected pair, got ${x}|]
+  x -> vmerror [i|unimplemented prim: #{p}|]
+  where
+    ret vs = pure $ vm & activeFrame . #locals <>:~ vs
+    ret1 v = ret [v]
+    arith_binop op (ObjImm (ImmInt x)) (ObjImm (ImmInt y)) =
+      ret1 $ ObjImm (ImmInt (op x y))
+    arith_binop _ x y = vmerror [i|bad arith: #{x}, #{y}|]
+
 
 
 popFrame :: (HasCallStack, Jalmot :> es) => Stack -> Eff es (Frame, Stack)
@@ -380,6 +387,11 @@ ppDoc p t =
         .syn-paren-2 { color: green; }
         .syn-paren-3 { color: navy; }
         .syn-paren-4 { color: purple; }
+        .stack-frame
+        { display: inline-flex
+        ; flex-direction: row
+        ; column-gap: 0.5em
+        }
       """
     body_ do
       details_ do
@@ -432,11 +444,26 @@ ppVM vm = do
     td_ do
       code_ curi
     td_ do
-      let xs = vm ^.. activeFrame . each . to ppDatum
-      sequence_ $ intersperse " | " xs
+      ppStack vm.stack
   where
     curi = vm ^?! failing (#code . _head . to ppDatum) (#tail . to ppDatum)
 
+ppStack :: Stack -> Html ()
+ppStack stk = do
+  span_ [class_ "stack"] do
+    stk ^.. each
+      & fmap ppFrame
+      & intersperse " | "
+      & sequence_
+
+ppFrame :: Frame -> Html ()
+ppFrame frm = do
+  span_ [class_ "stack-frame"] do
+    sequence_ $ frm ^.. #locals . each . to ppDatum
+
+ppData :: S.DataIso a => a -> Html ()
+ppData = htmlData . runJalmotUnsafe . S.toData S.dataIso
+
 ppDatum :: S.DatumIso a => a -> Html ()
 ppDatum = htmlDatum . runJalmotUnsafe . S.toDatum S.datumIso
 

test/Gyehoek/Test/CPS/Stackify.hs

@@ -20,42 +20,44 @@ test_stackify =
   , procedure
   ]
 
-evalsTo :: List Obj -> Sut.Exp -> Assertion
+evalsTo :: HasCallStack => List Obj -> Sut.Program -> Assertion
 evalsTo rs e = runJalmotUnsafe (Stk.eval e') @?= rs
   where
-    e' = e & CPS.MkLambda [] "_ktail"
-      & CPS.MkProgram
-      & Sut.stackifyProgram & runGenSym & runPureEff
+    e' = e & Sut.stackifyProgram & runGenSym & runPureEff
 
 trivialReturn = testGroup "trivial return"
   [ testCase "return int" do
       evalsTo [ObjImm (ImmInt 4)]
-        [cps|(continue halt 4)|]
+        [cps|(λ (ktail) (continue ktail 4))|]
   , testCase "return bool" do
       evalsTo [ObjImm (ImmBool True)]
-        [cps|(continue halt #t)|]
+        [cps|(λ (ktail) (continue ktail #t))|]
       evalsTo [ObjImm (ImmBool False)]
-        [cps|(continue halt #f)|]
+        [cps|(λ (ktail) (continue ktail #f))|]
   ]
 
 tailCall = testGroup "tail call"
   [ testCase "square" do
-      evalsTo [ObjImm (ImmInt 16)]
-        [cps|(letrec ((square (λ (x ktail)
-                                (prim (* x x)
-                                  (κ (x0) (continue ktail x0))))))
-               (square 4 halt))|]
+      evalsTo [ObjImm (ImmInt 16)] [cps|
+        (λ (ktail0)
+          (letrec ((square (λ (x ktail)
+                             (prim (* x x)
+                               (κ (x0) (continue ktail x0))))))
+            (square 4 halt)))
+      |]
   ]
 
 prim = testGroup "prim"
   [ testCase "multiply" do
       evalsTo [ObjImm (ImmInt 20)]
-        [cps|(prim (* 4 5)
-               (κ (x) (continue halt x)))|]
+        [cps|(λ (ktail0)
+            (prim (* 4 5)
+               (κ (x) (continue ktail0 x))))|]
   , testCase "add" do
       evalsTo [ObjImm (ImmInt 9)]
-        [cps|(prim (+ 4 5)
-               (κ (x) (continue halt x)))|]
+        [cps|(λ (ktail0)
+        (prim (+ 4 5)
+               (κ (x) (continue ktail0 x))))|]
   -- , testGroup "call/cc"
   --   [ testCase "trivial" do
   --       evalsTo [ObjImm (ImmInt 123)]
@@ -66,14 +68,17 @@ prim = testGroup "prim"
 
 condition = testCase "if" do
   evalsTo [ObjImm (ImmInt 123)]
-    [cps|(if #t (continue halt 123) (continue halt 456))|]
+    [cps|(λ (ktail0)
+      (if #t (continue ktail0 123) (continue ktail0 456)))|]
   evalsTo [ObjImm (ImmInt 456)]
-    [cps|(if #f (continue halt 123) (continue halt 456))|]
+    [cps|(λ (ktail0)
+      (if #f (continue ktail0 123) (continue ktail0 456)))|]
 
 procedure = testGroup "procedure"
   [ testCase "factorial" do
       evalsTo [ObjImm (ImmInt 720)]
-        [cps|(letrec ((fac (λ (n ktail)
+        [cps|(λ (ktail0)
+          (letrec ((fac (λ (n ktail)
                              (prim (zero? n)
                                (κ (x0)
                                  (if x0
@@ -86,5 +91,5 @@ procedure = testGroup "procedure"
                                                       (κ (x3)
                                                         (continue ktail x3))))))
                                          (fac x1 fac-k0))))))))))
-               (fac 6 halt))|]
+               (fac 6 halt)))|]
   ]

test/Gyehoek/Test/Stack/VM.hs

@@ -76,8 +76,7 @@ test_root = testGroup "stack machine"
           (tail-call 1))
         (define $square
           (pop! %x)
-          (prim %x2 (* %x %x))
-          (push! %x2)
+          (prim (* %x %x))
           (return 1))
       |]
   , testGroup "factorial"
@@ -90,20 +89,19 @@ test_root = testGroup "stack machine"
             (tail-call 1))
           (define $fac
             (load %n 0)
-            (prim %x0 (zero? %n))
+            (prim (zero? %n))
+            (pop! %x0)
             (if %x0
                 (then (push! 1)
                       (return 1))
-                (else (prim %x1 (- %n 1))
-                      (push! $fac-c0)
+                (else (push! $fac-c0)
                       (push! $fac)
-                      (push! %x1)
+                      (prim (- %n 1))
                       (call 1))))
           (define $fac-c0
             (pop! %x2)
             (pop! %n)
-            (prim %x3 (* %n %x2))
-            (push! %x3)
+            (prim (* %n %x2))
             (return 1))
       |]
         mkcase n = testCase [i|#{n}|] do