idk ^w^

src/Gyehoek/CPS/Lower.hs

@@ -28,6 +28,7 @@ import qualified Gyehoek.Sexp
 import Data.Text qualified as T
 import Data.Foldable (fold)
 import Gyehoek.Sexp (encodeOrShow, toSexp)
+import Debug.Pretty.Simple
 
 
 data Env = MkEnv

src/Gyehoek/CPS/Syntax.hs

@@ -2,6 +2,7 @@
 {-# LANGUAGE TemplateHaskell #-}
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE TypeFamilies #-}
 module Gyehoek.CPS.Syntax
   ( Val(..)
   , Kappa(..)
@@ -22,6 +23,8 @@ module Gyehoek.CPS.Syntax
   , pattern AbsLambda'
   , pattern AbsKappa'
   , Abs(..)
+  , free
+  , free'
   )
   where
 
@@ -41,6 +44,10 @@ import Language.Sexp.Located (Sexp)
 import qualified Data.InvertibleGrammar.Base as IG
 import qualified Gyehoek.Scheme.Syntax as Gyehoek
 import Data.InvertibleGrammar.Base (type (:-)((:-)))
+import Data.HashSet (HashSet)
+import qualified Data.HashSet as HS
+import Data.Hashable (Hashable)
+import Data.Monoid (Endo)
 
 -- Data types
 
@@ -97,7 +104,6 @@ makePrisms ''Exp
 
 instance S.SexpIso Val where
   sexpIso = match
-    --  $ With (. label)
     $ With (\var -> var . S.sexpIso)
     $ With (\lit -> lit . S.sexpIso)
     $ End
@@ -192,3 +198,56 @@ instance CPS Abs where toCPS = Gyehoek.Sexp.fromSexp
 
 cps :: QuasiQuoter
 cps = Gyehoek.Sexp.makeSx' [| toCPS |]
+
+
+
+deleteFrom :: (Foldable f, Hashable a) => f a -> HashSet a -> HashSet a
+deleteFrom = flip $ foldr HS.delete
+
+insertFrom :: (Foldable f, Hashable a) => f a -> HashSet a -> HashSet a
+insertFrom = flip $ foldr HS.insert
+
+toHashSetOf :: Hashable a => Getting (Endo (HashSet a)) s a -> s -> HashSet a
+toHashSetOf l = foldrOf l HS.insert mempty
+
+free :: Exp -> HashSet Name
+free = go where
+  gokap (MkKappa xs m) = go m & deleteFrom xs
+  golam (MkLambda xs k m) = go m & deleteFrom xs & sans k
+  goabs = \case
+    AbsKappa kap -> gokap kap
+    AbsLambda lam -> golam lam
+  go = \case
+    ExpPrim p k ->
+      p & toHashSetOf (folded . #ValVar)
+        & HS.union (gokap k)
+    ExpLetRec bs m ->
+      foldMapOf (each . _2) goabs bs <> go m
+      & deleteFrom (bs ^.. each . _1)
+    ExpContinue k xs -> HS.fromList $ k : xs ^.. each . #ValVar
+    ExpIf c t f -> toHashSetOf #ValVar c <> go t <> go f
+    ExpApply f xs k -> toHashSetOf (each . #ValVar) (f:xs) <> HS.singleton k
+
+-- | Free variables given in the order of their appearance.
+free' :: Exp -> List Name
+free' = go HS.empty where
+  gokap bound (MkKappa xs m) = go (bound & insertFrom xs) m
+  golam bound (MkLambda xs k m) = go (bound & insertFrom (k:xs)) m
+  goabs bound = \case
+    AbsKappa kap -> gokap bound kap
+    AbsLambda lam -> golam bound lam
+  go :: HashSet Name -> Exp -> List Name
+  go bound = \case
+    ExpPrim p k ->
+      p & toListOf (folded . #ValVar . filtered (`notElem` bound))
+        & (<> gokap bound k)
+    ExpLetRec bs m ->
+      foldMapOf (each . _2) (goabs bound') bs <> go bound' m
+      where bound' = bound & insertFrom (bs ^.. each . _1)
+    ExpContinue k xs -> filter (`notElem` bound) (k : xs ^.. each . #ValVar)
+    ExpIf c t f ->
+      (c ^.. #ValVar . filtered (`notElem` bound))
+      <> go bound t <> go bound f
+    ExpApply f xs k ->
+      (f:xs) ^.. (each . #ValVar . filtered (`notElem` bound))
+      <> (k ^.. filtered (`notElem` bound))