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/Simplify.hs | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

(limited to 'src/Simplify.hs')

diff --git a/src/Simplify.hs b/src/Simplify.hs
index b89d7f6..1889adc 100644
--- a/src/Simplify.hs
+++ b/src/Simplify.hs
@@ -200,12 +200,14 @@ simplify'Rec = \case
       EMaybe ext (ESnd ext e1) (ESnd ext e2) e3
 
   -- TODO: more array indexing
-  EIdx _ (EBuild _ _ e1 e2) e3 | not (hasAdds e1) -> acted $ simplify' $ elet e3 e2
+  EIdx _ (EBuild _ _ e1 e2) e3 | not (hasAdds e1), not (hasAdds e2) -> acted $ simplify' $ elet e3 e2
+  EIdx _ (EMap _ e1 e2) e3 | not (hasAdds e1) -> acted $ simplify' $ elet (EIdx ext e2 e3) e1
   EIdx _ (EReplicate1Inner _ e1 e2) e3 | not (hasAdds e1) -> acted $ simplify' $ EIdx ext e2 (EFst ext e3)
   EIdx _ (EUnit _ e1) e2 | not (hasAdds e2) -> acted $ simplify' $ e1
 
   -- TODO: more array shape
   EShape _ (EBuild _ _ e1 e2) | not (hasAdds e2) -> acted $ simplify' e1
+  EShape _ (EMap _ e1 e2) | not (hasAdds e1) -> acted $ simplify' (EShape ext e2)
 
   -- TODO: more constant folding
   EOp _ OIf (EConst _ STBool True) -> acted $ return (EInl ext STNil (ENil ext))
@@ -308,6 +310,7 @@ simplify'Rec = \case
   ELCase _ e a b c -> [simprec| ELCase ext *e *a *b *c |]
   EConstArr _ n t v -> pure $ EConstArr ext n t v
   EBuild _ n a b -> [simprec| EBuild ext n *a *b |]
+  EMap _ a b -> [simprec| EMap ext *a *b |]
   EFold1Inner _ cm a b c -> [simprec| EFold1Inner ext cm *a *b *c |]
   ESum1Inner _ e -> [simprec| ESum1Inner ext *e |]
   EUnit _ e -> [simprec| EUnit ext *e |]
@@ -315,6 +318,7 @@ simplify'Rec = \case
   EMaximum1Inner _ e -> [simprec| EMaximum1Inner ext *e |]
   EMinimum1Inner _ e -> [simprec| EMinimum1Inner ext *e |]
   EReshape _ n a b -> [simprec| EReshape ext n *a *b |]
+  EZip _ a b -> [simprec| EZip ext *a *b |]
   EFold1InnerD1 _ cm a b c -> [simprec| EFold1InnerD1 ext cm *a *b *c |]
   EFold1InnerD2 _ cm a b c -> [simprec| EFold1InnerD2 ext cm *a *b *c |]
   EConst _ t v -> pure $ EConst ext t v
@@ -364,6 +368,7 @@ hasAdds = \case
   ELCase _ e a b c -> hasAdds e || hasAdds a || hasAdds b || hasAdds c
   EConstArr _ _ _ _ -> False
   EBuild _ _ a b -> hasAdds a || hasAdds b
+  EMap _ a b -> hasAdds a || hasAdds b
   EFold1Inner _ _ a b c -> hasAdds a || hasAdds b || hasAdds c
   ESum1Inner _ e -> hasAdds e
   EUnit _ e -> hasAdds e
@@ -371,6 +376,7 @@ hasAdds = \case
   EMaximum1Inner _ e -> hasAdds e
   EMinimum1Inner _ e -> hasAdds e
   EReshape _ _ a b -> hasAdds a || hasAdds b
+  EZip _ a b -> hasAdds a || hasAdds b
   EFold1InnerD1 _ _ a b c -> hasAdds a || hasAdds b || hasAdds c
   EFold1InnerD2 _ _ a b c -> hasAdds a || hasAdds b || hasAdds c
   ECustom _ _ _ _ a b c d e -> hasAdds a || hasAdds b || hasAdds c || hasAdds d || hasAdds e
-- 
cgit v1.2.3-70-g09d2