Better fold D{1,2} primitives

1 files changed, 12 insertions, 10 deletions

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index afa62c6..587328d 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -245,21 +245,23 @@ ppExpr' d val expr = case expr of
     return $ ppParen (d > 10) $
       ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name1, ppString name2] a', b', c']
 
-  EFold1InnerD2 _ cm t2 ef ep ezi ebog ed -> do
-    let STArr _ (STPair t1 ttape) = typeOf ebog
-    name1 <- genNameIfUsedIn ttape (IS (IS (IS IZ))) ef
-    name2 <- genNameIfUsedIn t1 (IS (IS IZ)) ef
-    name3 <- genNameIfUsedIn t1 (IS IZ) ef
-    name4 <- genNameIfUsedIn (fromSMTy t2) IZ ef
-    ef' <- ppExpr' 0 (Const name4 `SCons` Const name3 `SCons` Const name2 `SCons` Const name1 `SCons` val) ef
-    ep' <- ppExpr' 11 val ep
-    ezi' <- ppExpr' 11 val ezi
+  EFold1InnerD2 _ cm ef ez eplus ebog ed -> do
+    let STArr _ tB = typeOf ebog
+        t2 = typeOf ez
+    namef1 <- genNameIfUsedIn tB (IS IZ) ef
+    namef2 <- genNameIfUsedIn t2 IZ ef
+    ef' <- ppExpr' 0 (Const namef2 `SCons` Const namef1 `SCons` val) ef
+    ez' <- ppExpr' 11 val ez
+    namep1 <- genNameIfUsedIn t2 (IS IZ) eplus
+    namep2 <- genNameIfUsedIn t2 IZ eplus
+    eplus' <- ppExpr' 0 (Const namep2 `SCons` Const namep1 `SCons` val) eplus
     ebog' <- ppExpr' 11 val ebog
     ed' <- ppExpr' 11 val ed
     let opname = "fold1iD2" ++ ppCommut cm
     return $ ppParen (d > 10) $
       ppApp (annotate AHighlight (ppString opname) <> ppX expr)
-            [ppLam [ppString name1, ppString name2, ppString name3, ppString name4] ef', ep', ezi', ebog', ed']
+            [ppLam [ppString namef1, ppString namef2] ef', ez'
+            ,ppLam [ppString namep1, ppString namep2] eplus', ebog', ed']
 
   EConst _ ty v
     | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr