WIP autoWeak

1 files changed, 73 insertions, 9 deletions

diff --git a/src/CHAD.hs b/src/CHAD.hs
index a6dd9ff..45d2d08 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -27,10 +27,12 @@ module CHAD (
 ) where
 
 import Data.Bifunctor (first, second)
+import Data.Functor.Const
 import Data.Kind (Type)
 import GHC.TypeLits (Symbol)
 
 import AST
+import AST.Count
 import Data
 import Lemmas
 
@@ -142,10 +144,6 @@ sreverse :: SList f ts -> SList f (Reverse ts)
 sreverse SNil = SNil
 sreverse (SCons t ts) = ssnoc (sreverse ts) t
 
-sappend :: SList f l1 -> SList f l2 -> SList f (Append l1 l2)
-sappend SNil l = l
-sappend (SCons x xs) l = SCons x (sappend xs l)
-
 stapeUnfoldings :: SList STy ts -> SList STy (TapeUnfoldings ts)
 stapeUnfoldings SNil = SNil
 stapeUnfoldings (SCons _ ts) = SCons (tapeTy ts) (stapeUnfoldings ts)
@@ -408,6 +406,63 @@ zeroTup :: SList STy env0 -> Ex env (Tup (D2E env0))
 zeroTup SNil = ENil ext
 zeroTup (SCons t env) = EPair ext (zeroTup env) (zero t)
 
+accumPromote :: forall t env sto proxy r.
+                proxy t
+             -> Descr env sto
+             -> OccEnv env
+             -> (forall stoRepl envPro.
+                    Descr env stoRepl
+                      -- ^ A revised environment description that switches
+                      -- arrays (used in the OccEnv) that are currently on
+                      -- "merge" storage, to "accum" storage.
+                 -> Subenv (Select env sto "merge") (Select env stoRepl "merge")
+                      -- ^ The new storage has fewer "merge"-storage entries.
+                 -> SList STy envPro
+                      -- ^ New entries on top of the original dual environment,
+                      -- that house the accumulators for the promoted arrays in
+                      -- the original environment.
+                 -> (forall shbinds.
+                            SList STy shbinds
+                         -> (D2 t : Append shbinds (D2AcE (Select env stoRepl "accum")))
+                            :> Append envPro (D2 t : Append shbinds (D2AcE (Select env sto "accum"))))
+                      -- ^ A weakening that converts a computation in the
+                      -- revised environment to one in the original environment
+                      -- extended with some accumulators.
+                 -> r)
+             -> r
+accumPromote _ DTop _ k = k DTop SETop SNil (\_ -> WId)
+accumPromote _ descr OccEnd k = k descr (subenvAll (select SMerge descr)) SNil (\_ -> WId)
+accumPromote pty (descr `DPush` (t, sto)) (occenv `OccPush` occ) k =
+  accumPromote pty descr occenv $ \(storepl :: Descr env1 stoRepl) mergesub (envpro :: SList _ envPro) wf ->
+    case (t, sto, occ) of
+      (STArr (arrn :: SNat arrn) (arrt :: STy arrt), SMerge, Occ _ c) | c > Zero ->
+        k (storepl `DPush` (t, SAccum))
+          (SENo mergesub)
+          (STAccum arrn arrt `SCons` envpro)
+          (\(shbinds :: SList _ shbinds) ->
+            let shbindsC = slistMap (\_ -> Const ()) shbinds
+            in
+            -- wf:
+            --                 D2 t : Append shbinds (D2AcE (Select envPro stoRepl "accum"))  :>                Append envPro (D2 t : Append shbinds (D2AcE (Select envPro sto1 "accum")))
+            -- WCopy wf:
+            --   TAccum n t3 : D2 t : Append shbinds (D2AcE (Select envPro stoRepl "accum"))  :>  TAccum n t3 : Append envPro (D2 t : Append shbinds (D2AcE (Select envPro sto1 "accum")))
+            --                       WPICK: ^                                                                 THESE TWO  ||
+            -- goal:                        |                                                                 ARE EQUAL  ||
+            --   D2 t : Append shbinds (TAccum n t3 : D2AcE (Select envPro stoRepl "accum"))  :>  TAccum n t3 : Append envPro (D2 t : Append shbinds (D2AcE (Select envPro sto1 "accum")))
+            WCopy (wf shbinds)
+            .> WPick @(TAccum arrn arrt) @(D2 t : shbinds) (Const () `SCons` shbindsC)
+                 (WId @(D2AcE (Select env1 stoRepl "accum"))))
+
+      (_, SAccum, _) ->
+        k (storepl `DPush` (t, SAccum))
+          mergesub
+          envpro
+          (\shbinds -> _ (wf shbinds))
+
+      _ -> _
+
+-- | @env'@ is a subset of @env@: each element of @env@ is either included in
+-- @env'@ ('SEYes') or not included in @env'@ ('SENo').
 data Subenv env env' where
   SETop :: Subenv '[] '[]
   SEYes :: Subenv env env' -> Subenv (t : env) (t : env')
