WIP: Store subset of D1 Gamma for recompute at binding sitesrecompute-primalstores

1 files changed, 22 insertions, 0 deletions

diff --git a/src/CHAD/AST/Weaken.hs b/src/CHAD/AST/Weaken.hs
index ac0d152..16403dd 100644
--- a/src/CHAD/AST/Weaken.hs
+++ b/src/CHAD/AST/Weaken.hs
@@ -109,6 +109,25 @@ infixr 3 .>
 (.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
 (.>) = flip WThen
 
+unweakenList :: as :> bs -> SList f bs -> SList f as
+unweakenList WId l = l
+unweakenList WSink (SCons _ l) = l
+unweakenList (WCopy w) (SCons x l) = SCons x (unweakenList w l)
+unweakenList (WPop w) l = let SCons _ l' = unweakenList w l in l'
+unweakenList (WThen w1 w2) l = unweakenList w1 (unweakenList w2 l)
+unweakenList WClosed _ = SNil
+unweakenList (WIdx i) l = SCons (slistIdx l i) l
+unweakenList (WPick @_ @_ @_ @env' pre w) (SCons x l) =
+  let (l1, l2) = sunappend @env' pre l
+  in sappend l1 (SCons x (unweakenList w l2))
+unweakenList (WSwap @env @as as bs) l =
+  let (l1, l23) = sunappend @(Append as env) bs l
+      (l2, l3) = sunappend @env as l23
+  in sappend l2 (sappend l1 l3)
+unweakenList (WStack @_ @env2 _ bs wab w12) l =
+  let (l1, l2) = sunappend @env2 bs l
+  in sappend (unweakenList wab l1) (unweakenList w12 l2)
+
 class KnownListSpine list where knownListSpine :: SList (Const ()) list
 instance KnownListSpine '[] where knownListSpine = SNil
 instance KnownListSpine list => KnownListSpine (t : list) where knownListSpine = SCons (Const ()) knownListSpine
@@ -133,6 +152,9 @@ wRaiseAbove :: SList f env1 -> proxy env -> env1 :> Append env1 env
 wRaiseAbove SNil _ = WClosed
 wRaiseAbove (SCons _ s) env = WCopy (wRaiseAbove s env)
 
+wUnder :: SList f as -> as :> bs -> Append as env :> Append bs env
+wUnder as w = WStack (slistMap (\_ -> Const ()) as) (subList) w WId
+
 wPops :: SList f bs -> Append bs env1 :> env2 -> env1 :> env2
 wPops SNil w = w
 wPops (_ `SCons` bs) w = wPops bs (WPop w)