1 files changed, 48 insertions, 0 deletions
diff --git a/haskell/app/Lift.hs b/haskell/app/Lift.hs
index d31eab1..fa22a17 100644
--- a/haskell/app/Lift.hs
+++ b/haskell/app/Lift.hs
@@ -39,3 +39,51 @@ findVarsAExp :: AExpF [String] -> [String]
 findVarsAExp (IdentF ns) = [ns]
 findVarsAExp (LamF args cexp) = converge (\\ args) (fold findVars cexp)
 findVarsAExp def = foldMap id def
+
+-- replacing free variables with a new one if it matches the argument
+
+compareNames n m = if n == m then n ++ "_" else n
+
+-- | replace bound variables in AExps
+replaceVarsAExp :: String -> AExpF a -> AExpF a
+replaceVarsAExp n (IdentF m) = IdentF $ compareNames m n
+replaceVarsAExp n (LamF args body) =
+    LamF
+        (fmap (\x -> if x == n then x ++ "_" else x) args)
+        (cata replaceVarsCExp body n)
+replaceVarsAExp _ rest = rest
+
+-- | replace bound variables in CExps
+replaceVarsCExp :: CExpF (String -> CExp) -> String -> CExp
+replaceVarsCExp (LetF name fc restf) = do
+    env <- id
+    rest <- restf
+    let newName = compareNames name env
+    return $ Let newName (replaceVarsFC env fc) rest
+replaceVarsCExp (IfF cond thenF elseF) = do
+    thenPart <- thenF
+    elsePart <- elseF
+    env <- id
+    return $ If (hoist (replaceVarsAExp env) cond) thenPart elsePart
+replaceVarsCExp (FCF fc) = do
+    env <- id
+    return $ FC (replaceVarsFC env fc)
+
+-- | replace bound variables in Function calls
+replaceVarsFC :: String -> Funcall -> Funcall
+replaceVarsFC n (Atom aexp) = Atom $ hoist (replaceVarsAExp n) aexp
+replaceVarsFC n (Call name args) =
+    Call
+        (compareNames name n)
+        $ fmap (hoist (replaceVarsAExp n)) args
+
+subVarsAExp :: ([String], AExpF a) -> AExpF ([String], a)
+subVarsAExp (env, LamF args body) = let
+  toReplace = intersect env args
+  newArgs = fmap (\x -> if x `elem` toReplace then x ++ "_" else x) args
+  newBody = foldl (cata replaceVarsCExp) body toReplace
+  in LamF newArgs (cotransverse subVarsCExp (newArgs ++ env, newBody))
+subVarsAExp (env, rest) = fmap (env,) rest
+
+subVarsCExp :: ([String], CExpF a) -> CExpF ([String], a)
+subVarsCExp = undefined