Merge branch 'sparse'HEADmaster

1 files changed, 19 insertions, 2 deletions

diff --git a/src/AST/Pretty.hs b/src/AST/Pretty.hs
index 41da656..fef9686 100644
--- a/src/AST/Pretty.hs
+++ b/src/AST/Pretty.hs
@@ -25,6 +25,7 @@ import System.IO.Unsafe (unsafePerformIO)
 
 import AST
 import AST.Count
+import AST.Sparse.Types
 import CHAD.Types
 import Data
 
@@ -304,18 +305,24 @@ ppExpr' d val expr = case expr of
            <> hardline <> e2')
         (ppApp (annotate AWith (ppString "with") <> ppX expr) [e1', ppLam [ppString name] e2'])
 
-  EAccum _ t prj e1 e2 e3 -> do
+  EAccum _ t prj e1 sp e2 e3 -> do
     e1' <- ppExpr' 11 val e1
     e2' <- ppExpr' 11 val e2
     e3' <- ppExpr' 11 val e3
     return $ ppParen (d > 10) $
-      ppApp (annotate AMonoid (ppString "accum") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (acPrjTy prj t)) [ppString (ppAcPrj t prj), e1', e2', e3']
+      ppApp (annotate AMonoid (ppString "accum") <> ppX expr <+> ppString "@" <> ppSMTy' 11 (applySparse sp (acPrjTy prj t)))
+            [ppString (ppAcPrj t prj), ppString (ppSparse (acPrjTy prj t) sp), e1', e2', e3']
 
   EZero _ t e1 -> do
     e1' <- ppExpr' 11 val e1
     return $ ppParen (d > 0) $
       annotate AMonoid (ppString "zero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1'
 
+  EDeepZero _ t e1 -> do
+    e1' <- ppExpr' 11 val e1
+    return $ ppParen (d > 0) $
+      annotate AMonoid (ppString "deepzero") <> ppX expr <+> ppString "@" <> ppSMTy' 11 t <+> e1'
+
   EPlus _ t a b -> do
     a' <- ppExpr' 11 val a
     b' <- ppExpr' 11 val b
@@ -368,6 +375,16 @@ ppAcPrj (SMTLEither _ t) (SAPRight prj) = "(|" ++ ppAcPrj t prj ++ ")"
 ppAcPrj (SMTMaybe t) (SAPJust prj) = "J" ++ ppAcPrj t prj
 ppAcPrj (SMTArr n t) (SAPArrIdx prj) = "[" ++ ppAcPrj t prj ++ "]" ++ intSubscript (fromSNat n)
 
+ppSparse :: SMTy a -> Sparse a b -> String
+ppSparse t sp | Just Refl <- isDense t sp = "D"
+ppSparse _ SpAbsent = "A"
+ppSparse t (SpSparse s) = "S" ++ ppSparse t s
+ppSparse (SMTPair t1 t2) (SpPair s1 s2) = "(" ++ ppSparse t1 s1 ++ "," ++ ppSparse t2 s2 ++ ")"
+ppSparse (SMTLEither t1 t2) (SpLEither s1 s2) = "(" ++ ppSparse t1 s1 ++ "|" ++ ppSparse t2 s2 ++ ")"
+ppSparse (SMTMaybe t) (SpMaybe s) = "M" ++ ppSparse t s
+ppSparse (SMTArr _ t) (SpArr s) = "A" ++ ppSparse t s
+ppSparse (SMTScal _) SpScal = "."
+
 ppX :: PrettyX x => Expr x env t -> ADoc
 ppX expr = annotate AExt $ ppString $ prettyXsuffix (extOf expr)