From 81d88dbc430ca6ec8390636f8b7162887b390873 Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Mon, 3 Nov 2025 23:09:37 +0100
Subject: WIP map + zip

---
 src/AST/Pretty.hs | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

(limited to 'src/AST/Pretty.hs')

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 68fc629..2c51b85 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -204,6 +204,14 @@ ppExpr' d val expr = case expr of
            <> hardline <> e')
         (ppApp (annotate AHighlight primName <> ppX expr) [a', ppLam [ppString name] e'])
 
+  EMap _ a b -> do
+    let STArr _ t1 = typeOf b
+    name <- genNameIfUsedIn' "i" t1 IZ a
+    a' <- ppExpr' 0 (Const name `SCons` val) a
+    b' <- ppExpr' 11 val b
+    return $ ppParen (d > 0) $
+      ppApp (annotate AHighlight (ppString "map") <> ppX expr) [ppLam [ppString name] a', b']
+
   EFold1Inner _ cm a b c -> do
     name1 <- genNameIfUsedIn (typeOf a) (IS IZ) a
     name2 <- genNameIfUsedIn (typeOf a) IZ a
@@ -238,7 +246,12 @@ ppExpr' d val expr = case expr of
   EReshape _ n esh e -> do
     esh' <- ppExpr' 11 val esh
     e' <- ppExpr' 11 val e
-    return $ ppParen (d > 10) $ ppString ("reshape" ++ intSubscript (fromSNat n)) <> ppX expr <+> esh' <+> e'
+    return $ ppParen (d > 10) $ ppApp (ppString ("reshape" ++ intSubscript (fromSNat n)) <> ppX expr) [esh', e']
+
+  EZip _ e1 e2 -> do
+    e1' <- ppExpr' 11 val e1
+    e2' <- ppExpr' 11 val e2
+    return $ ppParen (d > 10) $ ppApp (ppString "zip" <> ppX expr) [e1', e2']
 
   EFold1InnerD1 _ cm a b c -> do
     name1 <- genNameIfUsedIn (typeOf b) (IS IZ) a
-- 
cgit v1.2.3-70-g09d2