module Quox.Syntax.Term.Subst

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

%default total

namespace CanDSubst
  public export
  interface CanDSubst (0 tm : Nat -> Nat -> Type) where
    (//) : tm dfrom n -> Lazy (DSubst dfrom dto) -> tm dto n

||| 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
CanDSubst (Term q) where
  s          // Shift SZ = s
  TYPE l     // _        = TYPE l
  DCloT s ph // th       = DCloT s $ ph . th
  s          // th       = DCloT s th

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
CanDSubst (Elim q) where
  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 DSubst.ScopeTermN
  export %inline
  (//) : ScopeTermN s q dfrom n -> Lazy (DSubst dfrom dto) ->
         ScopeTermN s q dto n
  S ns (Y body) // th = S ns $ Y $ body // th
  S ns (N body) // th = S ns $ N $ body // th

namespace DSubst.DScopeTermN
  export %inline
  (//) : {s : Nat} ->
         DScopeTermN s q dfrom n -> Lazy (DSubst dfrom dto) ->
         DScopeTermN s q dto n
  S ns (Y body) // th = S ns $ Y $ body // pushN s th
  S ns (N body) // th = S ns $ N $ 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
CanSubstSelf (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

namespace CanTSubst
  public export
  interface CanTSubst (0 tm : TermLike) where
    (//) : tm q d from -> Lazy (TSubst q d from to) -> tm q d to

||| 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
CanTSubst Term 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

namespace ScopeTermN
  export %inline
  (//) : {s : Nat} ->
         ScopeTermN s q d from -> Lazy (TSubst q d from to) ->
         ScopeTermN s q d to
  S ns (Y body) // th = S ns $ Y $ body // pushN s th
  S ns (N body) // th = S ns $ N $ body // th

namespace DScopeTermN
  export %inline
  (//) : {s : Nat} ->
         DScopeTermN s q d from -> Lazy (TSubst q d from to) ->
         DScopeTermN s q d to
  S ns (Y body) // th = S ns $ Y $ body // map (// shift s) th
  S ns (N body) // th = S ns $ N $ body // th

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


public export %inline
dweakT : {by : Nat} -> Term q d n -> Term q (by + d) n
dweakT t = t // shift by

public export %inline
dweakE : {by : Nat} -> Elim q d n -> Elim q (by + d) n
dweakE t = t // shift by


public export %inline
weakT : {default 1 by : Nat} -> Term q d n -> Term q d (by + n)
weakT t = t // shift by

public export %inline
weakE : {default 1 by : Nat} -> Elim q d n -> Elim q d (by + n)
weakE t = t // shift by


parameters {s : Nat}
  namespace ScopeTermBody
    export %inline
    (.term) : ScopedBody s (Term q d) n -> Term q d (s + n)
    (Y b).term = b
    (N b).term = weakT b {by = s}

  namespace ScopeTermN
    export %inline
    (.term) : ScopeTermN s q d n -> Term q d (s + n)
    t.term = t.body.term

  namespace DScopeTermBody
    export %inline
    (.term) : ScopedBody s (\d => Term q d n) d -> Term q (s + d) n
    (Y b).term = b
    (N b).term = dweakT b {by = s}

  namespace DScopeTermN
    export %inline
    (.term) : DScopeTermN s q d n -> Term q (s + d) n
    t.term = t.body.term


export %inline
subN : ScopeTermN s q d n -> Vect s (Elim q d n) -> Term q d n
subN (S _ (Y body)) es = body // fromVect es
subN (S _ (N 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 (S _ (Y body)) ps = body // fromVect ps
dsubN (S _ (N 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