Unboxed NFA with Fixed-Width Unsigned Indices for Performance Optimization

Summary

Create a new UNFA (Unboxed NFA) type that restricts node indices to a configurable fixed-width unsigned integer (default: UInt32) instead of Nat, enabling performance optimizations by avoiding boxing/unboxing in hot loops. The goal is to implement:

1. UNFA type with StateId-indexed nodes
2. Translator from NFA to UNFA
3. Unboxed VM search function
4. Proof of correctness showing it reproduces the same transition paths as the original VM (when size constraints are satisfied)

Motivation

In Lean 4, fixed-width unsigned integers (e.g. UInt32) are represented as unboxed integer types at runtime, while Nat is boxed (by shifting the value to the left). The VM's hot loop frequently accesses NFA nodes using Fin nfa.nodes.size indices, which involves:

• Array indexing operations
• State set membership checks (SparseSet)
• Stack operations with state indices

By using fixed-width unsigned indices directly instead of Nat and Fin n, we can:

• Eliminate boxing/unboxing overhead in the hot loop
• Improve cache locality
• Reduce memory footprint (e.g. UInt32 is 4 bytes, UInt16 is 2 bytes vs Nat's 8 bytes on 64-bit systems)

while maintaining formal correctness through proofs.

Current Architecture

NFA Structure

inductive NFA.Node where
  | done
  | fail
  | epsilon (next : Nat)
  | anchor (anchor : Anchor) (next : Nat)
  | char (c : Char) (next : Nat)
  | split (next₁ next₂ : Nat)
  | save (offset : Nat) (next : Nat)
  | sparse (cs : Classes) (next : Nat)

structure NFA where
  nodes : Array NFA.Node
  start : Nat

VM Usage

The VM uses Fin nfa.nodes.size extensively:

• SearchState.states : SparseSet nfa.nodes.size (stores Fin n values)
• εStack : List (Update × Fin nfa.nodes.size)
• Vector σ.Update nfa.nodes.size for updates
• Array access: nfa[state] where state : Fin nfa.nodes.size

Proposed Implementation

0. Configurable Index Type (naming + easy experimentation)

In the current codebase, NFA node indices are consistently referred to as states in the VM (e.g. variables named state, state'). So we use StateId as the name for the index type.

-- Pick the unboxed index type here (experiment-friendly):
abbrev StateId := UInt32  -- could also try UInt16, UInt64, ...

-- Use the same type with a different name so that we have `UFin (n : StateCount)`
abbrev StateCount := StateId

1. UFin: Bounded StateId Index

First, we introduce UFin as the unboxed equivalent of Fin:

namespace Regex.UNFA

/-- A bounded `StateId` index, equivalent to `Fin n` but unboxed -/
structure UFin (n : StateCount) where
  val : StateId
  isLt : val < n

end Regex.UNFA

2. UNFA Type Definition

namespace Regex.UNFA

/-- Unboxed NFA node with embedded inBounds properties -/
inductive UNode (size : StateCount) where
  | done
  | fail
  | epsilon (next : UFin size)
  | anchor (anchor : Anchor) (next : UFin size)
  | char (c : Char) (next : UFin size)
  | split (next₁ next₂ : UFin size)
  | save (offset : StateId) (next : UFin size)
  | sparse (cs : Classes) (next : UFin size)

/-- Unboxed NFA with well-formedness embedded -/
structure UNFA (size : StateCount) where
  nodes : Vector (UNode size) size.toNat
  start : UFin size
  -- and other well-formedness properties from NFA.WellFormed, if necessary

end Regex.UNFA

Key design decisions:

• Use StateId as the runtime index representation (swap UInt32↔UInt16 easily)
• Also use StateCount := StateId so that sizes are carried unboxed, and converted via size.toNat only at the boundary (Vector length)
• All indices use UFin size, ensuring inBounds-ness at the type level
• UNFA embeds well-formedness. Since we construct it from a well-formed NFA, it shouldn't be hard to deal with well-formedness proofs.

2. Supporting Data Structures

USparseSet

structure USparseSet (n : StateCount) where
  count : StateId
  dense : Vector StateId n.toNat
  sparse : Vector StateId n.toNat
  -- Invariants: count ≤ n, and all dense/sparse values are < n

Array/Vector Access with USize

Lean 4 provides Array.uget and Vector.uget that take USize indices, which can be converted from StateId. We can also implement ufget and getElem for UFin indices.

3. NFA → UNFA Translator

def NFA.toUNFA (nfa : NFA) (wf : nfa.WellFormed) : Option (UNFA nfa.nodes.size.toStateCount) :=
  if h : nfa.nodes.size ≤ UInt32.size then
    -- Convert all Nat indices to UFin
    -- Convert NFA.Node to UNode with embedded proofs
    -- Well-formedness is preserved by construction
    some ⟨...⟩
  else
    .none  -- NFA too large for StateId

Conversion challenges:

• All Nat indices in nodes must be < StateId.size
• save offset values must also fit in StateId
• Need to prove that NFA.Node.inBounds implies UFin bounds
• Since UNode embeds size, conversion requires proving each node's indices are valid

4. Unboxed VM Search Function

namespace Regex.UNFA.VM

structure USearchState {s : String} (σ : Strategy s) {size : StateCount} (unfa : UNFA size) where
  states : USparseSet size
  updates : Vector σ.Update size.toNat

def captureNext {s : String} (σ : Strategy s) {size : StateCount} (unfa : UNFA size) (p : Pos s) : Option σ.Update :=
  -- Use Array.uget and Vector.uget with StateId.toUSize
  -- Direct unboxed access: unfa.nodes.uget (state.val.toUSize)

end Regex.UNFA.VM

Performance benefits:

• Array.uget / Vector.uget with USize avoids Fin conversions
• StateId.toUSize should stay unboxed (for UInt32/UInt16 etc.)

5. Correctness Proof

The key correctness property is along the lines of:

theorem UNFA.VM.captureNext_equiv {s : String} (nfa : NFA) (wf : nfa.WellFormed)
  (unfa : UNFA nfa.nodes.size.toStateCount) (h : nfa.toUNFA wf = some unfa) (p : Pos s) :
  VM.captureNext σ nfa wf p = UNFA.VM.captureNext σ unfa p :=
  -- Proof that transition paths are identical
  -- when size constraints are satisfied
  -- Well-formedness is embedded in UNFA, so no separate wf' parameter needed

6. Regex Integration

Update Regex to support both boxed and unboxed representations:

inductive Regex where
  | boxed (nfa : NFA) (useBacktracker : Bool)
  | unboxed (size : StateCount) (unfa : UNFA size)

def Regex.withUNFA (self : Regex) : Regex :=
  match self with
  | .boxed nfa useBacktracker =>
    if useBacktracker then
      self  -- Backtracker not supported, keep boxed
    else
      match nfa.toUNFA nfa.wf with
      | some unfa => .unboxed nfa.nodes.size.toUInt32 unfa
      | none => self  -- Fallback to boxed if too large
  | .unboxed _ _ => self  -- Already unboxed

This allows:

• Automatic conversion from boxed to unboxed when possible
• Fallback to boxed NFA when size exceeds StateId.size
• Backtracker continues to use boxed NFA (not optimized)

Implementation Plan

Phase 1: Core Types

• Introduce StateId := UIntXX knob (e.g. UInt32 by default)
• Define UFin (n : StateCount) structure (stores StateId + < n proof)
• Define UNode (size : StateCount) inductive type with embedded UFin indices
• Define UNFA (size : StateCount) structure with embedded well-formedness
• Implement USparseSet with StateId indices
• Prove basic properties (UFin bounds, UNode inBounds, etc.)

Phase 2: Translator

• Implement NFA.toUNFA conversion
• Prove well-formedness preservation
• Handle size overflow cases gracefully

Phase 3: Unboxed VM

• Implement UNFA.VM.captureNext using Array.uget/Vector.uget (via StateId.toUSize)
• Implement UNFA.VM.captureNextBuf
• Use StateId.toUSize for array indexing in hot paths
• Optimize hot paths (εClosure, stepChar) with direct unboxed access

Phase 4: Correctness Proofs

• Prove captureNext_equiv theorem
• Prove transition path equivalence
• Verify capture group correctness

Phase 5: Integration

• Update Regex to inductive with boxed/unboxed cases
• Implement Regex.withUNFA conversion
• Update Regex.captureNextBuf to dispatch to appropriate VM
• Benchmark performance improvements

Design Notes

Embedding Well-Formedness

Since we only deal with an NFA of fixed size, embedding well-formedness into UNFA works well:

• start : UFin ensures start < nodes.size at the type level
• UNode's UFin parameters ensure all indices are inBounds
• No need for separate WellFormed predicate or runtime checks

Conversion Proofs

When converting NFA.Node to UNode, we need to prove:

• NFA.Node.inBounds nfa.nodes.size → UFin bounds for each index
• Well-formedness preservation: NFA.WellFormed → UNFA (embedded)
• These proofs may require careful handling of the size constraint

Open Questions

1. Size limit handling: Fallback to boxed NFA when size exceeds StateId.size (handled by withUNFA)
2. Benchmarking: How to measure actual performance gains? (See CodSpeed integration issue)
3. UFin vs direct StateId: Should we use UFin everywhere, or allow direct StateId in some contexts?

References

• Lean 4 FFI and Unboxed Types
• Current NFA implementation: regex/Regex/NFA/Basic.lean
• Current VM implementation: regex/Regex/VM/Basic.lean