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};
use super::types_impls::{FindImplementationsParams, FindImplementationsResponse, ImplementationHit};
pub(super) fn run_find_implementations(
idx: Option<&crate::index::IndexDb>,
params: FindImplementationsParams,
cache: &super::MapCache,
) -> Result<CallToolResult, McpError> {
let limit = params.limit.unwrap_or(SEARCH_LIMIT_DEFAULT).min(SEARCH_LIMIT_MAX) as usize;
let Some(idx) = idx else {
return find_implementations_in_ram(cache, params, limit);
};
let cursor_bytes = params.cursor.as_ref().map(|c| c.decode_fjall()).transpose()?;
let finder = memchr::memmem::Finder::new(params.trait_name.as_bytes());
let lower: Bound<Vec<u8>> = match cursor_bytes.as_deref() {
Some(k) => Bound::Excluded(k.to_vec()),
None => Bound::Unbounded,
};
let scan_cap = limit.saturating_mul(8).max(2_000);
let mut hits: Vec<ImplementationHit> = Vec::with_capacity(limit.min(64));
let mut hit_keys: Vec<Vec<u8>> = Vec::with_capacity(limit.min(64));
let mut total: usize = 0;
let mut total_is_partial = false;
let mut last_emitted_key: Option<Vec<u8>> = None;
let mut has_more = false;
let mut matched: usize = 0;
for guard in idx
.implementations_by_trait
.range::<Vec<u8>, _>((lower, Bound::Unbounded))
{
let (k, _) = guard
.into_inner()
.map_err(|e| McpError::internal_error(format!("impl index iter: {e}"), None))?;
let Some((trait_name, impl_type, rel, start_byte)) = crate::index::keys::parse_impl_by_trait(&k) else {
continue;
};
if finder.find(trait_name.as_bytes()).is_none() {
continue;
}
if let Some(lang_filter) = params.language.as_deref() {
let l1_lang = cache.by_path.get(&rel).map(|l1| l1.language.as_str());
if l1_lang != Some(lang_filter) {
continue;
}
}
total += 1;
matched += 1;
if hits.len() < limit {
let (start_row, start_col) = resolve_impl_row_col(cache, &rel, start_byte);
hits.push(ImplementationHit {
path: rel,
trait_name,
impl_type,
start_row,
start_col,
});
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
};
let budget = super::budget::apply_budget(hits, params.max_tokens);
let (hits, budgeted, next_cursor) = if budget.budgeted {
let kept = budget.items.len();
let cursor = hit_keys.get(kept - 1).map(|k| Cursor::encode_fjall(k));
(budget.items, true, cursor)
} else {
(budget.items, false, next_cursor)
};
json_result(&FindImplementationsResponse {
trait_name: params.trait_name,
total,
total_is_partial,
budgeted,
hits,
next_cursor,
})
}
fn resolve_impl_row_col(cache: &super::MapCache, rel: &crate::path::RelPath, start_byte: u32) -> (u32, u32) {
let Some(l1) = cache.by_path.get(rel) else {
return (0, 0);
};
if let Some(imp) = l1.implementations.iter().find(|i| i.start_byte == start_byte) {
(imp.start_row + 1, imp.start_col)
} else {
(0, 0)
}
}
fn find_implementations_in_ram(
cache: &super::MapCache,
params: FindImplementationsParams,
limit: usize,
) -> Result<CallToolResult, McpError> {
let Some(index) = cache.impls.as_ref() else {
return json_result(&FindImplementationsResponse {
trait_name: params.trait_name,
total: 0,
total_is_partial: false,
budgeted: false,
hits: Vec::new(),
next_cursor: None,
});
};
let cursor_bytes = params.cursor.as_ref().map(|c| c.decode_fjall()).transpose()?;
let finder = memchr::memmem::Finder::new(params.trait_name.as_bytes());
let start = match cursor_bytes.as_deref() {
Some(c) => index.entries.partition_point(|e| e.key.as_slice() <= c),
None => 0,
};
let scan_cap = limit.saturating_mul(8).max(2_000);
let mut hits: Vec<ImplementationHit> = Vec::with_capacity(limit.min(64));
let mut hit_keys: Vec<Vec<u8>> = Vec::with_capacity(limit.min(64));
let mut total: usize = 0;
let mut total_is_partial = false;
let mut has_more = false;
let mut matched: usize = 0;
for entry in &index.entries[start..] {
if finder.find(entry.trait_name.as_bytes()).is_none() {
continue;
}
if let Some(lang) = params.language.as_deref()
&& cache.by_path.get(&entry.rel).map(|l1| l1.language.as_str()) != Some(lang)
{
continue;
}
total += 1;
matched += 1;
if hits.len() < limit {
hits.push(ImplementationHit {
path: entry.rel.clone(),
trait_name: entry.trait_name.clone(),
impl_type: entry.impl_type.clone(),
start_row: entry.start_row,
start_col: entry.start_col,
});
hit_keys.push(entry.key.clone());
} 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
};
let budget = super::budget::apply_budget(hits, params.max_tokens);
let (hits, budgeted, next_cursor) = if budget.budgeted {
let kept = budget.items.len();
(
budget.items,
true,
hit_keys.get(kept - 1).map(|k| Cursor::encode_fjall(k)),
)
} else {
(budget.items, false, next_cursor)
};
json_result(&FindImplementationsResponse {
trait_name: params.trait_name,
total,
total_is_partial,
budgeted,
hits,
next_cursor,
})
}
pub(crate) struct InRamImplIndex {
entries: Vec<InRamImpl>,
}
struct InRamImpl {
key: Vec<u8>,
trait_name: String,
impl_type: String,
rel: crate::path::RelPath,
start_row: u32,
start_col: u32,
}
impl InRamImplIndex {
pub(crate) fn build(by_path: &std::collections::BTreeMap<crate::path::RelPath, crate::extract::FileMapL1>) -> Self {
let mut entries: Vec<InRamImpl> = Vec::new();
for (rel, l1) in by_path {
for imp in &l1.implementations {
if let Some(key) =
crate::index::keys::impl_by_trait(&imp.trait_name, &imp.impl_type, rel, imp.start_byte)
{
entries.push(InRamImpl {
key,
trait_name: imp.trait_name.clone(),
impl_type: imp.impl_type.clone(),
rel: rel.clone(),
start_row: imp.start_row + 1,
start_col: imp.start_col,
});
}
}
}
entries.sort_unstable_by(|a, b| a.key.cmp(&b.key));
Self { entries }
}
}