feat(collect): fused per-read SMEM collection (passes 1+2) + FFI - #32
Merged
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (3)
📝 WalkthroughWalkthroughAdds fused per-read SMEM collection: a Rust zigzag implementation (pass-1 + pass-2), ABI-compatible types, a public LearnedIndex::collect_smems API, C-FFI entrypoint ChangesPer-Read SMEM Collection Engine
sequenceDiagram
participant Caller as C Caller
participant FFI as prmi_collect_smems
participant RustAPI as LearnedIndex::collect_smems
participant Pass1 as zz_step1
participant Pass2 as zz_step1_reseed
participant Output as Output Buffer
Caller->>FFI: call with read, opts, packed pac, out buffer/cap
FFI->>FFI: validate pointers/lengths, pack pac, reject max_mem_intv != 0
FFI->>RustAPI: invoke collect_smems with decoded opts and pac
RustAPI->>Pass1: run pass-1 zigzag collection
Pass1->>RustAPI: emit SMEMs
RustAPI->>Pass2: select pivots and run pass-2 reseed
Pass2->>RustAPI: emit SMEMs
RustAPI->>RustAPI: sort SMEMs within read
RustAPI->>Output: copy SMEMs into caller buffer (or return needed size)
RustAPI-->>FFI: return result (count or needed)
FFI->>FFI: write *out_n, map to return codes (0, -4, -5, -3)
FFI-->>Caller: return status
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
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. Comment |
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
prmi/src/index/collect.rs (3)
155-158: 📐 Maintainability & Code Quality | ⚡ Quick winAdd a regression test for the direct-Rust pass-3 guard.
This public entrypoint intentionally panics on
max_mem_intv != 0, but that caller path is not locked down by a focused unit test in this file. A small#[should_panic]/catch_unwindtest would keep the “fail hard instead of silently under-seeding” contract from regressing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@prmi/src/index/collect.rs` around lines 155 - 158, Add a focused regression test that verifies the public entrypoint panics when max_mem_intv != 0: call the function collect_smems (the guarded code path) with an options struct where opts.max_mem_intv is set > 0 and assert the panic using either #[should_panic] on the test or std::panic::catch_unwind; ensure the test constructs the same options type used by collect_smems and exercises the public call path so the panic guard remains covered by CI.
159-163: 🚀 Performance & Scalability | ⚡ Quick winReturn
Err(needed)before sorting the retry path.When
outis too small, this still pays the full within-read sort even though the caller has to retry anyway. The overflow contract only needs the count, so the capacity check can happen immediately aftercollect_smems_unsorted().♻️ Proposed change
- let mut smems = self.collect_smems_unsorted(read, rid, opts, pac, enc); - Self::sort_within_read(&mut smems); - if smems.len() > out.len() { + let mut smems = self.collect_smems_unsorted(read, rid, opts, pac, enc); + if smems.len() > out.len() { return Err(smems.len()); } + Self::sort_within_read(&mut smems); out[..smems.len()].copy_from_slice(&smems); Ok(smems.len())🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@prmi/src/index/collect.rs` around lines 159 - 163, The current code calls Self::sort_within_read(&mut smems) before checking capacity and returning Err(smems.len()), wasting work when the caller must retry; move the capacity check immediately after calling collect_smems_unsorted(read, rid, opts, pac, enc) and return Err(smems.len()) if smems.len() > out.len() so sort_within_read is only invoked when the output buffer is large enough (i.e., call collect_smems_unsorted -> if smems.len() > out.len() return Err(smems.len()) -> then call Self::sort_within_read(&mut smems) and proceed).
197-220: 🚀 Performance & Scalability | ⚡ Quick winPre-reserve the per-read vectors on this hot path.
collect_smems_unsortedruns once per read, but bothsmemsandreseedsstart at capacity 0 and grow via repeated reallocations. The method docs already advertise a2 * read.len()safe output bound, and pass-2 can reservenum1, so reserving here would remove avoidable allocator churn from the path this PR is optimizing.♻️ Proposed change
- let mut smems: Vec<Smem> = Vec::new(); + let mut smems: Vec<Smem> = Vec::with_capacity(rlen.saturating_mul(2)); @@ - let mut reseeds: Vec<(usize, i64)> = Vec::new(); + let mut reseeds: Vec<(usize, i64)> = Vec::with_capacity(num1);🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@prmi/src/index/collect.rs` around lines 197 - 220, Reserve the per-read vectors to avoid reallocations: before the zigzag pass, call smems.reserve(2 * rlen) (smems is Vec<Smem>) since the method docs guarantee a 2 * read.len() bound; after the first pass, once num1 = smems.len() is known, call reseeds.reserve(num1) (reseeds is Vec<(usize, i64)>) so the reseed collection does not reallocate while being filled.prmi-sys/tests/spectrum_ffi.rs (1)
1962-2020: 📐 Maintainability & Code Quality | ⚡ Quick winExercise the documented
out == NULL && out_cap == 0sizing path.The new API explicitly special-cases a null output buffer when capacity is zero, but this test only checks overflow with a real buffer. Add a dry-run call on the non-empty
readwithout = ptr::null_mut()andout_cap = 0, then assertrc == -4andout_n == rust_nso regressions in that branch are caught.Suggested test addition
+ // Dry-run sizing path: NULL out is allowed when out_cap == 0. + let mut dry_run_n: i32 = -1; + let rc_dry_run = unsafe { + prmi_collect_smems( + handle, + read.as_ptr(), + read.len() as i32, + 9, + &copts, + pac_packed.as_ptr(), + pnb, + ptr::null_mut(), + 0, + &mut dry_run_n, + ) + }; + assert_eq!(rc_dry_run, -4, "NULL/0 dry-run must report required capacity"); + assert_eq!(dry_run_n as usize, rust_n, "dry-run must report the exact SMEM count"); + // Overflow: out_cap one short -> -4, out_n set to the required count.Based on the new FFI contract in
prmi-sys/src/lib.rs,outmay be null iffout_cap == 0.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@prmi-sys/tests/spectrum_ffi.rs` around lines 1962 - 2020, Add a dry-run invocation of prmi_collect_smems that exercises the documented "out == NULL && out_cap == 0" sizing path: call prmi_collect_smems with the existing non-empty read buffer (read.as_ptr(), read.len() as i32), copts, pac_packed.as_ptr(), pnb, out = ptr::null_mut(), out_cap = 0 and a new out_n (e.g., dry_n) initialized to -1; then assert the return code is -4 and that dry_n equals rust_n. Place this before or alongside the existing overflow test so regressions in prmi_collect_smems’s null-out sizing branch are caught (refer to prmi_collect_smems, read, rust_n, and the overflow assertions for placement and naming).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@prmi-sys/src/lib.rs`:
- Around line 1877-1884: The entrypoint prmi_collect_smems currently rejects any
NULL `read` pointer unconditionally; change the null-pointer validation so
`read` is allowed to be NULL when `read_len == 0` (i.e. treat NULL+0 as a valid
empty slice). Concretely, replace the `read.is_null()` requirement with a check
that fails only if `read.is_null() && read_len != 0`, keeping the other null
checks for `handle`, `opts`, `pac`, and `out_n` and preserving the existing
negative-length/out-cap validations (`read_len < 0 || out_cap < 0`).
---
Nitpick comments:
In `@prmi-sys/tests/spectrum_ffi.rs`:
- Around line 1962-2020: Add a dry-run invocation of prmi_collect_smems that
exercises the documented "out == NULL && out_cap == 0" sizing path: call
prmi_collect_smems with the existing non-empty read buffer (read.as_ptr(),
read.len() as i32), copts, pac_packed.as_ptr(), pnb, out = ptr::null_mut(),
out_cap = 0 and a new out_n (e.g., dry_n) initialized to -1; then assert the
return code is -4 and that dry_n equals rust_n. Place this before or alongside
the existing overflow test so regressions in prmi_collect_smems’s null-out
sizing branch are caught (refer to prmi_collect_smems, read, rust_n, and the
overflow assertions for placement and naming).
In `@prmi/src/index/collect.rs`:
- Around line 155-158: Add a focused regression test that verifies the public
entrypoint panics when max_mem_intv != 0: call the function collect_smems (the
guarded code path) with an options struct where opts.max_mem_intv is set > 0 and
assert the panic using either #[should_panic] on the test or
std::panic::catch_unwind; ensure the test constructs the same options type used
by collect_smems and exercises the public call path so the panic guard remains
covered by CI.
- Around line 159-163: The current code calls Self::sort_within_read(&mut smems)
before checking capacity and returning Err(smems.len()), wasting work when the
caller must retry; move the capacity check immediately after calling
collect_smems_unsorted(read, rid, opts, pac, enc) and return Err(smems.len()) if
smems.len() > out.len() so sort_within_read is only invoked when the output
buffer is large enough (i.e., call collect_smems_unsorted -> if smems.len() >
out.len() return Err(smems.len()) -> then call Self::sort_within_read(&mut
smems) and proceed).
- Around line 197-220: Reserve the per-read vectors to avoid reallocations:
before the zigzag pass, call smems.reserve(2 * rlen) (smems is Vec<Smem>) since
the method docs guarantee a 2 * read.len() bound; after the first pass, once
num1 = smems.len() is known, call reseeds.reserve(num1) (reseeds is Vec<(usize,
i64)>) so the reseed collection does not reallocate while being filled.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 64928c79-89ef-429f-87bf-93fb009f05f4
📒 Files selected for processing (4)
prmi-sys/src/lib.rsprmi-sys/tests/spectrum_ffi.rsprmi/src/index/collect.rsprmi/src/index/mod.rs
Add `LearnedIndex::collect_smems` and the `prmi_collect_smems` FFI: the native port of the consumer's zigzag SMEM-collection driver, collapsing the per-read walk (91-155 stateless C->Rust FFI crossings) into ONE call. The per-call search cost is already at parity with bwa-meme; the residual is call COUNT, and running the zigzag in-process lets the SA-interval state cross zz_left_span -> zz_right_emit -> next-step in-register (the thing the FFI boundary forces cold). Built entirely on the merged, byte-identity-tested per-call primitives (mem_search, mem_search_backward, mem_search_backward_truncated_span_rc, forward_truncate_below_maximal) -- no new search math. Pass 1 (zigzag, min_intv=1 == the definitional MEM set), reseed selection, and pass 2 (per-pivot reseed walk), then the within-read two-stage sort. `#[repr(C)] Smem` mirrors the consumer's SMEM field-for-field (layout-pinned). Adaptations to main: `zz_left_span` uses `mem_search_backward(0,1,1,..)` (the byte-identical match_len of the perf `mem_search_backward_impl(..,false)`, one extra discarded gallop); `zz_left_span_reseed` uses main's 7-arg `_span_rc`. Pass 3 (max_mem_intv > 0) is DEFERRED -- it needs the still-unmerged forward routing. The FFI rejects max_mem_intv != 0 with rc -5 (passes 1+2 are byte-identical for max_mem_intv == 0); `collect_smems` documents the precondition. Byte-identity: pass-1 == independent MEM-set oracle; pass-2 == transcribed reseed selection + zz_step1_reseed; full pipeline == composed oracle; sort/overflow/ determinism/thread-safety covered; the FFI returns SMEMs byte-for-byte equal to the Rust path (+ empty-read, pass-3-reject -5, overflow -4 contracts).
nh13
force-pushed
the
feat/v0.3-collect-core
branch
from
June 13, 2026 03:54
5b1e836 to
d97bfe9
Compare
Author
|
@coderabbitai review |
✅ Action performedReview finished.
|
This was referenced Jun 13, 2026
This was referenced Jun 13, 2026
nh13
added a commit
that referenced
this pull request
Jun 14, 2026
Extend the cpp_caller FFI smoke with a prmi_collect_smems block: it covers the
fused per-read SMEM collector's shape (no prior coverage) AND demonstrates the
idiomatic C calling pattern for a consumer adopting it —
- prmi_collect_opts_t setup (min_seed_len/split_len/split_width/max_mem_intv,
with max_mem_intv > 0 enabling pass 3);
- the out_cap=0 size-probe (rc -4 with *out_n = required count);
- the grow-on-(-4) retry loop (resize to *out_n, retry) — the buffer-sizing
contract a C caller must implement;
- iterating the returned prmi_smem_t[] and validating each (rid/span/occ).
Built against the cbindgen-generated prmi.h (the collect symbols land via the
#32 FFI). run_smoke.sh passes: size-probe rc=-4 needed=1, then one well-formed
SMEM spanning the reference-derived read. -Wall -Wextra clean.
This was referenced Jun 20, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
v0.3 perf series — PR-D2: the fused
collect_smemsmodule (the central deliverable; built on #31'sforward_truncate_below_maximaland the merged TRUNC_IV primitives).What
Adds
LearnedIndex::collect_smems+ theprmi_collect_smemsFFI — the native port of the consumer's zigzag SMEM-collection driver. Today the consumer crosses the C→Rust FFI boundary 91–155 times per read (stateless: eachzz_left_span/zz_right_emitre-issues a coldprmi_mem_search). This collapses the whole per-read walk into one call, carrying the SA-interval state acrosszz_left_span → zz_right_emit → next-stepin-register — the thing the FFI boundary forces cold.The per-call search cost is already at parity with bwa-meme (MODE2: prmi 7.0 vs 7.76 probes/call); the entire residual was call count (bwa-meme ~53). Fusing the walk in-process is what collapses it.
How
Built entirely on the merged, byte-identity-tested per-call primitives —
mem_search,mem_search_backward,mem_search_backward_truncated_span_rc(#27),forward_truncate_below_maximal(#31). No new search math. Pass 1 (zigzag,min_intv=1= the definitional MEM set), reseed selection, pass 2 (per-pivot reseed walk), then the within-read two-stage sort.#[repr(C)] Smemmirrors the consumer'sSMEMfield-for-field (layout-pinned by a Rust offset test + an FFIconst _size/align assert).Byte-identity (the safety net)
zz_step1_reseed;zz_left_span/zz_right_emiteach vs their own oracle; sort / overflow / determinism / thread-safety covered;Scope — Pass 3 deferred
max_mem_intv > 0(the pass-3max_mem_intvstrategy) needs the still-unmerged forward-routing primitive, so it is deferred: the FFI rejectsmax_mem_intv != 0with rc-5, andcollect_smemspanics on the precondition violation (the FFI pre-checks, so the panic is unreachable on the consumer path). Passes 1+2 (max_mem_intv == 0) are byte-identical to FMI seeding today. Pass 3 lands in a follow-up (PR-D′) once the forward routing merges.CAVEAT for the box-gate: if the consumer config sets
max_mem_intv > 0(bwa-mem ships a nonzero long-MEM-reseed default), full-SAM e2e parity is gated on the whole D→C→D′ chain — this PR's correctness rests on the passes-1+2 byte-identity proptests, not the box-gate.Adaptations to main
zz_left_spanusesmem_search_backward(0,1,1,..)— the byte-identicalmatch_lenof the perfmem_search_backward_impl(..,false)(one extra discarded interval-recovery gallop; a deferrable micro-opt);zz_left_span_reseeduses main's 7-arg_span_rc(perf(spectrum): min_intv-truncated backward reseed span via RC-downward scan #27 droppedsa_start/est_hint).Self-review
Pre-PR
/coderabbitai-review: 1 actionable + 1 nitpick, both fixed — (1) upgraded the deferred-pass-3 guard fromdebug_assertto a hardassertso a direct Rust caller can't silently get an under-seeded result in release; (2) clarified theattribprofiling-bucket doc (slot 4 reserved for pass 3).Summary by CodeRabbit
New Features
Tests