1 files changed, 18 insertions, 6 deletions
diff --git a/src/AST/Env.hs b/src/AST/Env.hs
index bc2b9e0..422f0f7 100644
--- a/src/AST/Env.hs
+++ b/src/AST/Env.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE PatternSynonyms #-}
 {-# LANGUAGE PolyKinds #-}
 {-# LANGUAGE QuantifiedConstraints #-}
+{-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TypeOperators #-}
 module AST.Env where
@@ -12,6 +13,7 @@ import Data.Type.Equality
 
 import AST.Sparse
 import AST.Weaken
+import CHAD.Types
 import Data
 
 
@@ -38,18 +40,18 @@ subList SNil SETop = SNil
 subList (SCons x xs) (SEYes s sub) = SCons (subtApply s x) (subList xs sub)
 subList (SCons _ xs) (SENo sub) = subList xs sub
 
-subenvAll :: IsSubType s => SList f env -> Subenv' s env env
+subenvAll :: (IsSubType s, IsSubTypeSubject s f) => SList f env -> Subenv' s env env
 subenvAll SNil = SETop
-subenvAll (SCons _ env) = SEYes subtFull (subenvAll env)
+subenvAll (SCons t env) = SEYes (subtFull t) (subenvAll env)
 
 subenvNone :: SList f env -> Subenv' s env '[]
 subenvNone SNil = SETop
 subenvNone (SCons _ env) = SENo (subenvNone env)
 
-subenvOnehot :: IsSubType s => SList f env -> Idx env t -> Subenv' s env '[t]
-subenvOnehot (SCons _ env) IZ = SEYes subtFull (subenvNone env)
-subenvOnehot (SCons _ env) (IS i) = SENo (subenvOnehot env i)
-subenvOnehot SNil i = case i of {}
+subenvOnehot :: SList f env -> Idx env t -> s t t' -> Subenv' s env '[t']
+subenvOnehot (SCons _ env) IZ sp = SEYes sp (subenvNone env)
+subenvOnehot (SCons _ env) (IS i) sp = SENo (subenvOnehot env i sp)
+subenvOnehot SNil i _ = case i of {}
 
 subenvCompose :: IsSubType s => Subenv' s env1 env2 -> Subenv' s env2 env3 -> Subenv' s env1 env3
 subenvCompose SETop SETop = SETop
@@ -71,3 +73,13 @@ wUndoSubenv :: Subenv' (:~:) env env' -> env' :> env
 wUndoSubenv SETop = WId
 wUndoSubenv (SEYes Refl sub) = WCopy (wUndoSubenv sub)
 wUndoSubenv (SENo sub) = WSink .> wUndoSubenv sub
+
+subenvMap :: (forall a a'. f a -> s a a' -> s' a a') -> SList f env -> Subenv' s env env' -> Subenv' s' env env'
+subenvMap _ SNil SETop = SETop
+subenvMap f (t `SCons` l) (SEYes s sub) = SEYes (f t s) (subenvMap f l sub)
+subenvMap f (_ `SCons` l) (SENo sub) = SENo (subenvMap f l sub)
+
+subenvD2E :: Subenv env env' -> Subenv (D2E env) (D2E env')
+subenvD2E SETop = SETop
+subenvD2E (SEYesR sub) = SEYesR (subenvD2E sub)
+subenvD2E (SENo sub) = SENo (subenvD2E sub)