WIP preserve only subset of D0 bindings in dual (...)

The point of this is to ensure that when an expression occurs in a
Build, then the parts of D0 that are only there to make sharing work out
for D1 are not laboriously taped in an array and preserved for D2, only
for D2 to ignore them.

However, while the subtape machinery is a good first step, this is not
everything: the current Build translation makes a Build for the
(elementwise) tape and separately a build for the primal. Because the
primal _does_ generally need the subtaped-away stuff, we can't just not
tape those.

TODO: figure out how to resolve this / what the next step is.

2 files changed, 24 insertions, 8 deletions

diff --git a/src/AST/Count.hs b/src/AST/Count.hs
index 31720a5..a928743 100644
--- a/src/AST/Count.hs
+++ b/src/AST/Count.hs
@@ -150,12 +150,12 @@ deleteUnused (_ `SCons` env) (OccPush occenv (Occ _ count)) k =
 
 unsafeWeakenWithSubenv :: Subenv env env' -> Expr x env t -> Expr x env' t
 unsafeWeakenWithSubenv = \sub ->
-  subst (\x t i -> case sinkWithSubenv i sub of
+  subst (\x t i -> case sinkViaSubenv i sub of
                      Just i' -> EVar x t i'
                      Nothing -> error "unsafeWeakenWithSubenv: Index occurred that was subenv'd away")
   where
-    sinkWithSubenv :: Idx env t -> Subenv env env' -> Maybe (Idx env' t)
-    sinkWithSubenv IZ (SEYes _) = Just IZ
-    sinkWithSubenv IZ (SENo _) = Nothing
-    sinkWithSubenv (IS i) (SEYes sub) = IS <$> sinkWithSubenv i sub
-    sinkWithSubenv (IS i) (SENo sub) = sinkWithSubenv i sub
+    sinkViaSubenv :: Idx env t -> Subenv env env' -> Maybe (Idx env' t)
+    sinkViaSubenv IZ (SEYes _) = Just IZ
+    sinkViaSubenv IZ (SENo _) = Nothing
+    sinkViaSubenv (IS i) (SEYes sub) = IS <$> sinkViaSubenv i sub
+    sinkViaSubenv (IS i) (SENo sub) = sinkViaSubenv i sub
diff --git a/src/AST/Env.hs b/src/AST/Env.hs
index c33bad3..4f34166 100644
--- a/src/AST/Env.hs
+++ b/src/AST/Env.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE EmptyCase #-}
+{-# LANGUAGE ExplicitForAll #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE StandaloneDeriving #-}
@@ -14,8 +15,8 @@ import Data
 -- @env'@ ('SEYes') or not included in @env'@ ('SENo').
 data Subenv env env' where
   SETop :: Subenv '[] '[]
-  SEYes :: Subenv env env' -> Subenv (t : env) (t : env')
-  SENo  :: Subenv env env' -> Subenv (t : env) env'
+  SEYes :: forall t env env'. Subenv env env' -> Subenv (t : env) (t : env')
+  SENo  :: forall t env env'. Subenv env env' -> Subenv (t : env) env'
 deriving instance Show (Subenv env env')
 
 subList :: SList f env -> Subenv env env' -> SList f env'
@@ -41,3 +42,18 @@ subenvCompose SETop SETop = SETop
 subenvCompose (SEYes sub1) (SEYes sub2) = SEYes (subenvCompose sub1 sub2)
 subenvCompose (SEYes sub1) (SENo sub2) = SENo (subenvCompose sub1 sub2)
 subenvCompose (SENo sub1) sub2 = SENo (subenvCompose sub1 sub2)
+
+subenvConcat :: Subenv env1 env1' -> Subenv env2 env2' -> Subenv (Append env2 env1) (Append env2' env1')
+subenvConcat sub1 SETop = sub1
+subenvConcat sub1 (SEYes sub2) = SEYes (subenvConcat sub1 sub2)
+subenvConcat sub1 (SENo sub2) = SENo (subenvConcat sub1 sub2)
+
+sinkWithSubenv :: Subenv env env' -> env0 :> Append env' env0
+sinkWithSubenv SETop = WId
+sinkWithSubenv (SEYes sub) = WSink .> sinkWithSubenv sub
+sinkWithSubenv (SENo sub) = sinkWithSubenv sub
+
+wUndoSubenv :: Subenv env env' -> env' :> env
+wUndoSubenv SETop = WId
+wUndoSubenv (SEYes sub) = WCopy (wUndoSubenv sub)
+wUndoSubenv (SENo sub) = WSink .> wUndoSubenv sub