WIP map + zip

1 files changed, 19 insertions, 0 deletions

diff --git a/src/Compile.hs b/src/Compile.hs
index bf7817a..d6ad7ec 100644
--- a/src/Compile.hs
+++ b/src/Compile.hs
@@ -802,6 +802,15 @@ compile' env = \case
 
     return (CELit arrname)
 
+  -- TODO: actually generate decent code here
+  EMap _ e1 e2 -> do
+    let STArr n _ = typeOf e2
+    compile' env $
+      elet e2 $
+        EBuild ext n (EShape ext (evar IZ)) $
+          elet (EIdx ext (evar (IS IZ)) (EVar ext (tTup (sreplicate n tIx)) IZ)) $
+            weakenExpr (WCopy (WSink .> WSink)) e1
+
   EFold1Inner _ commut efun ex0 earr -> do
     let STArr (SS n) t = typeOf earr
 
@@ -951,6 +960,16 @@ compile' env = \case
               [("buf", CEProj (CELit arrname) "buf")
               ,("sh", CELit ("{" ++ intercalate ", " [printCExpr 0 e "" | e <- indexTupleComponents dim shname] ++ "}"))])
 
+  -- TODO: actually generate decent code here
+  EZip _ e1 e2 -> do
+    let STArr n _ = typeOf e1
+    compile' env $
+      elet e1 $
+      elet (weakenExpr WSink e2) $
+        EBuild ext n (EShape ext (evar (IS IZ))) $
+          EPair ext (EIdx ext (evar (IS (IS IZ))) (EVar ext (tTup (sreplicate n tIx)) IZ))
+                    (EIdx ext (evar (IS IZ)) (EVar ext (tTup (sreplicate n tIx)) IZ))
+
   EFold1InnerD1 _ commut efun ex0 earr -> do
     let STArr (SS n) t = typeOf earr
         STPair _ bty = typeOf efun