Tests pass, should check if output is sensible

7 files changed, 275 insertions, 136 deletions

diff --git a/src/AST/Accum.hs b/src/AST/Accum.hs
index 158b4d9..619c2b1 100644
--- a/src/AST/Accum.hs
+++ b/src/AST/Accum.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TypeData #-}
 {-# LANGUAGE TypeFamilies #-}
@@ -33,35 +34,38 @@ data SAcPrj (p :: AcPrj) (a :: Ty) (b :: Ty) where
   -- SAPArrSlice :: SNat m -> SAcPrj (APArrSlice m) (TArr n t) (TArr (n - m) t)
 deriving instance Show (SAcPrj p a b)
 
-type data StillDense = AI_D | AI_S
-data SStillDense dense where
-  SAI_D :: SStillDense AI_D
-  SAI_S :: SStillDense AI_S
-deriving instance Show (SStillDense dense)
+type data AIDense = AID | AIS
 
-type family AcIdx dense p t where
-  AcIdx dense APHere t = TNil
-  AcIdx AI_D (APFst p) (TPair a b) = AcIdx AI_D p a
-  AcIdx AI_D (APSnd p) (TPair a b) = AcIdx AI_D p b
-  AcIdx AI_S (APFst p) (TPair a b) = TPair (AcIdx AI_S p a) (ZeroInfo b)
-  AcIdx AI_S (APSnd p) (TPair a b) = TPair (ZeroInfo a) (AcIdx AI_S p b)
-  AcIdx dense (APLeft p) (TLEither a b) = AcIdx AI_S p a
-  AcIdx dense (APRight p) (TLEither a b) = AcIdx AI_S p b
-  AcIdx dense (APJust p) (TMaybe a) = AcIdx AI_S p a
-  AcIdx AI_D (APArrIdx p) (TArr n a) = TPair (Tup (Replicate n TIx)) (AcIdx AI_D p a)
-  AcIdx AI_S (APArrIdx p) (TArr n a) =
-    -- ((index, shapes info), recursive info)
+data SAIDense d where
+  SAID :: SAIDense AID
+  SAIS :: SAIDense AIS
+deriving instance Show (SAIDense d)
+
+type family AcIdx d p t where
+  AcIdx d APHere t = TNil
+  AcIdx AID (APFst p) (TPair a b) = AcIdx AID p a
+  AcIdx AID (APSnd p) (TPair a b) = AcIdx AID p b
+  AcIdx AIS (APFst p) (TPair a b) = TPair (AcIdx AIS p a) (ZeroInfo b)
+  AcIdx AIS (APSnd p) (TPair a b) = TPair (ZeroInfo a) (AcIdx AIS p b)
+  AcIdx d (APLeft p) (TLEither a b) = AcIdx d p a
+  AcIdx d (APRight p) (TLEither a b) = AcIdx d p b
+  AcIdx d (APJust p) (TMaybe a) = AcIdx d p a
+  AcIdx AID (APArrIdx p) (TArr n a) =
+    -- (index, recursive info)
+    TPair (Tup (Replicate n TIx)) (AcIdx AID p a)
+  AcIdx AIS (APArrIdx p) (TArr n a) =
+    -- ((index, shape info), recursive info)
     TPair (TPair (Tup (Replicate n TIx)) (ZeroInfo (TArr n a)))
-          (AcIdx AI_S p a)
-  -- AcIdx AI_D (APArrSlice m) (TArr n a) =
+          (AcIdx AIS p a)
+  -- AcIdx AID (APArrSlice m) (TArr n a) =
   --   -- index
   --   Tup (Replicate m TIx)
-  -- AcIdx AI_S (APArrSlice m) (TArr n a) =
+  -- AcIdx AIS (APArrSlice m) (TArr n a) =
   --   -- (index, array shape)
   --   TPair (Tup (Replicate m TIx)) (Tup (Replicate n TIx))
 
-type AcIdxD p t = AcIdx AI_D p t
-type AcIdxS p t = AcIdx AI_S p t
+type AcIdxD p t = AcIdx AID p t
+type AcIdxS p t = AcIdx AIS p t
 
 acPrjTy :: SAcPrj p a b -> SMTy a -> SMTy b
 acPrjTy SAPHere t = t
@@ -88,6 +92,16 @@ tZeroInfo (SMTMaybe _) = STNil
 tZeroInfo (SMTArr n t) = STArr n (tZeroInfo t)
 tZeroInfo (SMTScal _) = STNil
 
+-- | Info needed to create a zero-valued deep accumulator for a monoid type.
+-- Should be constructable from a D1.
+type family DeepZeroInfo t where
+  DeepZeroInfo TNil = TNil
+  DeepZeroInfo (TPair a b) = TPair (DeepZeroInfo a) (DeepZeroInfo b)
+  DeepZeroInfo (TLEither a b) = TLEither (DeepZeroInfo a) (DeepZeroInfo b)
+  DeepZeroInfo (TMaybe a) = TMaybe (DeepZeroInfo a)
+  DeepZeroInfo (TArr n a) = TArr n (DeepZeroInfo a)
+  DeepZeroInfo (TScal t) = TNil
+
 -- -- | Additional info needed for accumulation. This is empty unless there is
 -- -- sparsity in the monoid.
 -- type family AccumInfo t where
diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index 03a36f6..05be524 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -134,7 +134,7 @@ occCountGeneral onehot unpush alter many = go WId
       ECustom _ _ _ _ _ _ _ a b -> re a <> re b
       ERecompute _ e -> re e
       EWith _ _ a b -> re a <> re1 b
-      EAccum _ _ _ a b e -> re a <> re b <> re e
+      EAccum _ _ _ a _ b e -> re a <> re b <> re e
       EZero _ _ e -> re e
       EPlus _ _ a b -> re a <> re b
       EOneHot _ _ _ a b -> re a <> re b
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)
 
