module Quox.Syntax.Term.Subst

import Quox.Syntax.Term.Base

%default total


export FromVar (Elim d) where fromVar = B
export FromVar (Term d) where fromVar = E . fromVar

||| does the minimal reasonable work:
||| - deletes the closure around a free name since it doesn't do anything
||| - deletes an identity substitution
||| - composes (lazily) with an existing top-level closure
||| - immediately looks up a bound variable
||| - otherwise, wraps in a new closure
export
CanSubst (Elim d) (Elim d) where
  F x       // _  = F x
  B i       // th = th !! i
  CloE e ph // th = assert_total CloE e $ ph . th
  e         // th = case force th of
                         Shift SZ => e
                         th       => CloE e th

||| does the minimal reasonable work:
||| - deletes the closure around an atomic constant like `TYPE`
||| - deletes an identity substitution
||| - composes (lazily) with an existing top-level closure
||| - goes inside `E` in case it is a simple variable or something
||| - otherwise, wraps in a new closure
export
CanSubst (Elim d) (Term d) where
  TYPE l    // _  = TYPE l
  E e       // th = E $ e // th
  CloT s ph // th = CloT s $ ph . th
  s         // th = case force th of
                         Shift SZ => s
                         th       => CloT s th

export
CanSubst (Elim d) (ScopeTerm d) where
  TUsed   body // th = TUsed   $ body // push th
  TUnused body // th = TUnused $ body // th

export CanSubst Var (Term      d) where s // th = s // map (B {d}) th
export CanSubst Var (Elim      d) where e // th = e // map (B {d}) th
export CanSubst Var (ScopeTerm d) where s // th = s // map (B {d}) th


infixl 8 //., ///
mutual
  namespace Term
    ||| applies a term substitution with a less ambiguous type
    export
    (//.) : Term d from -> TSubst d from to -> Term d to
    t //. th = t // th

    ||| applies a dimension substitution with the same behaviour as `(//)`
    ||| above
    export
    (///) : Term dfrom n -> DSubst dfrom dto -> Term dto n
    TYPE l     /// _        = TYPE l
    E e        /// th       = E $ e /// th
    DCloT s ph /// th       = DCloT s $ ph . th
    s          /// Shift SZ = s
    s          /// th       = DCloT s th

    ||| applies a term and dimension substitution
    public export %inline
    subs : Term dfrom from -> DSubst dfrom dto -> TSubst dto from to ->
           Term dto to
    subs s th ph = s /// th // ph

  namespace Elim
    ||| applies a term substitution with a less ambiguous type
    export
    (//.) : Elim d from -> TSubst d from to -> Elim d to
    e //. th = e // th

    ||| applies a dimension substitution with the same behaviour as `(//)`
    ||| above
    export
    (///) : Elim dfrom n -> DSubst dfrom dto -> Elim dto n
    F x        /// _        = F x
    B i        /// _        = B i
    DCloE e ph /// th       = DCloE e $ ph . th
    e          /// Shift SZ = e
    e          /// th       = DCloE e th

    ||| applies a term and dimension substitution
    public export %inline
    subs : Elim dfrom from -> DSubst dfrom dto -> TSubst dto from to ->
           Elim dto to
    subs e th ph = e /// th // ph

  namespace ScopeTerm
    ||| applies a term substitution with a less ambiguous type
    export
    (//.) : ScopeTerm d from -> TSubst d from to -> ScopeTerm d to
    body //. th = body // th

    ||| applies a dimension substitution with the same behaviour as `(//)`
    ||| above
    export
    (///) : ScopeTerm dfrom n -> DSubst dfrom dto -> ScopeTerm dto n
    TUsed   body /// th = TUsed   $ body /// th
    TUnused body /// th = TUnused $ body /// th

    ||| applies a term and dimension substitution
    public export %inline
    subs : ScopeTerm dfrom from -> DSubst dfrom dto -> TSubst dto from to ->
           ScopeTerm dto to
    subs body th ph = body /// th // ph

export CanShift (Term      d) where s // by = s //. Shift by
export CanShift (Elim      d) where e // by = e //. Shift by
export CanShift (ScopeTerm d) where s // by = s //. Shift by


export %inline
comp' : DSubst dfrom dto -> TSubst dfrom from mid -> TSubst dto mid to ->
        TSubst dto from to
comp' th ps ph = map (/// th) ps . ph


export
fromDScopeTerm : DScopeTerm d n -> Term (S d) n
fromDScopeTerm (DUsed   body) = body
fromDScopeTerm (DUnused body) = body /// shift 1

export
fromScopeTerm : ScopeTerm d n -> Term d (S n)
fromScopeTerm (TUsed   body) = body
fromScopeTerm (TUnused body) = body //. shift 1