Deal with functions using an argument more than once

1 files changed, 12 insertions, 3 deletions

diff --git a/pf.icl b/pf.icl
index 0e8ff6a..7055898 100644
--- a/pf.icl
+++ b/pf.icl
@@ -47,17 +47,25 @@ optim :: Expr -> Expr
 optim (Lambda id e)
 | hasIdent id e = case optim e of
 	Ident id` -> if (id == id`) (Ident "id") (Ident "const" @ e)
-	f @ x -> case moveOutside id (f @ x) of
-		a=:(f @ Ident id`) -> if (id == id`) f (Ident "const" @ a)
-		a                  -> Ident "const" @ a
+	e=:(f @ x) -> optim (optApp e)
 	e -> Lambda id e
 | otherwise = Ident "const" @ optim e
+where
+	optApp :: Expr -> Expr
+	optApp (f @ x)
+	| hasIdent id f && hasIdent id x = Ident "ap" @ Lambda id f @ Lambda id x
+	| otherwise = case moveOutside id (f @ x) of
+		a=:(f @ Ident id`) -> if (id == id`) f (Ident "const" @ a)
+		a                  -> Ident "const" @ a
+optim (Ident "flip" @ Ident "ap" @ Ident "id") = Ident "join"
+optim (f @ x) = optim f @ optim x
 optim e = e
 
 moveOutside :: Ident Expr -> Expr
 moveOutside _ (Ident id) = Ident id
 moveOutside _ (Literal l) = Literal l
 moveOutside id (f @ x)
+| hasIdent id f && hasIdent id x = abort "moveOutside on @ with id in both sides"
 | hasIdent id f = case moveOutside id f of
 	Ident id -> Ident "flip" @ Ident "id" @ x @ Ident id
 	g @ Ident id -> Ident "flip" @ g @ x @ Ident id
@@ -81,6 +89,7 @@ Start = map do
 	, Lambda "x" (Lambda "y" (Ident "y" @ Ident "x"))
 	, Lambda "x" (Lambda "y" (Ident "y" @ Literal (LitInt 10)))
 	, Lambda "f" (Lambda "a" (Lambda "b" (Lambda "c" (Lambda "d" (Ident "f" @ Ident "b" @ Ident "c" @ Ident "d" @ Ident "a")))))
+	, Lambda "f" (Lambda "x" (Ident "f" @ Ident "x" @ Ident "x"))
 	]
 where
 	do e = (print e, " ==> ", print (optim e), "\n")