quox/lib/Quox/Context.idr

module Quox.Context

import Quox.Syntax.Shift
import Quox.Pretty
import Quox.Name
import Quox.NatExtra
import Quox.PrettyValExtra

import Data.DPair
import Data.Nat
import Quox.SingletonExtra
import Data.SnocList
import Data.SnocVect
import Data.Vect
import Control.Monad.Identity
import Derive.Prelude

%language ElabReflection
%default total


||| a sequence of bindings under an existing context. each successive element
||| has one more bound variable, which correspond to all previous elements
||| as well as the surrounding context.
public export
data Telescope : (tm : Nat -> Type) -> (from, to : Nat) -> Type where
  Lin  : Telescope tm from from
  (:<) : Telescope tm from to -> tm to -> Telescope tm from (S to)
%name Telescope tel

public export
Telescope' : (a : Type) -> (from, to : Nat) -> Type
Telescope' a = Telescope (\_ => a)

||| a top level context is actually just a telescope over no existing bindings
public export
Context : (tm : Nat -> Type) -> (len : Nat) -> Type
Context tm len = Telescope tm 0 len

public export
Context' : (a : Type) -> (len : Nat) -> Type
Context' a = Context (\_ => a)

public export
NContext : Nat -> Type
NContext = Context' BaseName

public export
BContext : Nat -> Type
BContext = Context' BindName

public export
BTelescope : Nat -> Nat -> Type
BTelescope = Telescope' BindName


public export
unsnoc : Context tm (S n) -> (Context tm n, tm n)
unsnoc (tel :< x) = (tel, x)

public export
head : Context tm (S n) -> tm n
head = snd . unsnoc

public export
tail : Context tm (S n) -> Context tm n
tail = fst . unsnoc


parameters {0 tm : Nat -> Type} (f : forall n. tm n -> a)
  export
  toSnocListWith : Telescope tm _ _ -> SnocList a
  toSnocListWith [<]        = [<]
  toSnocListWith (tel :< t) = toSnocListWith tel :< f t

  export
  toListWith : Telescope tm _ _ -> List a
  toListWith tel = toListAcc tel [] where
    toListAcc : Telescope tm _ _ -> List a -> List a
    toListAcc [<]        acc = acc
    toListAcc (tel :< t) acc = toListAcc tel (f t :: acc)

  export
  toSnocVectWith : Context tm n -> SnocVect n a
  toSnocVectWith [<]        = [<]
  toSnocVectWith (tel :< t) = toSnocVectWith tel :< f t

export %inline
toSnocList : Telescope tm _ _ -> SnocList (Exists tm)
toSnocList = toSnocListWith (Evidence _)

export %inline
toSnocList' : Telescope' a _ _ -> SnocList a
toSnocList' = toSnocListWith id

export %inline
toList : Telescope tm _ _ -> List (Exists tm)
toList = toListWith (Evidence _)

export %inline
toList' : Telescope' a _ _ -> List a
toList' = toListWith id

export %inline
toSnocVect : Context tm n -> SnocVect n (Exists tm)
toSnocVect = toSnocVectWith (Evidence _)

export %inline
toSnocVect' : Context' a n -> SnocVect n a
toSnocVect' = toSnocVectWith id



export
fromSnocVect : SnocVect n a -> Context' a n
fromSnocVect [<]       = [<]
fromSnocVect (sx :< x) = fromSnocVect sx :< x


export
tabulateVar : (n : Nat) -> (Var n -> a) -> Context' a n
tabulateVar 0     f = [<]
tabulateVar (S k) f = tabulateVar k (f . VS) :< f VZ

export
allVars : (n : Nat) -> Context' (Var n) n
allVars n = tabulateVar n id

public export
replicate : (n : Nat) -> a -> Context' a n
replicate n x = tabulateVar n $ const x


public export
replicateLTE : from `LTE'` to => a -> Telescope' a from to
replicateLTE @{LTERefl}    x = [<]
replicateLTE @{LTESuccR _} x = replicateLTE x :< x

