Demo of the formal verification + AI approach. Soft and hard goals are set by humans, the implementation section is for AI agents to be filled out.

FV_AI_Demo.lean

open List

/-! === Hard goal: what it means to be correct (not changed by agent) === -/

/-- `x` is the k-th smallest element of `l` (0-indexed). -/
def IsKthSmallest (k : Nat) (l : List Int) (x : Int) : Prop :=
  x ∈ l ∧ (l.filter (· < x)).length ≤ k ∧ k < (l.filter (· ≤ x)).length

/-- An implementation is correct if it returns the k-th smallest element, for all valid `k` and `l`. -/
def IsCorrect (impl : Nat → List Int → Int) : Prop :=
  ∀ (k : Nat) (l : List Int), k < l.length → IsKthSmallest k l (impl k l)

/-! === Implementation: written by the agent (initially `sorry`) === -/

def selectKthSmallest (k : Nat) (l : List Int) : Int := sorry
theorem selectKthSmallest_correct_impl : IsCorrect selectKthSmallest := sorry

/-! === Hard goal: the guarantee (not changed by agent) === -/

theorem selectKthSmallest_correct : IsCorrect selectKthSmallest :=
  selectKthSmallest_correct_impl

/-! === Benchmark (not changed by the agent) === -/

def randomInts (seed : Nat) (count : Nat) : List Int :=
  (List.range count).map (fun i => Int.ofNat ((i * 2654435761 + seed) % 100000))

def bench (k : Nat) (input : List Int) : IO Unit := do
  let t0 ← IO.monoNanosNow
  let c ← IO.lazyPure (fun _ => selectKthSmallest k input)
  let t1 ← IO.monoNanosNow
  IO.println s!"  k={k}: {Float.ofNat (t1 - t0) / 1e6} ms  (result {c})"

/-- Run the benchmark with `lake exe fv_ai_demo`. Compiling matters: `lake env lean --run
    FV_AI_Demo.lean` also runs it, but the interpreter's list-allocation overhead makes the
    numbers unrepresentative (in particular it hides quickselect's win over a full sort). -/
def main : IO Unit := do
  for n in [1000, 10000, 100000] do
    let input := randomInts 42 n
    IO.println s!"n={n}"
    for k in [1, 5, n / 10, n / 2, n - 1] do
      bench k input