From 79e072eddf0ec2a97ca455c27cb5ff6f2132bbab Mon Sep 17 00:00:00 2001
From: Tom Smeding <tom@tomsmeding.com>
Date: Tue, 22 Oct 2024 22:02:06 +0200
Subject: Differentiate Replicate

---
 src/AST/Count.hs  | 2 +-
 src/AST/Pretty.hs | 5 +++--
 2 files changed, 4 insertions(+), 3 deletions(-)

(limited to 'src/AST')

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index ad68685..dbec446 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -116,7 +116,7 @@ occCountGeneral onehot unpush alter many = go WId
       EConstArr{} -> mempty
       EBuild1 _ a b -> re a <> many (re1 b)
       EBuild _ _ a b -> re a <> many (re1 b)
-      EFold1Inner _ a b -> many (unpush (unpush (go (WSink .> WSink .> w) a))) <> re b
+      EFold1Inner _ a b c -> many (unpush (unpush (go (WSink .> WSink .> w) a))) <> re b <> re c
       ESum1Inner _ e -> re e
       EUnit _ e -> re e
       EReplicate1Inner _ a b -> re a <> re b
diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 8f1fe67..d811912 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -124,14 +124,15 @@ ppExpr' d val = \case
     return $ showParen (d > 10) $
       showString "build " . a' . showString (" (\\" ++ name ++ " -> ") . e' . showString ")"
 
-  EFold1Inner _ a b -> do
+  EFold1Inner _ a b c -> do
     name1 <- genNameIfUsedIn (typeOf a) (IS IZ) a
     name2 <- genNameIfUsedIn (typeOf a) IZ a
     a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a
     b' <- ppExpr' 11 val b
+    c' <- ppExpr' 11 val c
     return $ showParen (d > 10) $
       showString ("fold1i (\\" ++ name1 ++ " " ++ name2 ++ " -> ") . a'
-      . showString ") " . b'
+      . showString ") " . b' . showString " " . c'
 
   ESum1Inner _ e -> do
     e' <- ppExpr' 11 val e
-- 
cgit v1.2.3-70-g09d2