{-# LANGUAGE DataKinds #-}
{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PatternSynonyms #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeOperators #-}
module AST.Env where

import Data.Type.Equality

import AST.Sparse
import AST.Weaken
import CHAD.Types
import Data


-- | @env'@ is a subset of @env@: each element of @env@ is either included in
-- @env'@ ('SEYes') or not included in @env'@ ('SENo').
data Subenv' s env env' where
  SETop :: Subenv' s '[] '[]
  SEYes :: forall t t' env env' s. s t t' -> Subenv' s env env' -> Subenv' s (t : env) (t' : env')
  SENo  :: forall t env env' s. Subenv' s env env' -> Subenv' s (t : env) env'
deriving instance (forall t t'. Show (s t t')) => Show (Subenv' s env env')

type Subenv = Subenv' (:~:)
type SubenvS = Subenv' Sparse

pattern SEYesR :: forall tenv tenv'. ()
               => forall t env env'. (tenv ~ t : env, tenv' ~ t : env')
               => Subenv env env' -> Subenv tenv tenv'
pattern SEYesR s = SEYes Refl s

{-# COMPLETE SETop, SEYesR, SENo #-}

subList :: (IsSubType s, IsSubTypeSubject s f) => SList f env -> Subenv' s env env' -> SList f env'
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, IsSubTypeSubject s f) => SList f env -> Subenv' s env env
subenvAll SNil = SETop
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 :: 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
subenvCompose (SEYes s1 sub1) (SEYes s2 sub2) = SEYes (subtTrans s1 s2) (subenvCompose sub1 sub2)
subenvCompose (SEYes _ sub1) (SENo sub2) = SENo (subenvCompose sub1 sub2)
subenvCompose (SENo sub1) sub2 = SENo (subenvCompose sub1 sub2)

subenvConcat :: Subenv' s env1 env1' -> Subenv' s env2 env2' -> Subenv' s (Append env2 env1) (Append env2' env1')
subenvConcat sub1 SETop = sub1
subenvConcat sub1 (SEYes s sub2) = SEYes s (subenvConcat sub1 sub2)
subenvConcat sub1 (SENo sub2) = SENo (subenvConcat sub1 sub2)

sinkWithSubenv :: Subenv' s 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 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)