Autoweak!

1 files changed, 38 insertions, 15 deletions

diff --git a/src/CHAD.hs b/src/CHAD.hs
index 45d2d08..97632c7 100644
--- a/src/CHAD.hs
+++ b/src/CHAD.hs
@@ -585,19 +585,20 @@ weakenRets w (Rets binds list) =
   in Rets binds' (slistMap (weakenRetPair (bindingsBinds binds) w) list)
 
 rebaseRetPair :: forall env b1 b2 env0 sto t f.
-                 SList f env0 -> SList f b1 -> SList f b2
+                 Descr env0 sto -> 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)
+rebaseRetPair descr b1 b2 (RetPair p sub d)
   | Refl <- lemAppendAssoc @b2 @b1 @env =
-      RetPair p sub (weakenExpr (autoWeak @['("d", '[D2 t]), '("b2", b2), '("b1", b1), '("tl", D2AcE (Select env0 sto "accum"))]
+      RetPair p sub (weakenExpr (autoWeak
+                                  (Seg' @"d" @'[D2 t] $.. Seg @"b2" b2 $.. Seg @"b1" b1 $.. Seg @"tl" (d2ace (select SAccum descr)))
                                   (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
-retConcat (SCons (Ret (b :: Bindings _ _ shbinds) p sub d) list)
-  | Rets binds1 pairs1 <- retConcat list
+retConcat :: forall env0 sto list. Descr env0 sto -> SList (Ret env0 sto) list -> Rets env0 sto (D1E env0) list
+retConcat _ SNil = Rets BTop SNil
+retConcat descr (SCons (Ret (b :: Bindings _ _ shbinds) p sub d) list)
+  | Rets binds1 pairs1 <- retConcat descr list
   , Rets (binds :: Bindings _ _ shbinds2) pairs <- weakenRets (sinkWithBindings b) (Rets binds1 pairs1)
   , Refl <- lemAppendAssoc @shbinds2 @shbinds @(D1E env0)
   , Refl <- lemAppendAssoc @shbinds2 @shbinds @(D2AcE (Select env0 sto "accum"))
@@ -605,7 +606,7 @@ retConcat (SCons (Ret (b :: Bindings _ _ shbinds) p sub d) list)
          (SCons (RetPair (weakenExpr (sinkWithBindings binds) p)
                          sub
                          (weakenExpr (WCopy (sinkWithBindings binds)) d))
-                (slistMap (rebaseRetPair (bindingsBinds b) (bindingsBinds binds)) pairs))
+                (slistMap (rebaseRetPair descr (bindingsBinds b) (bindingsBinds binds)) pairs))
 
 d1op :: SOp a t -> Ex env (D1 a) -> Ex env (D1 t)
 d1op (OAdd t) e = EOp ext (OAdd t) e
@@ -731,7 +732,12 @@ drev des = \case
         (weakenExpr wbody0' body1)
         subBoth
         (ELet ext
-           (weakenExpr (WCopy (wStack @(D2AcE (Select env sto "accum")) (wRaiseAbove (bindingsBinds body0) (SCons (typeOf rhs1) (bindingsBinds rhs0)))))
+           (weakenExpr (autoWeak (Seg' @"d" @'[D2 t]
+                                  $.. Seg @"body" (bindingsBinds body0)
+                                  $.. Seg @"rhs" (SCons (typeOf rhs1) (bindingsBinds rhs0))
+                                  $.. Seg @"tl" (d2ace (select SAccum des)))
+                                 (LSeg @"d" :++: LSeg @"body" :++: LSeg @"tl")
+                                 (LSeg @"d" :++: LSeg @"body" :++: LSeg @"rhs" :++: LSeg @"tl"))
                        body2') $
          ELet ext
            (ELet ext (ESnd ext (EVar ext bodyResType IZ)) $
@@ -742,7 +748,7 @@ drev des = \case
 
   EPair _ a b
     | Rets binds (RetPair a1 subA a2 `SCons` RetPair b1 subB b2 `SCons` SNil)
-        <- retConcat $ drev des a `SCons` drev des b `SCons` SNil
+        <- retConcat des $ drev des a `SCons` drev des b `SCons` SNil
     , let dt = STPair (d2 (typeOf a)) (d2 (typeOf b)) ->
     subenvPlus (select SMerge des) subA subB $ \subBoth _ _ plus_A_B ->
     Ret binds
@@ -834,8 +840,14 @@ drev des = \case
                     ELet ext
                       (EVar ext (d2 (typeOf a)) (wSinks @(Tape rhs_a_binds : D2 t : t_primal_ty : Append e_binds (D2AcE (Select env sto "accum"))) (sappend (bindingsBinds a0) prerebinds) @> IS IZ)) $
                     ELet ext
-                      (weakenExpr (wStack @(D2AcE (Select env sto "accum")) $
-                                     WCopy (wRaiseAbove (sappend (bindingsBinds a0) prerebinds) (tapeA `SCons` d2 (typeOf a) `SCons` tPrimal `SCons` bindingsBinds e0) .> wRaiseAbove (bindingsBinds a0) prerebinds))
+                      (weakenExpr (autoWeak (Seg' @"d" @'[D2 t]
+                                             $.. Seg @"a0" (bindingsBinds a0)
+                                             $.. Seg @"prea0" prerebinds
+                                             $.. Seg @"recon" (tapeA `SCons` d2 (typeOf a) `SCons` SNil)
+                                             $.. Seg @"binds" (tPrimal `SCons` bindingsBinds e0)
+                                             $.. Seg @"tl" (d2ace (select SAccum des)))
+                                            (LSeg @"d" :++: LSeg @"a0" :++: LSeg @"tl")
+                                            (LSeg @"d" :++: (LSeg @"a0" :++: LSeg @"prea0") :++: LSeg @"recon" :++: LSeg @"binds" :++: LSeg @"tl"))
                                   a2') $
                     EPair ext
                      (expandSubenvZeros (subList (select SMerge des) subAB) sAB_A $
@@ -847,8 +859,14 @@ drev des = \case
                     ELet ext
                       (EVar ext (d2 (typeOf a)) (wSinks @(Tape rhs_b_binds : D2 t : t_primal_ty : Append e_binds (D2AcE (Select env sto "accum"))) (sappend (bindingsBinds b0) prerebinds) @> IS IZ)) $
                     ELet ext
-                      (weakenExpr (wStack @(D2AcE (Select env sto "accum")) $
-                                     WCopy (wRaiseAbove (sappend (bindingsBinds b0) prerebinds) (tapeB `SCons` d2 (typeOf a) `SCons` tPrimal `SCons` bindingsBinds e0) .> wRaiseAbove (bindingsBinds b0) prerebinds))
+                      (weakenExpr (autoWeak (Seg' @"d" @'[D2 t]
+                                             $.. Seg @"b0" (bindingsBinds b0)
+                                             $.. Seg @"preb0" prerebinds
+                                             $.. Seg @"recon" (tapeB `SCons` d2 (typeOf a) `SCons` SNil)
+                                             $.. Seg @"binds" (tPrimal `SCons` bindingsBinds e0)
+                                             $.. Seg @"tl" (d2ace (select SAccum des)))
+                                            (LSeg @"d" :++: LSeg @"b0" :++: LSeg @"tl")
+                                            (LSeg @"d" :++: (LSeg @"b0" :++: LSeg @"preb0") :++: LSeg @"recon" :++: LSeg @"binds" :++: LSeg @"tl"))
                                   b2') $
                     EPair ext
                       (expandSubenvZeros (subList (select SMerge des) subAB) sAB_B $
@@ -900,7 +918,12 @@ drev des = \case
     Ret (bconcat (ne0 `BPush` (tIx, ne1))
                  (fst (weakenBindings weakenExpr (WCopy (wSinks (bindingsBinds ne0))) ve0)))
         (EBuild1 ext
-           (weakenExpr (wStack @(D1E env) (wSinks (bindingsBinds ve0) .> WSink @TIx @ne_binds))
+           (weakenExpr (autoWeak (Seg @"ve0" (bindingsBinds ve0)
+                                  $.. Seg' @"i" @'[TIx]
+                                  $.. Seg @"ne0" (bindingsBinds ne0)
+                                  $.. Seg @"tl" (sD1eEnv des))
+                                 (LSeg @"ne0" :++: LSeg @"tl")
+                                 ((LSeg @"ve0" :++: LSeg @"i" :++: LSeg @"ne0") :++: LSeg @"tl"))
                        ne1)
            (subst (\_ t i -> case splitIdx @(TIx : D1E env) (bindingsBinds e0) i of
                                Left ibind ->