1 files changed, 28 insertions, 7 deletions
diff --git a/haskell/app/Lift.hs b/haskell/app/Lift.hs
index fa22a17..bf2dcb1 100644
--- a/haskell/app/Lift.hs
+++ b/haskell/app/Lift.hs
@@ -45,6 +45,8 @@ findVarsAExp def = foldMap id def
 compareNames n m = if n == m then n ++ "_" else n
 
 -- | replace bound variables in AExps
+
+-- >>> hoist (replaceVarsAExp "n") factorial
 replaceVarsAExp :: String -> AExpF a -> AExpF a
 replaceVarsAExp n (IdentF m) = IdentF $ compareNames m n
 replaceVarsAExp n (LamF args body) =
@@ -78,12 +80,31 @@ replaceVarsFC n (Call name args) =
         $ 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, 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
+{- | cotransverse of a complex expression. The first string list represents
+| bound variables, the second one a queue of variables to be replaced
+-}
+subVarsCExp :: ([String], [String], CExpF a) -> CExpF ([String], [String], a)
+subVarsCExp (env, queue, LetF name fc rest) =
+    let
+        (newName, newFC) =
+            foldr
+                (\m (n, f) -> if n == m then (n ++ "_", replaceVarsFC n f) else (n, f)) -- repeatedly replace variables in the function call
+                (name, fc)
+                queue -- make sure to do oldest first!
+     in
+        LetF newName newFC (if name `elem` env then (env, newName : queue, rest) else (name : env, name : queue, rest))
+subVarsCExp (env, queue, IfF cond thenPart elsePart) =
+    IfF
+        (foldr (\x c -> hoist (replaceVarsAExp x) c) cond queue)
+        (env, queue, thenPart)
+        (env, queue, elsePart)
+subVarsCExp (env, queue, FCF fc) = FCF $ foldr replaceVarsFC fc queue