Problem: kmp

Your task is to implement the Knuth-Morris-Pratt (KMP) algorithm to efficiently find all occurrences of a needle (pattern) within a haystack (text) and verify its functional correctness in Lean using Mathlib. Preconditions: You can assume the input text and pattern are arrays of UInt8 (representing bytes) and their lengths are at most 1,000,000. Requirements: Implement the method kmpSearch which returns an Array Nat. (1) Preprocessing: You must implement a helper to compute the Longest Prefix Suffix (LPS) array (or "next" array), where lps[i] stores the length of the longest proper prefix of the sub-pattern that is also a suffix of that sub-pattern. (2) Search: The main loop must use the LPS array to skip unnecessary comparisons in the text when a mismatch occurs, achieving linear time complexity. (3) Functional Correctness: Prove that despite the optimization, the result is identical to a naive search (Soundness and Completeness). Verification Challenges: (1) LPS Property: Proving that your preprocessing logic actually satisfies the mathematical definition of the LPS table. (2) Safe Skipping: Proving that the "jump" logic (moving the pattern index based on LPS) never skips a valid match. (3) Complex Invariants: The main loop invariant must relate the current state of the pattern index to the prefix of the pattern matched against the text so far.


import Mathlib

-- Precondition definitions
@[reducible, simp]
def kmpSearch_precond (haystack : Array UInt8) (needle : Array UInt8) : Prop :=
  -- !benchmark @start precond
  haystack.size < 1000000 ∧ 
  needle.size < 1000000
  -- !benchmark @end precond

-- !benchmark @start auxcode
-- !benchmark @end auxcode

-- Main function definition
def kmpSearch (haystack : Array UInt8) (needle : Array UInt8)
    (_ : kmpSearch_precond haystack needle) : Array Nat :=
  -- !benchmark @start code
  sorry
  -- !benchmark @end code

-- Postcondition auxiliary definitions
def matchesAt (haystack needle : Array UInt8) (i : Nat) : Prop :=
  i + needle.size ≤ haystack.size ∧
  (∀ j, j < needle.size → haystack.getD (i + j) 0 = needle.getD j 0)

-- Postcondition definitions
@[reducible, simp]
def kmpSearch_postcond (haystack : Array UInt8) (needle : Array UInt8)
    (result : Array Nat) (_ : kmpSearch_precond haystack needle) : Prop :=
  -- !benchmark @start postcond
  (∀ i, i < result.size → matchesAt haystack needle (result.getD i 0)) ∧ -- Soundness
  (∀ i, matchesAt haystack needle i → ∃ k, k < result.size ∧ result.getD k 0 = i) -- Completeness
  -- !benchmark @end postcond

-- !benchmark @start lemma
-- !benchmark @end lemma

-- Proof content
theorem kmpSearch_postcond_satisfied (haystack : Array UInt8) (needle : Array UInt8)
    (h_precond : kmpSearch_precond haystack needle) :
    kmpSearch_postcond haystack needle (kmpSearch haystack needle h_precond) h_precond := by
  -- !benchmark @start proof
  sorry
  -- !benchmark @end proof