pretty printing errors

lib/Quox/Typing/Context.idr

@@ -2,10 +2,15 @@ module Quox.Typing.Context
 
 import Quox.Syntax
 import Quox.Context
+import Quox.Pretty
 
 %default total
 
 
+public export
+QContext : Type -> Nat -> Type
+QContext = Context'
+
 public export
 TContext : Type -> Nat -> Nat -> Type
 TContext q d = Context (Term q d)
@@ -14,10 +19,6 @@ public export
 QOutput : Type -> Nat -> Type
 QOutput = Context'
 
-public export
-NContext : Nat -> Type
-NContext = Context' BaseName
-
 public export
 DimAssign : Nat -> Type
 DimAssign = Context' DimConst
@@ -30,17 +31,19 @@ record TyContext q d n where
   dnames : NContext d
   tctx   : TContext q d n
   tnames : NContext n
+  qtys   : QContext q n -- only used for printing
 %name TyContext ctx
 
 
 public export
 record EqContext q n where
   constructor MkEqContext
-  -- (only used for errors; not needed by the actual equality test)
-  dassign : DimAssign dimLen
-  dnames  : NContext dimLen
+  {dimLen : Nat}
+  dassign : DimAssign dimLen -- only used for printing
+  dnames  : NContext dimLen -- only used for printing
   tctx    : TContext q 0 n
   tnames  : NContext n
+  qtys    : QContext q n -- only used for printing
 %name EqContext ctx
 
 
@@ -56,17 +59,20 @@ namespace TContext
 namespace TyContext
   public export %inline
   empty : TyContext q 0 0
-  empty = MkTyContext {dctx = new, dnames = [<], tctx = [<], tnames = [<]}
+  empty =
+    MkTyContext {dctx = new, dnames = [<], tctx = [<], tnames = [<], qtys = [<]}
 
   export %inline
-  extendTyN : Telescope (\n => (BaseName, Term q d n)) from to ->
+  extendTyN : Telescope (\n => (q, BaseName, Term q d n)) from to ->
               TyContext q d from -> TyContext q d to
   extendTyN xss ctx =
-    let (xs, ss) = unzip xss in {tctx $= (. ss), tnames $= (. xs)} ctx
+    let (qs, xs, ss) = unzip3 xss in
+    {qtys $= (. qs), tctx $= (. ss), tnames $= (. xs)} ctx
 
   export %inline
-  extendTy : BaseName -> Term q d n -> TyContext q d n -> TyContext q d (S n)
-  extendTy x s = extendTyN [< (x, s)]
+  extendTy : q -> BaseName -> Term q d n -> TyContext q d n ->
+             TyContext q d (S n)
+  extendTy q x s = extendTyN [< (q, x, s)]
 
   export %inline
   extendDim : BaseName -> TyContext q d n -> TyContext q (S d) n
@@ -98,41 +104,96 @@ makeDAssign (Shift SZ)   = [<]
 makeDAssign (K e ::: th) = makeDAssign th :< e
 
 export
