UnMonoid

1 files changed, 17 insertions, 10 deletions

diff --git a/src/AST.hs b/src/AST.hs
index 9ad0d4d..fff290a 100644
--- a/src/AST.hs
+++ b/src/AST.hs
@@ -396,17 +396,24 @@ emap f arr =
                             (EVar ext (tTup (sreplicate n tIx)) IZ)) $
            weakenExpr (WCopy (WSink .> WSink)) f
 
-ezip :: Ex env (TArr n a) -> Ex env (TArr n b) -> Ex env (TArr n (TPair a b))
-ezip a b =
-  let STArr n t1 = typeOf a
-      STArr _ t2 = typeOf b
-  in ELet ext a $
-     ELet ext (weakenExpr WSink b) $
+ezipWith :: Ex (b : a : env) c -> Ex env (TArr n a) -> Ex env (TArr n b) -> Ex env (TArr n c)
+ezipWith f arr1 arr2 =
+  let STArr n t1 = typeOf arr1
+      STArr _ t2 = typeOf arr2
+  in ELet ext arr1 $
+     ELet ext (weakenExpr WSink arr2) $
        EBuild ext n (EShape ext (EVar ext (STArr n t1) (IS IZ))) $
-         EPair ext (EIdx ext (EVar ext (STArr n t1) (IS (IS IZ)))
-                             (EVar ext (tTup (sreplicate n tIx)) IZ))
-                   (EIdx ext (EVar ext (STArr n t2) (IS IZ))
-                             (EVar ext (tTup (sreplicate n tIx)) IZ))
+         ELet ext (EIdx ext (EVar ext (STArr n t1) (IS (IS IZ)))
+                            (EVar ext (tTup (sreplicate n tIx)) IZ)) $
+         ELet ext (EIdx ext (EVar ext (STArr n t2) (IS (IS IZ)))
+                            (EVar ext (tTup (sreplicate n tIx)) (IS IZ))) $
+           weakenExpr (WCopy (WCopy (WSink .> WSink .> WSink))) f
+
+ezip :: Ex env (TArr n a) -> Ex env (TArr n b) -> Ex env (TArr n (TPair a b))
+ezip arr1 arr2 =
+  let STArr _ t1 = typeOf arr1
+      STArr _ t2 = typeOf arr2
+  in ezipWith (EPair ext (EVar ext t1 (IS IZ)) (EVar ext t2 IZ)) arr1 arr2
 
 eif :: Ex env (TScal TBool) -> Ex env a -> Ex env a -> Ex env a
 eif a b c = ECase ext (EOp ext OIf a) (weakenExpr WSink b) (weakenExpr WSink c)