1 files changed, 20 insertions, 6 deletions
diff --git a/src/Analysis/Identity.hs b/src/Analysis/Identity.hs
index 6301dc1..7b896a3 100644
--- a/src/Analysis/Identity.hs
+++ b/src/Analysis/Identity.hs
@@ -202,11 +202,19 @@ idana env expr = case expr of
     res <- VIArr <$> genId <*> shidsToVec dim shids
     pure (res, EBuild res dim e1' e2')
 
+  EMap _ e1 e2 -> do
+    let STArr _ t = typeOf e2
+    x1 <- genIds t
+    (_, e1') <- idana (x1 `SCons` env) e1
+    (v2, e2') <- idana env e2
+    let VIArr _ sh = v2
+    res <- VIArr <$> genId <*> pure sh
+    pure (res, EMap res e1' e2')
+
   EFold1Inner _ cm e1 e2 e3 -> do
     let t1 = typeOf e1
-    x1 <- genIds t1
-    x2 <- genIds t1
-    (_, e1') <- idana (x1 `SCons` x2 `SCons` env) e1
+    x1 <- genIds (STPair t1 t1)
+    (_, e1') <- idana (x1 `SCons` env) e1
     (_, e2') <- idana env e2
     (v3, e3') <- idana env e3
     let VIArr _ (_ :< sh) = v3
@@ -250,11 +258,17 @@ idana env expr = case expr of
     res <- VIArr <$> genId <*> shidsToVec dim v1
     pure (res, EReshape res dim e1' e2')
 
+  EZip _ e1 e2 -> do
+    (v1, e1') <- idana env e1
+    (_, e2') <- idana env e2
+    let VIArr _ sh = v1
+    res <- VIArr <$> genId <*> pure sh
+    pure (res, EZip res e1' e2')
+
   EFold1InnerD1 _ cm e1 e2 e3 -> do
     let t1 = typeOf e2
-    x1 <- genIds t1
-    x2 <- genIds t1
-    (_, e1') <- idana (x1 `SCons` x2 `SCons` env) e1
+    x1 <- genIds (STPair t1 t1)
+    (_, e1') <- idana (x1 `SCons` env) e1
     (_, e2') <- idana env e2
     (v3, e3') <- idana env e3
     let VIArr _ sh'@(_ :< sh) = v3