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

{-# LANGUAGE PartialTypeSignatures #-}
{-# OPTIONS -Wno-partial-type-signatures #-}

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

module AST.Weaken (module AST.Weaken, Append) where

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

import Data
import Lemmas


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)

splitIdx :: forall env2 env1 t f. SList f env1 -> Idx (Append env1 env2) t -> Either (Idx env1 t) (Idx env2 t)
splitIdx SNil i = Right i
splitIdx (SCons _ _) IZ = Left IZ
splitIdx (SCons _ l) (IS i) = first IS (splitIdx l i)

data env :> env' where
  WId :: env :> env
  WSink :: forall t env. env :> (t : env)
  WCopy :: forall t env env'. 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
  WPick :: forall t pre env env'. SList (Const ()) pre -> env :> env'
                               -> Append pre (t : env) :> t : Append pre env'
  WSwap :: forall env as bs. SList (Const ()) as -> SList (Const ()) bs
        -> Append as (Append bs env) :> Append bs (Append as env)
  WStack :: forall env1 env2 as bs. SList (Const ()) as -> SList (Const ()) bs
         -> as :> bs -> env1 :> env2
         -> Append as env1 :> Append bs env2
deriving instance Show (env :> env')
infix 4 :>

infixr 2 @>
(@>) :: 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
WPick SNil w @> i = WCopy w @> i
WPick (_ `SCons` _) _ @> IZ = IS IZ
WPick @t (_ `SCons` pre) w @> IS i = WCopy WSink .> WPick @t pre w @> i
WSwap @env (as :: SList _ as) (bs :: SList _ bs) @> i =
  case splitIdx @(Append bs env) as i of
    Left i' -> indexSinks bs (indexRaiseAbove @env as i')
    Right i' -> case splitIdx @env bs i' of
                  Left j -> indexRaiseAbove @(Append as env) bs j
                  Right j -> indexSinks bs (indexSinks as j)
WStack @env1 @env2 as bs wlo whi @> i =
  case splitIdx @env1 as i of
    Left i' -> indexRaiseAbove @env2 bs (wlo @> i')
    Right i' -> indexSinks bs (whi @> i')

indexSinks :: SList f as -> Idx bs t -> Idx (Append as bs) t
indexSinks SNil j = j
indexSinks (_ `SCons` bs') j = IS (indexSinks bs' j)

indexRaiseAbove :: forall env as t f. SList f as -> Idx as t -> Idx (Append as env) t
indexRaiseAbove = flip go
  where
    go :: forall as'. Idx as' t -> SList f as' -> Idx (Append as' env) t
    go IZ (_ `SCons` _) = IZ
    go (IS i) (_ `SCons` as) = IS (go i as)

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

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

wSinksAnd :: forall env env' bs f. SList f bs -> env :> env' -> env :> Append bs env'
wSinksAnd SNil w = w
wSinksAnd (SCons _ spine) w = WSink .> wSinksAnd spine w

wCopies :: SList f bs -> env1 :> env2 -> Append bs env1 :> Append bs env2
wCopies bs w =
  let bs' = slistMap (\_ -> Const ()) bs
  in WStack bs' bs' WId w

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)