diff --git a/src/AST/Sparse.hs b/src/AST/Sparse.hs
index 369d395..f0a1f2a 100644
--- a/src/AST/Sparse.hs
+++ b/src/AST/Sparse.hs
@@ -1,116 +1,19 @@
+{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
-{-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE StandaloneDeriving #-}
-{-# LANGUAGE TypeFamilies #-}
-{-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE RankNTypes #-}
-{-# LANGUAGE DataKinds #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 
 {-# OPTIONS_GHC -fmax-pmcheck-models=80 #-}
-module AST.Sparse where
+module AST.Sparse (module AST.Sparse, module AST.Sparse.Types) where
 
-import Data.Kind (Constraint, Type)
 import Data.Type.Equality
 
 import AST
+import AST.Sparse.Types
+import Data (SBool(..))
 
 
-data Sparse t t' where
-  SpSparse :: Sparse t t' -> Sparse t (TMaybe t')
-  SpAbsent :: Sparse t TNil
-
-  SpPair :: Sparse a a' -> Sparse b b' -> Sparse (TPair a b) (TPair a' b')
-  SpLEither :: Sparse a a' -> Sparse b b' -> Sparse (TLEither a b) (TLEither a' b')
-  SpMaybe :: Sparse t t' -> Sparse (TMaybe t) (TMaybe t')
-  SpArr :: Sparse t t' -> Sparse (TArr n t) (TArr n t')
-  SpScal :: Sparse (TScal t) (TScal t)
-deriving instance Show (Sparse t t')
-
-class ApplySparse f where
-  applySparse :: Sparse t t' -> f t -> f t'
-
-instance ApplySparse STy where
-  applySparse (SpSparse s) t = STMaybe (applySparse s t)
-  applySparse SpAbsent _ = STNil
-  applySparse (SpPair s1 s2) (STPair t1 t2) = STPair (applySparse s1 t1) (applySparse s2 t2)
-  applySparse (SpLEither s1 s2) (STLEither t1 t2) = STLEither (applySparse s1 t1) (applySparse s2 t2)
-  applySparse (SpMaybe s) (STMaybe t) = STMaybe (applySparse s t)
-  applySparse (SpArr s) (STArr n t) = STArr n (applySparse s t)
-  applySparse SpScal t = t
-
-instance ApplySparse SMTy where
-  applySparse (SpSparse s) t = SMTMaybe (applySparse s t)
-  applySparse SpAbsent _ = SMTNil
-  applySparse (SpPair s1 s2) (SMTPair t1 t2) = SMTPair (applySparse s1 t1) (applySparse s2 t2)
-  applySparse (SpLEither s1 s2) (SMTLEither t1 t2) = SMTLEither (applySparse s1 t1) (applySparse s2 t2)
-  applySparse (SpMaybe s) (SMTMaybe t) = SMTMaybe (applySparse s t)
-  applySparse (SpArr s) (SMTArr n t) = SMTArr n (applySparse s t)
-  applySparse SpScal t = t
-
-
-class IsSubType s where
-  type IsSubTypeSubject (s :: k -> k -> Type) (f :: k -> Type) :: Constraint
-  subtApply :: IsSubTypeSubject s f => s t t' -> f t -> f t'
-  subtTrans :: s a b -> s b c -> s a c
-  subtFull :: IsSubTypeSubject s f => f t -> s t t
-
-instance IsSubType (:~:) where
-  type IsSubTypeSubject (:~:) f = ()
-  subtApply = gcastWith
-  subtTrans = trans
-  subtFull _ = Refl
-
-instance IsSubType Sparse where
-  type IsSubTypeSubject Sparse f = f ~ SMTy
-  subtApply = applySparse
-
-  subtTrans s1 (SpSparse s2) = SpSparse (subtTrans s1 s2)
-  subtTrans _ SpAbsent = SpAbsent
-  subtTrans (SpPair s1a s1b) (SpPair s2a s2b) = SpPair (subtTrans s1a s2a) (subtTrans s1b s2b)
-  subtTrans (SpLEither s1a s1b) (SpLEither s2a s2b) = SpLEither (subtTrans s1a s2a) (subtTrans s1b s2b)
-  subtTrans (SpSparse s1) (SpMaybe s2) = SpSparse (subtTrans s1 s2)
-  subtTrans (SpMaybe s1) (SpMaybe s2) = SpMaybe (subtTrans s1 s2)
-  subtTrans (SpArr s1) (SpArr s2) = SpArr (subtTrans s1 s2)
-  subtTrans SpScal SpScal = SpScal
-
-  subtFull = spDense
-
-spDense :: SMTy t -> Sparse t t
-spDense SMTNil = SpAbsent
-spDense (SMTPair t1 t2) = SpPair (spDense t1) (spDense t2)
-spDense (SMTLEither t1 t2) = SpLEither (spDense t1) (spDense t2)
-spDense (SMTMaybe t) = SpMaybe (spDense t)
-spDense (SMTArr _ t) = SpArr (spDense t)
-spDense (SMTScal _) = SpScal
-
-isDense :: SMTy t -> Sparse t t' -> Maybe (t :~: t')
-isDense SMTNil SpAbsent = Just Refl
-isDense _ SpSparse{} = Nothing
-isDense _ SpAbsent = Nothing
-isDense (SMTPair t1 t2) (SpPair s1 s2)
-  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
-  | otherwise = Nothing
-isDense (SMTLEither t1 t2) (SpLEither s1 s2)
-  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
-  | otherwise = Nothing
-isDense (SMTMaybe t) (SpMaybe s)
-  | Just Refl <- isDense t s = Just Refl
-  | otherwise = Nothing
-isDense (SMTArr _ t) (SpArr s)
-  | Just Refl <- isDense t s = Just Refl
-  | otherwise = Nothing
-isDense (SMTScal _) SpScal = Just Refl
-
-isAbsent :: Sparse t t' -> Bool
-isAbsent (SpSparse s) = isAbsent s
-isAbsent SpAbsent = True
-isAbsent (SpPair s1 s2) = isAbsent s1 && isAbsent s2
-isAbsent (SpLEither s1 s2) = isAbsent s1 && isAbsent s2
-isAbsent (SpMaybe s) = isAbsent s
-isAbsent (SpArr s) = isAbsent s
-isAbsent SpScal = False
-
 sparsePlus :: SMTy t -> Sparse t t' -> Ex env t' -> Ex env t' -> Ex env t'
 sparsePlus _ SpAbsent _ _ = ENil ext
 sparsePlus t sp e1 e2 | Just Refl <- isDense t sp = EPlus ext t e1 e2
@@ -143,11 +46,6 @@ sparsePlus (SMTArr _ t) (SpArr sp) e1 e2 = ezipWith (sparsePlus t sp (evar (IS I
 sparsePlus t@SMTScal{} SpScal e1 e2 = EPlus ext t e1 e2
 
 
-data SBool b where
-  SF :: SBool False
-  ST :: SBool True
-deriving instance Show (SBool b)
-
 data Injection sp a b where
   -- | 'Inj' is purposefully also allowed when @sp@ is @False@ so that
   -- 'sparsePlusS' can provide injections even if the caller doesn't require
diff --git a/src/AST/Sparse/Types.hs b/src/AST/Sparse/Types.hs
new file mode 100644
index 0000000..10cac4e
--- /dev/null
+++ b/src/AST/Sparse/Types.hs
@@ -0,0 +1,107 @@
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+module AST.Sparse.Types where
+
+import AST.Types
+
+import Data.Kind (Type, Constraint)
+import Data.Type.Equality
+
+
+data Sparse t t' where
+  SpSparse :: Sparse t t' -> Sparse t (TMaybe t')
+  SpAbsent :: Sparse t TNil
+
+  SpPair :: Sparse a a' -> Sparse b b' -> Sparse (TPair a b) (TPair a' b')
+  SpLEither :: Sparse a a' -> Sparse b b' -> Sparse (TLEither a b) (TLEither a' b')
+  SpMaybe :: Sparse t t' -> Sparse (TMaybe t) (TMaybe t')
+  SpArr :: Sparse t t' -> Sparse (TArr n t) (TArr n t')
+  SpScal :: Sparse (TScal t) (TScal t)
+deriving instance Show (Sparse t t')
+
+class ApplySparse f where
+  applySparse :: Sparse t t' -> f t -> f t'
+
+instance ApplySparse STy where
+  applySparse (SpSparse s) t = STMaybe (applySparse s t)
+  applySparse SpAbsent _ = STNil
+  applySparse (SpPair s1 s2) (STPair t1 t2) = STPair (applySparse s1 t1) (applySparse s2 t2)
+  applySparse (SpLEither s1 s2) (STLEither t1 t2) = STLEither (applySparse s1 t1) (applySparse s2 t2)
+  applySparse (SpMaybe s) (STMaybe t) = STMaybe (applySparse s t)
+  applySparse (SpArr s) (STArr n t) = STArr n (applySparse s t)
+  applySparse SpScal t = t
+
+instance ApplySparse SMTy where
+  applySparse (SpSparse s) t = SMTMaybe (applySparse s t)
+  applySparse SpAbsent _ = SMTNil
+  applySparse (SpPair s1 s2) (SMTPair t1 t2) = SMTPair (applySparse s1 t1) (applySparse s2 t2)
+  applySparse (SpLEither s1 s2) (SMTLEither t1 t2) = SMTLEither (applySparse s1 t1) (applySparse s2 t2)
+  applySparse (SpMaybe s) (SMTMaybe t) = SMTMaybe (applySparse s t)
+  applySparse (SpArr s) (SMTArr n t) = SMTArr n (applySparse s t)
+  applySparse SpScal t = t
+
+
+class IsSubType s where
+  type IsSubTypeSubject (s :: k -> k -> Type) (f :: k -> Type) :: Constraint
+  subtApply :: IsSubTypeSubject s f => s t t' -> f t -> f t'
+  subtTrans :: s a b -> s b c -> s a c
+  subtFull :: IsSubTypeSubject s f => f t -> s t t
+
+instance IsSubType (:~:) where
+  type IsSubTypeSubject (:~:) f = ()
+  subtApply = gcastWith
+  subtTrans = trans
+  subtFull _ = Refl
+
+instance IsSubType Sparse where
+  type IsSubTypeSubject Sparse f = f ~ SMTy
+  subtApply = applySparse
+
+  subtTrans s1 (SpSparse s2) = SpSparse (subtTrans s1 s2)
+  subtTrans _ SpAbsent = SpAbsent
+  subtTrans (SpPair s1a s1b) (SpPair s2a s2b) = SpPair (subtTrans s1a s2a) (subtTrans s1b s2b)
+  subtTrans (SpLEither s1a s1b) (SpLEither s2a s2b) = SpLEither (subtTrans s1a s2a) (subtTrans s1b s2b)
+  subtTrans (SpSparse s1) (SpMaybe s2) = SpSparse (subtTrans s1 s2)
+  subtTrans (SpMaybe s1) (SpMaybe s2) = SpMaybe (subtTrans s1 s2)
+  subtTrans (SpArr s1) (SpArr s2) = SpArr (subtTrans s1 s2)
+  subtTrans SpScal SpScal = SpScal
+
+  subtFull = spDense
+
+spDense :: SMTy t -> Sparse t t
+spDense SMTNil = SpAbsent
+spDense (SMTPair t1 t2) = SpPair (spDense t1) (spDense t2)
+spDense (SMTLEither t1 t2) = SpLEither (spDense t1) (spDense t2)
+spDense (SMTMaybe t) = SpMaybe (spDense t)
+spDense (SMTArr _ t) = SpArr (spDense t)
+spDense (SMTScal _) = SpScal
+
+isDense :: SMTy t -> Sparse t t' -> Maybe (t :~: t')
+isDense SMTNil SpAbsent = Just Refl
+isDense _ SpSparse{} = Nothing
+isDense _ SpAbsent = Nothing
+isDense (SMTPair t1 t2) (SpPair s1 s2)
+  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
+  | otherwise = Nothing
+isDense (SMTLEither t1 t2) (SpLEither s1 s2)
+  | Just Refl <- isDense t1 s1, Just Refl <- isDense t2 s2 = Just Refl
+  | otherwise = Nothing
+isDense (SMTMaybe t) (SpMaybe s)
+  | Just Refl <- isDense t s = Just Refl
+  | otherwise = Nothing
+isDense (SMTArr _ t) (SpArr s)
+  | Just Refl <- isDense t s = Just Refl
+  | otherwise = Nothing
+isDense (SMTScal _) SpScal = Just Refl
+
+isAbsent :: Sparse t t' -> Bool
+isAbsent (SpSparse s) = isAbsent s
+isAbsent SpAbsent = True
+isAbsent (SpPair s1 s2) = isAbsent s1 && isAbsent s2
+isAbsent (SpLEither s1 s2) = isAbsent s1 && isAbsent s2
+isAbsent (SpMaybe s) = isAbsent s
+isAbsent (SpArr s) = isAbsent s
+isAbsent SpScal = False
diff --git a/src/AST/SplitLets.hs b/src/AST/SplitLets.hs
index 3c353d4..2dad17a 100644
--- a/src/AST/SplitLets.hs
+++ b/src/AST/SplitLets.hs
@@ -63,7 +63,7 @@ splitLets' = \sub -> \case
   ECustom x s t p a b c e1 e2 -> ECustom x s t p a b c (splitLets' sub e1) (splitLets' sub e2)
   ERecompute x e -> ERecompute x (splitLets' sub e)
   EWith x t e1 e2 -> EWith x t (splitLets' sub e1) (splitLets' (sinkF sub) e2)
-  EAccum x t p e1 e2 e3 -> EAccum x t p (splitLets' sub e1) (splitLets' sub e2) (splitLets' sub e3)
+  EAccum x t p e1 sp e2 e3 -> EAccum x t p (splitLets' sub e1) sp (splitLets' sub e2) (splitLets' sub e3)
   EZero x t ezi -> EZero x t (splitLets' sub ezi)
   EPlus x t a b -> EPlus x t (splitLets' sub a) (splitLets' sub b)
   EOneHot x t p a b -> EOneHot x t p (splitLets' sub a) (splitLets' sub b)
diff --git a/src/AST/UnMonoid.hs b/src/AST/UnMonoid.hs
index 389dd5a..d498aaa 100644
--- a/src/AST/UnMonoid.hs
+++ b/src/AST/UnMonoid.hs
@@ -1,18 +1,22 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE TypeOperators #-}
-module AST.UnMonoid (unMonoid, zero, plus) where
+module AST.UnMonoid (unMonoid, zero, plus, acPrjCompose) where
 
 import AST
+import AST.Sparse.Types
 import Data
 
 
--- | Remove 'EZero', 'EPlus' and 'EOneHot' from the program by expanding them
--- into their concrete implementations.
+-- | Remove 'EZero', 'EDeepZero', 'EPlus' and 'EOneHot' from the program by
+-- expanding them into their concrete implementations. Also ensure that
+-- 'EAccum' has a dense sparsity.
 unMonoid :: Ex env t -> Ex env t
 unMonoid = \case
   EZero _ t e -> zero t e
+  EDeepZero _ t e -> deepZero t e
   EPlus _ t a b -> plus t (unMonoid a) (unMonoid b)
   EOneHot _ t p a b -> onehot t p (unMonoid a) (unMonoid b)
 
@@ -49,7 +53,10 @@ unMonoid = \case
   ECustom _ t1 t2 t3 a b c e1 e2 -> ECustom ext t1 t2 t3 (unMonoid a) (unMonoid b) (unMonoid c) (unMonoid e1) (unMonoid e2)
   ERecompute _ e -> ERecompute ext (unMonoid e)
   EWith _ t a b -> EWith ext t (unMonoid a) (unMonoid b)
-  EAccum _ t p a b e -> EAccum ext t p (unMonoid a) (unMonoid b) (unMonoid e)
+  EAccum _ t p eidx sp eval eacc ->
+    accumulateSparse (acPrjTy p t) sp eval $ \w prj2 idx2 val2 ->
+    acPrjCompose SAID p (weakenExpr w eidx) prj2 idx2 $ \prj' idx' ->
+      EAccum ext t prj' (unMonoid idx') (spDense (acPrjTy prj' t)) (unMonoid val2) (weakenExpr w (unMonoid eacc))
   EError _ t s -> EError ext t s
 
 zero :: SMTy t -> Ex env (ZeroInfo t) -> Ex env t
@@ -66,6 +73,27 @@ zero (SMTScal t) _ = case t of
   STF32 -> EConst ext STF32 0.0
   STF64 -> EConst ext STF64 0.0
 
+deepZero :: SMTy t -> Ex env (DeepZeroInfo t) -> Ex env t
+deepZero SMTNil _ = ENil ext
+deepZero (SMTPair t1 t2) e =
+  ELet ext e $ EPair ext (deepZero t1 (EFst ext (EVar ext (typeOf e) IZ)))
+                         (deepZero t2 (ESnd ext (EVar ext (typeOf e) IZ)))
+deepZero (SMTLEither t1 t2) e =
+  elcase e
+    (ELNil ext (fromSMTy t1) (fromSMTy t2))
+    (ELInl ext (fromSMTy t2) (deepZero t1 (evar IZ)))
+    (ELInr ext (fromSMTy t1) (deepZero t2 (evar IZ)))
+deepZero (SMTMaybe t) e =
+  emaybe e
+    (ENothing ext (fromSMTy t))
+    (EJust ext (deepZero t (evar IZ)))
+deepZero (SMTArr _ t) e = emap (deepZero t (evar IZ)) e
+deepZero (SMTScal t) _ = case t of
+  STI32 -> EConst ext STI32 0
+  STI64 -> EConst ext STI64 0
+  STF32 -> EConst ext STF32 0.0
+  STF64 -> EConst ext STF64 0.0
+
 plus :: SMTy t -> Ex env t -> Ex env t -> Ex env t
 plus SMTNil _ _ = ENil ext
 plus (SMTPair t1 t2) a b =
@@ -143,3 +171,78 @@ onehot typ topprj idx arg = case (typ, topprj) of
              (onehot t1 prj (ESnd ext (EVar ext (typeOf idx) (IS IZ))) (weakenExpr (WSink .> WSink) arg))
              (ELet ext (EIdx ext (ESnd ext (EFst ext (EVar ext (typeOf idx) (IS IZ)))) (EVar ext tidx IZ)) $
                 zero t1 (EVar ext (tZeroInfo t1) IZ))
+
+accumulateSparse
+  :: SMTy t -> Sparse t t' -> Ex env t'
+  -> (forall p b env'. env :> env' -> SAcPrj p t b -> Ex env' (AcIdxD p t) -> Ex env' b -> Ex env' TNil)
+  -> Ex env TNil
+accumulateSparse topty topsp arg accum = case (topty, topsp) of
+  (_, s) | Just Refl <- isDense topty s ->
+    accum WId SAPHere (ENil ext) arg
+  (SMTScal _, SpScal) ->
+    accum WId SAPHere (ENil ext) arg  -- should be handled by isDense already, but meh
+  (_, SpSparse s) ->
+    emaybe arg
+      (ENil ext)
+      (accumulateSparse topty s (evar IZ) (\w -> accum (WPop w)))
+  (_, SpAbsent) ->
+    ENil ext
+  (SMTPair t1 t2, SpPair s1 s2) ->
+    eunPair arg $ \w1 e1 e2 ->
+      elet (accumulateSparse t1 s1 e1 (\w prj -> accum (w .> w1) (SAPFst prj))) $
+        accumulateSparse t2 s2 (weakenExpr WSink e2) (\w prj -> accum (w .> WSink .> w1) (SAPSnd prj))
+  (SMTLEither t1 t2, SpLEither s1 s2) ->
+    elcase arg
+      (ENil ext)
+      (accumulateSparse t1 s1 (evar IZ) (\w prj -> accum (WPop w) (SAPLeft prj)))
+      (accumulateSparse t2 s2 (evar IZ) (\w prj -> accum (WPop w) (SAPRight prj)))
+  (SMTMaybe t, SpMaybe s) ->
+    emaybe arg
+      (ENil ext)
+      (accumulateSparse t s (evar IZ) (\w prj -> accum (WPop w) (SAPJust prj)))
+  (SMTArr n t, SpArr s) ->
+    let tn = tTup (sreplicate n tIx) in
+    elet arg $
+    elet (EBuild ext n (EShape ext (evar IZ)) $
+            accumulateSparse t s
+              (EIdx ext (evar (IS IZ)) (EVar ext tn IZ))
+              (\w prj idx val -> accum (WPop (WPop w)) (SAPArrIdx prj) (EPair ext (EVar ext tn (w @> IZ)) idx) val)) $
+      ENil ext
+
+acPrjCompose
+  :: SAIDense dense
+  -> SAcPrj p1 a b -> Ex env (AcIdx dense p1 a)
+  -> SAcPrj p2 b c -> Ex env (AcIdx dense p2 b)
+  -> (forall p'. SAcPrj p' a c -> Ex env (AcIdx dense p' a) -> r) -> r
+acPrjCompose _ SAPHere _ p2 idx2 k = k p2 idx2
+acPrjCompose SAID (SAPFst p1) idx1 p2 idx2 k =
+  acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPFst p') idx'
+acPrjCompose SAID (SAPSnd p1) idx1 p2 idx2 k =
+  acPrjCompose SAID p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPSnd p') idx'
+acPrjCompose SAIS (SAPFst p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (efst (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPFst p') (elet idx1 $ EPair ext idx' (esnd (evar IZ)))
+acPrjCompose SAIS (SAPSnd p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPSnd p') (elet idx1 $ EPair ext (efst (evar IZ)) idx')
+acPrjCompose d (SAPLeft p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPLeft p') idx'
+acPrjCompose d (SAPRight p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPRight p') idx'
+acPrjCompose d (SAPJust p1) idx1 p2 idx2 k =
+  acPrjCompose d p1 idx1 p2 idx2 $ \p' idx' ->
+    k (SAPJust p') idx'
+acPrjCompose SAID (SAPArrIdx p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAID p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx')
+acPrjCompose SAIS (SAPArrIdx p1) idx1 p2 idx2 k
+  | Dict <- styKnown (typeOf idx1) =
+  acPrjCompose SAIS p1 (esnd (evar IZ)) p2 (weakenExpr WSink idx2) $ \p' idx' ->
+    k (SAPArrIdx p') (elet idx1 $ EPair ext (EFst ext (EVar ext (typeOf idx1) IZ)) idx')