src/Gyehoek/CPS/Close.hs

@@ -11,28 +11,24 @@ import Gyehoek.Prelude
 
 close :: forall es. GenSym :> es => Exp -> Eff es Exp
 close = transformM \case
-  ExpLetRec [(f, AbsLambda lam@(MkLambda bs kb m))] e -> do
+  ExpLetRec [(f, ab)] e -> do
     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.
-    let frees = nub $ free' lam
-    m' <- ifoldrM @_ @_ @(Eff es)
-             (\n x q -> do
-               q_l <- gensym' @Name "env-cont"
-               let p = if x == f then PrimEnv @Val else PrimEnvRef n
-               pure [cps|
-                 (letrec ((#{q_l} (κ (#{x}) #{q})))
-                   (prim #{p}
-                     #{q_l}))
-               |])
-             m frees
+    let frees = nub $ freeWithBound' [f] ab
+    env_cont_l <- gensym' @Name "env-cont"
+    let m = ab ^. absBody
+    let m' = [cps|
+      (letrec ((#{env_cont_l} (κ #{frees} #{m})))
+        (prim (get-env) #{env_cont_l}))
+    |]
     e_l <- gensym' @Name "make-closure-cont"
+    let ab' = ab & absBody .~ m'
     pure [cps|
-      (letrec ((#{f_code} (λ (##{bs} #{kb})
-                            #{m'})))
+      (letrec ((#{f_code} #{ab'}))
         (letrec ((#{e_l} (κ (#{f}) #{e})))
-          (prim (make-closure #{f_code} ##{frees})
+          (prim (make-closure ($ #{f_code}) ##{frees})
             #{e_l})))
     |]
 

src/Gyehoek/CPS/Contify.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE ApplicativeDo #-}
 module Gyehoek.CPS.Contify
   ( contifyProgram
   ) where
@@ -7,31 +8,43 @@ import Gyehoek.CPS.Syntax
 import Gyehoek.Prelude
 import qualified Data.HashSet as HS
 import Control.Lens.Unsound (adjoin)
+import Debug.Pretty.Simple
+import qualified Data.HashMap.Strict as H
+import Control.Monad.Writer.Lazy
+import Control.Monad.Trans.Tardis (liftTardisT)
 
 
 -- | ain't no way...
-type T es = TardisT (HashSet Name) (HashSet Name) (Eff es)
+-- type T = WriterT (HashSet Name) (Tardis (HashSet Name) (HashSet Name))
+type T = TardisT (HashSet Name) (HashSet Name) (Writer (HashSet Name))
 
-runT :: T es a -> Eff es a
-runT = (`evalTardisT` (mempty,mempty))
+evalT :: T a -> a
+-- evalT = (`evalTardis` (mempty,mempty)) . fmap fst . runWriterT
+evalT = fst . runWriter . (`evalTardisT` (mempty,mempty))
+
+runT :: T a -> (a, HashSet Name)
+-- runT = (`evalTardis` (mempty,mempty)) . runWriterT
+runT = runWriter . (`evalTardisT` (mempty,mempty))
 
 -- | inline function if it hasn't been used in the past, and won't
 -- be used in the future.
-tryInline :: Name -> Kappa -> T es Kexp
+tryInline :: Name -> Kappa -> T Kexp
 tryInline kname kap = do
   modifyBackwards (HS.insert kname)
   p <- getsPast (HS.member kname)
   modifyForwards (HS.insert kname)
   q <- getsFuture (HS.member kname)
-  pure $ if p || q
+  let c = p || q
+  liftTardisT . tell $ if c then HS.singleton kname else mempty
+  pure $ if c
     then KexpVar kname
     else KexpKappa kap
 
 getKap :: HashMap Name Abs -> Name -> Maybe Kappa
 getKap g kname = g ^? ix kname . #AbsKappa
 
-contify :: HashMap Name Abs -> Exp -> T es Exp
-contify g = \case
+contify :: HashMap Name Abs -> Exp -> T Exp
+contify g = transformM \case
   ExpApply f xs (KexpVar kname) | Just kap <- getKap g kname
    -> ExpApply f xs <$> tryInline kname kap
   ExpPrim p (KexpVar kname) | Just kap <- getKap g kname
@@ -41,63 +54,14 @@ contify g = \case
 contifyProgram :: HoistedProgram -> Eff es HoistedProgram
 contifyProgram p = do
   let g = p.bindings
-  runT $
-    traverseOf
-      (adjoin
-        (#bindings . each . body)
-        (#body . body))
-      (contify g)
-      p
-
-
-
-type BidirectionalState = Tardis (HashSet Text) (HashSet Text)
-
-runBidirectionalState = (`evalTardis` (mempty,mempty))
-
-markAsUsed :: Hashable a => a -> HashSet a -> HashSet a
-markAsUsed = HS.insert
-
-isMarkedAsUsed :: Hashable a => a -> HashSet a -> Bool
-isMarkedAsUsed = HS.member
-
--- | our bidirectional state monad. updates can be sent forwards or
--- backwards.
-runBidirectionalState :: BidirectionalState a -> a
-
-data InlineResult
-  -- | represents the successful inlining of a function.
-  = Inlined Text
-  -- | represents a function that couldn't be inlined.
-  | NotInlined Text
-  deriving (Show)
-
-
--- | try inlining two functions, `x` and `y`. each is used exactly
--- once, so both should be successfully inlined.
--- 
--- >>> runBidirectionalState example1
--- ( Inlined "x"
--- , Inlined "y"
--- )
-example1 :: BidirectionalState (InlineResult, InlineResult)
-example1 = do
-  x <- tryInline "x"
-  y <- tryInline "y"
-  pure (x,y)
-
--- | try inlining two functions, `x` and `y`. `x` is used twice, so
--- can't be inlined. but `y` is used once, and thus ought to be
--- inlined.
---
--- >>> runBidirectionalState example2
--- ( NotInlined "x"
--- , Inlined "y"
--- , NotInlined "x"
--- )
-example2 :: BidirectionalState (InlineResult, InlineResult, InlineResult)
-example2 = do
-  x <- tryInline "x"
-  y <- tryInline "y"
-  x' <- tryInline "x"
-  pure (x,y,x')
+  let (p',contifiedVars) =
+        runT $
+          traverseOf
+            (adjoin
+              (#bindings . each . body)
+              (#body . body))
+            (contify g)
+            p
+  pTraceShowM contifiedVars
+  -- pure $ p' & #bindings %~ H.filterWithKey \k _ -> HS.member k contifiedVars
+  pure p'

src/Gyehoek/CPS/Stackify.hs

@@ -54,28 +54,27 @@ emitRoutine :: Stackify :> es => Stk.Routine -> Eff es ()
 emitRoutine rt = tell [rt]
 
 stackify
-  :: (GenSym :> es, Stackify :> es)
+  :: forall es. (GenSym :> es, Stackify :> es)
   => Env -> Exp -> Eff es BlockBuilder
 
-stackify g (ExpLetRec [(f, AbsKappa kap)] e) = do
-  kap' <- stackifyKappa g kap
-  emitRoutine . Stk.MkRoutine (MkLabel f) . buildBlock $ kap'
-  stackify g e
-
-stackify g (ExpLetRec [(f, AbsLambda lam)] e) = do
-  lam' <- stackifyLambda g (MkLabel f) lam
-  emitRoutine lam'
+stackify g (ExpLetRec bs e) = do
+  for_ bs \(f,a) ->
+    emitRoutine =<< case a of
+      AbsKappa kap -> stackifyKappa g (MkLabel f) kap
+      AbsLambda lam -> stackifyLambda g (MkLabel f) lam
   stackify g e
 
 stackify g (ExpIf c t f) = do
   let c' = stackifyVal g c
-  t' <- buildBlock <$> stackify g t
-  f' <- buildBlock <$> stackify g f
-  pure . Tail $ Stk.If c' t' f'
+  let jump l =
+       Stk.MkBlock
+         [Stk.Push . Stk.ValLabel . MkLabel $ l]
+         (Stk.TailCall 0)
+  pure . Tail $ Stk.If c' (jump t) (jump f)
 
 stackify g (ExpApply f xs ktail) = do
   pure $
-    Code [ Stk.Push $ stackifyVal g (ValVar ktail)
+    Code [ Stk.Push $ stackifyVal g (ValVar $ ktail ^?! #KexpVar)
          , Stk.Push $ stackifyVal g f
          ] $
     Code (pushArgs g xs) $
@@ -91,23 +90,24 @@ stackify g e@(ExpContinue k xs)
     Tail (Stk.Return (length xs))
 
 stackify g (ExpPrim (PrimCallCC withcc) cc) = do
-  cc' <- stackifyKappa g cc
+  let cc' = cc ^?! #KexpVar . to MkLabel
   cc_l <- gensym' @Label "cc"
-  emitRoutine . Stk.MkRoutine cc_l . buildBlock $ cc'
   pure $
     Code [ Stk.Push $ stackifyVal g withcc
-         , Stk.Push $ stackifyVal g (ValLabel cc_l)
+         , Stk.Push $ stackifyVal g (ValLabel cc')
          ] $
     Tail Stk.CallCC
 
-stackify g (ExpPrim p (MkKappa rs e)) = do
-  _
-  pure $ Code [ Stk.Prim (stackifyVal g <$> p) ] _
+stackify g (ExpPrim p cc) = pure $
+  Code [ Stk.Push (Stk.ValLabel . MkLabel $ cc ^?! #KexpVar)
+       , Stk.Prim (stackifyVal g <$> p)
+       ] $
+  Tail $ Stk.TailCall 1
 
 stackify _ e = error [i|unimplemented exp: #{e}|]
 
-loadArgs :: Free a => Env -> a -> List Name -> List Stk.Instr
-loadArgs g e = imapOf itraversed \n x -> load g e (MkReg x) n
+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 ]
@@ -118,12 +118,20 @@ _ValName = failing #_ValVar (#_ValImm . #_ImmLabel . #_MkLabel)
 
 stackifyKappa
   :: (Stackify :> es, GenSym :> es)
-  => Env -> Kappa
-  -> Eff es BlockBuilder
-stackifyKappa g (MkKappa xs m) = do
+  => Env -> Label -> Kappa
+  -> Eff es Stk.Routine
+stackifyKappa g kname (MkKappa xs m) = do
   let g' = g & #bound <>:~ xs
-  Code [ _ | x <- g'.bound `intersect` free' m ]
-    <$> stackify g' m
+  m' <- stackify g' m
+  pure $
+    Stk.MkRoutine kname . buildBlock $
+    Code (loadArgs xs) $
+    Code [ Stk.Load (MkReg r) j
+         | v <- g ^.. #liveness . ix kname . each
+         , (j,r) <- itoListOf (#bound . itraversed) g
+         , r == v
+         ] $
+    m'
 
 stackifyLambda
   :: (Stackify :> es, GenSym :> es)
@@ -134,7 +142,8 @@ stackifyLambda g name (MkLambda xs k m) = do
   pure $
     Stk.MkRoutine name . buildBlock $
     Code (loadArgs xs) $
-    Code [Stk.Load (MkReg k) (length xs + 1)] m'
+    -- Code [Stk.Load (MkReg k) (length xs + 1)] $
+    m'
 
 stackifyVal :: Env -> Val -> Stk.Val
 stackifyVal g = \case
@@ -171,40 +180,80 @@ emptyEnv = MkEnv
 
 
 
-stackifyProgram :: GenSym :> es => Program -> Eff es Stk.Program
-stackifyProgram (MkProgram lam) = do
-  let g = emptyEnv
-  (_,p) <- runStackify $ emitRoutine =<< stackifyLambda g "start" lam
-  pure p
+stackifyProgram :: GenSym :> es => HoistedProgram -> Eff es Stk.Program
+stackifyProgram p = do
+  let liveness = p & foldMapOf
+        (#bindings . itraversed . withIndex . aside #AbsKappa)
+        \(kname,kap) -> H.singleton
+          (MkLabel kname)
+          (nub $ freeWithBound' (H.keysSet p.bindings) kap)
+  let g = MkEnv
+        { bound = mempty
+        , liveness
+        , tail = p.body.ktail }
+  let e = p.body & #body %~ ExpLetRec (H.toList p.bindings)
+  (_,p') <- runStackify $ emitRoutine =<< stackifyLambda g "start" e
+  pure p'
 
-letfn :: Program
-letfn = [cps|
-(λ (start-ktail0)
-  (letrec ((lambda-body1
-            (λ (x lambda-tail2)
-              (prim (* x x) (κ (r3) (continue lambda-tail2 r3))))))
-    (letrec ((let-body6
-              (κ (square)
-                (letrec ((r4 (κ (x5) (continue start-ktail0 x5))))
-                  (square 4 r4)))))
-      (continue let-body6 lambda-body1))))
-|]
-
-blah :: Program
+blah :: HoistedProgram
 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)))
+(letrec ((prim-k3 (κ (r2) (if r2 truthy-cont4 falsey-cont5)))
+         (prim-k7 (κ (r6) (fac r6 r8)))
+         (make-closure-cont15 (κ (fac) (fac 20 r12)))
+         (falsey-cont5 (κ () (prim (- n 1) prim-k7)))
+         (r12 (κ (x13) (continue start-ktail0 x13)))
+         (truthy-cont4 (κ () (continue lambda-tail1 1)))
+         (prim-k11 (κ (r10) (continue lambda-tail1 r10)))
+         (r8 (κ (x9) (prim (* n x9) prim-k11)))
+         (fac-code14 (λ (n lambda-tail1) (prim (zero? n) prim-k3))))
+  (λ (start-ktail0)
+    (prim (make-closure $fac-code14) make-closure-cont15)))
+|]
+
+p :: HoistedProgram
+p = [cps|
+(letrec ((prim-k3-code26 (κ (r2) (prim (env-ref 0) env-cont29)))
+         (env-cont29 (κ (lambda-tail1) (prim (env-ref 1) env-cont28)))
+         (env-cont18
+          (κ (lambda-tail1)
+            (prim
+             (make-closure $prim-k11-code14 lambda-tail1)
+             make-closure-cont16)))
+         (make-closure-cont20 (κ (r8) (fac r6 r8)))
+         (env-cont15 (κ (lambda-tail1) (continue lambda-tail1 r10)))
+         (falsey-cont5
+          (κ ()
+            (prim
+             (make-closure $prim-k7-code21 fac n lambda-tail1)
+             make-closure-cont25)))
+         (prim-k7-code21 (κ (r6) (prim (env-ref 0) env-cont24)))
+         (env-cont19 (κ (n) (prim (env-ref 1) env-cont18)))
+         (truthy-cont4 (κ () (continue lambda-tail1 1)))
+         (env-cont28 (κ (n) (prim (env-ref 2) env-cont27)))
+         (fac-code34
+          (λ (n lambda-tail1)
+            (prim
+             (make-closure $prim-k3-code26 lambda-tail1 n fac)
+             make-closure-cont30)))
+         (make-closure-cont16 (κ (prim-k11) (prim (* n x9) prim-k11)))
+         (env-cont32 (κ (start-ktail0) (continue start-ktail0 x13)))
+         (prim-k11-code14 (κ (r10) (prim (env-ref 0) env-cont15)))
+         (env-cont22
+          (κ (lambda-tail1)
+            (prim
+             (make-closure $r8-code17 n lambda-tail1)
+             make-closure-cont20)))
+         (make-closure-cont25 (κ (prim-k7) (prim (- n 1) prim-k7)))
+         (env-cont23 (κ (n) (prim (env-ref 2) env-cont22)))
+         (r12-code31 (κ (x13) (prim (env-ref 0) env-cont32)))
+         (make-closure-cont33 (κ (r12) (fac 20 r12)))
+         (make-closure-cont35
+          (κ (fac)
+            (prim (make-closure $r12-code31 start-ktail0) make-closure-cont33)))
+         (env-cont27 (κ (fac) (if r2 truthy-cont4 falsey-cont5)))
+         (make-closure-cont30 (κ (prim-k3) (prim (zero? n) prim-k3)))
+         (r8-code17 (κ (x9) (prim (env-ref 0) env-cont19)))
+         (env-cont24 (κ (fac) (prim (env-ref 1) env-cont23))))
+  (λ (start-ktail0)
+    (prim (make-closure $fac-code34) make-closure-cont35)))
 |]

src/Gyehoek/CPS/Syntax.hs

@@ -41,6 +41,7 @@ module Gyehoek.CPS.Syntax
   , Free(..)
   , pattern ValLabel
   , pattern ObjLabel
+  , absBody
   )
   where
 
@@ -60,6 +61,7 @@ import Gyehoek.GenSym (Gen)
 import Data.String (IsString)
 import Control.Applicative
 import qualified Data.HashMap.Strict as H
+import GHC.Records (HasField (..))
 
 -- Data types
 
@@ -194,6 +196,16 @@ _AbsLambda' = prism'
 
 instance Plated Exp where plate = uniplate
 
+
+absBody :: Lens' Abs Exp
+absBody = lens
+  (\case
+    AbsLambda lam -> lam.body
+    AbsKappa kap -> kap.body)
+  (\cases
+    (AbsLambda lam) b -> AbsLambda $ lam & #body .~ b
+    (AbsKappa kap) b -> AbsKappa $ kap & #body .~ b)
+
 
 -- DatumIso instances
 
@@ -354,6 +366,7 @@ instance CPS Kappa where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Lambda where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Abs where toCPS = S.fromDatumUnsafe S.datumIso
 instance CPS Program where toCPS = S.fromDatumUnsafe S.datumIso
+instance CPS HoistedProgram where toCPS = S.fromDatumUnsafe S.datumIso
 
 cps :: S.QuasiQuoter
 cps = S.makeSx' [| toCPS |]

src/Gyehoek/Driver.hs

@@ -127,12 +127,12 @@ driver opts = do
   hoistedCps <- hoistProgram closedCps
   when opts.dumpHoisted do
     hPutStrLn FS.stdout =<< S.encodeWith S.datumIso hoistedCps
-  contifiedCps <- contifyProgram hoistedCps
-  when opts.dumpContified do
-    hPutStrLn FS.stdout =<< S.encodeWith S.datumIso contifiedCps
+  -- contifiedCps <- contifyProgram hoistedCps
+  -- when opts.dumpContified do
+  --   hPutStrLn FS.stdout =<< S.encodeWith S.datumIso contifiedCps
   let rt_is p = is (_Just . p) opts.runtime
   dumpOrRun opts.dumpStackified (rt_is #Stackify)
-    (stackifyProgram closedCps)
+    (stackifyProgram hoistedCps)
     (hPutStrLn FS.stdout <=< S.encodeDataWith S.dataIso)
     (eval >=> fmap writeObj
          >>> T.unwords
@@ -148,7 +148,7 @@ driver opts = do
   --   inspectWasm
   --   (\wat -> withFile opts.output FS.WriteMode \h -> hPutStrLn h wat)
   when opts.traceStackified do
-    stackifyProgram closedCps >>= traceEval
+    stackifyProgram hoistedCps >>= traceEval
 
 parse_e2e :: FilePath -> IO Scm.Program
 parse_e2e = runJalmotIO . runFileSystem . readScm

src/Gyehoek/Scheme/Syntax.hs

@@ -81,6 +81,7 @@ data Prim e
   | PrimZeroP e
   | PrimNewline
   | PrimMakeClosure { code :: e, env :: List e }
+  | PrimGetEnv
   | PrimEnv
   | PrimEnvRef Int
   | PrimCallCC e
@@ -168,6 +169,7 @@ primDatumIso namefn a = S.match
     $ S.With (. ht1 "zero?")
     $ S.With (. ht0 "newline")
     $ S.With (. ht1' "make-closure")
+    $ S.With (. ht0 "get-env")
     $ S.With (. ht0 "env")
     $ S.With (. S.headTagged1 (namefn "env-ref") S.int)
     $ S.With (. ht1 "call/cc")