module DeBruijn.Sub (
    -- * Substitution
    Sub (..),
    unSub,
    substIdx,
    emptySub,
    snocSub,
    keepSub,
    weakenSub,
    nameMe,
    -- ** Category
    idSub,
    compSub,
    -- * Classes
    Var (..),
    Subst (..),
) where

import Data.Coerce (coerce)
import Data.Kind   (Constraint, Type)
import Data.Proxy  (Proxy (..))

import DeBruijn.Ctx
import DeBruijn.Env
import DeBruijn.Idx
import DeBruijn.Ren
import DeBruijn.Size
import DeBruijn.Wk

-- | Substitutions.
type Sub :: (Ctx -> Type) -> Ctx -> Ctx -> Type
newtype Sub t ctx ctx' = MkSub (Env ctx (t ctx'))

unSub :: Sub t ctx ctx' -> Env ctx (t ctx')
unSub :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Sub t ctx ctx' -> Env ctx (t ctx')
unSub = Sub t ctx ctx' -> Env ctx (t ctx')
forall a b. Coercible a b => a -> b
coerce

-- | Substitute index.
substIdx :: Sub t ctx ctx' -> Idx ctx -> t ctx'
substIdx :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Sub t ctx ctx' -> Idx ctx -> t ctx'
substIdx (MkSub Env ctx (t ctx')
ts) Idx ctx
i = Idx ctx -> Env ctx (t ctx') -> t ctx'
forall (ctx :: Ctx) a. Idx ctx -> Env ctx a -> a
lookupEnv Idx ctx
i Env ctx (t ctx')
ts

-- | Substitution from empty context.
emptySub :: Sub t EmptyCtx ctx'
emptySub :: forall (t :: Ctx -> *) (ctx' :: Ctx). Sub t EmptyCtx ctx'
emptySub = Env EmptyCtx (t ctx') -> Sub t EmptyCtx ctx'
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub Env EmptyCtx (t ctx')
forall a. Env EmptyCtx a
EmptyEnv

snocSub :: Sub t ctx ctx' -> t ctx' -> Sub t (S ctx) ctx'
snocSub :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Sub t ctx ctx' -> t ctx' -> Sub t (S ctx) ctx'
snocSub (MkSub Env ctx (t ctx')
s) t ctx'
t = Env (S ctx) (t ctx') -> Sub t (S ctx) ctx'
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub (Env ctx (t ctx')
s Env ctx (t ctx') -> t ctx' -> Env (S ctx) (t ctx')
forall (ctx1 :: Ctx) a. Env ctx1 a -> a -> Env ('S ctx1) a
:> t ctx'
t)

keepSub :: (Renamable t, Var t) => Sub t ctx ctx' -> Sub t (S ctx) (S ctx')
keepSub :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
(Renamable t, Var t) =>
Sub t ctx ctx' -> Sub t (S ctx) (S ctx')
keepSub (MkSub Env ctx (t ctx')
ts) = Env (S ctx) (t (S ctx')) -> Sub t (S ctx) (S ctx')
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub ((t ctx' -> t (S ctx')) -> Env ctx (t ctx') -> Env ctx (t (S ctx'))
forall a b. (a -> b) -> Env ctx a -> Env ctx b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Wk ctx' (S ctx') -> t ctx' -> t (S ctx')
forall (n :: Ctx) (m :: Ctx). Wk n m -> t n -> t m
forall (t :: Ctx -> *) (n :: Ctx) (m :: Ctx).
Renamable t =>
Wk n m -> t n -> t m
weaken Wk ctx' (S ctx')
forall (ctx :: Ctx). Wk ctx (S ctx)
wk1) Env ctx (t ctx')
ts Env ctx (t (S ctx')) -> t (S ctx') -> Env (S ctx) (t (S ctx'))
forall (ctx1 :: Ctx) a. Env ctx1 a -> a -> Env ('S ctx1) a
:> Idx (S ctx') -> t (S ctx')
forall (ctx :: Ctx). Idx ctx -> t ctx
forall (t :: Ctx -> *) (ctx :: Ctx). Var t => Idx ctx -> t ctx
var Idx (S ctx')
forall (n1 :: Ctx). Idx ('S n1)
IZ)

-- | Precompose 'Sub' with weakening.
weakenSub :: Wk ctx ctx' -> Sub t ctx' ctx'' -> Sub t ctx ctx''
weakenSub :: forall (ctx :: Ctx) (ctx' :: Ctx) (t :: Ctx -> *) (ctx'' :: Ctx).
Wk ctx ctx' -> Sub t ctx' ctx'' -> Sub t ctx ctx''
weakenSub Wk ctx ctx'
w (MkSub Env ctx' (t ctx'')
ts) = Env ctx (t ctx'') -> Sub t ctx ctx''
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub (Wk ctx ctx' -> Env ctx' (t ctx'') -> Env ctx (t ctx'')
forall (ctx :: Ctx) (ctx' :: Ctx) a.
Wk ctx ctx' -> Env ctx' a -> Env ctx a
weakenEnv Wk ctx ctx'
w Env ctx' (t ctx'')
ts)

-- TODO:
nameMe :: Renamable t => Sub t ctx ctx' -> Wk ctx' ctx'' -> Sub t ctx ctx''
nameMe :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx) (ctx'' :: Ctx).
Renamable t =>
Sub t ctx ctx' -> Wk ctx' ctx'' -> Sub t ctx ctx''
nameMe (MkSub Env ctx (t ctx')
ts) Wk ctx' ctx''
w = Env ctx (t ctx'') -> Sub t ctx ctx''
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub ((t ctx' -> t ctx'') -> Env ctx (t ctx') -> Env ctx (t ctx'')
forall a b. (a -> b) -> Env ctx a -> Env ctx b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Wk ctx' ctx'' -> t ctx' -> t ctx''
forall (n :: Ctx) (m :: Ctx). Wk n m -> t n -> t m
forall (t :: Ctx -> *) (n :: Ctx) (m :: Ctx).
Renamable t =>
Wk n m -> t n -> t m
weaken Wk ctx' ctx''
w) Env ctx (t ctx')
ts)

-------------------------------------------------------------------------------
-- Classes
-------------------------------------------------------------------------------

-- | Terms which contain indices as variables.
type Var :: (Ctx -> Type) -> Constraint
class Var t where
    var :: Idx ctx -> t ctx

-- | Terms which we can subsitute into.
type Subst :: (Ctx -> Type) -> Constraint
class Var t => Subst t where
    subst :: Sub t ctx ctx' -> t ctx -> t ctx'

instance Var Proxy where
    var :: forall (ctx :: Ctx). Idx ctx -> Proxy ctx
var Idx ctx
_ = Proxy ctx
forall {k} (t :: k). Proxy t
Proxy

instance Subst Proxy where
    subst :: forall (ctx :: Ctx) (ctx' :: Ctx).
Sub Proxy ctx ctx' -> Proxy ctx -> Proxy ctx'
subst Sub Proxy ctx ctx'
_ Proxy ctx
_ = Proxy ctx'
forall {k} (t :: k). Proxy t
Proxy

instance Var Idx where
    var :: forall (ctx :: Ctx). Idx ctx -> Idx ctx
var = Idx ctx -> Idx ctx
forall a. a -> a
id

instance Subst Idx where
    subst :: forall (ctx :: Ctx) (ctx' :: Ctx).
Sub Idx ctx ctx' -> Idx ctx -> Idx ctx'
subst = Sub Idx ctx ctx' -> Idx ctx -> Idx ctx'
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Sub t ctx ctx' -> Idx ctx -> t ctx'
substIdx

-------------------------------------------------------------------------------
-- Category
-------------------------------------------------------------------------------

-- | Identity substitution
idSub :: Var t => Size ctx -> Sub t ctx ctx
idSub :: forall (t :: Ctx -> *) (ctx :: Ctx).
Var t =>
Size ctx -> Sub t ctx ctx
idSub Size ctx
s = Env ctx (t ctx) -> Sub t ctx ctx
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub (Env ctx (t ctx) -> Sub t ctx ctx)
-> Env ctx (t ctx) -> Sub t ctx ctx
forall a b. (a -> b) -> a -> b
$ Size ctx -> (Idx ctx -> t ctx) -> Env ctx (t ctx)
forall (ctx :: Ctx) a. Size ctx -> (Idx ctx -> a) -> Env ctx a
tabulateEnv Size ctx
s Idx ctx -> t ctx
forall (ctx :: Ctx). Idx ctx -> t ctx
forall (t :: Ctx -> *) (ctx :: Ctx). Var t => Idx ctx -> t ctx
var

-- | Substitution composition.
compSub :: Subst t => Sub t ctx ctx' -> Sub t ctx' ctx'' -> Sub t ctx ctx''
compSub :: forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx) (ctx'' :: Ctx).
Subst t =>
Sub t ctx ctx' -> Sub t ctx' ctx'' -> Sub t ctx ctx''
compSub (MkSub Env ctx (t ctx')
s) Sub t ctx' ctx''
s' = Env ctx (t ctx'') -> Sub t ctx ctx''
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Env ctx (t ctx') -> Sub t ctx ctx'
MkSub ((t ctx' -> t ctx'') -> Env ctx (t ctx') -> Env ctx (t ctx'')
forall a b. (a -> b) -> Env ctx a -> Env ctx b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Sub t ctx' ctx'' -> t ctx' -> t ctx''
forall (ctx :: Ctx) (ctx' :: Ctx).
Sub t ctx ctx' -> t ctx -> t ctx'
forall (t :: Ctx -> *) (ctx :: Ctx) (ctx' :: Ctx).
Subst t =>
Sub t ctx ctx' -> t ctx -> t ctx'
subst Sub t ctx' ctx''
s') Env ctx (t ctx')
s)