wip: stackify

2026-08-18 00:16:55 -06:00
parent c4bcf38374
commit d91e059a84
7 changed files with 185 additions and 71 deletions
@@ -55,6 +55,7 @@ library
    Gyehoek.CPS.Close
    Gyehoek.CPS.Convert
    Gyehoek.CPS.Lower
+    Gyehoek.CPS.Stackify
    Gyehoek.CPS.Syntax
    Gyehoek.Driver
    Gyehoek.GenSym
@@ -102,9 +103,8 @@ test-suite test
  type:             exitcode-stdio-1.0
  hs-source-dirs:   test
  main-is:          Main.hs
-
-  -- cabal-fmt: expand test
  other-modules:
+    Gyehoek.Test.CPS.Stackify
    Gyehoek.Test.CPS.Syntax
    Gyehoek.Test.Golden
    Gyehoek.Test.Sexp
@@ -113,6 +113,7 @@ test-suite test
  build-depends:
    , base
    , directory
+    , effectful
    , filepath
    , generic-lens
    , gyehoek
@@ -0,0 +1,87 @@
+{-# LANGUAGE OverloadedLists #-}
+module Gyehoek.CPS.Stackify
+  ( stackifyExp
+  , stackifyProgram
+  , module Gyehoek.CPS.Syntax
+  ) where
+
+import Gyehoek.CPS.Syntax
+import Gyehoek.Stack.Syntax qualified as Stk
+import Data.Sequence (Seq)
+import Effectful
+import Gyehoek.GenSym
+import Effectful.Writer.Static.Shared
+import Control.Lens
+import Data.String.Interpolate
+import Gyehoek.Stack.Syntax (Imm(..))
+import Data.HashSet (HashSet)
+import qualified Data.HashSet as HS
+import GHC.Generics (Generic)
+import Data.Foldable
+import Data.HashMap.Strict (HashMap)
+import qualified Data.HashMap.Strict as H
+
+
+type Stackify = Writer Stk.Program
+
+runStackify :: Eff (Stackify : es) a -> Eff es (a, Stk.Program)
+runStackify = runWriter
+
+stackify
+  :: (GenSym :> es, Stackify :> es)
+  => Env -> Exp -> Eff es (Seq Stk.Instr)
+
+stackify g (ExpLetRec [(f, AbsLambda' xs k m)] e) = do
+  let xs' = (k:xs) <&> \x -> (x, Stk.ValReg x)
+  m' <- stackify (g & #bound .~ H.fromList xs') m
+  lam_body <- gensym' "lambda-body"
+  tell [Stk.MkBlock lam_body xs . toList $
+        [Stk.PopCont "ktail"] <> m']
+  stackify (g & #bound . at f ?~ Stk.ValLabel lam_body) e
+
+stackify g (ExpApply f xs ktail) = do
+  pure [ Stk.PushCont ktail
+       , Stk.Call (stackifyVal g f) (stackifyVal g <$> xs)
+       ]
+
+stackify g (ExpContinue k xs) = pure
+  [ Stk.Call (var g k) (stackifyVal g <$> xs)
+  ]
+
+stackify g (ExpPrim p (MkKappa [x] e)) = do
+  e' <- stackify (g & #bound . at x ?~ Stk.ValReg x) e
+  pure $ [ Stk.Prim x (stackifyVal g <$> p) ] <> e'
+
+stackify _ e = error [i|unimplemented exp: #{e}|]
+
+stackifyVal :: Env -> Val -> Stk.Val
+stackifyVal g = \case
+  ValLit (LitInt n) -> Stk.ValImm (ImmInt n)
+  ValLit (LitBool b) -> Stk.ValImm (ImmBool b)
+  ValVar v -> var g 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 v
+
+
+
+data Env = MkEnv
+  { bound :: HashMap Name Stk.Val
+  }
+  deriving (Show, Generic)
+
+emptyEnv :: Env
+emptyEnv = MkEnv mempty
+
+
+
+stackifyExp :: GenSym :> es => Name -> Exp -> Eff es Stk.Program
+stackifyExp lbl e = do
+  (code,p) <- runStackify $ stackify emptyEnv e
+  pure $ p <> Stk.MkProgram [ Stk.MkBlock lbl [] (code ^.. each) ]
+
+stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
+stackifyProgram (MkProgram e) = stackifyExp "main" e
@@ -10,6 +10,7 @@ module Gyehoek.Stack.Syntax
  , Imm(..)
  , Prim(..)
  , Name
+  , pattern ValLabel
  ) where

 import Control.Lens
@@ -57,23 +58,24 @@ data Instr
  | PopCont Name
  | PushCont Name
  | Prim Name (Prim Val)
-  | CallLabel Name (List Val)
-  | CallReg Name (List Val)
+  | Call Val (List Val)
  | If Val (List Instr) (List Instr)
  deriving stock (Show, Generic, Data)

 data Val
-  = ValLabel Name
-  | ValReg Name
+  = ValReg Name
  | ValImm Imm
  deriving stock (Show, Generic, Data, Eq)

+pattern ValLabel :: Name -> Val
+pattern ValLabel x = ValImm (ImmLabel x)
+
 data Imm
  = ImmInt Int
  | ImmBool Bool
+  | ImmLabel Name
  deriving stock (Show, Generic, Data, Eq)

 data Obj
  = ObjImm Imm
-  | ObjLabel Name
  deriving (Show, Generic, Data, Eq)
@@ -1,8 +1,10 @@
+{-# LANGUAGE ViewPatterns #-}
 module Gyehoek.Stack.VM
  ( VM(..)
  , Env(..)
  , eval
  , trace
+  , module Gyehoek.Stack.Syntax
  ) where

 import Gyehoek.Stack.Syntax
@@ -68,23 +70,18 @@ stepI e vm (Pop r) = case vm ^. #stack of

 stepI e vm (PopCont r) = case vm ^. #kstack of
  [] -> error "empty stack"
-  (x:xs) -> vm & #registers . at r ?~ ObjLabel x
+  (x:xs) -> vm & #registers . at r ?~ ObjImm (ImmLabel x)
               & #kstack .~ xs

-stepI e vm (CallReg r xs) = stepI e vm (CallLabel l xs)
-  where l = case vm ^? #registers . at r . _Just . #ObjLabel of
+stepI e vm (Call v xs) =
+  case evalToLabel e vm v of
+    "halt" -> vm & #result ?~ fmap (evalVal e vm) xs
+    l -> vm & #code .~ b.code
+            & #registers .~ fmap (evalVal e vm) (H.fromList $ b.params `zip` xs)
+      where
+        b = case e ^. #blocks . at l of
          Just x -> x
-          Nothing -> error [i|tried to call undefined register: #{r}|]
-
-stepI e vm (CallLabel "halt" xs) = vm & #result ?~ fmap (evalVal e vm) xs
-
-stepI e vm (CallLabel l xs) =
-  vm & #code .~ b.code
-     & #registers .~ fmap (evalVal e vm) (H.fromList $ b.params `zip` xs)
-  where
-    b = case e ^. #blocks . at l of
-      Just x -> x
-      Nothing -> error [i|undefined label: #{l}|]
+          Nothing -> error [i|undefined label: #{l}|]

 stepI e vm (If c t f) =
  case evalVal e vm c of
@@ -93,6 +90,11 @@ stepI e vm (If c t f) =

 stepI e vm ins = error [i|unimplemented instruction: #{ins}|]

+evalToLabel e vm v =
+  case evalVal e vm v of
+    ObjImm (ImmLabel x) -> x
+    x -> error [i|not a label: #{x}|]
+
 evalVal :: Env -> VM -> Val -> Obj
 evalVal e vm = \case
  ValImm imm -> ObjImm imm
@@ -104,7 +106,7 @@ initialVM :: VM
 initialVM = MkVM
  { stack = []
  , kstack = ["halt"]
-  , code = [CallLabel "main" []]
+  , code = [Call (ValImm $ ImmLabel "main") []]
  , registers = mempty
  , stdout = ""
  , result = Nothing
@@ -131,39 +133,3 @@ trace p = initialVM & unfoldr \vm ->
      Just _ -> Nothing
      Nothing -> Just (vm, step e vm)
  where e = initialEnv p
-
-fac =
-      [ MkBlock "fac" ["n"]
-        [ Prim "x0" $ PrimZeroP (ValReg "n")
-        , If (ValReg "x0")
-          [ PopCont "ktail"
-          , CallReg "ktail" [ValImm (ImmInt 1)]
-          ]
-          [ Push (ValReg "n")
-          , Prim "x1" $ PrimSub (ValReg "n") (ValImm (ImmInt 1))
-          , PushCont "fac-k0"
-          , CallLabel "fac" [ValReg "x1"]
-          ]
-        ]
-      , MkBlock "fac-k0" ["x2"]
-        [ Pop "n"
-        , Prim "x3" $ PrimMul (ValReg "x2") (ValReg "n")
-        , PopCont "ktail"
-        , CallReg "ktail" [ValReg "x3"]
-        ]
-      ]
-
-fac3 = MkProgram $
-  [ MkBlock "main" []
-    [ CallLabel "fac" [ValImm (ImmInt 3)]
-    ]
-  ] <> fac
-
-fac3_trace = trace fac3
-
-showVM :: VM -> Text
-showVM vm = [i|(#{instr}) ; #{stk} #{kstk}|]
-  where
-    instr = vm ^?! #code . _head
-    stk = vm.stack
-    kstk = vm.kstack
@@ -0,0 +1,54 @@
+module Gyehoek.Test.CPS.Stackify (root) where
+
+import Test.Tasty (TestTree, testGroup)
+import Test.Tasty.HUnit
+import qualified Gyehoek.CPS.Stackify as Sut
+import Gyehoek.Stack.VM as Stk
+import Data.List (List)
+import Gyehoek.CPS.Syntax (cps)
+import Gyehoek.GenSym (runGenSym)
+import Effectful
+
+
+root :: IO TestTree
+root = pure . testGroup "stackify" $
+  [ trivialReturn
+  , tailCall
+  , prim
+  ]
+
+evalsTo :: List Obj -> Sut.Exp -> Assertion
+evalsTo rs e =
+  Stk.eval e' @?= rs
+  where e' = runPureEff . runGenSym $ Sut.stackifyExp "main" e
+
+trivialReturn = testGroup "trivial return"
+  [ testCase "return int" do
+      evalsTo [ObjImm (ImmInt 4)]
+        [cps|(continue halt 4)|]
+  , testCase "return bool" do
+      evalsTo [ObjImm (ImmBool True)]
+        [cps|(continue halt #t)|]
+      evalsTo [ObjImm (ImmBool False)]
+        [cps|(continue halt #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))|]
+  ]
+
+prim = testGroup "prim"
+  [ testCase "multiply" do
+      evalsTo [ObjImm (ImmInt 20)]
+        [cps|(prim (* 4 5)
+               (κ (x) (continue halt x)))|]
+  , testCase "add" do
+      evalsTo [ObjImm (ImmInt 9)]
+        [cps|(prim (+ 4 5)
+               (κ (x) (continue halt x)))|]
+  ]
@@ -27,40 +27,42 @@ lit_int = testCase "lit int" do
  evalsTo [ObjImm (ImmInt 3)]
    [ MkBlock "main" []
      [ PopCont "ktail"
-      , CallReg "ktail" [ValImm (ImmInt 3)]
+      , Call (ValReg "ktail") [ValImm (ImmInt 3)]
      ]
    ]

+vlb = ValImm . ImmLabel
+
 procedure = testGroup "procedure"
  [ testCase "return constant" do
      evalsTo [ObjImm (ImmInt 123)]
        [ MkBlock "main" []
-          [ CallLabel "silly" []
+          [ Call (ValLabel "silly") []
          ]
        , MkBlock "silly" []
          [ PopCont "ktail"
-          , CallReg "ktail" [ValImm (ImmInt 123)]
+          , Call (ValReg "ktail") [ValImm (ImmInt 123)]
          ]
        ]
  , testCase "identity function" do
      evalsTo [ObjImm (ImmInt 45)]
        [ MkBlock "main" []
-          [ CallLabel "id" [ValImm (ImmInt 45)]
+          [ Call (ValLabel "id") [ValImm (ImmInt 45)]
          ]
        , MkBlock "id" ["x"]
          [ PopCont "ktail"
-          , CallReg "ktail" [ValReg "x"]
+          , Call (ValReg "ktail") [ValReg "x"]
          ]
        ]
  , testCase "square" do
      evalsTo [ObjImm (ImmInt 16)]
        [ MkBlock "main" []
-          [ CallLabel "square" [ValImm (ImmInt 4)]
+          [ Call (ValLabel "square") [ValImm (ImmInt 4)]
          ]
        , MkBlock "square" ["x"]
          [ PopCont "ktail"
          , Prim "x2" $ PrimMul (ValReg "x") (ValReg "x")
-          , CallReg "ktail" [ValReg "x2"]
+          , Call (ValReg "ktail") [ValReg "x2"]
          ]
        ]
  , testCase "factorial" do
@@ -69,29 +71,29 @@ procedure = testGroup "procedure"
              [ Prim "x0" $ PrimZeroP (ValReg "n")
              , If (ValReg "x0")
                [ PopCont "ktail"
-                , CallReg "ktail" [ValImm (ImmInt 1)]
+                , Call (ValReg "ktail") [ValImm (ImmInt 1)]
                ]
                [ Push (ValReg "n")
                , Prim "x1" $ PrimSub (ValReg "n") (ValImm (ImmInt 1))
                , PushCont "fac-k0"
-                , CallLabel "fac" [ValReg "x1"]
+                , Call (ValLabel "fac") [ValReg "x1"]
                ]
              ]
            , MkBlock "fac-k0" ["x2"]
              [ Pop "n"
              , Prim "x3" $ PrimMul (ValReg "x2") (ValReg "n")
              , PopCont "ktail"
-              , CallReg "ktail" [ValReg "x3"]
+              , Call (ValReg "ktail") [ValReg "x3"]
              ]
            ]
      evalsTo [ObjImm (ImmInt 1)] $
        [ MkBlock "main" []
-          [ CallLabel "fac" [ValImm (ImmInt 0)]
+          [ Call (ValLabel "fac") [ValImm (ImmInt 0)]
          ]
        ] ++ fac
      evalsTo [ObjImm (ImmInt 720)] $
        [ MkBlock "main" []
-          [ CallLabel "fac" [ValImm (ImmInt 6)]
+          [ Call (ValLabel "fac") [ValImm (ImmInt 6)]
          ]
        ] ++ fac
  ]
@@ -110,7 +112,7 @@ trivialPrimTest rs p =
  [ MkBlock "main" []
    [ PopCont "ktail"
    , Prim "x1" p
-    , CallReg "ktail" [ValReg "x1"]
+    , Call (ValReg "ktail") [ValReg "x1"]
    ]
  ]

@@ -6,6 +6,7 @@ import qualified Gyehoek.Test.Golden
 import qualified Gyehoek.Test.Sexp
 import qualified Gyehoek.Test.CPS.Syntax
 import qualified Gyehoek.Test.Stack.VM
+import qualified Gyehoek.Test.CPS.Stackify


 main :: IO ()
@@ -17,5 +18,6 @@ root = testGroup "test" <$> sequenceA
  ,-} Gyehoek.Test.Sexp.root
  , Gyehoek.Test.CPS.Syntax.root
  , Gyehoek.Test.Stack.VM.root
+  , Gyehoek.Test.CPS.Stackify.root
  ]