use std::ops::Bound;
use rmcp::ErrorData as McpError;
use rmcp::model::CallToolResult;
use super::cursor::Cursor;
use super::helpers::{
SEARCH_LIMIT_DEFAULT, SEARCH_LIMIT_MAX, json_result, kind_to_str, parse_kind,
};
use super::types::ReferenceHit;
pub(super) fn resolve_call_line_col(
idx: &crate::index::IndexDb,
rel: &crate::path::RelPath,
start_byte: u32,
) -> (u32, u32) {
let key = crate::index::keys::call_by_path(rel, start_byte);
let value = match idx.calls_by_path.get(key) {
Ok(Some(v)) => v,
_ => return (0, 0),
};
let call: crate::extract::Call = match rmp_serde::from_slice(&value) {
Ok(c) => c,
Err(_) => return (0, 0),
};
(call.start_row + 1, call.start_col)
}
pub(super) fn run_find_references(
idx: Option<&crate::index::IndexDb>,
params: super::types::FindReferencesParams,
) -> Result<CallToolResult, McpError> {
use super::types::FindReferencesResponse;
let format = super::toon::ResponseFormat::parse(params.format.as_deref());
let limit = params
.limit
.unwrap_or(SEARCH_LIMIT_DEFAULT)
.min(SEARCH_LIMIT_MAX) as usize;
let Some(idx) = idx else {
return super::toon::format_result(
&FindReferencesResponse {
name: params.name,
total: 0,
total_is_partial: false,
budgeted: false,
hits: Vec::new(),
next_cursor: None,
},
format,
);
};
let cursor_bytes = params
.cursor
.as_ref()
.map(|c| c.decode_fjall())
.transpose()?;
let scan = scan_calls_by_name(idx, ¶ms.name, limit, cursor_bytes.as_deref())?;
let total = scan.total;
let total_is_partial = scan.total_is_partial;
let budgeted = budget_call_page(scan, params.max_tokens);
super::toon::format_result(
&FindReferencesResponse {
name: params.name,
total,
total_is_partial,
budgeted: budgeted.budgeted,
hits: budgeted.hits,
next_cursor: budgeted.next_cursor,
},
format,
)
}
pub(super) fn run_find_callers(
idx: Option<&crate::index::IndexDb>,
params: super::types::FindCallersParams,
cache: &super::MapCache,
) -> Result<CallToolResult, McpError> {
use super::types::{DefinitionView, FindCallersResponse};
let limit = params
.limit
.unwrap_or(SEARCH_LIMIT_DEFAULT)
.min(SEARCH_LIMIT_MAX) as usize;
let kind_filter = params.kind.as_deref().map(parse_kind).transpose()?;
let Some(idx) = idx else {
return json_result(&FindCallersResponse {
definition: None,
total: 0,
total_is_partial: false,
budgeted: false,
hits: Vec::new(),
next_cursor: None,
});
};
let definition: Option<DefinitionView> = cache
.by_path
.get(¶ms.path)
.and_then(|l1| {
l1.symbols
.iter()
.find(|s| s.name == params.name && kind_filter.is_none_or(|k| s.kind == k))
})
.map(|sym| DefinitionView {
path: params.path.clone(),
name: sym.name.clone(),
kind: kind_to_str(sym.kind),
start_row: sym.start_row,
start_col: sym.start_col,
});
let cursor_bytes = params
.cursor
.as_ref()
.map(|c| c.decode_fjall())
.transpose()?;
let scan = scan_calls_by_name(idx, ¶ms.name, limit, cursor_bytes.as_deref())?;
let total = scan.total;
let total_is_partial = scan.total_is_partial;
let budgeted = budget_call_page(scan, params.max_tokens);
json_result(&FindCallersResponse {
definition,
total,
total_is_partial,
budgeted: budgeted.budgeted,
hits: budgeted.hits,
next_cursor: budgeted.next_cursor,
})
}
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(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 total: u32 = 0;
let mut total_is_partial = false;
let scan_cap = limit.saturating_mul(8).max(2_000);
let mut last_emitted_key: Option<Vec<u8>> = None;
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,
});
hit_keys.push(k.to_vec());
last_emitted_key = Some(k.to_vec());
} else {
has_more = true;
}
if matched >= scan_cap {
total_is_partial = true;
break;
}
}
let next_cursor = if has_more {
last_emitted_key.as_deref().map(Cursor::encode_fjall)
} else {
None
};
Ok(CallScanPage {
total,
total_is_partial,
hits,
next_cursor,
hit_keys,
})
}