WIP accum top-level args

1 files changed, 34 insertions, 115 deletions

diff --git a/src/CHAD.hs b/src/CHAD.hs
index 9f58f73..3da083b 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -32,18 +32,17 @@ module CHAD (
 ) where
 
 import Data.Functor.Const
-import Data.Kind (Type)
 import Data.Type.Bool (If)
 import Data.Type.Equality (type (==))
 import GHC.Stack (HasCallStack)
-import GHC.TypeLits (Symbol)
 
 import AST
 import AST.Bindings
 import AST.Count
 import AST.Env
 import AST.Weaken.Auto
-import CHAD.Heuristics
+import CHAD.Accum
+import CHAD.EnvDescr
 import CHAD.Types
 import Data
 import Lemmas
@@ -197,66 +196,6 @@ reconstructBindings binds tape =
      ,sreverse (stapeUnfoldings binds))
 
 
----------------------- ENVIRONMENT DESCRIPTION AND STORAGE ---------------------
-
-type Storage :: Symbol -> Type
-data Storage s where
-  SAccum :: Storage "accum"  -- ^ in the monad state as a mutable accumulator
-  SMerge :: Storage "merge"  -- ^ just return and merge
-  SDiscr :: Storage "discr"  -- ^ we happen to know this is a discrete type and won't need any contributions
-deriving instance Show (Storage s)
-
--- | Environment description
-data Descr env sto where
-  DTop :: Descr '[] '[]
-  DPush :: Descr env sto -> (STy t, Storage s) -> Descr (t : env) (s : sto)
-deriving instance Show (Descr env sto)
-
-descrList :: Descr env sto -> SList STy env
-descrList DTop = SNil
-descrList (des `DPush` (t, _)) = t `SCons` descrList des
-
--- | This could have more precise typing on the output storage.
-subDescr :: Descr env sto -> Subenv env env'
-         -> (forall sto'. Descr env' sto'
-                       -> Subenv (Select env sto "merge") (Select env' sto' "merge")
-                       -> Subenv (D2AcE (Select env sto "accum")) (D2AcE (Select env' sto' "accum"))
-                       -> Subenv (D1E env) (D1E env')
-                       -> r)
-         -> r
-subDescr DTop SETop k = k DTop SETop SETop SETop
-subDescr (des `DPush` (t, sto)) (SEYes sub) k =
-  subDescr des sub $ \des' submerge subaccum subd1e ->
-    case sto of
-      SMerge -> k (des' `DPush` (t, sto)) (SEYes submerge) subaccum (SEYes subd1e)
-      SAccum -> k (des' `DPush` (t, sto)) submerge (SEYes subaccum) (SEYes subd1e)
-      SDiscr -> k (des' `DPush` (t, sto)) submerge subaccum (SEYes subd1e)
-subDescr (des `DPush` (_, sto)) (SENo sub) k =
-  subDescr des sub $ \des' submerge subaccum subd1e ->
-    case sto of
-      SMerge -> k des' (SENo submerge) subaccum (SENo subd1e)
-      SAccum -> k des' submerge (SENo subaccum) (SENo subd1e)
-      SDiscr -> k des' submerge subaccum (SENo subd1e)
-
--- | Select only the types from the environment that have the specified storage
-type family Select env sto s where
-  Select '[] '[] _ = '[]
-  Select (t : ts) (s : sto) s = t : Select ts sto s
-  Select (_ : ts) (_ : sto) s = Select ts sto s
-
-select :: Storage s -> Descr env sto -> SList STy (Select env sto s)
-select _ DTop = SNil
-select s@SAccum (DPush des (t, SAccum)) = SCons t (select s des)
-select s@SMerge (DPush des (_, SAccum)) = select s des
-select s@SDiscr (DPush des (_, SAccum)) = select s des
-select s@SAccum (DPush des (_, SMerge)) = select s des
-select s@SMerge (DPush des (t, SMerge)) = SCons t (select s des)
-select s@SDiscr (DPush des (_, SMerge)) = select s des
-select s@SAccum (DPush des (_, SDiscr)) = select s des
-select s@SMerge (DPush des (_, SDiscr)) = select s des
-select s@SDiscr (DPush des (t, SDiscr)) = SCons t (select s des)
-
-
 ---------------------------------- DERIVATIVES ---------------------------------
 
 d1op :: SOp a t -> Ex env (D1 a) -> Ex env (D1 t)
@@ -283,24 +222,24 @@ data D2Op a t = Linear (forall env. Ex env (D2 t) -> Ex env (D2 a))
 
 d2op :: SOp a t -> D2Op a t
 d2op op = case op of
-  OAdd t -> d2opBinArrangeInt t $ Linear $ \d -> EInr ext STNil (EPair ext d d)
+  OAdd t -> d2opBinArrangeInt t $ Linear $ \d -> EJust ext (EPair ext d d)
   OMul t -> d2opBinArrangeInt t $ Nonlinear $ \e d ->
-    EInr ext STNil (EPair ext (EOp ext (OMul t) (EPair ext (ESnd ext e) d))
-                              (EOp ext (OMul t) (EPair ext (EFst ext e) d)))
+    EJust ext (EPair ext (EOp ext (OMul t) (EPair ext (ESnd ext e) d))
+                         (EOp ext (OMul t) (EPair ext (EFst ext e) d)))
   ONeg t -> d2opUnArrangeInt t $ Linear $ \d -> EOp ext (ONeg t) d
-  OLt t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
-  OLe t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
-  OEq t -> Linear $ \_ -> EInl ext (STPair (d2 (STScal t)) (d2 (STScal t))) (ENil ext)
+  OLt t -> Linear $ \_ -> ENothing ext (STPair (d2 (STScal t)) (d2 (STScal t)))
+  OLe t -> Linear $ \_ -> ENothing ext (STPair (d2 (STScal t)) (d2 (STScal t)))
+  OEq t -> Linear $ \_ -> ENothing ext (STPair (d2 (STScal t)) (d2 (STScal t)))
   ONot -> Linear $ \_ -> ENil ext
-  OAnd -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
-  OOr -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
+  OAnd -> Linear $ \_ -> ENothing ext (STPair STNil STNil)
+  OOr -> Linear $ \_ -> ENothing ext (STPair STNil STNil)
   OIf -> Linear $ \_ -> ENil ext
   ORound64 -> Linear $ \_ -> EConst ext STF64 0.0
   OToFl64 -> Linear $ \_ -> ENil ext
   ORecip t -> floatingD2 t $ Nonlinear $ \e d -> EOp ext (OMul t) (EPair ext (EOp ext (ONeg t) (EOp ext (ORecip t) (EOp ext (OMul t) (EPair ext e e)))) d)
   OExp t -> floatingD2 t $ Nonlinear $ \e d -> EOp ext (OMul t) (EPair ext (EOp ext (OExp t) e) d)
   OLog t -> floatingD2 t $ Nonlinear $ \e d -> EOp ext (OMul t) (EPair ext (EOp ext (ORecip t) e) d)
-  OIDiv t -> integralD2 t $ Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
+  OIDiv t -> integralD2 t $ Linear $ \_ -> ENothing ext (STPair STNil STNil)
   where
     d2opUnArrangeInt :: SScalTy a
                      -> (D2s a ~ TScal a => D2Op (TScal a) t)
@@ -316,11 +255,11 @@ d2op op = case op of
                       -> (D2s a ~ TScal a => D2Op (TPair (TScal a) (TScal a)) t)
                       -> D2Op (TPair (TScal a) (TScal a)) t
     d2opBinArrangeInt ty float = case ty of
-      STI32 -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
-      STI64 -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
+      STI32 -> Linear $ \_ -> ENothing ext (STPair STNil STNil)
+      STI64 -> Linear $ \_ -> ENothing ext (STPair STNil STNil)
       STF32 -> float
       STF64 -> float
-      STBool -> Linear $ \_ -> EInl ext (STPair STNil STNil) (ENil ext)
+      STBool -> Linear $ \_ -> ENothing ext (STPair STNil STNil)
 
     floatingD2 :: ScalIsFloating a ~ True
                => SScalTy a -> ((D2s a ~ TScal a, ScalIsNumeric a ~ True) => r) -> r
@@ -332,13 +271,8 @@ d2op op = case op of
     integralD2 STI32 k = k
     integralD2 STI64 k = k
 
-sD1eEnv :: Descr env sto -> SList STy (D1E env)
-sD1eEnv DTop = SNil
-sD1eEnv (DPush d (t, _)) = SCons (d1 t) (sD1eEnv d)
-
-d2ace :: SList STy env -> SList STy (D2AcE env)
-d2ace SNil = SNil
-d2ace (SCons t ts) = SCons (STAccum (d2 t)) (d2ace ts)
+desD1E :: Descr env sto -> SList STy (D1E env)
+desD1E = d1e . descrList
 
 -- d1W :: env :> env' -> D1E env :> D1E env'
 -- d1W WId = WId
@@ -610,24 +544,6 @@ accumPromote pdty (descr `DPush` (t :: STy t, sto)) k =
     --   STScal{} -> False
     --   STAccum{} -> error "An accumulator in merge storage?"
 
-makeAccumulators :: SList STy envPro -> Ex (Append (D2AcE envPro) env) t -> Ex env (InvTup t (D2E envPro))
-makeAccumulators SNil e = e
-makeAccumulators (STArr n t `SCons` envpro) e =
-  makeAccumulators envpro $
-    EWith (zero (STArr n t)) e
-makeAccumulators (t `SCons` _) _ = error $ "makeAccumulators: Not only arrays in envpro: " ++ show t
-
-uninvertTup :: SList STy list -> STy core -> Ex env (InvTup core list) -> Ex env (TPair core (Tup list))
-uninvertTup SNil _ e = EPair ext e (ENil ext)
-uninvertTup (t `SCons` list) tcore e =
-  ELet ext (uninvertTup list (STPair tcore t) e) $
-    let recT = STPair (STPair tcore t) (tTup list)  -- type of the RHS of that let binding
-    in EPair ext
-         (EFst ext (EFst ext (EVar ext recT IZ)))
-         (EPair ext
-            (ESnd ext (EVar ext recT IZ))
-            (ESnd ext (EFst ext (EVar ext recT IZ))))
-
 
 ---------------------------- RETURN TRIPLE FROM CHAD ---------------------------
 
@@ -701,7 +617,7 @@ freezeRet :: Descr env sto
           -> Ex (D2 t : Append (D2AcE (Select env sto "accum")) (D1E env)) (TPair (D1 t) (Tup (D2E (Select env sto "merge"))))
 freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) =
   let (e0', wInsertD2Ac) = weakenBindings weakenExpr (WSink .> wSinks (d2ace (select SAccum descr))) e0
-      e2' = weakenExpr (WCopy (wCopies (subList (bindingsBinds e0) subtape) (wRaiseAbove (d2ace (select SAccum descr)) (sD1eEnv descr)))) e2
+      e2' = weakenExpr (WCopy (wCopies (subList (bindingsBinds e0) subtape) (wRaiseAbove (d2ace (select SAccum descr)) (desD1E descr)))) e2
   in letBinds e0' $
        EPair ext
          (weakenExpr wInsertD2Ac e1)
@@ -709,7 +625,7 @@ freezeRet descr (Ret e0 subtape e1 sub e2 :: Ret _ _ t) =
                                           &. #tape (subList (bindingsBinds e0) subtape)
                                           &. #shbinds (bindingsBinds e0)
                                           &. #d2ace (d2ace (select SAccum descr))
-                                          &. #tl (sD1eEnv descr))
+                                          &. #tl (desD1E descr))
                                          (#d :++: LPreW #tape #shbinds (wUndoSubenv subtape) :++: #d2ace :++: #tl)
                                          (#shbinds :++: #d :++: #d2ace :++: #tl))
                       e2') $
@@ -782,7 +698,7 @@ drev des = \case
         subtape
         (EPair ext a1 b1)
         subBoth
-        (ECase ext (EVar ext (STEither STNil (STPair (d2 (typeOf a)) (d2 (typeOf b)))) IZ)
+        (EMaybe ext
            (zeroTup (subList (select SMerge des) subBoth))
            (ELet ext (ELet ext (EFst ext (EVar ext dt IZ))
                       (weakenExpr (WCopy (wSinks' @[_,_])) a2)) $
@@ -790,7 +706,8 @@ drev des = \case
                       (weakenExpr (WCopy (wSinks' @[_,_,_])) b2)) $
             plus_A_B
              (EVar ext (tTup (d2e (subList (select SMerge des) subA))) (IS IZ))
-             (EVar ext (tTup (d2e (subList (select SMerge des) subB))) IZ)))
+             (EVar ext (tTup (d2e (subList (select SMerge des) subB))) IZ))
+           (EVar ext (STMaybe (STPair (d2 (typeOf a)) (d2 (typeOf b)))) IZ))
 
   EFst _ e
     | Ret e0 subtape e1 sub e2 <- drev des e
@@ -799,7 +716,7 @@ drev des = \case
         subtape
         (EFst ext e1)
         sub
-        (ELet ext (EInr ext STNil (EPair ext (EVar ext (d2 t1) IZ) (zero t2))) $
+        (ELet ext (EJust ext (EPair ext (EVar ext (d2 t1) IZ) (zero t2))) $
            weakenExpr (WCopy WSink) e2)
 
   ESnd _ e
@@ -809,7 +726,7 @@ drev des = \case
         subtape
         (ESnd ext e1)
         sub
-        (ELet ext (EInr ext STNil (EPair ext (zero t1) (EVar ext (d2 t2) IZ))) $
+        (ELet ext (EJust ext (EPair ext (zero t1) (EVar ext (d2 t2) IZ))) $
            weakenExpr (WCopy WSink) e2)
 
   ENil _ -> Ret BTop SETop (ENil ext) (subenvNone (select SMerge des)) (ENil ext)
@@ -820,11 +737,12 @@ drev des = \case
         subtape
         (EInl ext (d1 t2) e1)
         sub
-        (ECase ext (EVar ext (STEither STNil (STEither (d2 (typeOf e)) (d2 t2))) IZ)
+        (EMaybe ext
            (zeroTup (subList (select SMerge des) sub))
            (ECase ext (EVar ext (STEither (d2 (typeOf e)) (d2 t2)) IZ)
               (weakenExpr (WCopy (wSinks' @[_,_])) e2)
-              (EError (tTup (d2e (subList (select SMerge des) sub))) "inl<-dinr")))
+              (EError (tTup (d2e (subList (select SMerge des) sub))) "inl<-dinr"))
+           (EVar ext (STMaybe (STEither (d2 (typeOf e)) (d2 t2))) IZ))
 
   EInr _ t1 e
     | Ret e0 subtape e1 sub e2 <- drev des e ->
@@ -832,11 +750,12 @@ drev des = \case
         subtape
         (EInr ext (d1 t1) e1)
         sub
-        (ECase ext (EVar ext (STEither STNil (STEither (d2 t1) (d2 (typeOf e)))) IZ)
+        (EMaybe ext
            (zeroTup (subList (select SMerge des) sub))
            (ECase ext (EVar ext (STEither (d2 t1) (d2 (typeOf e))) IZ)
               (EError (tTup (d2e (subList (select SMerge des) sub))) "inr<-dinl")
-              (weakenExpr (WCopy (wSinks' @[_,_])) e2)))
+              (weakenExpr (WCopy (wSinks' @[_,_])) e2))
+           (EVar ext (STMaybe (STEither (d2 t1) (d2 (typeOf e)))) IZ))
 
   ECase _ e (a :: Ex _ t) b
     | STEither t1 t2 <- typeOf e
@@ -907,7 +826,7 @@ drev des = \case
                       (EInr ext (d2 t1)
                         (ESnd ext (EVar ext (STPair (tTup (d2e (subList (select SMerge des) subB))) (d2 t2)) IZ))))) $
          ELet ext
-           (ELet ext (EInr ext STNil (ESnd ext (EVar ext tCaseRet IZ))) $
+           (ELet ext (EJust ext (ESnd ext (EVar ext tCaseRet IZ))) $
               weakenExpr (WCopy (wSinks' @[_,_,_])) e2) $
          plus_AB_E
            (EFst ext (EVar ext tCaseRet (IS IZ)))
@@ -986,16 +905,16 @@ drev des = \case
                          (EVar ext shty IZ)
                          (letBinds (fst (weakenBindings weakenExpr (autoWeak (#ix (shty `SCons` SNil)
                                                                               &. #sh (shty `SCons` SNil)
-                                                                              &. #d1env (sD1eEnv des)
-                                                                              &. #d1env' (sD1eEnv usedDes))
+                                                                              &. #d1env (desD1E des)
+                                                                              &. #d1env' (desD1E usedDes))
                                                                              (#ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
                                                                              (#ix :++: #sh :++: #d1env))
                                                                    e0)) $
                             let w = autoWeak (#ix (shty `SCons` SNil)
                                                    &. #sh (shty `SCons` SNil)
                                                    &. #e0 (bindingsBinds e0)
-                                                   &. #d1env (sD1eEnv des)
-                                                   &. #d1env' (sD1eEnv usedDes))
+                                                   &. #d1env (desD1E des)
+                                                   &. #d1env' (desD1E usedDes))
                                                   (#e0 :++: #ix :++: LPreW #d1env' #d1env (wUndoSubenv subD1eUsed))
                                                   (#e0 :++: #ix :++: #sh :++: #d1env)
                             in EPair ext (weakenExpr w e1) (collectexpr w)))