@@ -440,11 +495,15 @@ subList SNil SETop = SNil
 subList (SCons x xs) (SEYes sub) = SCons x (subList xs sub)
 subList (SCons _ xs) (SENo sub) = subList xs sub
 
-subenvNone :: SList STy env -> Subenv env '[]
+subenvNone :: SList f env -> Subenv env '[]
 subenvNone SNil = SETop
 subenvNone (SCons _ env) = SENo (subenvNone env)
 
-subenvOnehot :: SList STy env -> Idx env t -> Subenv env '[t]
+subenvAll :: SList f env -> Subenv env env
+subenvAll SNil = SETop
+subenvAll (SCons _ env) = SEYes (subenvAll env)
+
+subenvOnehot :: SList f env -> Idx env t -> Subenv env '[t]
 subenvOnehot (SCons _ env) IZ = SEYes (subenvNone env)
 subenvOnehot (SCons _ env) (IS i) = SENo (subenvOnehot env i)
 subenvOnehot SNil i = case i of {}
@@ -525,10 +584,15 @@ weakenRets w (Rets binds list) =
   let (binds', _) = weakenBindings weakenExpr w binds
   in Rets binds' (slistMap (weakenRetPair (bindingsBinds binds) w) list)
 
-rebaseRetPair :: forall env b1 b2 env0 sto t f. SList f b1 -> SList f b2 -> RetPair env0 sto (Append b1 env) b2 t -> RetPair env0 sto env (Append b2 b1) t
-rebaseRetPair b1 b2 (RetPair p sub d)
+rebaseRetPair :: forall env b1 b2 env0 sto t f.
+                 SList f env0 -> SList f b1 -> SList f b2
+              -> RetPair env0 sto (Append b1 env) b2 t -> RetPair env0 sto env (Append b2 b1) t
+rebaseRetPair env b1 b2 (RetPair p sub d)
   | Refl <- lemAppendAssoc @b2 @b1 @env =
-      RetPair p sub (weakenExpr (WCopy (wStack @(D2AcE (Select env0 sto "accum")) (wRaiseAbove b2 b1))) d)
+      RetPair p sub (weakenExpr (autoWeak @['("d", '[D2 t]), '("b2", b2), '("b1", b1), '("tl", D2AcE (Select env0 sto "accum"))]
+                                  (LSeg @"d" :++: (LSeg @"b2" :++: LSeg @"tl"))
+                                  (LSeg @"d" :++: ((LSeg @"b2" :++: LSeg @"b1") :++: LSeg @"tl")))
+                                d)
 
 retConcat :: forall env0 sto list. SList (Ret env0 sto) list -> Rets env0 sto (D1E env0) list
 retConcat SNil = Rets BTop SNil