Only rewrite _free_ variables in rewall

1 files changed, 19 insertions, 0 deletions

diff --git a/src/Haskell/AST.hs b/src/Haskell/AST.hs
index 072fd97..2238b6d 100644
--- a/src/Haskell/AST.hs
+++ b/src/Haskell/AST.hs
@@ -1,6 +1,7 @@
 module Haskell.AST where
 
 import Data.List
+import qualified Data.Set as Set
 import Pretty
 
 
@@ -134,3 +135,21 @@ instance AllRefs Expr where
 
 instance AllRefs Inst where
     allRefs (Inst _ _ ds) = nub $ concatMap allRefs ds
+
+
+boundVars :: Pat -> Set.Set Name
+boundVars PatAny = mempty
+boundVars (PatVar n) = Set.singleton n
+boundVars (PatCon _ ps) = Set.unions (map boundVars ps)
+boundVars (PatTup ps) = Set.unions (map boundVars ps)
+
+freeVariables :: Expr -> Set.Set Name
+freeVariables (App e es) = freeVariables e <> Set.unions (map freeVariables es)
+freeVariables (Ref n) = Set.singleton n
+freeVariables (Con _) = mempty
+freeVariables (Num _) = mempty
+freeVariables (Tup es) = Set.unions (map freeVariables es)
+freeVariables (Lam ns e) = freeVariables e Set.\\ Set.fromList ns
+freeVariables (Case e pairs) =
+    freeVariables e <> Set.unions [freeVariables e' Set.\\ boundVars p
+                                  | (p, e') <- pairs]