diff options
| author | Tom Smeding <tom@tomsmeding.com> | 2024-10-22 22:02:06 +0200 |
|---|---|---|
| committer | Tom Smeding <tom@tomsmeding.com> | 2024-10-22 22:02:06 +0200 |
| commit | 79e072eddf0ec2a97ca455c27cb5ff6f2132bbab (patch) | |
| tree | 2099dc7e9d9a1109d844bca73277ca82983a02c2 /src/AST | |
| parent | e7d7ac0fd8b81c1d6fae9ab7c1e4654133c631ea (diff) | |
Differentiate Replicate
Diffstat (limited to 'src/AST')
| -rw-r--r-- | src/AST/Count.hs | 2 | ||||
| -rw-r--r-- | src/AST/Pretty.hs | 5 |
2 files changed, 4 insertions, 3 deletions
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 |