module Quox.Syntax.Term.Subst

import Quox.Syntax.Term.Base
import Data.Vect

%default total


infixl 8 ///
mutual
  namespace Term
    ||| 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 dim-closure
    ||| - otherwise, wraps in a new closure
    export
    (///) : Term q dfrom n -> DSubst dfrom dto -> Term q dto n
    s          /// Shift SZ = s
    TYPE l     /// _        = TYPE l
    DCloT s ph /// th       = DCloT s $ ph . th
    s          /// th       = DCloT s th

  namespace Elim
    private
    subDArgs : Elim q dfrom n -> DSubst dfrom dto -> Elim q dto n
    subDArgs (f :% d) th = subDArgs f th :% (d // th)
    subDArgs e        th = DCloE e th

    ||| does the minimal reasonable work:
    ||| - deletes the closure around a term variable
    ||| - deletes an identity substitution
    ||| - composes (lazily) with an existing top-level dim-closure
    ||| - immediately looks up bound variables in a
    |||   top-level sequence of dimension applications
    ||| - otherwise, wraps in a new closure
    export
    (///) : Elim q dfrom n -> DSubst dfrom dto -> Elim q dto n
    e          /// Shift SZ = e
    F x        /// _        = F x
    B i        /// _        = B i
    f :% d     /// th       = subDArgs (f :% d) th
    DCloE e ph /// th       = DCloE e $ ph . th
    e          /// th       = DCloE e th

  namespace ScopeTermN
    export
    (///) : ScopeTermN s q dfrom n -> DSubst dfrom dto -> ScopeTermN s q dto n
    TUsed   body /// th = TUsed   $ body /// th
    TUnused body /// th = TUnused $ body /// th

  namespace DScopeTermN
    export
    (///) : {s : Nat} ->
            DScopeTermN s q dfrom n -> DSubst dfrom dto ->
            DScopeTermN s q dto n
    DUsed   body /// th = DUsed   $ body /// pushN s th
    DUnused body /// th = DUnused $ body /// th


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

||| does the minimal reasonable work:
||| - deletes the closure around a *free* name
||| - 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 q d) (Elim q 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 q d) (Term q 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 %inline
{s : Nat} -> CanSubst (Elim q d) (ScopeTermN s q d) where
  TUsed   body // th = TUsed   $ body // pushN s th
  TUnused body // th = TUnused $ body // th

export %inline
{s : Nat} -> CanSubst (Elim q d) (DScopeTermN s q d) where
  DUsed   body // th = DUsed   $ body // map (/// shift s) th
  DUnused body // th = DUnused $ body // th

export %inline CanSubst Var (Term q d) where s // th = s // map (B {q, d}) th
export %inline CanSubst Var (Elim q d) where e // th = e // map (B {q, d}) th

export %inline
{s : Nat} -> CanSubst Var (ScopeTermN s q d) where
  b // th = b // map (B {q, d}) th

export %inline
{s : Nat} -> CanSubst Var (DScopeTermN s q d) where
  b // th = b // map (B {q, d}) th


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

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

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

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

  namespace ScopeTermN
    ||| applies a term substitution with a less ambiguous type
    export %inline
    (//.) : {s : Nat} ->
            ScopeTermN s q d from -> TSubst q d from to ->
            ScopeTermN s q d to
    body //. th = body // th

    ||| applies a term and dimension substitution
    public export %inline
    subs : {s : Nat} ->
           ScopeTermN s q dfrom from ->
           DSubst dfrom dto -> TSubst q dto from to ->
           ScopeTermN s q dto to
    subs body th ph = body /// th // ph

  namespace DScopeTermN
    ||| applies a term substitution with a less ambiguous type
    export %inline
    (//.) : {s : Nat} -> DScopeTermN s q d from -> TSubst q d from to ->
            DScopeTermN s q d to
    body //. th = body // th

    ||| applies a term and dimension substitution
    public export %inline
    subs : {s : Nat} ->
           DScopeTermN s q dfrom from ->
           DSubst dfrom dto -> TSubst q dto from to ->
           DScopeTermN s q dto to
    subs body th ph = body /// th // ph

export %inline CanShift (Term q d) where s // by = s //. Shift by
export %inline CanShift (Elim q d) where e // by = e //. Shift by

export %inline
{s : Nat} -> CanShift (ScopeTermN s q d) where
  b // by = b //. Shift by


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


namespace ScopeTermN
  export %inline
  (.term) : {s : Nat} -> ScopeTermN s q d n -> Term q d (s + n)
  (TUsed   body).term = body
  (TUnused body).term = body //. shift s

namespace DScopeTermN
  export %inline
  (.term) : {s : Nat} -> DScopeTermN s q d n -> Term q (s + d) n
  (DUsed   body).term = body
  (DUnused body).term = body /// shift s


export %inline
subN : ScopeTermN s q d n -> Vect s (Elim q d n) -> Term q d n
subN (TUsed   body) es = body // fromVect es
subN (TUnused body) _  = body

export %inline
sub1 : ScopeTerm q d n -> Elim q d n -> Term q d n
sub1 t e = subN t [e]

export %inline
dsubN : DScopeTermN s q d n -> Vect s (Dim d) -> Term q d n
dsubN (DUsed   body) ps = body /// fromVect ps
dsubN (DUnused body) _  = body

export %inline
dsub1 : DScopeTerm q d n -> Dim d -> Term q d n
dsub1 t p = dsubN t [p]


public export %inline
(.zero) : DScopeTerm q d n -> Term q d n
body.zero = dsub1 body $ K Zero

public export %inline
(.one) : DScopeTerm q d n -> Term q d n
body.one = dsub1 body $ K One