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

Tests.lean

@@ -168,6 +168,41 @@ def runNeutralRepresentationChecks : IO Bool := do
     IO.println "fail  stuck eliminators stay in the neutral fragment"
     pure false
 
+def runContextInvariantChecks : IO Bool := do
+  let cxt :=
+    (Cxt.empty.bind "A" (.univ 0)).define "x" .zero .nat
+  let wfOk := Cxt.isWellFormed cxt
+  let lookupOk :=
+    match cxt.lookup "x", cxt.lookup "A" with
+    | some (ix, tx), some (iA, tA) =>
+        ix.val == 0 &&
+        iA.val == 1 &&
+        (match tx with | .nat => true | _ => false) &&
+        (match tA with | .univ 0 => true | _ => false)
+    | _, _ =>
+        false
+  let indexedOk :=
+    let latest : Fin cxt.lvl := ⟨0, by decide⟩
+    let outer : Fin cxt.lvl := ⟨1, by decide⟩
+    (match cxt.typeAt latest with
+     | ("x", .nat) => true
+     | _ => false) &&
+    (match cxt.typeAt outer with
+     | ("A", .univ 0) => true
+     | _ => false) &&
+    (match cxt.valueAt latest with
+     | .zero => true
+     | _ => false) &&
+    (match cxt.valueAt outer with
+     | .neu (.var 0) => true
+     | _ => false)
+  if wfOk && lookupOk && indexedOk then
+    IO.println "pass  context lookup is indexed and scope invariants hold"
+    pure true
+  else
+    IO.println "fail  context lookup is indexed and scope invariants hold"
+    pure false
+
 def runLevelSolverChecks : IO Bool := do
   let satOk :=
     match solveLevelConstraints [(.max 0 1, 1), (1, .succ 1)] with
@@ -212,9 +247,10 @@ def main : IO UInt32 := do
   let results ← cases.mapM runCase
   let safetyOk ← runInternalSafetyChecks
   let neutralOk ← runNeutralRepresentationChecks
+  let contextOk ← runContextInvariantChecks
   let levelOk ← runLevelSolverChecks
   let prettyOk ← runPrettyPrinterChecks
-  let allResults := results ++ [safetyOk, neutralOk, levelOk, prettyOk]
+  let allResults := results ++ [safetyOk, neutralOk, contextOk, levelOk, prettyOk]
   let failed := allResults.countP (· == false)
   if failed == 0 then
     IO.println s!"\n{allResults.length} passed, 0 failed"