public export
replicateTel : (from, to : Nat) -> a ->
               Either (from `GT` to) (Telescope' a from to)
replicateTel from to x = case isLTE' from to of
  Yes lte  => Right $ replicateLTE x
  No  nlte => Left  $ notLTEImpliesGT $ nlte . fromLTE


public export
(.) : Telescope tm from mid -> Telescope tm mid to -> Telescope tm from to
tel1 . [<]         = tel1
tel1 . (tel2 :< s) = (tel1 . tel2) :< s

export
(<><) : Telescope' a from to -> Vect n a -> Telescope' a from (n + to)
(<><) tel [] = tel
(<><) tel (x :: xs) {n = S n} =
  rewrite plusSuccRightSucc n to in
  (tel :< x) <>< xs

export
(++) : Telescope' a from to -> SnocVect n a -> Telescope' a from (n + to)
tel ++ [<]       = tel
tel ++ (sx :< x) = (tel ++ sx) :< x

public export
getShiftWith : (forall from, to. tm from -> Shift from to -> tm to) ->
               Shift len out -> Context tm len -> Var len -> tm out
getShiftWith shft by (ctx :< t) VZ     = t `shft` ssDown by
getShiftWith shft by (ctx :< t) (VS i) = getShiftWith shft (ssDown by) ctx i

public export %inline
getShift : CanShift tm => Shift len out -> Context tm len -> Var len -> tm out
getShift = getShiftWith (//)

public export %inline
getWith : (forall from, to. tm from -> Shift from to -> tm to) ->
          Context tm len -> Var len -> tm len
getWith shft = getShiftWith shft SZ

export infixl 8 !!
public export %inline
(!!) : CanShift tm => Context tm len -> Var len -> tm len
(!!) = getWith (//)

export infixl 8 !!!
public export %inline
(!!!) : Context' tm len -> Var len -> tm
(!!!) = getWith const

public export
find : Alternative f =>
       (forall n. tm n -> Bool) -> Context tm len -> f (Var len)
find p [<]        = empty
find p (ctx :< x) = (guard (p x) $> VZ) <|> (VS <$> find p ctx)


export
0 telescopeLTE : Telescope _ from to -> from `LTE` to
telescopeLTE [<]        = reflexive
telescopeLTE (tel :< _) = lteSuccRight $ telescopeLTE tel

export
(gt : from `GT` to) => Uninhabited (Telescope _ from to) where
  uninhabited tel = void $ LTEImpliesNotGT (telescopeLTE tel) gt

export %hint
0 succGT : S n `GT` n
succGT = LTESucc reflexive


public export
drop : (m : Nat) -> Context term (m + n) -> Context term n
drop 0     ctx        = ctx
drop (S m) (ctx :< _) = drop m ctx


parameters {auto _ : Applicative f}
  export
  traverse : (forall n. tm1 n -> f (tm2 n)) ->
             Telescope tm1 from to -> f (Telescope tm2 from to)
  traverse f [<]        = pure [<]
  traverse f (tel :< x) = [|traverse f tel :< f x|]

  export %inline
  traverse' : (a -> f b) -> Telescope' a from to -> f (Telescope' b from to)
  traverse' f = traverse f

  export infixl 3 `app`
  ||| like `(<*>)` but with effects
  export
  app : Telescope (\n => tm1 n -> f (tm2 n)) from to ->
        Telescope tm1 from to -> f (Telescope tm2 from to)
  app [<]         [<]         = pure [<]
  app (ftel :< f) (xtel :< x) = [|app ftel xtel :< f x|]
  app [<]         (xtel :< _) = void $ uninhabited xtel
  app (ftel :< _) [<]         = void $ uninhabited ftel

  export %inline
  sequence : Telescope (f . tm) from to -> f (Telescope tm from to)
  sequence = traverse id


parameters {0 tm1, tm2 : Nat -> Type}
           (f : forall n. tm1 n -> tm2 n)
  export %inline
  map : Telescope tm1 from to -> Telescope tm2 from to
  map = runIdentity . traverse (pure . f)

  export %inline
  (<$>) : Telescope tm1 from to -> Telescope tm2 from to
  (<$>) = map



export %inline
(<*>) : Telescope (\n => tm1 n -> tm2 n) from to ->
        Telescope tm1 from to -> Telescope tm2 from to
ftel <*> xtel = runIdentity $ (pure .) <$> ftel `app` xtel
-- ...but can't write pure without `from,to` being ω, so no idiom brackets ☹

export %inline
(<$) : (forall n. tm1 n) -> Telescope tm2 from to -> Telescope tm1 from to
x <$ tel = const x <$> tel


export %inline
zipWith : (forall n. tm1 n -> tm2 n -> tm3 n) ->
          Telescope tm1 from to -> Telescope tm2 from to ->
          Telescope tm3 from to
zipWith f tel1 tel2 = f <$> tel1 <*> tel2

export %inline
zipWithLazy : forall tm1, tm2, tm3.
              (forall n. tm1 n -> tm2 n -> tm3 n) ->
              Telescope tm1 from to -> Telescope tm2 from to ->
              Telescope (\n => Lazy (tm3 n)) from to
zipWithLazy f = zipWith $ delay .: f

export
unzip : Telescope (\n => (tm1 n, tm2 n)) from to ->
        (Telescope tm1 from to, Telescope tm2 from to)
unzip [<]             = ([<], [<])
unzip (tel :< (x, y)) = let (xs, ys) = unzip tel in (xs :< x, ys :< y)

export
unzip3 : Telescope (\n => (tm1 n, tm2 n, tm3 n)) from to ->
         (Telescope tm1 from to, Telescope tm2 from to, Telescope tm3 from to)
unzip3 [<] = ([<], [<], [<])
unzip3 (tel :< (x, y, z)) =
  let (xs, ys, zs) = unzip3 tel in (xs :< x, ys :< y, zs :< z)



public export
lengthPrf : Telescope _ from to -> Subset Nat (\len => len + from = to)
lengthPrf [<]        = Element 0 Refl
lengthPrf (tel :< _) =
  let len = lengthPrf tel in Element (S len.fst) (cong S len.snd)

public export
lengthPrfSing : Telescope _ from to -> Singleton from -> Singleton to
lengthPrfSing tel (Val from) =
  let Element len prf = lengthPrf tel in
  rewrite sym prf in
  Val (len + from)

public export
lengthPrf0 : Context _ to -> Singleton to
lengthPrf0 ctx = lengthPrfSing ctx (Val 0)

public export %inline
length : Telescope {} -> Nat
length = fst . lengthPrf

public export
null : Telescope _ from to -> Bool
null [<] = True
null _   = False


export
foldl : {0 acc : Nat -> Type} ->
        (f : forall n. acc n -> tm (n + from) -> acc (S n)) ->
        (z : acc 0) -> (tel : Telescope tm from to) -> acc (length tel)
foldl f z [<]        = z
foldl f z (tel :< t) = f (foldl f z tel) (rewrite (lengthPrf tel).snd in t)

export %inline
foldl_ : (acc -> tm -> acc) -> acc -> Telescope' tm from to -> acc
foldl_ f z tel = foldl f z tel

export %inline
foldMap : Monoid a => (forall n. tm n -> a) -> Telescope tm from to -> a
foldMap f = foldl (\acc, tm => acc <+> f tm) neutral

export %inline
fold : Monoid a => Telescope' a from to -> a
fold = foldMap id

||| like `fold` but calculate the elements only when actually appending
export %inline
foldLazy : Monoid a => Telescope' (Lazy a) from to -> a
foldLazy = foldMap force


export %inline
all, any : (forall n. tm n -> Bool) -> Telescope tm from to -> Bool
all p = foldMap @{All} p
any p = foldMap @{Any} p


export %inline
(forall n. Eq (tm n)) => Eq (Telescope tm from to) where
  (==) = foldLazy @{All} .: zipWithLazy (==)

export %inline
(forall n. DecEq (tm n)) => DecEq (Telescope tm from to) where
  decEq [<] [<] = Yes Refl
  decEq (xs :< x) (ys :< y) = do
    let Yes Refl = decEq xs ys
      | No n => No $ \case Refl => n Refl
    let Yes Refl = decEq x y
      | No n => No $ \case Refl => n Refl
    Yes Refl
  decEq [<]      (_ :< _) = No $ \case _ impossible
  decEq (_ :< _) [<]      = No $ \case _ impossible

export %inline
(forall n. Ord (tm n)) => Ord (Telescope tm from to) where
  compare = foldLazy .: zipWithLazy compare

export %inline
(forall n. Show (tm n)) => Show (Telescope tm from to) where
  showPrec d = showPrec d . toSnocList where
    Show (Exists tm) where showPrec d t = showPrec d t.snd

export %inline
(forall n. PrettyVal (tm n)) => PrettyVal (Telescope tm from to) where
  prettyVal xs = prettyVal $ toSnocList' $ map prettyVal xs


parameters {opts : LayoutOpts} {0 tm : Nat -> Type}
           (nameHL : HL)
           (pterm : forall n. BContext n -> tm n -> Eff Pretty (Doc opts))
  private
  prettyOne : BindName -> BContext to -> tm to -> Eff Pretty (Doc opts)
  prettyOne x xs tm = hsep <$> sequence
    [hl nameHL $ prettyBind' x, hl Delim $ text ":", pterm xs tm]

  private
  prettyEach : BContext to -> Telescope tm from to ->
               Eff Pretty (Telescope' (Doc opts) from to)
  prettyEach _         [<]       = pure [<]
  prettyEach (xs :< x) (ts :< t) = [|prettyEach xs ts :< prettyOne x xs t|]

  export
  prettyTel : BContext to -> Telescope tm from to -> Eff Pretty (Doc opts)
  prettyTel names tel = do
    docs <- prettyEach names tel
    comma <- hl Delim $ text ","
    pure $ separateTight comma $ toList' docs


namespace BContext
  export
  toNames : BTelescope {} -> SnocList BaseName
  toNames = foldl (\xs, x => xs :< x.val) [<]


public export
record NameContexts q d n where
  constructor MkNameContexts
  {auto qtyLen  : Singleton q}
  {auto dimLen  : Singleton d}
  {auto termLen : Singleton n}
  qnames : BContext q
  dnames : BContext d
  tnames : BContext n
%runElab deriveIndexed "NameContexts" [Show]

namespace NameContexts
  export
  MkNameContexts' : (qnames : BContext q) -> (dnames : BContext d) ->
                    (tnames : BContext n) -> NameContexts q d n
  MkNameContexts' {qnames, dnames, tnames} =
    MkNameContexts {
      qnames, dnames, tnames,
      qtyLen  = lengthPrf0 qnames,
      dimLen  = lengthPrf0 dnames,
      termLen = lengthPrf0 tnames
    }

  export
  fromQNames : BContext q -> NameContexts q 0 0
  fromQNames qnames =
    MkNameContexts qnames [<] [<] {qtyLen = lengthPrf0 qnames}

  export %inline
  empty : NameContexts 0 0 0
  empty = fromQNames [<]

  export
  extendQtyN' : BTelescope q q' -> NameContexts q d n -> NameContexts q' d n
  extendQtyN' xs = {qnames $= (. xs), qtyLen $= lengthPrfSing xs}

  export
  extendDimN' : BTelescope d d' -> NameContexts q d n -> NameContexts q d' n
  extendDimN' xs = {dnames $= (. xs), dimLen $= lengthPrfSing xs}

  export
  extendDimN : {s : Nat} -> BContext s ->
               NameContexts q d n -> NameContexts q (s + d) n
  extendDimN xs = {dnames $= (++ toSnocVect' xs), dimLen $= map (s +)}

  export
  extendDim : BindName -> NameContexts q d n -> NameContexts q (S d) n
  extendDim i = extendDimN [< i]

  export
  extendTermN' : BTelescope n n' -> NameContexts q d n -> NameContexts q d n'
  extendTermN' xs = {tnames $= (. xs), termLen $= lengthPrfSing xs}

  export
  extendTermN : {s : Nat} -> BContext s ->
                NameContexts q d n -> NameContexts q d (s + n)
  extendTermN xs = {tnames $= (++ toSnocVect' xs), termLen $= map (s +)}

  export
  extendTerm : BindName -> NameContexts q d n -> NameContexts q d (S n)
  extendTerm x = extendTermN [< x]