1 files changed, 21 insertions, 5 deletions
diff --git a/src/CHAD/Types.hs b/src/CHAD/Types.hs
index 130493d..6662cbf 100644
--- a/src/CHAD/Types.hs
+++ b/src/CHAD/Types.hs
@@ -17,8 +17,8 @@ type family D1 t where
 
 type family D2 t where
   D2 TNil = TNil
-  D2 (TPair a b) = TEither TNil (TPair (D2 a) (D2 b))
-  D2 (TEither a b) = TEither TNil (TEither (D2 a) (D2 b))
+  D2 (TPair a b) = TMaybe (TPair (D2 a) (D2 b))
+  D2 (TEither a b) = TMaybe (TEither (D2 a) (D2 b))
   D2 (TMaybe t) = TMaybe (D2 t)
   D2 (TArr n t) = TArr n (D2 t)
   D2 (TScal t) = D2s t
@@ -51,10 +51,14 @@ d1 (STArr n t) = STArr n (d1 t)
 d1 (STScal t) = STScal t
 d1 STAccum{} = error "Accumulators not allowed in input program"
 
+d1e :: SList STy env -> SList STy (D1E env)
+d1e SNil = SNil
+d1e (t `SCons` env) = d1 t `SCons` d1e env
+
 d2 :: STy t -> STy (D2 t)
 d2 STNil = STNil
-d2 (STPair a b) = STEither STNil (STPair (d2 a) (d2 b))
-d2 (STEither a b) = STEither STNil (STEither (d2 a) (d2 b))
+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) = STArr n (d2 t)
 d2 (STScal t) = case t of
@@ -67,7 +71,11 @@ d2 STAccum{} = error "Accumulators not allowed in input program"
 
 d2e :: SList STy env -> SList STy (D2E env)
 d2e SNil = SNil
-d2e (SCons t ts) = SCons (d2 t) (d2e ts)
+d2e (t `SCons` ts) = d2 t `SCons` d2e ts
+
+d2ace :: SList STy env -> SList STy (D2AcE env)
+d2ace SNil = SNil
+d2ace (t `SCons` ts) = STAccum (d2 t) `SCons` d2ace ts
 
 
 data CHADConfig = CHADConfig
@@ -75,10 +83,18 @@ data CHADConfig = CHADConfig
     chcLetArrayAccum :: Bool
   , -- | D[case] will bind variables containing arrays in accumulator mode.
     chcCaseArrayAccum :: Bool
+  , -- | Introduce top-level arguments containing arrays in accumulator mode.
+    chcArgArrayAccum :: Bool
   }
 
 defaultConfig :: CHADConfig
 defaultConfig = CHADConfig
   { chcLetArrayAccum = False
   , chcCaseArrayAccum = False
+  , chcArgArrayAccum = False
   }
+
+chcSetAccum :: CHADConfig -> CHADConfig
+chcSetAccum c = c { chcLetArrayAccum = True
+                  , chcCaseArrayAccum = True
+                  , chcArgArrayAccum = True }