-makeEqContext : TyContext q d n -> DSubst d 0 -> EqContext q n
-makeEqContext ctx th = MkEqContext {
+makeEqContext' : {d : Nat} -> TyContext q d n -> DSubst d 0 -> EqContext q n
+makeEqContext' ctx th = MkEqContext {
   dassign = makeDAssign th,
   dnames  = ctx.dnames,
   tctx    = map (// th) ctx.tctx,
-  tnames  = ctx.tnames
+  tnames  = ctx.tnames,
+  qtys    = ctx.qtys
 }
 
+export
+makeEqContext : TyContext q d n -> DSubst d 0 -> EqContext q n
+makeEqContext ctx@(MkTyContext {dnames, _}) th =
+  let (d' ** Refl) = lengthPrf0 dnames in makeEqContext' ctx th
+
 namespace EqContext
   public export %inline
   empty : EqContext q 0
-  empty = MkEqContext [<] [<] [<] [<]
+  empty = MkEqContext {
+    dassign = [<], dnames = [<], tctx = [<], tnames = [<], qtys = [<]
+  }
 
   export %inline
-  extendTyN : Telescope (\n => (BaseName, Term q 0 n)) from to ->
+  extendTyN : Telescope (\n => (q, BaseName, Term q 0 n)) from to ->
               EqContext q from -> EqContext q to
   extendTyN tel ctx =
-    let (xs, ss) = unzip tel in {tctx $= (. ss), tnames $= (. xs)} ctx
+    let (qs, xs, ss) = unzip3 tel in
+    {qtys $= (. qs), tctx $= (. ss), tnames $= (. xs)} ctx
 
   export %inline
-  extendTy : BaseName -> Term q 0 n -> EqContext q n -> EqContext q (S n)
-  extendTy x s = extendTyN [< (x, s)]
+  extendTy : q -> BaseName -> Term q 0 n -> EqContext q n -> EqContext q (S n)
+  extendTy q x s = extendTyN [< (q, x, s)]
 
   export %inline
   extendDim : BaseName -> DimConst -> EqContext q n -> EqContext q n
   extendDim x e ctx = {dassign $= (:< e), dnames $= (:< x)} ctx
 
   export
-  toTyContext : (e : EqContext q n) -> (d ** TyContext q d n)
-  toTyContext (MkEqContext {dassign, dnames, tctx, tnames})
-      with (lengthPrf0 dnames)
-    _ | Element d eq =
-      (d ** MkTyContext {
-        dctx   = rewrite eq in fromGround dassign,
-        dnames = rewrite eq in dnames,
-        tctx   = map (// shift0 d) tctx,
-        tnames
-      })
+  toTyContext : (e : EqContext q n) -> TyContext q e.dimLen n
+  toTyContext (MkEqContext {dimLen, dassign, dnames, tctx, tnames, qtys}) =
+    MkTyContext {
+      dctx = fromGround dassign,
+      tctx = map (// shift0 dimLen) tctx,
+      dnames, tnames, qtys
+    }
+
+
+parameters {auto _ : (Eq q, PrettyHL q)} (unicode : Bool)
+  export covering
+  prettyTContext : QContext q n -> NContext n -> TContext q d n -> Doc HL
+  prettyTContext qs ns ctx = separate comma $ toList $ go qs ns ctx where
+    go : QContext q m -> NContext m -> TContext q d m -> SnocList (Doc HL)
+    go [<]       [<]       [<] = [<]
+    go (qs :< q) (xs :< x) (ctx :< t) =
+      go qs xs ctx :< runPretty unicode (prettyBind [q] x t)
+
+  private
+  prettyDVars : NContext d -> Doc HL
+  prettyDVars ds = hseparate comma $ map (pretty0 unicode . DV) $ toList' ds
+
+  export
+  prettyDimEq1 : (PrettyHL a, PrettyHL b) => NContext d -> a -> b -> Doc HL
+  prettyDimEq1 ds p q = runPretty unicode $
+    local {dnames := toSnocList' ds} $ do
+      p <- pretty0M p; q <- pretty0M q
+      pure $ hsep [p, hl Syntax "=", q]
+
+  export
+  prettyDimEqCons : NContext d -> DimEq' d -> Doc HL
+  prettyDimEqCons ds eqs = hseparate comma $ toList $ go ds eqs
+  where
+    go : NContext e -> Context (Maybe . Dim) e -> SnocList (Doc HL)
+    go [<]       [<]              = [<]
+    go (ds :< _) (eqs :< Nothing) = go ds eqs
+    go (ds :< x) (eqs :< Just d)  = go ds eqs :< prettyDimEq1 ds (DV x) d
+
+  export
+  prettyDimEq : NContext d -> DimEq d -> Doc HL
+  prettyDimEq ds ZeroIsOne =
+    prettyDVars ds <+> comma <++> prettyDimEq1 [<] Zero One
+  prettyDimEq ds (C eqs) =
+    prettyDVars ds <+> comma <%%> prettyDimEqCons ds eqs
+
+  export covering
+  prettyTyContext : TyContext q d n -> Doc HL
+  prettyTyContext (MkTyContext dctx dnames tctx tnames qtys) =
+    sep [prettyDimEq dnames dctx <++> hl Syntax "|",
+         prettyTContext qtys tnames tctx]
+
+  export covering
+  prettyEqContext : EqContext q n -> Doc HL
+  prettyEqContext (MkEqContext dassign dnames tctx tnames qtys) =
+    sep [prettyDimEqCons dnames (fromGround' dassign) <++> hl Syntax "|",
+         prettyTContext qtys tnames tctx]