{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE StandaloneKindSignatures #-}
{-# LANGUAGE EmptyCase #-}

-- The reason why this is a separate module with "little" in it:
{-# LANGUAGE AllowAmbiguousTypes #-}

module AST.Weaken where

import Data.Functor.Const
import Data.Kind (Type)

import Data


type Idx :: [k] -> k -> Type
data Idx env t where
  IZ :: Idx (t : env) t
  IS :: Idx env t -> Idx (a : env) t
deriving instance Show (Idx env t)

data env :> env' where
  WId :: env :> env
  WSink :: env :> (t : env)
  WCopy :: env :> env' -> (t : env) :> (t : env')
  WPop :: (t : env) :> env' -> env :> env'
  WThen :: env1 :> env2 -> env2 :> env3 -> env1 :> env3
  WClosed :: SList (Const ()) env -> '[] :> env
  WIdx :: Idx env t -> (t : env) :> env
deriving instance Show (env :> env')

infixr @>
(@>) :: env :> env' -> Idx env t -> Idx env' t
WId @> i = i
WSink @> i = IS i
WCopy _ @> IZ = IZ
WCopy w @> (IS i) = IS (w @> i)
WPop w @> i = w @> IS i
WThen w1 w2 @> i = w2 @> w1 @> i
WClosed _ @> i = case i of {}
WIdx j @> IZ = j
WIdx _ @> IS i = i

(.>) :: env2 :> env3 -> env1 :> env2 -> env1 :> env3
(.>) = flip WThen

type family Append a b where
  Append '[] l = l
  Append (x : xs) l = x : Append xs l

class KnownListSpine list where knownListSpine :: SList (Const ()) list
instance KnownListSpine '[] where knownListSpine = SNil
instance KnownListSpine list => KnownListSpine (t : list) where knownListSpine = SCons (Const ()) knownListSpine

wSinks' :: forall list env. KnownListSpine list => env :> Append list env
wSinks' = wSinks (knownListSpine :: SList (Const ()) list)

wSinks :: forall env bs f. SList f bs -> env :> Append bs env
wSinks SNil = WId
wSinks (SCons _ spine) = WSink .> wSinks spine

wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2
wCopies SNil w = w
wCopies (SCons _ spine) w = WCopy (wCopies spine w)

-- wStack :: forall env b1 b2. b1 :> b2 -> Append b1 env :> Append b2 env
-- wStack WId = WId
-- wStack WSink = WSink
-- wStack (WCopy w) = WCopy (wStack @env w)
-- wStack (WPop w) = WPop (wStack @env w)
-- wStack (WThen w1 w2) = WThen (wStack @env w1) (wStack @env w2)
-- wStack (WClosed s) = wSinks s
-- wStack (WIdx i) = WIdx (_ i)

wRaiseAbove :: SList f env1 -> SList g env -> env1 :> Append env1 env
wRaiseAbove SNil env = WClosed (slistMap (\_ -> Const ()) env)
wRaiseAbove (SCons _ s) env = WCopy (wRaiseAbove s env)