module Quox.Syntax.Subst

import Quox.Syntax.Shift
import Quox.Syntax.Var
import Quox.Name
import Quox.Pretty

import Data.List

%default total


infixr 5 :::
public export
data Subst : (Nat -> Type) -> Nat -> Nat -> Type where
  Shift : Shift from to -> Subst env from to
  (:::) : (t : Lazy (env to)) -> Subst env from to -> Subst env (S from) to
%name Subst th, ph, ps


private
Repr : (Nat -> Type) -> Nat -> Type
Repr f to = (List (f to), Nat)

private
repr : Subst f from to -> Repr f to
repr (Shift by) = ([], by.nat)
repr (t ::: th) = let (ts, i) = repr th in (t::ts, i)


export Eq  (f to) => Eq  (Subst f from to) where (==) = (==) `on` repr
export Ord (f to) => Ord (Subst f from to) where compare = compare `on` repr


infixl 8 //
public export
interface FromVar env => CanSubst env term where
  (//) : term from -> Subst env from to -> term to

public export
CanSubst1 : (Nat -> Type) -> Type
CanSubst1 f = CanSubst f f


infixl 8 !!
public export
(!!) : FromVar term => Subst term from to -> Var from -> term to
(Shift by) !! i      = fromVar $ shift by i
(t ::: th) !! VZ     = t
(t ::: th) !! (VS i) = th !! i


public export
CanSubst Var Var where
  i    // Shift by   = shift by i
  VZ   // (t ::: th) = t
  VS i // (t ::: th) = i // th


public export %inline
shift : (by : Nat) -> Subst env from (by + from)
shift by = Shift $ fromNat by


infixl 9 .
public export
(.) : CanSubst1 f => Subst f from mid -> Subst f mid to -> Subst f from to
Shift SZ      . ph         = ph
Shift (SS by) . Shift bz   = Shift $ SS by . bz
Shift (SS by) . (t ::: th) = Shift by . th
(t ::: th)    . ph         = (t // ph) ::: (th . ph)

public export %inline
id : Subst f n n
id = shift 0

public export
map : (f to -> g to) -> Subst f from to -> Subst g from to
map f (Shift by) = Shift by
map f (t ::: th) = f t ::: map f th


public export %inline
push : CanSubst1 f => Subst f from to -> Subst f (S from) (S to)
push th = fromVar VZ ::: (th . shift 1)


||| `prettySubst pr bnd op cl unicode th` pretty-prints the substitution `th`,
||| with the following arguments:
|||
||| * `th : Subst f from to`
||| * `pr : (unicode : Bool) -> f to -> m (Doc HL)` prints a single element
||| * `names : List Name` is a list of known bound var names
||| * `bnd : HL` is the highlight to use for bound variables being subsituted
||| * `op, cl : Doc HL` are the opening and closing brackets
||| * `unicode : Bool` is whether to use unicode characters in the output
|||   (also passed into `pr`)
export
prettySubstM : Pretty.HasEnv m =>
               (pr : f to -> m (Doc HL)) ->
               (names : List Name) -> (bnd : HL) -> (op, cl : Doc HL) ->
               Subst f from to -> m (Doc HL)
prettySubstM pr names bnd op cl th =
  encloseSep (hl Delim op) (hl Delim cl) (hl Delim "; ") <$>
    withPrec Outer (go 0 th)
 where
  go1 : Nat -> f to -> m (Doc HL)
  go1 i t = pure $ hang 2 $ sep
    [hsep [!(prettyVar' bnd bnd names i),
           hl Delim !(ifUnicode "≔" ":=")],
     !(pr t)]

  go : forall from. Nat -> Subst f from to -> m (List (Doc HL))
  go _ (Shift SZ) = pure []
  go _ (Shift by) = [|pure (prettyShift bnd by)|]
  go i (t ::: th) = [|go1 i t :: go (S i) th|]

||| prints with [square brackets] and the `TVar` highlight for variables
export
PrettyHL (f to) => PrettyHL (Subst f from to) where
  prettyM th = prettySubstM prettyM (!ask).tnames TVar "[" "]" th