WIP map + zip

1 files changed, 10 insertions, 0 deletions

diff --git a/src/Interpreter.hs b/src/Interpreter.hs
index 9e3d2a6..d982261 100644
--- a/src/Interpreter.hs
+++ b/src/Interpreter.hs
@@ -116,6 +116,9 @@ interpret'Rec env = \case
   EBuild _ dim a b -> do
     sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env a
     arrayGenerateM sh (\idx -> interpret' (V (tTup (sreplicate dim tIx)) (tupRepIdx ixUncons dim idx) `SCons` env) b)
+  EMap _ a b -> do
+    let STArr _ t = typeOf b
+    arrayMapM (\x -> interpret' (V t x `SCons` env) a) =<< interpret' env b
   EFold1Inner _ _ a b c -> do
     let t = typeOf b
     let f = \x y -> interpret' (V t y `SCons` V t x `SCons` env) a
@@ -150,6 +153,13 @@ interpret'Rec env = \case
     sh <- unTupRepIdx ShNil ShCons dim <$> interpret' env esh
     arr <- interpret' env e
     return $ arrayReshape sh arr
+  EZip _ a b -> do
+    arr1 <- interpret' env a
+    arr2 <- interpret' env b
+    let sh = arrayShape arr1
+    when (sh /= arrayShape arr2) $
+      error "Interpreter: mismatched shapes in EZip"
+    return $ arrayGenerateLin sh (\i -> (arr1 `arrayIndexLinear` i, arr2 `arrayIndexLinear` i))
   EFold1InnerD1 _ _ a b c -> do
     let t = typeOf b
     let f = \x y -> interpret' (V t y `SCons` V t x `SCons` env) a