use std::ops::Bound;
use rmcp::ErrorData as McpError;
use super::cursor::Cursor;
use super::helpers_calls::resolve_call_line_col;
use super::types::ReferenceHit;
pub(super) struct CallScanPage {
pub total: u32,
pub total_is_partial: bool,
pub hits: Vec<ReferenceHit>,
pub next_cursor: Option<Cursor>,
pub hit_keys: Vec<Vec<u8>>,
pub hit_starts: Vec<u32>,
}
pub(super) struct BudgetedCallPage {
pub hits: Vec<ReferenceHit>,
pub next_cursor: Option<Cursor>,
pub budgeted: bool,
}
pub(super) fn budget_call_page(page: CallScanPage, max_tokens: Option<u32>) -> BudgetedCallPage {
if max_tokens.is_none() {
return BudgetedCallPage {
hits: page.hits,
next_cursor: page.next_cursor,
budgeted: false,
};
}
let budget = super::budget::apply_budget(page.hits, max_tokens);
if !budget.budgeted {
return BudgetedCallPage {
hits: budget.items,
next_cursor: page.next_cursor,
budgeted: false,
};
}
let kept = budget.items.len();
let next_cursor = page.hit_keys.get(kept - 1).map(|k| Cursor::encode_fjall(k));
BudgetedCallPage {
hits: budget.items,
next_cursor,
budgeted: true,
}
}
fn scan_calls_by_name(
idx: &crate::index::IndexDb,
name: &str,
limit: usize,
cursor_after: Option<&[u8]>,
) -> Result<CallScanPage, McpError> {
let finder = memchr::memmem::Finder::new(name.as_bytes());
let lower: Bound<Vec<u8>> = match cursor_after {
Some(k) => Bound::Excluded(k.to_vec()),
None => Bound::Unbounded,
};
let mut hits: Vec<ReferenceHit> = Vec::with_capacity(limit.min(64));
let mut hit_keys: Vec<Vec<u8>> = Vec::with_capacity(limit.min(64));
let mut hit_starts: Vec<u32> = Vec::with_capacity(limit.min(64));
let mut total: u32 = 0;
let mut total_is_partial = false;
let scan_cap = limit.saturating_mul(8).max(2_000);
let mut has_more = false;
let mut matched: usize = 0;
for guard in idx.calls_by_callee.range::<Vec<u8>, _>((lower, Bound::Unbounded)) {
let (k, _) = guard
.into_inner()
.map_err(|e| McpError::internal_error(format!("index iter: {e}"), None))?;
let Some((callee, rel, start)) = crate::index::keys::parse_call_by_callee(&k) else {
continue;
};
if finder.find(callee.as_bytes()).is_none() {
continue;
}
total += 1;
matched += 1;
if hits.len() < limit {
let (line, column) = resolve_call_line_col(idx, &rel, start);
hits.push(ReferenceHit {
path: rel,
line,
column,
callee,
resolved: None,
});
hit_keys.push(k.to_vec());
hit_starts.push(start);
} else {
has_more = true;
}
if matched >= scan_cap {
total_is_partial = true;
break;
}
}
let next_cursor = if has_more {
hit_keys.last().map(|k| Cursor::encode_fjall(k))
} else {
None
};
Ok(CallScanPage {
total,
total_is_partial,
hits,
next_cursor,
hit_keys,
hit_starts,
})
}
pub(super) fn scan_calls(
idx: Option<&crate::index::IndexDb>,
cache: &super::MapCache,
name: &str,
limit: usize,
cursor_after: Option<&[u8]>,
) -> Result<CallScanPage, McpError> {
match idx {
Some(idx) => scan_calls_by_name(idx, name, limit, cursor_after),
None => Ok(match cache.calls.as_ref() {
Some(calls) => scan_calls_in_ram(calls, name, limit, cursor_after),
None => empty_call_page(),
}),
}
}
fn empty_call_page() -> CallScanPage {
CallScanPage {
total: 0,
total_is_partial: false,
hits: Vec::new(),
next_cursor: None,
hit_keys: Vec::new(),
hit_starts: Vec::new(),
}
}
pub(super) fn scan_calls_in_ram(
index: &InRamCallIndex,
name: &str,
limit: usize,
cursor_after: Option<&[u8]>,
) -> CallScanPage {
let finder = memchr::memmem::Finder::new(name.as_bytes());
let start = match cursor_after {
Some(cursor) => index.entries.partition_point(|e| e.key.as_slice() <= cursor),
None => 0,
};
let mut hits: Vec<ReferenceHit> = Vec::with_capacity(limit.min(64));
let mut hit_keys: Vec<Vec<u8>> = Vec::with_capacity(limit.min(64));
let mut hit_starts: Vec<u32> = Vec::with_capacity(limit.min(64));
let mut total: u32 = 0;
let mut total_is_partial = false;
let scan_cap = limit.saturating_mul(8).max(2_000);
let mut has_more = false;
let mut matched: usize = 0;
for entry in &index.entries[start..] {
if finder.find(entry.callee.as_bytes()).is_none() {
continue;
}
total += 1;
matched += 1;
if hits.len() < limit {
hits.push(ReferenceHit {
path: entry.rel.clone(),
line: entry.line,
column: entry.column,
callee: entry.callee.clone(),
resolved: None,
});
hit_keys.push(entry.key.clone());
hit_starts.push(entry.start_byte);
} else {
has_more = true;
}
if matched >= scan_cap {
total_is_partial = true;
break;
}
}
let next_cursor = if has_more {
hit_keys.last().map(|k| Cursor::encode_fjall(k))
} else {
None
};
CallScanPage {
total,
total_is_partial,
hits,
next_cursor,
hit_keys,
hit_starts,
}
}
pub(crate) struct InRamCallIndex {
entries: Vec<InRamCall>,
by_path: ahash::AHashMap<crate::path::RelPath, Vec<CallRef>>,
}
struct InRamCall {
key: Vec<u8>,
callee: String,
rel: crate::path::RelPath,
start_byte: u32,
line: u32,
column: u32,
}
pub(crate) struct CallRef {
pub callee: String,
pub start_byte: u32,
pub line: u32,
pub column: u32,
}
impl InRamCallIndex {
pub(crate) fn build(store: &crate::store::Store) -> Self {
use rayon::prelude::*;
let per_file: Vec<(crate::path::RelPath, Vec<crate::extract::Call>)> = store
.index
.files
.par_iter()
.filter_map(|(rel, entry)| {
let calls = store.read_l2_by_hex(&entry.hash_hex).ok().flatten()?.calls;
Some((rel.clone(), calls))
})
.collect();
let mut entries: Vec<InRamCall> = Vec::new();
let mut by_path: ahash::AHashMap<crate::path::RelPath, Vec<CallRef>> =
ahash::AHashMap::with_capacity(per_file.len());
for (rel, calls) in per_file {
let mut refs: Vec<CallRef> = Vec::with_capacity(calls.len());
for call in calls {
if let Some(key) = crate::index::keys::call_by_callee(&call.callee, &rel, call.start_byte) {
entries.push(InRamCall {
key,
callee: call.callee.clone(),
rel: rel.clone(),
start_byte: call.start_byte,
line: call.start_row + 1,
column: call.start_col,
});
}
refs.push(CallRef {
line: call.start_row + 1,
column: call.start_col,
callee: call.callee,
start_byte: call.start_byte,
});
}
by_path.insert(rel, refs);
}
entries.sort_unstable_by(|a, b| a.key.cmp(&b.key));
Self { entries, by_path }
}
pub(crate) fn callers_of<'a>(
&'a self,
name: &'a str,
) -> impl Iterator<Item = (&'a crate::path::RelPath, u32)> + 'a {
self.entries
.iter()
.filter(move |c| c.callee == name)
.map(|c| (&c.rel, c.start_byte))
}
pub(crate) fn calls_in_file(&self, rel: &crate::path::RelPath) -> &[CallRef] {
self.by_path.get(rel).map_or(&[], Vec::as_slice)
}
}