WIP map + zip

1 files changed, 16 insertions, 0 deletions

diff --git a/src/Analysis/Identity.hs b/src/Analysis/Identity.hs
index 6301dc1..71da793 100644
--- a/src/Analysis/Identity.hs
+++ b/src/Analysis/Identity.hs
@@ -202,6 +202,15 @@ 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
@@ -250,6 +259,13 @@ 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