WIP fold: everything but Compile (slow, but should be sound)

1 files changed, 32 insertions, 2 deletions

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 9018602..afa62c6 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -72,6 +72,7 @@ genNameIfUsedIn' prefix ty idx ex
                                            _ -> return "_"
   | otherwise                 = genName' prefix
 
+-- TODO: let this return a type-tagged thing so that name environments are more typed than Const
 genNameIfUsedIn :: STy a -> Idx env a -> Expr x env t -> M String
 genNameIfUsedIn = \t -> genNameIfUsedIn' (nameBaseForType t) t
 
@@ -209,8 +210,7 @@ ppExpr' d val expr = case expr of
     a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a
     b' <- ppExpr' 11 val b
     c' <- ppExpr' 11 val c
-    let opname = case cm of Commut -> "fold1i(C)"
-                            Noncommut -> "fold1i"
+    let opname = "fold1i" ++ ppCommut cm
     return $ ppParen (d > 10) $
       ppApp (annotate AHighlight (ppString opname) <> ppX expr) [ppLam [ppString name1, ppString name2] a', b', c']
 
@@ -235,6 +235,32 @@ ppExpr' d val expr = case expr of
     e' <- ppExpr' 11 val e
     return $ ppParen (d > 10) $ ppString "minimum1i" <> ppX expr <+> e'
 
+  EFold1InnerD1 _ cm a b c -> do
+    name1 <- genNameIfUsedIn (typeOf b) (IS IZ) a
+    name2 <- genNameIfUsedIn (typeOf b) IZ a
+    a' <- ppExpr' 0 (Const name2 `SCons` Const name1 `SCons` val) a
+    b' <- ppExpr' 11 val b
+    c' <- ppExpr' 11 val c
+    let opname = "fold1iD1" ++ ppCommut cm
+    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
+    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']
+
   EConst _ ty v
     | Dict <- scalRepIsShow ty -> return $ ppString (showsPrec d v "") <> ppX expr
 
@@ -386,6 +412,10 @@ ppSparse (SMTMaybe t) (SpMaybe s) = "M" ++ ppSparse t s
 ppSparse (SMTArr _ t) (SpArr s) = "A" ++ ppSparse t s
 ppSparse (SMTScal _) SpScal = "."
 
+ppCommut :: Commutative -> String
+ppCommut Commut = "(C)"
+ppCommut Noncommut = ""
+
 ppX :: PrettyX x => Expr x env t -> ADoc
 ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr)