WIP revamp accumulators again: explicit monoid types

No more D2 in accumulators! Paving the way for configurable sparsity of
products and arrays. The idea is to make separate monoid types for a
"product cotangent" and an "array cotangent" that can be lowered to
either a sparse monoid or a non-sparse monoid. Downsides of this
approach: lots of API duplication.

1 files changed, 49 insertions, 2 deletions

diff --git a/src/AST/Types.hs b/src/AST/Types.hs
index b20fc2d..c8515fc 100644
--- a/src/AST/Types.hs
+++ b/src/AST/Types.hs
@@ -27,6 +27,8 @@ type data Ty
   | TArr Nat Ty  -- ^ rank, element type
   | TScal ScalTy
   | TAccum Ty  -- ^ contained type must be a monoid type
+  -- sparse monoid types
+  | TLEither Ty Ty
 
 type data ScalTy = TI32 | TI64 | TF32 | TF64 | TBool
 
@@ -38,7 +40,9 @@ data STy t where
   STMaybe :: STy a -> STy (TMaybe a)
   STArr :: SNat n -> STy t -> STy (TArr n t)
   STScal :: SScalTy t -> STy (TScal t)
-  STAccum :: STy t -> STy (TAccum t)
+  STAccum :: SMTy t -> STy (TAccum t)
+  -- sparse monoid types
+  STLEither :: STy a -> STy b -> STy (TLEither a b)
 deriving instance Show (STy t)
 
 instance GCompare STy where
@@ -56,12 +60,54 @@ instance GCompare STy where
     (STScal t) (STScal t') -> gorderingLift1 (gcompare t t')
     STScal{} _ -> GLT ; _ STScal{} -> GGT
     (STAccum t) (STAccum t') -> gorderingLift1 (gcompare t t')
-    -- STAccum{} _ -> GLT ; _ STAccum{} -> GGT
+    STAccum{} _ -> GLT ; _ STAccum{} -> GGT
+    (STLEither a b) (STLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
+    -- STLEither{} _ -> GLT ; _ STLEither{} -> GGT
 
 instance TestEquality STy where testEquality = geq
 instance GEq STy where geq = defaultGeq
 instance GShow STy where gshowsPrec = defaultGshowsPrec
 
+-- | Monoid types
+type SMTy :: Ty -> Type
+data SMTy t where
+  SMTNil :: SMTy TNil
+  SMTPair :: SMTy a -> SMTy b -> SMTy (TPair a b)
+  -- TODO: call this SMTLEither
+  SMTLEither :: SMTy a -> SMTy b -> SMTy (TLEither a b)
+  SMTMaybe :: SMTy a -> SMTy (TMaybe a)
+  SMTArr :: SNat n -> SMTy t -> SMTy (TArr n t)
+  SMTScal :: ScalIsNumeric t ~ True => SScalTy t -> SMTy (TScal t)
+deriving instance Show (SMTy t)
+
+instance GCompare SMTy where
+  gcompare = \cases
+    SMTNil SMTNil -> GEQ
+    SMTNil _ -> GLT ; _ SMTNil -> GGT
+    (SMTPair a b) (SMTPair a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
+    SMTPair{} _ -> GLT ; _ SMTPair{} -> GGT
+    (SMTLEither a b) (SMTLEither a' b') -> gorderingLift2 (gcompare a a') (gcompare b b')
+    SMTLEither{} _ -> GLT ; _ SMTLEither{} -> GGT
+    (SMTMaybe a) (SMTMaybe a') -> gorderingLift1 (gcompare a a')
+    SMTMaybe{} _ -> GLT ; _ SMTMaybe{} -> GGT
+    (SMTArr n t) (SMTArr n' t') -> gorderingLift2 (gcompare n n') (gcompare t t')
+    SMTArr{} _ -> GLT ; _ SMTArr{} -> GGT
+    (SMTScal t) (SMTScal t') -> gorderingLift1 (gcompare t t')
+    -- SMTScal{} _ -> GLT ; _ SMTScal{} -> GGT
+
+instance TestEquality SMTy where testEquality = geq
+instance GEq SMTy where geq = defaultGeq
+instance GShow SMTy where gshowsPrec = defaultGshowsPrec
+
+fromSMTy :: SMTy t -> STy t
+fromSMTy = \case
+  SMTNil -> STNil
+  SMTPair t1 t2 -> STPair (fromSMTy t1) (fromSMTy t2)
+  SMTLEither t1 t2 -> STLEither (fromSMTy t1) (fromSMTy t2)
+  SMTMaybe t -> STMaybe (fromSMTy t)
+  SMTArr n t -> STArr n (fromSMTy t)
+  SMTScal sty -> STScal sty
+
 data SScalTy t where
   STI32 :: SScalTy TI32
   STI64 :: SScalTy TI64
@@ -136,6 +182,7 @@ hasArrays (STMaybe t) = hasArrays t
 hasArrays STArr{} = True
 hasArrays STScal{} = False
 hasArrays STAccum{} = True
+hasArrays (STLEither a b) = hasArrays a || hasArrays b
 
 type family Tup env where
   Tup '[] = TNil