quox/src/Quox/Context.idr

module Quox.Context

import Quox.Syntax.Shift
import Quox.Pretty

import Data.DPair
import Data.Nat
import Data.SnocList
import Control.Monad.Identity

%default total


infixl 5 :<
||| 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 global 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' : Type -> (from, to : Nat) -> Type
Telescope' a = Telescope (\_ => a)

||| a global 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' : Type -> (len : Nat) -> Type
Context' a = Context (\_ => a)


export
toSnocList : Telescope tm _ _ -> SnocList (Exists tm)
toSnocList [<]        = [<]
toSnocList (tel :< t) = toSnocList tel :< Evidence _ t

private
toListAcc : Telescope tm _ _ -> List (Exists tm) -> List (Exists tm)
toListAcc [<]        acc = acc
toListAcc (tel :< t) acc = toListAcc tel (Evidence _ t :: acc)

export %inline
toList : Telescope tm _ _ -> List (Exists tm)
toList tel = toListAcc tel []

export %inline
toSnocList' : Telescope' a _ _ -> SnocList a
toSnocList' = map snd . toSnocList

export %inline
toList' : Telescope' a _ _ -> List a
toList' = map snd . toList


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


export
getWith : CanShift tm => Context tm len -> Var len -> Shift len out -> tm out
getWith (ctx :< t) VZ     th = t // drop1 th
getWith (ctx :< t) (VS i) th = getWith ctx i (drop1 th)

infixl 8 !!
export %inline
(!!) : CanShift tm => Context tm len -> Var len -> tm len
ctx !! i = getWith ctx i SZ


||| a triangle of bindings. each type binding in a context counts the ues of
||| others in its type, and all of these together form a triangle.
public export
Triangle : (tm : Nat -> Type) -> (len : Nat) -> Type
Triangle = Context . Context

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

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

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

export %hint
0 succGT : S n `GT` n
succGT = LTESucc $ reflexive {rel=LTE}


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|]

  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

export %inline
map : (forall n. tm1 n -> tm2 n) ->
      Telescope tm1 from to -> Telescope tm2 from to
map f = runIdentity . traverse (pure . f)

infixr 4 <$>
export %inline
(<$>) : (forall n. tm1 n -> tm2 n) ->
        Telescope tm1 from to -> Telescope tm2 from to
(<$>) = map

infixl 3 <*>
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 relevant,
-- so no idiom brackets ☹


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
zipWith3 : (forall n. tm1 n -> tm2 n -> tm3 n -> tm4 n) ->
           Telescope tm1 from to ->
           Telescope tm2 from to ->
           Telescope tm3 from to ->
           Telescope tm4 from to
zipWith3 f tel1 tel2 tel3 = f <$> tel1 <*> tel2 <*> tel3

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 %inline
zipWith3Lazy : forall tm1, tm2, tm3, tm4.
               (forall n. tm1 n -> tm2 n -> tm3 n -> tm4 n) ->
               Telescope tm1 from to ->
               Telescope tm2 from to ->
               Telescope tm3 from to ->
               Telescope (\n => Lazy (tm4 n)) from to
zipWith3Lazy f = zipWith3 $ \x, y, z => delay $ f x y z


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 %inline
length : Telescope {} -> Nat
length = fst . lengthPrf


parameters {0 acc : Nat -> Type}
  export
  foldl : (forall m, n. acc m -> tm n -> acc (S m)) ->
          acc 0 -> (tel : Telescope tm from to) -> acc (length tel)
  foldl f z [<]        = z
  foldl f z (tel :< t) = f (foldl f z tel) t

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

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

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


export %inline
and : Telescope' (Lazy Bool) _ _ -> Bool
and = force . fold @{All}

export %inline
all : (forall n. tm n -> Bool) -> Telescope tm _ _ -> Bool
all p = and . map (delay . p)

export %inline
all2 : (forall n. tm n -> tm2 n -> Bool) ->
       Telescope tm from to -> Telescope tm2 from to -> Bool
all2 p = and .: zipWithLazy p


export %inline
or : Telescope' (Lazy Bool) _ _ -> Bool
or = force . fold @{Any}

export %inline
any : (forall n. tm n -> Bool) -> Telescope tm _ _ -> Bool
any p = or . map (delay . p)

export %inline
any2 : (forall n. tm1 n -> tm2 n -> Bool) ->
       Telescope tm1 from to -> Telescope tm2 from to -> Bool
any2 p = or .: zipWithLazy p


export %inline
(forall n. Eq (tm n)) => Eq (Telescope tm from to) where
  (==) = all2 (==)

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

export %inline
(forall n. PrettyHL (tm n)) => PrettyHL (Telescope tm from to) where
  prettyM tel = separate (hl Delim ";") <$> traverse prettyM (toList tel)