1 files changed, 930 insertions, 0 deletions

diff --git a/src/CHAD/AST/Count.hs b/src/CHAD/AST/Count.hs
new file mode 100644
index 0000000..133093a
--- /dev/null
+++ b/src/CHAD/AST/Count.hs
@@ -0,0 +1,930 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DeriveGeneric #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE EmptyCase #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE PolyKinds #-}
+{-# LANGUAGE QuantifiedConstraints #-}
+{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE TypeOperators #-}
+{-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE PatternSynonyms #-}
+module CHAD.AST.Count where
+
+import Data.Functor.Product
+import Data.Some
+import Data.Type.Equality
+import GHC.Generics (Generic, Generically(..))
+
+import CHAD.Array
+import CHAD.AST
+import CHAD.AST.Env
+import CHAD.Data
+
+
+-- | The monoid operation combines assuming that /both/ branches are taken.
+class Monoid a => Occurrence a where
+  -- | One of the two branches is taken
+  (<||>) :: a -> a -> a
+  -- | This code is executed many times
+  scaleMany :: a -> a
+
+
+data Count = Zero | One | Many
+  deriving (Show, Eq, Ord)
+
+instance Semigroup Count where
+  Zero <> n = n
+  n <> Zero = n
+  _ <> _ = Many
+instance Monoid Count where
+  mempty = Zero
+instance Occurrence Count where
+  (<||>) = max
+  scaleMany Zero = Zero
+  scaleMany _ = Many
+
+data Occ = Occ { _occLexical :: Count
+               , _occRuntime :: Count }
+  deriving (Eq, Generic)
+  deriving (Semigroup, Monoid) via Generically Occ
+
+instance Show Occ where
+  showsPrec d (Occ l r) = showParen (d > 10) $
+    showString "Occ " . showsPrec 11 l . showString " " . showsPrec 11 r
+
+instance Occurrence Occ where
+  Occ l1 r1 <||> Occ l2 r2 = Occ (l1 <> l2) (r1 <||> r2)
+  scaleMany (Occ l c) = Occ l (scaleMany c)
+
+
+data Substruc t t' where
+  -- If you add constructors here, do not forget to update the COMPLETE pragmas of any pattern synonyms below
+  SsFull :: Substruc t t
+  SsNone :: Substruc t TNil
+  SsPair :: Substruc a a' -> Substruc b b' -> Substruc (TPair a b) (TPair a' b')
+  SsEither :: Substruc a a' -> Substruc b b' -> Substruc (TEither a b) (TEither a' b')
+  SsLEither :: Substruc a a' -> Substruc b b' -> Substruc (TLEither a b) (TLEither a' b')
+  SsMaybe :: Substruc a a' -> Substruc (TMaybe a) (TMaybe a')
+  SsArr :: Substruc a a' -> Substruc (TArr n a) (TArr n a')  -- ^ union of usages of all array elements
+  SsAccum :: Substruc a a' -> Substruc (TAccum a) (TAccum a')
+
+pattern SsPair' :: forall a b t'. forall a' b'. t' ~ TPair a' b' => Substruc a a' -> Substruc b b' -> Substruc (TPair a b) t'
+pattern SsPair' s1 s2 <- ((\case { SsFull -> SsPair SsFull SsFull ; s -> s }) -> SsPair s1 s2)
+  where SsPair' = SsPair
+{-# COMPLETE SsNone, SsPair', SsEither, SsLEither, SsMaybe, SsArr, SsAccum #-}
+
+pattern SsArr' :: forall n a t'. forall a'. t' ~ TArr n a' => Substruc a a' -> Substruc (TArr n a) t'
+pattern SsArr' s <- ((\case { SsFull -> SsArr SsFull ; s -> s }) -> SsArr s)
+  where SsArr' = SsArr
+{-# COMPLETE SsNone, SsPair, SsEither, SsLEither, SsMaybe, SsArr', SsAccum #-}
+
+instance Semigroup (Some (Substruc t)) where
+  Some SsFull <> _ = Some SsFull
+  _ <> Some SsFull = Some SsFull
+  Some SsNone <> s = s
+  s <> Some SsNone = s
+  Some (SsPair a b) <> Some (SsPair a' b') = withSome (Some a <> Some a') $ \a2 -> withSome (Some b <> Some b') $ \b2 -> Some (SsPair a2 b2)
+  Some (SsEither a b) <> Some (SsEither a' b') = withSome (Some a <> Some a') $ \a2 -> withSome (Some b <> Some b') $ \b2 -> Some (SsEither a2 b2)
+  Some (SsLEither a b) <> Some (SsLEither a' b') = withSome (Some a <> Some a') $ \a2 -> withSome (Some b <> Some b') $ \b2 -> Some (SsLEither a2 b2)
+  Some (SsMaybe a) <> Some (SsMaybe a') = withSome (Some a <> Some a') $ \a2 -> Some (SsMaybe a2)
+  Some (SsArr a) <> Some (SsArr a') = withSome (Some a <> Some a') $ \a2 -> Some (SsArr a2)
+  Some (SsAccum a) <> Some (SsAccum a') = withSome (Some a <> Some a') $ \a2 -> Some (SsAccum a2)
+instance Monoid (Some (Substruc t)) where
+  mempty = Some SsNone
+
+instance TestEquality (Substruc t) where
+  testEquality SsFull s = isFull s
+  testEquality s SsFull = sym <$> isFull s
+  testEquality SsNone SsNone = Just Refl
+  testEquality SsNone _ = Nothing
+  testEquality _ SsNone = Nothing
+  testEquality (SsPair a b) (SsPair a' b') | Just Refl <- testEquality a a', Just Refl <- testEquality b b' = Just Refl | otherwise = Nothing
+  testEquality (SsEither a b) (SsEither a' b') | Just Refl <- testEquality a a', Just Refl <- testEquality b b' = Just Refl | otherwise = Nothing
+  testEquality (SsLEither a b) (SsLEither a' b') | Just Refl <- testEquality a a', Just Refl <- testEquality b b' = Just Refl | otherwise = Nothing
+  testEquality (SsMaybe s) (SsMaybe s') | Just Refl <- testEquality s s' = Just Refl | otherwise = Nothing
+  testEquality (SsArr s) (SsArr s') | Just Refl <- testEquality s s' = Just Refl | otherwise = Nothing
+  testEquality (SsAccum s) (SsAccum s') | Just Refl <- testEquality s s' = Just Refl | otherwise = Nothing
+
+isFull :: Substruc t t' -> Maybe (t :~: t')
+isFull SsFull = Just Refl
+isFull SsNone = Nothing  -- TODO: nil?
+isFull (SsPair a b) | Just Refl <- isFull a, Just Refl <- isFull b = Just Refl | otherwise = Nothing
+isFull (SsEither a b) | Just Refl <- isFull a, Just Refl <- isFull b = Just Refl | otherwise = Nothing
+isFull (SsLEither a b) | Just Refl <- isFull a, Just Refl <- isFull b = Just Refl | otherwise = Nothing
+isFull (SsMaybe s) | Just Refl <- isFull s = Just Refl | otherwise = Nothing
+isFull (SsArr s) | Just Refl <- isFull s = Just Refl | otherwise = Nothing
+isFull (SsAccum s) | Just Refl <- isFull s = Just Refl | otherwise = Nothing
+
+applySubstruc :: Substruc t t' -> STy t -> STy t'
+applySubstruc SsFull t = t
+applySubstruc SsNone _ = STNil
+applySubstruc (SsPair s1 s2) (STPair a b) = STPair (applySubstruc s1 a) (applySubstruc s2 b)
+applySubstruc (SsEither s1 s2) (STEither a b) = STEither (applySubstruc s1 a) (applySubstruc s2 b)
+applySubstruc (SsLEither s1 s2) (STLEither a b) = STLEither (applySubstruc s1 a) (applySubstruc s2 b)
+applySubstruc (SsMaybe s) (STMaybe t) = STMaybe (applySubstruc s t)
+applySubstruc (SsArr s) (STArr n t) = STArr n (applySubstruc s t)
+applySubstruc (SsAccum s) (STAccum t) = STAccum (applySubstrucM s t)
+
+applySubstrucM :: Substruc t t' -> SMTy t -> SMTy t'
+applySubstrucM SsFull t = t
+applySubstrucM SsNone _ = SMTNil
+applySubstrucM (SsPair s1 s2) (SMTPair a b) = SMTPair (applySubstrucM s1 a) (applySubstrucM s2 b)
+applySubstrucM (SsLEither s1 s2) (SMTLEither a b) = SMTLEither (applySubstrucM s1 a) (applySubstrucM s2 b)
+applySubstrucM (SsMaybe s) (SMTMaybe t) = SMTMaybe (applySubstrucM s t)
+applySubstrucM (SsArr s) (SMTArr n t) = SMTArr n (applySubstrucM s t)
+applySubstrucM _ t = case t of {}
+
+data ExMap a b = ExMap (forall env. Ex env a -> Ex env b)
+               | a ~ b => ExMapId
+
+fromExMap :: ExMap a b -> Ex env a -> Ex env b
+fromExMap (ExMap f) = f
+fromExMap ExMapId = id
+
+simplifySubstruc :: STy t -> Substruc t t' -> Substruc t t'
+simplifySubstruc STNil SsNone = SsFull
+
+simplifySubstruc _ SsFull = SsFull
+simplifySubstruc _ SsNone = SsNone
+simplifySubstruc (STPair t1 t2) (SsPair s1 s2) = SsPair (simplifySubstruc t1 s1) (simplifySubstruc t2 s2)
+simplifySubstruc (STEither t1 t2) (SsEither s1 s2) = SsEither (simplifySubstruc t1 s1) (simplifySubstruc t2 s2)
+simplifySubstruc (STLEither t1 t2) (SsLEither s1 s2) = SsLEither (simplifySubstruc t1 s1) (simplifySubstruc t2 s2)
+simplifySubstruc (STMaybe t) (SsMaybe s) = SsMaybe (simplifySubstruc t s)
+simplifySubstruc (STArr _ t) (SsArr s) = SsArr (simplifySubstruc t s)
+simplifySubstruc (STAccum t) (SsAccum s) = SsAccum (simplifySubstruc (fromSMTy t) s)
+
+-- simplifySubstruc' :: Substruc t t'
+--                  -> (forall t'2. Substruc t t'2 -> ExMap t'2 t' -> r) -> r
+-- simplifySubstruc' SsFull k = k SsFull ExMapId
+-- simplifySubstruc' SsNone k = k SsNone ExMapId
+-- simplifySubstruc' (SsPair s1 s2) k =
+--   simplifySubstruc' s1 $ \s1' f1 ->
+--   simplifySubstruc' s2 $ \s2' f2 ->
+--   case (s1', s2') of
+--     (SsFull, SsFull) ->
+--       k SsFull (case (f1, f2) of
+--                   (ExMapId, ExMapId) -> ExMapId
+--                   _ -> ExMap (\e -> eunPair e $ \_ e1 e2 ->
+--                                       EPair ext (fromExMap f1 e1) (fromExMap f2 e2)))
+--     (SsNone, SsNone) -> k SsNone (ExMap (\_ -> EPair ext (fromExMap f1 (ENil ext)) (fromExMap f2 (ENil ext))))
+--     _ -> k (SsPair s1' s2') (ExMap (\e -> elet e $ EPair ext (fromExMap f1 (EFst ext (evar IZ))) (fromExMap f2 (ESnd ext (evar IZ)))))
+-- simplifySubstruc' _ _ = _
+
+-- ssUnpair :: Substruc (TPair a b) -> (Substruc a, Substruc b)
+-- ssUnpair SsFull = (SsFull, SsFull)
+-- ssUnpair SsNone = (SsNone, SsNone)
+-- ssUnpair (SsPair a b) = (a, b)
+
+-- ssUnleft :: Substruc (TEither a b) -> Substruc a
+-- ssUnleft SsFull = SsFull
+-- ssUnleft SsNone = SsNone
+-- ssUnleft (SsEither a _) = a
+
+-- ssUnright :: Substruc (TEither a b) -> Substruc b
+-- ssUnright SsFull = SsFull
+-- ssUnright SsNone = SsNone
+-- ssUnright (SsEither _ b) = b
+
+-- ssUnlleft :: Substruc (TLEither a b) -> Substruc a
+-- ssUnlleft SsFull = SsFull
+-- ssUnlleft SsNone = SsNone
+-- ssUnlleft (SsLEither a _) = a
+
+-- ssUnlright :: Substruc (TLEither a b) -> Substruc b
+-- ssUnlright SsFull = SsFull
+-- ssUnlright SsNone = SsNone
+-- ssUnlright (SsLEither _ b) = b
+
+-- ssUnjust :: Substruc (TMaybe a) -> Substruc a
+-- ssUnjust SsFull = SsFull
+-- ssUnjust SsNone = SsNone
+-- ssUnjust (SsMaybe a) = a
+
+-- ssUnarr :: Substruc (TArr n a) -> Substruc a
+-- ssUnarr SsFull = SsFull
+-- ssUnarr SsNone = SsNone
+-- ssUnarr (SsArr a) = a
+
+-- ssUnaccum :: Substruc (TAccum a) -> Substruc a
+-- ssUnaccum SsFull = SsFull
+-- ssUnaccum SsNone = SsNone
+-- ssUnaccum (SsAccum a) = a
+
+
+type family MapEmpty env where
+  MapEmpty '[] = '[]
+  MapEmpty (t : env) = TNil : MapEmpty env
+
+data OccEnv a env env' where
+  OccEnd :: OccEnv a env (MapEmpty env)  -- not necessarily top!
+  OccPush :: OccEnv a env env' -> a -> Substruc t t' -> OccEnv a (t : env) (t' : env')
+
+instance Semigroup a => Semigroup (Some (OccEnv a env)) where
+  Some OccEnd <> e = e
+  e <> Some OccEnd = e
+  Some (OccPush e o s) <> Some (OccPush e' o' s') = withSome (Some e <> Some e') $ \e2 -> withSome (Some s <> Some s') $ \s2 -> Some (OccPush e2 (o <> o') s2)
+
+instance Semigroup a => Monoid (Some (OccEnv a env)) where
+  mempty = Some OccEnd
+
+instance Occurrence a => Occurrence (Some (OccEnv a env)) where
+  Some OccEnd <||> e = e
+  e <||> Some OccEnd = e
+  Some (OccPush e o s) <||> Some (OccPush e' o' s') = withSome (Some e <||> Some e') $ \e2 -> withSome (Some s <> Some s') $ \s2 -> Some (OccPush e2 (o <||> o') s2)
+
+  scaleMany (Some OccEnd) = Some OccEnd
+  scaleMany (Some (OccPush e o s)) = withSome (scaleMany (Some e)) $ \e2 -> Some (OccPush e2 (scaleMany o) s)
+
+onehotOccEnv :: Monoid a => Idx env t -> a -> Substruc t t' -> Some (OccEnv a env)
+onehotOccEnv IZ v s = Some (OccPush OccEnd v s)
+onehotOccEnv (IS i) v s
+  | Some env' <- onehotOccEnv i v s
+  = Some (OccPush env' mempty SsNone)
+
+occEnvPop :: OccEnv a (t : env) (t' : env') -> (OccEnv a env env', Substruc t t')
+occEnvPop (OccPush e _ s) = (e, s)
+occEnvPop OccEnd = (OccEnd, SsNone)
+
+occEnvPop' :: OccEnv a (t : env) env' -> (forall t' env''. env' ~ t' : env'' => OccEnv a env env'' -> Substruc t t' -> r) -> r
+occEnvPop' (OccPush e _ s) k = k e s
+occEnvPop' OccEnd k = k OccEnd SsNone
+
+occEnvPopSome :: Some (OccEnv a (t : env)) -> Some (OccEnv a env)
+occEnvPopSome (Some (OccPush e _ _)) = Some e
+occEnvPopSome (Some OccEnd) = Some OccEnd
+
+occEnvPrj :: Monoid a => OccEnv a env env' -> Idx env t -> (a, Some (Substruc t))
+occEnvPrj OccEnd _ = mempty
+occEnvPrj (OccPush _ o s) IZ = (o, Some s)
+occEnvPrj (OccPush e _ _) (IS i) = occEnvPrj e i
+
+occEnvPrjS :: OccEnv a env env' -> Idx env t -> Some (Product (Substruc t) (Idx env'))
+occEnvPrjS OccEnd IZ = Some (Pair SsNone IZ)
+occEnvPrjS OccEnd (IS i) | Some (Pair s i') <- occEnvPrjS OccEnd i = Some (Pair s (IS i'))
+occEnvPrjS (OccPush _ _ s) IZ = Some (Pair s IZ)
+occEnvPrjS (OccPush e _ _) (IS i)
+  | Some (Pair s' i') <- occEnvPrjS e i
+  = Some (Pair s' (IS i'))
+
+projectSmallerSubstruc :: Substruc t t'big -> Substruc t t'small -> Ex env t'big -> Ex env t'small
+projectSmallerSubstruc topsbig topssmall ex = case (topsbig, topssmall) of
+  _ | Just Refl <- testEquality topsbig topssmall -> ex
+
+  (SsFull, SsFull) -> ex
+  (SsNone, SsNone) -> ex
+  (SsNone, _) -> error "projectSmallerSubstruc: smaller substructure not smaller"
+  (_, SsNone) ->
+    case typeOf ex of
+      STNil -> ex
+      _ -> use ex $ ENil ext
+
+  (SsPair s1 s2, SsPair s1' s2') ->
+    eunPair ex $ \_ e1 e2 ->
+      EPair ext (projectSmallerSubstruc s1 s1' e1) (projectSmallerSubstruc s2 s2' e2)
+  (s@SsPair{}, SsFull) -> projectSmallerSubstruc s (SsPair SsFull SsFull) ex
+  (SsFull, s@SsPair{}) -> projectSmallerSubstruc (SsPair SsFull SsFull) s ex
+
+  (SsEither s1 s2, SsEither s1' s2')
+    | STEither t1 t2 <- typeOf ex ->
+    let e1 = projectSmallerSubstruc s1 s1' (EVar ext t1 IZ)
+        e2 = projectSmallerSubstruc s2 s2' (EVar ext t2 IZ)
+    in ecase ex
+         (EInl ext (typeOf e2) e1)
+         (EInr ext (typeOf e1) e2)
+  (s@SsEither{}, SsFull) -> projectSmallerSubstruc s (SsEither SsFull SsFull) ex
+  (SsFull, s@SsEither{}) -> projectSmallerSubstruc (SsEither SsFull SsFull) s ex
+
+  (SsLEither s1 s2, SsLEither s1' s2')
+    | STLEither t1 t2 <- typeOf ex ->
+    let e1 = projectSmallerSubstruc s1 s1' (EVar ext t1 IZ)
+        e2 = projectSmallerSubstruc s2 s2' (EVar ext t2 IZ)
+    in elcase ex
+         (ELNil ext (typeOf e1) (typeOf e2))
+         (ELInl ext (typeOf e2) e1)
+         (ELInr ext (typeOf e1) e2)
+  (s@SsLEither{}, SsFull) -> projectSmallerSubstruc s (SsLEither SsFull SsFull) ex
+  (SsFull, s@SsLEither{}) -> projectSmallerSubstruc (SsLEither SsFull SsFull) s ex
+
+  (SsMaybe s1, SsMaybe s1')
+    | STMaybe t1 <- typeOf ex ->
+    let e1 = projectSmallerSubstruc s1 s1' (EVar ext t1 IZ)
+    in emaybe ex
+         (ENothing ext (typeOf e1))
+         (EJust ext e1)
+  (s@SsMaybe{}, SsFull) -> projectSmallerSubstruc s (SsMaybe SsFull) ex
+  (SsFull, s@SsMaybe{}) -> projectSmallerSubstruc (SsMaybe SsFull) s ex
+
+  (SsArr s1, SsArr s2) -> emap (projectSmallerSubstruc s1 s2 (evar IZ)) ex
+  (s@SsArr{}, SsFull) -> projectSmallerSubstruc s (SsArr SsFull) ex
+  (SsFull, s@SsArr{}) -> projectSmallerSubstruc (SsArr SsFull) s ex
+
+  (SsAccum _, SsAccum _) -> error "TODO smaller ssaccum"
+  (s@SsAccum{}, SsFull) -> projectSmallerSubstruc s (SsAccum SsFull) ex
+  (SsFull, s@SsAccum{}) -> projectSmallerSubstruc (SsAccum SsFull) s ex
+
+
+-- | A boolean for each entry in the environment, with the ability to uniformly
+-- mask the top part above a certain index.
+data EnvMask env where
+  EMRest :: Bool -> EnvMask env
+  EMPush :: EnvMask env -> Bool -> EnvMask (t : env)
+
+envMaskPrj :: EnvMask env -> Idx env t -> Bool
+envMaskPrj (EMRest b) _ = b
+envMaskPrj (_ `EMPush` b) IZ = b
+envMaskPrj (env `EMPush` _) (IS i) = envMaskPrj env i
+
+occCount :: Idx env a -> Expr x env t -> Occ
+occCount idx ex
+  | Some env <- occCountAll ex
+  = fst (occEnvPrj env idx)
+
+occCountAll :: Expr x env t -> Some (OccEnv Occ env)
+occCountAll ex = occCountX SsFull ex $ \env _ -> Some env
+
+pruneExpr :: SList f env -> Expr x env t -> Ex env t
+pruneExpr env ex = occCountX SsFull ex $ \_ mkex -> mkex (fullOccEnv env)
+  where
+    fullOccEnv :: SList f env -> OccEnv () env env
+    fullOccEnv SNil = OccEnd
+    fullOccEnv (_ `SCons` e) = OccPush (fullOccEnv e) () SsFull
+
+-- In one traversal, count occurrences of variables and determine what parts of
+-- expressions are actually used. These two results are computed independently:
+-- even if (almost) nothing of a particular term is actually used, variable
+-- references in that term still count as usual.
+--
+-- In @occCountX s t k@:
+-- * s: how much of the result of this term is required
+-- * t: the term to analyse
+-- * k: is passed the actual environment usage of this expression, including
+--   occurrence counts. The callback reconstructs a new expression in an
+--   updated "response" environment. The response must be at least as large as
+--   the computed usages.
+occCountX :: forall env t t' x r. Substruc t t' -> Expr x env t
+          -> (forall env'. OccEnv Occ env env'
+                           -- response OccEnv must be at least as large as the OccEnv returned above
+                        -> (forall env''. OccEnv () env env'' -> Ex env'' t')
+                        -> r)
+          -> r
+occCountX initialS topexpr k = case topexpr of
+  EVar _ t i ->
+    withSome (onehotOccEnv i (Occ One One) s) $ \env ->
+    k env $ \env' ->
+    withSome (occEnvPrjS env' i) $ \(Pair s' i') ->
+      projectSmallerSubstruc s' s (EVar ext (applySubstruc s' t) i')
+  ELet _ rhs body ->
+    occCountX s body $ \envB mkbody ->
+    occEnvPop' envB $ \envB' s1 ->
+    occCountX s1 rhs $ \envR mkrhs ->
+    withSome (Some envB' <> Some envR) $ \env ->
+    k env $ \env' ->
+      ELet ext (mkrhs env') (mkbody (OccPush env' () s1))
+  EPair _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      SsPair' s1 s2 ->
+        occCountX s1 a $ \env1 mka ->
+        occCountX s2 b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EPair ext (mka env') (mkb env')
+  EFst _ e ->
+    occCountX (SsPair s SsNone) e $ \env1 mke ->
+    k env1 $ \env' ->
+      EFst ext (mke env')
+  ESnd _ e ->
+    occCountX (SsPair SsNone s) e $ \env1 mke ->
+    k env1 $ \env' ->
+      ESnd ext (mke env')
+  ENil _ ->
+    case s of
+      SsFull -> k OccEnd (\_ -> ENil ext)
+      SsNone -> k OccEnd (\_ -> ENil ext)
+  EInl _ t e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      SsEither s1 s2 ->
+        occCountX s1 e $ \env1 mke ->
+        k env1 $ \env' ->
+          EInl ext (applySubstruc s2 t) (mke env')
+      SsFull -> occCountX (SsEither SsFull SsFull) topexpr k
+  EInr _ t e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      SsEither s1 s2 ->
+        occCountX s2 e $ \env1 mke ->
+        k env1 $ \env' ->
+          EInr ext (applySubstruc s1 t) (mke env')
+      SsFull -> occCountX (SsEither SsFull SsFull) topexpr k
+  ECase _ e a b ->
+    occCountX s a $ \env1' mka ->
+    occCountX s b $ \env2' mkb ->
+    occEnvPop' env1' $ \env1 s1 ->
+    occEnvPop' env2' $ \env2 s2 ->
+    occCountX (SsEither s1 s2) e $ \env0 mke ->
+    withSome (Some env0 <> (Some env1 <||> Some env2)) $ \env ->
+    k env $ \env' ->
+      ECase ext (mke env') (mka (OccPush env' () s1)) (mkb (OccPush env' () s2))
+  ENothing _ t ->
+    case s of
+      SsNone -> k OccEnd (\_ -> ENil ext)
+      SsMaybe s' -> k OccEnd (\_ -> ENothing ext (applySubstruc s' t))
+      SsFull -> occCountX (SsMaybe SsFull) topexpr k
+  EJust _ e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      SsMaybe s' ->
+        occCountX s' e $ \env1 mke ->
+        k env1 $ \env' ->
+          EJust ext (mke env')
+      SsFull -> occCountX (SsMaybe SsFull) topexpr k
+  EMaybe _ a b e ->
+    occCountX s a $ \env1 mka ->
+    occCountX s b $ \env2' mkb ->
+    occEnvPop' env2' $ \env2 s2 ->
+    occCountX (SsMaybe s2) e $ \env0 mke ->
+    withSome (Some env0 <> (Some env1 <||> Some env2)) $ \env ->
+    k env $ \env' ->
+      EMaybe ext (mka env') (mkb (OccPush env' () s2)) (mke env')
+  ELNil _ t1 t2 ->
+    case s of
+      SsNone -> k OccEnd (\_ -> ENil ext)
+      SsLEither s1 s2 -> k OccEnd (\_ -> ELNil ext (applySubstruc s1 t1) (applySubstruc s2 t2))
+      SsFull -> occCountX (SsLEither SsFull SsFull) topexpr k
+  ELInl _ t e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      SsLEither s1 s2 ->
+        occCountX s1 e $ \env1 mke ->
+        k env1 $ \env' ->
+          ELInl ext (applySubstruc s2 t) (mke env')
+      SsFull -> occCountX (SsLEither SsFull SsFull) topexpr k
+  ELInr _ t e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      SsLEither s1 s2 ->
+        occCountX s2 e $ \env1 mke ->
+        k env1 $ \env' ->
+          ELInr ext (applySubstruc s1 t) (mke env')
+      SsFull -> occCountX (SsLEither SsFull SsFull) topexpr k
+  ELCase _ e a b c ->
+    occCountX s a $ \env1 mka ->
+    occCountX s b $ \env2' mkb ->
+    occCountX s c $ \env3' mkc ->
+    occEnvPop' env2' $ \env2 s1 ->
+    occEnvPop' env3' $ \env3 s2 ->
+    occCountX (SsLEither s1 s2) e $ \env0 mke ->
+    withSome (Some env0 <> (Some env1 <||> Some env2 <||> Some env3)) $ \env ->
+    k env $ \env' ->
+      ELCase ext (mke env') (mka env') (mkb (OccPush env' () s1)) (mkc (OccPush env' () s2))
+
+  EConstArr _ n t x ->
+    case s of
+      SsNone -> k OccEnd (\_ -> ENil ext)
+      SsArr' SsNone -> k OccEnd (\_ -> EBuild ext n (eshapeConst (arrayShape x)) (ENil ext))
+      SsArr' SsFull -> k OccEnd (\_ -> EConstArr ext n t x)
+
+  EBuild _ n a b ->
+    case s of
+      SsNone ->
+        occCountX SsFull a $ \env1 mka ->
+        occCountX SsNone b $ \env2'' mkb ->
+        occEnvPop' env2'' $ \env2' s2 ->
+        withSome (Some env1 <> scaleMany (Some env2')) $ \env ->
+        k env $ \env' ->
+          use (EBuild ext n (mka env') $
+                 use (elet (projectSmallerSubstruc SsFull s2 (EVar ext (tTup (sreplicate n tIx)) IZ)) $
+                        weakenExpr (WCopy WSink) (mkb (OccPush env' () s2))) $
+                   ENil ext) $
+            ENil ext
+      SsArr' s' ->
+        occCountX SsFull a $ \env1 mka ->
+        occCountX s' b $ \env2'' mkb ->
+        occEnvPop' env2'' $ \env2' s2 ->
+        withSome (Some env1 <> scaleMany (Some env2')) $ \env ->
+        k env $ \env' ->
+          EBuild ext n (mka env') $
+            elet (projectSmallerSubstruc SsFull s2 (EVar ext (tTup (sreplicate n tIx)) IZ)) $
+              weakenExpr (WCopy WSink) (mkb (OccPush env' () s2))
+
+  EMap _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1'' mka ->
+        occEnvPop' env1'' $ \env1' s1 ->
+        occCountX (SsArr s1) b $ \env2 mkb ->
+        withSome (scaleMany (Some env1') <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (EMap ext (mka (OccPush env' () s1)) (mkb env')) $
+            ENil ext
+      SsArr' s' ->
+        occCountX s' a $ \env1'' mka ->
+        occEnvPop' env1'' $ \env1' s1 ->
+        occCountX (SsArr s1) b $ \env2 mkb ->
+        withSome (scaleMany (Some env1') <> Some env2) $ \env ->
+        k env $ \env' ->
+          EMap ext (mka (OccPush env' () s1)) (mkb env')
+
+  EFold1Inner _ commut a b c ->
+    occCountX SsFull a $ \env1'' mka ->
+    occEnvPop' env1'' $ \env1' s1' ->
+    let s1 = case s1' of
+               SsNone -> Some SsNone
+               SsPair' s1'a s1'b -> Some s1'a <> Some s1'b
+        s0 = case s of
+               SsNone -> Some SsNone
+               SsArr' s' -> Some s' in
+    withSome (s1 <> s0) $ \sElt ->
+    occCountX sElt b $ \env2 mkb ->
+    occCountX (SsArr sElt) c $ \env3 mkc ->
+    withSome (scaleMany (Some env1') <> Some env2 <> Some env3) $ \env ->
+    k env $ \env' ->
+      projectSmallerSubstruc (SsArr sElt) s $
+        EFold1Inner ext commut
+           (projectSmallerSubstruc SsFull sElt $
+             mka (OccPush env' () (SsPair sElt sElt)))
+           (mkb env') (mkc env')
+
+  ESum1Inner _ e -> handleReduction (ESum1Inner ext) e
+
+  EUnit _ e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env mke ->
+        k env $ \env' ->
+          use (mke env') $ ENil ext
+      SsArr' s' ->
+        occCountX s' e $ \env mke ->
+        k env $ \env' ->
+          EUnit ext (mke env')
+
+  EReplicate1Inner _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      SsArr' s' ->
+        occCountX SsFull a $ \env1 mka ->
+        occCountX (SsArr s') b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EReplicate1Inner ext (mka env') (mkb env')
+
+  EMaximum1Inner _ e -> handleReduction (EMaximum1Inner ext) e
+  EMinimum1Inner _ e -> handleReduction (EMinimum1Inner ext) e
+
+  EReshape _ n esh e ->
+    case s of
+      SsNone ->
+        occCountX SsNone esh $ \env1 mkesh ->
+        occCountX SsNone e $ \env2 mke ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mkesh env') $ use (mke env') $ ENil ext
+      SsArr' s' ->
+        occCountX SsFull esh $ \env1 mkesh ->
+        occCountX (SsArr s') e $ \env2 mke ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EReshape ext n (mkesh env') (mke env')
+
+  EZip _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      SsArr' SsNone ->
+        occCountX (SsArr SsNone) a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mkb env') $ mka env'
+      SsArr' (SsPair' SsNone s2) ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX (SsArr s2) b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $
+            emap (EPair ext (ENil ext) (evar IZ)) (mkb env')
+      SsArr' (SsPair' s1 SsNone) ->
+        occCountX (SsArr s1) a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mkb env') $
+            emap (EPair ext (evar IZ) (ENil ext)) (mka env')
+      SsArr' (SsPair' s1 s2) ->
+        occCountX (SsArr s1) a $ \env1 mka ->
+        occCountX (SsArr s2) b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EZip ext (mka env') (mkb env')
+
+  EFold1InnerD1 _ cm e1 e2 e3 ->
+    case s of
+      -- If nothing is necessary, we can execute a fold and then proceed to ignore it
+      SsNone ->
+        let foldex = EFold1Inner ext cm (EFst ext (mapExt (\_ -> ext) e1))
+                                        (mapExt (\_ -> ext) e2) (mapExt (\_ -> ext) e3)
+        in occCountX SsNone foldex $ \env1 mkfoldex -> k env1 mkfoldex
+      -- If we don't need the stores, still a fold suffices
+      SsPair' sP SsNone ->
+        let foldex = EFold1Inner ext cm (EFst ext (mapExt (\_ -> ext) e1))
+                                        (mapExt (\_ -> ext) e2) (mapExt (\_ -> ext) e3)
+        in occCountX sP foldex $ \env1 mkfoldex -> k env1 $ \env' -> EPair ext (mkfoldex env') (ENil ext)
+      -- If for whatever reason the additional stores themselves are
+      -- unnecessary but the shape of the array is, then oblige
+      SsPair' sP (SsArr' SsNone) ->
+        let STArr sn _ = typeOf e3
+            foldex =
+              elet (mapExt (\_ -> ext) e3) $
+                EPair ext
+                  (EShape ext (evar IZ))
+                  (EFold1Inner ext cm (EFst ext (mapExt (\_ -> ext) (weakenExpr (WCopy WSink) e1)))
+                                      (mapExt (\_ -> ext) (weakenExpr WSink e2))
+                                      (evar IZ))
+        in occCountX (SsPair SsFull sP) foldex $ \env1 mkfoldex ->
+           k env1 $ \env' ->
+             eunPair (mkfoldex env') $ \_ eshape earr ->
+               EPair ext earr (EBuild ext sn eshape (ENil ext))
+      -- If at least some of the additional stores are required, we need to keep this a mapAccum
+      SsPair' _ (SsArr' sB) ->
+        -- TODO: propagate usage of primals
+        occCountX (SsPair SsFull sB) e1 $ \env1_1' mka ->
+        occEnvPop' env1_1' $ \env1' _ ->
+        occCountX SsFull e2 $ \env2 mkb ->
+        occCountX SsFull e3 $ \env3 mkc ->
+        withSome (scaleMany (Some env1') <> Some env2 <> Some env3) $ \env ->
+        k env $ \env' ->
+          projectSmallerSubstruc (SsPair SsFull (SsArr sB)) s $
+            EFold1InnerD1 ext cm (mka (OccPush env' () SsFull))
+                                 (mkb env') (mkc env')
+
+  EFold1InnerD2 _ cm ef ebog ed ->
+    -- TODO: propagate usage of duals
+    occCountX SsFull ef $ \env1_2' mkef ->
+    occEnvPop' env1_2' $ \env1_1' _ ->
+    occEnvPop' env1_1' $ \env1' sB ->
+    occCountX (SsArr sB) ebog $ \env2 mkebog ->
+    occCountX SsFull ed $ \env3 mked ->
+    withSome (scaleMany (Some env1') <> Some env2 <> Some env3) $ \env ->
+    k env $ \env' ->
+      projectSmallerSubstruc SsFull s $
+        EFold1InnerD2 ext cm
+          (mkef (OccPush (OccPush env' () sB) () SsFull))
+          (mkebog env') (mked env')
+
+  EConst _ t x ->
+    k OccEnd $ \_ ->
+      case s of
+        SsNone -> ENil ext
+        SsFull -> EConst ext t x
+
+  EIdx0 _ e ->
+    occCountX (SsArr s) e $ \env1 mke ->
+    k env1 $ \env' ->
+      EIdx0 ext (mke env')
+
+  EIdx1 _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      SsArr' s' ->
+        occCountX (SsArr s') a $ \env1 mka ->
+        occCountX SsFull b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EIdx1 ext (mka env') (mkb env')
+
+  EIdx _ a b ->
+    case s of
+      SsNone ->
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      _ ->
+        occCountX (SsArr s) a $ \env1 mka ->
+        occCountX SsFull b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          EIdx ext (mka env') (mkb env')
+
+  EShape _ e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      _ ->
+        occCountX (SsArr SsNone) e $ \env1 mke ->
+        k env1 $ \env' ->
+          projectSmallerSubstruc SsFull s $ EShape ext (mke env')
+
+  EOp _ op e ->
+    case s of
+      SsNone ->
+        occCountX SsNone e $ \env1 mke ->
+        k env1 $ \env' ->
+          use (mke env') $ ENil ext
+      _ ->
+        occCountX SsFull e $ \env1 mke ->
+        k env1 $ \env' ->
+          projectSmallerSubstruc SsFull s $ EOp ext op (mke env')
+
+  ECustom _ t1 t2 t3 e1 e2 e3 a b
+    | typeHasAccums t1 || typeHasAccums t2 || typeHasAccums t3 ->
+        error "Accumulators not allowed in input/output/tape of an ECustom"
+    | otherwise ->
+    case s of
+      SsNone ->
+        -- Allowed to ignore e1/e2/e3 here because no accumulators are
+        -- communicated, and hence no relevant effects exist
+        occCountX SsNone a $ \env1 mka ->
+        occCountX SsNone b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (mka env') $ use (mkb env') $ ENil ext
+      s' ->  -- Let's be pessimistic for safety
+        occCountX SsFull a $ \env1 mka ->
+        occCountX SsFull b $ \env2 mkb ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          projectSmallerSubstruc SsFull s' $
+            ECustom ext t1 t2 t3 (mapExt (const ext) e1) (mapExt (const ext) e2) (mapExt (const ext) e3) (mka env') (mkb env')
+
+  ERecompute _ e ->
+    occCountX s e $ \env1 mke ->
+    k env1 $ \env' ->
+      ERecompute ext (mke env')
+
+  EWith _ t a b ->
+    case s of
+      SsNone ->  -- TODO: simplifier should remove accumulations to an unused with, and then remove the with
+        occCountX SsNone b $ \env2' mkb ->
+        occEnvPop' env2' $ \env2 s1 ->
+        withSome (case s1 of
+                    SsFull -> Some SsFull
+                    SsAccum s' -> Some s'
+                    SsNone -> Some SsNone) $ \s1' ->
+        occCountX s1' a $ \env1 mka ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          use (EWith ext (applySubstrucM s1' t) (mka env') (mkb (OccPush env' () (SsAccum s1')))) $
+            ENil ext
+      SsPair sB sA ->
+        occCountX sB b $ \env2' mkb ->
+        occEnvPop' env2' $ \env2 s1 ->
+        let s1' = case s1 of
+                    SsFull -> Some SsFull
+                    SsAccum s' -> Some s'
+                    SsNone -> Some SsNone in
+        withSome (Some sA <> s1') $ \sA' ->
+        occCountX sA' a $ \env1 mka ->
+        withSome (Some env1 <> Some env2) $ \env ->
+        k env $ \env' ->
+          projectSmallerSubstruc (SsPair sB sA') (SsPair sB sA) $
+            EWith ext (applySubstrucM sA' t) (mka env') (mkb (OccPush env' () (SsAccum sA')))
+      SsFull -> occCountX (SsPair SsFull SsFull) topexpr k
+
+  EAccum _ t p a sp b e ->
+    -- TODO: do better!
+    occCountX SsFull a $ \env1 mka ->
+    occCountX SsFull b $ \env2 mkb ->
+    occCountX SsFull e $ \env3 mke ->
+    withSome (Some env1 <> Some env2) $ \env12 ->
+    withSome (Some env12 <> Some env3) $ \env ->
+    k env $ \env' ->
+      case s of {SsFull -> id; SsNone -> id} $
+        EAccum ext t p (mka env') sp (mkb env') (mke env')
+
+  EZero _ t e ->
+    occCountX (subZeroInfo s) e $ \env1 mke ->
+    k env1 $ \env' ->
+      EZero ext (applySubstrucM s t) (mke env')
+    where
+      subZeroInfo :: Substruc t1 t2 -> Substruc (ZeroInfo t1) (ZeroInfo t2)
+      subZeroInfo SsFull = SsFull
+      subZeroInfo SsNone = SsNone
+      subZeroInfo (SsPair s1 s2) = SsPair (subZeroInfo s1) (subZeroInfo s2)
+      subZeroInfo SsEither{} = error "Either is not a monoid"
+      subZeroInfo SsLEither{} = SsNone
+      subZeroInfo SsMaybe{} = SsNone
+      subZeroInfo (SsArr s') = SsArr (subZeroInfo s')
+      subZeroInfo SsAccum{} = error "Accum is not a monoid"
+
+  EDeepZero _ t e ->
+    occCountX (subDeepZeroInfo s) e $ \env1 mke ->
+    k env1 $ \env' ->
+      EDeepZero ext (applySubstrucM s t) (mke env')
+    where
+      subDeepZeroInfo :: Substruc t1 t2 -> Substruc (DeepZeroInfo t1) (DeepZeroInfo t2)
+      subDeepZeroInfo SsFull = SsFull
+      subDeepZeroInfo SsNone = SsNone
+      subDeepZeroInfo (SsPair s1 s2) = SsPair (subDeepZeroInfo s1) (subDeepZeroInfo s2)
+      subDeepZeroInfo SsEither{} = error "Either is not a monoid"
+      subDeepZeroInfo (SsLEither s1 s2) = SsLEither (subDeepZeroInfo s1) (subDeepZeroInfo s2)
+      subDeepZeroInfo (SsMaybe s') = SsMaybe (subDeepZeroInfo s')
+      subDeepZeroInfo (SsArr s') = SsArr (subDeepZeroInfo s')
+      subDeepZeroInfo SsAccum{} = error "Accum is not a monoid"
+
+  EPlus _ t a b ->
+    occCountX s a $ \env1 mka ->
+    occCountX s b $ \env2 mkb ->
+    withSome (Some env1 <> Some env2) $ \env ->
+    k env $ \env' ->
+      EPlus ext (applySubstrucM s t) (mka env') (mkb env')
+
+  EOneHot _ t p a b ->
+    occCountX SsFull a $ \env1 mka ->
+    occCountX SsFull b $ \env2 mkb ->  -- TODO: do better
+    withSome (Some env1 <> Some env2) $ \env ->
+    k env $ \env' ->
+      projectSmallerSubstruc SsFull s $ EOneHot ext t p (mka env') (mkb env')
+
+  EError _ t msg ->
+    k OccEnd $ \_ -> EError ext (applySubstruc s t) msg
+  where
+    s = simplifySubstruc (typeOf topexpr) initialS
+
+    handleReduction :: t ~ TArr n (TScal t2)
+                    => (forall env2. Ex env2 (TArr (S n) (TScal t2)) -> Ex env2 (TArr n (TScal t2)))
+                    -> Expr x env (TArr (S n) (TScal t2))
+                    -> r
+    handleReduction reduce e
+      | STArr (SS n) _ <- typeOf e =
+      case s of
+        SsNone ->
+          occCountX SsNone e $ \env mke ->
+          k env $ \env' ->
+            use (mke env') $ ENil ext
+        SsArr' SsNone ->
+          occCountX (SsArr SsNone) e $ \env mke ->
+          k env $ \env' ->
+            elet (mke env') $
+              EBuild ext n (EFst ext (EShape ext (evar IZ))) (ENil ext)
+        SsArr' SsFull ->
+          occCountX (SsArr SsFull) e $ \env mke ->
+          k env $ \env' ->
+            reduce (mke env')
+
+
+deleteUnused :: SList f env -> Some (OccEnv Occ env) -> (forall env'. Subenv env env' -> r) -> r
+deleteUnused SNil (Some OccEnd) k = k SETop
+deleteUnused (_ `SCons` env) (Some OccEnd) k =
+  deleteUnused env (Some OccEnd) $ \sub -> k (SENo sub)
+deleteUnused (_ `SCons` env) (Some (OccPush occenv (Occ _ count) _)) k =
+  deleteUnused env (Some occenv) $ \sub ->
+    case count of Zero -> k (SENo sub)
+                  _    -> k (SEYesR sub)
+
+unsafeWeakenWithSubenv :: Subenv env env' -> Expr x env t -> Expr x env' t
+unsafeWeakenWithSubenv = \sub ->
+  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
+    sinkViaSubenv :: Idx env t -> Subenv env env' -> Maybe (Idx env' t)
+    sinkViaSubenv IZ (SEYesR _) = Just IZ
+    sinkViaSubenv IZ (SENo _) = Nothing
+    sinkViaSubenv (IS i) (SEYesR sub) = IS <$> sinkViaSubenv i sub
+    sinkViaSubenv (IS i) (SENo sub) = sinkViaSubenv i sub