Replace the listy context API with indexed lookups and explicit scope invariants

BidirTT/Check.lean

@@ -46,7 +46,7 @@ mutual
   partial def infer (cxt : Cxt) : Raw → TCM (Tm × Val)
     | .var x =>
         match cxt.lookup x with
-        | some (i, a) => pure (.var i, a)
+        | some (i, a) => pure (.var i.val, a)
         | none        => throw s!"unknown variable {x}"
     | .nat => pure (.nat, .univ 0)
     | .zero => pure (.zero, .nat)

BidirTT/Context.lean

@@ -2,29 +2,62 @@ import BidirTT.Value
 
 namespace BidirTT
 
+abbrev CxtEntry := Name × Val
+
 structure Cxt where
   env   : Env
-  types : List (Name × Val)
+  types : List CxtEntry
   lvl   : Lvl
-  deriving Inhabited
+  envWf : env.length = lvl
+  typesWf : types.length = lvl
 
-def Cxt.empty : Cxt := ⟨[], [], 0⟩
+def Cxt.empty : Cxt := ⟨[], [], 0, rfl, rfl⟩
+
+instance : Inhabited Cxt := ⟨Cxt.empty⟩
 
 def Cxt.bind (cxt : Cxt) (x : Name) (a : Val) : Cxt :=
   { env   := .neu (.var cxt.lvl) :: cxt.env
   , types := (x, a) :: cxt.types
-  , lvl   := cxt.lvl + 1 }
+  , lvl   := cxt.lvl + 1
+  , envWf := by simp [cxt.envWf]
+  , typesWf := by simp [cxt.typesWf] }
 
 def Cxt.define (cxt : Cxt) (x : Name) (v a : Val) : Cxt :=
   { env   := v :: cxt.env
   , types := (x, a) :: cxt.types
-  , lvl   := cxt.lvl + 1 }
+  , lvl   := cxt.lvl + 1
+  , envWf := by simp [cxt.envWf]
+  , typesWf := by simp [cxt.typesWf] }
 
-private def lookupGo : Name → List (Name × Val) → Nat → Option (Nat × Val)
-  | _, [],              _ => none
-  | x, (y, a) :: rest,  i => if x == y then some (i, a) else lookupGo x rest (i+1)
+private def castFin {n m : Nat} (h : n = m) (i : Fin n) : Fin m :=
+  ⟨i.val, h ▸ i.isLt⟩
 
-def Cxt.lookup (cxt : Cxt) (x : Name) : Option (Nat × Val) :=
-  lookupGo x cxt.types 0
+private def lookupGo (x : Name) : (entries : List CxtEntry) → Option (Fin entries.length × Val)
+  | [] => none
+  | (y, a) :: rest =>
+      if x == y then
+        some (⟨0, by simp⟩, a)
+      else
+        match lookupGo x rest with
+        | some (i, a) =>
+            some (⟨i.val + 1, by simp [i.isLt]⟩, a)
+        | none =>
+            none
+
+def Cxt.isWellFormed (cxt : Cxt) : Bool :=
+  cxt.env.length == cxt.lvl && cxt.types.length == cxt.lvl
+
+def Cxt.typeAt (cxt : Cxt) (idx : Fin cxt.lvl) : CxtEntry :=
+  cxt.types.get (castFin cxt.typesWf.symm idx)
+
+def Cxt.valueAt (cxt : Cxt) (idx : Fin cxt.lvl) : Val :=
+  cxt.env.get (castFin cxt.envWf.symm idx)
+
+def Cxt.lookup (cxt : Cxt) (x : Name) : Option (Fin cxt.lvl × Val) :=
+  match lookupGo x cxt.types with
+  | some (idx, a) =>
+      some (castFin cxt.typesWf idx, a)
+  | none =>
+      none
 
 end BidirTT