wip: stackify

gyehoek.cabal

@@ -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

src/Gyehoek/CPS/Stackify.hs

@@ -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

src/Gyehoek/Stack/Syntax.hs

@@ -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)
+  | Call Val (List Val)
-  | CallReg Name (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)

src/Gyehoek/Stack/VM.hs

@@ -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)
+stepI e vm (Call v xs) =
-  where l = case vm ^? #registers . at r . _Just . #ObjLabel of
+  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}|]
+          Nothing -> error [i|undefined label: #{l}|]
-
-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}|]
 
 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

test/Gyehoek/Test/CPS/Stackify.hs

@@ -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)))|]
+  ]

test/Gyehoek/Test/Stack/VM.hs

@@ -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"]
     ]
   ]
 

test/Main.hs

@@ -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
   ]