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, 27 insertions, 17 deletions

diff --git a/src/CHAD/Types.hs b/src/CHAD/Types.hs
index 7f49cef..74e7dbd 100644
--- a/src/CHAD/Types.hs
+++ b/src/CHAD/Types.hs
@@ -14,14 +14,16 @@ type family D1 t where
   D1 (TMaybe a) = TMaybe (D1 a)
   D1 (TArr n t) = TArr n (D1 t)
   D1 (TScal t) = TScal t
+  D1 (TLEither a b) = TLEither (D1 a) (D1 b)
 
 type family D2 t where
   D2 TNil = TNil
   D2 (TPair a b) = TMaybe (TPair (D2 a) (D2 b))
-  D2 (TEither a b) = TMaybe (TEither (D2 a) (D2 b))
+  D2 (TEither a b) = TLEither (D2 a) (D2 b)
   D2 (TMaybe t) = TMaybe (D2 t)
   D2 (TArr n t) = TMaybe (TArr n (D2 t))
   D2 (TScal t) = D2s t
+  D2 (TLEither a b) = TLEither (D2 a) (D2 b)
 
 type family D2s t where
   D2s TI32 = TNil
@@ -40,7 +42,7 @@ type family D2E env where
 
 type family D2AcE env where
   D2AcE '[] = '[]
-  D2AcE (t : env) = TAccum t : D2AcE env
+  D2AcE (t : env) = TAccum (D2 t) : D2AcE env
 
 d1 :: STy t -> STy (D1 t)
 d1 STNil = STNil
@@ -50,32 +52,40 @@ d1 (STMaybe t) = STMaybe (d1 t)
 d1 (STArr n t) = STArr n (d1 t)
 d1 (STScal t) = STScal t
 d1 STAccum{} = error "Accumulators not allowed in input program"
+d1 (STLEither a b) = STLEither (d1 a) (d1 b)
 
 d1e :: SList STy env -> SList STy (D1E env)
 d1e SNil = SNil
 d1e (t `SCons` env) = d1 t `SCons` d1e env
 
+d2M :: STy t -> SMTy (D2 t)
+d2M STNil = SMTNil
+d2M (STPair a b) = SMTMaybe (SMTPair (d2M a) (d2M b))
+d2M (STEither a b) = SMTLEither (d2M a) (d2M b)
+d2M (STMaybe t) = SMTMaybe (d2M t)
+d2M (STArr n t) = SMTMaybe (SMTArr n (d2M t))
+d2M (STScal t) = case t of
+  STI32 -> SMTNil
+  STI64 -> SMTNil
+  STF32 -> SMTScal STF32
+  STF64 -> SMTScal STF64
+  STBool -> SMTNil
+d2M STAccum{} = error "Accumulators not allowed in input program"
+d2M (STLEither a b) = SMTLEither (d2M a) (d2M b)
+
 d2 :: STy t -> STy (D2 t)
-d2 STNil = STNil
-d2 (STPair a b) = STMaybe (STPair (d2 a) (d2 b))
-d2 (STEither a b) = STMaybe (STEither (d2 a) (d2 b))
-d2 (STMaybe t) = STMaybe (d2 t)
-d2 (STArr n t) = STMaybe (STArr n (d2 t))
-d2 (STScal t) = case t of
-  STI32 -> STNil
-  STI64 -> STNil
-  STF32 -> STScal STF32
-  STF64 -> STScal STF64
-  STBool -> STNil
-d2 STAccum{} = error "Accumulators not allowed in input program"
+d2 = fromSMTy . d2M
+
+d2eM :: SList STy env -> SList SMTy (D2E env)
+d2eM SNil = SNil
+d2eM (t `SCons` ts) = d2M t `SCons` d2eM ts
 
 d2e :: SList STy env -> SList STy (D2E env)
-d2e SNil = SNil
-d2e (t `SCons` ts) = d2 t `SCons` d2e ts
+d2e = slistMap fromSMTy . d2eM
 
 d2ace :: SList STy env -> SList STy (D2AcE env)
 d2ace SNil = SNil
-d2ace (t `SCons` ts) = STAccum t `SCons` d2ace ts
+d2ace (t `SCons` ts) = STAccum (d2M t) `SCons` d2ace ts
 
 
 data CHADConfig = CHADConfig