perf(collect): reuse per-read SMEM/reseed buffers via CollectScratch

[nh13]

LearnedIndex::collect_smems allocated and freed two Vecs on every read — the emitted SMEM buffer and the pass-2 reseed work-list. bwa-mem3 runs the collector once per read over millions of reads, so those two allocate/free pairs recur on the compute-bound seeding path.

This adds a reusable CollectScratch and a collect_smems_into entry point that clears and reuses those buffers instead of allocating fresh ones:

• The walk moves into collect_smems_unsorted_into, which clears the caller-held buffers on entry and drains (rather than consumes) the reseed list so its capacity is retained between reads.
• collect_smems stays a thin wrapper that allocates a throwaway CollectScratch, so every existing caller, the FFI slice contract, and the output are unchanged.
• The C entry point prmi_collect_smems now holds a thread-local CollectScratch, so the amortization reaches the consumer without any ABI change while preserving per-thread isolation.

Byte-identical by construction — the buffers are cleared before use. Verified by:

• A new collect_smems_into_equals_collect_smems proptest that runs a reused, pre-dirtied scratch against the fresh-allocating path across the full opts sweep (long-then-short reads, to catch stale-tail bugs).
• The existing oracle / ISA byte-identity proptests, unchanged.
• An end-to-end collect_gate SMEM-dump comparison vs main: 0 differences across a 24-point k/split_len/split_width/max_mem_intv grid on a 1 Mbp reference.
• cargo clippy --all-features --all-targets -D warnings and cargo +nightly fmt --check clean; full cargo test workspace suite passes.

Performance (secondary): a new wall_gate example measures the seeding wall (single-thread, warm, load-excluded). The scratch path is a small, consistent win over the allocating path — min/median/mean ns/read all lower (paired runs ~0.5–3%) — with identical output. This is a hygiene / allocation-pressure win, not a large speedup.

Summary by CodeRabbit

• New Features
  • Introduced CollectScratch API for SMEM collection operations.
  • Added collect_smems_into method on LearnedIndex for configurable collection workflows.
  • Added new example demonstrating SMEM collection benchmarking with performance metrics.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: ce63df7e-0097-42c8-aea0-127b5fd0924b

📥 Commits

Reviewing files that changed from the base of the PR and between 01a2627 and 13311c3.

📒 Files selected for processing (3)

• prmi-sys/src/lib.rs
• prmi/examples/wall_gate.rs
• prmi/src/index/collect.rs

---

Walkthrough

Introduces CollectScratch to eliminate per-call Vec allocations in the SMEM collection pipeline. Refactors all three collection passes to reuse caller-provided scratch buffers, exposes collect_smems_into as a public API, wires a thread-local CollectScratch into the C FFI, adds proptest coverage for byte identity, and adds a wall_gate benchmark example.

Changes

CollectScratch reuse for SMEM collection

Layer / File(s)	Summary
CollectScratch struct and collect_smems_into public API prmi/src/index/collect.rs	Defines CollectScratch with smems/reseeds buffers. Refactors collect_smems into a thin wrapper allocating a throwaway scratch and delegating to new collect_smems_into. Marks the old allocating unsorted helper #[cfg(test)].
Pipeline refactor: pass-1/2/3 reuse scratch buffers prmi/src/index/collect.rs	collect_smems_unsorted_into clears and reuses scratch.smems and scratch.reseeds. Pass-1 emits into scratch.smems; reseed selection accumulates into scratch.reseeds instead of a fresh Vec; pass-2 drains scratch.reseeds to avoid borrow conflicts; pass-3 appends into the same scratch.smems.
Proptest: collect_smems_into equals collect_smems prmi/src/index/collect.rs	New proptest reuses one CollectScratch across a batch and asserts byte-identical count and SMEM contents against the fresh-allocating path.
FFI thread-local scratch wiring prmi-sys/src/lib.rs	prmi_collect_smems defines a thread_local! SCRATCH (RefCell<CollectScratch>), borrows it mutably inside catch_unwind, and calls collect_smems_into instead of collect_smems.
wall_gate benchmark example prmi/examples/wall_gate.rs	Env-var–driven harness loading an on-disk index and PAC, iterating FASTQ reads with scratch-reuse or alloc path selection, buffer-resize retry on Err(needed), and throughput/latency output.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• fg-labs/prmi#32: Introduced the prmi_collect_smems FFI entry point and the collect_smems call path that this PR replaces with collect_smems_into.
• fg-labs/prmi#34: Implements the pass-3 max_mem_intv forward-narrowing logic that this PR routes through the reused scratch.smems buffer.
• fg-labs/prmi#40: Modifies pass-2 reseed logic and threads ReseedHint through the same reseed call paths now refactored to drain scratch.reseeds.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'perf(collect): reuse per-read SMEM/reseed buffers via CollectScratch' accurately captures the main optimization: introducing buffer reuse via the new CollectScratch structure to avoid repeated allocations in the SMEM collection path.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch perf/native-collect-buffer-reuse

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[nh13]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[nh13]

@coderabbitai review

[coderabbitai]

Action performed

Review triggered.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.

[nh13]

@coderabbitai review

[coderabbitai]

✅ Action performed

Review finished.

Note: CodeRabbit is an incremental review system and does not re-review already reviewed commits. This command is applicable only when automatic reviews are paused.