use anyhow::Result;
use serde::Serialize;
use crate::{Filter, Hit, Prefer, fleet::Fleet, store::EdgeEnd};
pub const DEFAULT_BUDGET: i64 = 4000;
pub const SPANS: usize = 8;
const EDGES_PER_SYMBOL: usize = 8;
const TEXT_SPANS: usize = 1;
#[derive(Debug, Clone, Serialize)]
pub struct Brief {
pub question: String,
pub budget: i64,
pub tokens: i64,
pub counted_with: &'static str,
pub floor: i64,
pub spans: Vec<Span>,
pub symbols: Vec<Symbol>,
pub cut: Cut,
}
#[derive(Debug, Clone, Serialize)]
pub struct Span {
#[serde(serialize_with = "crate::serialize_plain")]
pub path: String,
pub start_line: u32,
pub end_line: u32,
#[serde(skip_serializing_if = "Vec::is_empty")]
pub lists: Vec<&'static str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol_kind: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub symbol_line: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
pub fresh: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub text: Option<String>,
#[serde(skip_serializing_if = "std::ops::Not::not")]
pub over_budget: bool,
}
#[derive(Debug, Clone, Serialize)]
pub struct Symbol {
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
pub found_by: &'static str,
pub callers: Vec<EdgeEnd>,
pub callees: Vec<EdgeEnd>,
#[serde(skip_serializing_if = "is_zero")]
pub hidden: usize,
}
#[derive(Debug, Clone, Default, Serialize)]
pub struct Cut {
#[serde(skip_serializing_if = "is_zero")]
pub spans: usize,
#[serde(skip_serializing_if = "is_zero")]
pub span_text: usize,
#[serde(skip_serializing_if = "is_zero")]
pub symbols: usize,
#[serde(skip_serializing_if = "is_zero")]
pub edges: usize,
}
fn is_zero(n: &usize) -> bool {
*n == 0
}
impl Cut {
pub fn is_empty(&self) -> bool {
self.spans == 0 && self.span_text == 0 && self.symbols == 0 && self.edges == 0
}
}
pub fn brief(
fleet: &mut Fleet,
only: Option<&[String]>,
question: &str,
budget: i64,
filter: &Filter,
prefer: Prefer,
) -> Result<Brief> {
brief_leaning(fleet, only, question, budget, filter, Some(prefer))
}
pub fn brief_leaning(
fleet: &mut Fleet,
only: Option<&[String]>,
question: &str,
budget: i64,
filter: &Filter,
prefer: Option<Prefer>,
) -> Result<Brief> {
let hits = fleet.search_leaning(only, question, SPANS, filter, prefer)?;
let counter = fleet.counter();
let counted_with = counter.label();
let mut spans: Vec<Span> = Vec::new();
let mut kept: Vec<usize> = Vec::new();
let mut tokens: i64 = 0;
let mut floor: i64 = 0;
let mut cut = Cut::default();
for (i, hit) in hits.iter().enumerate() {
let span = Span::locating(hit);
let cost = counter.count(&span.locator());
if spans.is_empty() {
floor = cost;
} else if tokens + cost > budget {
cut.spans += 1;
continue;
}
tokens += cost;
spans.push(span);
kept.push(i);
}
let code = crate::code_shaped(question);
let texted: usize = if code {
let product =
|i: usize| crate::filter::is_code(&hits[i].path) && !crate::is_test_path(&hits[i].path);
let about_tests = question.to_lowercase().contains("test");
kept.iter()
.position(|&i| product(i) || (about_tests && crate::filter::is_code(&hits[i].path)))
.or_else(|| {
kept.iter()
.position(|&i| crate::filter::is_code(&hits[i].path))
})
.unwrap_or(0)
} else {
0
};
let navigational = |kind: &str| crate::graph::NAVIGATIONAL_KINDS.contains(&kind);
let mut symbols: Vec<Symbol> = Vec::new();
let mut seen: Vec<String> = Vec::new();
for hit in kept.iter().skip(texted).take(TEXT_SPANS).map(|&i| &hits[i]) {
let Some(name) = hit.symbol.as_deref() else {
continue;
};
if code && hit.symbol_kind.as_deref().is_some_and(navigational) {
continue;
}
if seen.iter().any(|s| s == name) {
continue;
}
seen.push(name.to_string());
let own = hit.store.as_ref().map(std::slice::from_ref);
let mut found = fleet.neighbours_in(own.or(only), name, &[], false)?;
if code {
found.callers.retain(|e| !navigational(&e.symbol.kind));
found.callees.retain(|e| !navigational(&e.symbol.kind));
}
let (callers, hidden_callers) = cap(found.callers);
let (callees, hidden_callees) = cap(found.callees);
if callers.is_empty() && callees.is_empty() {
continue;
}
let symbol = Symbol {
name: name.to_string(),
store: hit.store.clone(),
found_by: "graph",
callers,
callees,
hidden: hidden_callers + hidden_callees,
};
let cost = counter.count(&symbol.rendered());
if tokens + cost > budget {
cut.symbols += 1;
cut.edges += symbol.callers.len() + symbol.callees.len();
continue;
}
tokens += cost;
symbols.push(symbol);
}
for (i, (span, hit)) in spans
.iter_mut()
.zip(kept.iter().map(|&i| &hits[i]))
.enumerate()
{
if i < texted || i >= texted + TEXT_SPANS {
continue;
}
let cost = counter.count(&hit.text);
if tokens + cost > budget {
cut.span_text += 1;
span.over_budget = true;
continue;
}
tokens += cost;
span.text = Some(hit.text.clone());
}
Ok(Brief {
question: question.to_string(),
budget,
tokens,
counted_with,
floor,
spans,
symbols,
cut,
})
}
fn cap(mut edges: Vec<EdgeEnd>) -> (Vec<EdgeEnd>, usize) {
let hidden = edges.len().saturating_sub(EDGES_PER_SYMBOL);
edges.truncate(EDGES_PER_SYMBOL);
(edges, hidden)
}
impl Span {
fn locating(hit: &Hit) -> Self {
Self {
path: hit.path.clone(),
start_line: hit.start_line,
end_line: hit.end_line,
lists: hit.lists.clone(),
symbol: hit.symbol.clone(),
symbol_kind: hit.symbol_kind.clone(),
symbol_line: hit.symbol_line,
over_budget: false,
store: hit.store.clone(),
fresh: hit.fresh,
text: None,
}
}
fn locator(&self) -> String {
let mut line = format!("{}:{}-{}", self.path, self.start_line, self.end_line);
if let Some(symbol) = &self.symbol {
line.push(' ');
line.push_str(symbol);
}
if let Some(kind) = &self.symbol_kind {
line.push(' ');
line.push_str(kind);
}
for list in &self.lists {
line.push(' ');
line.push_str(list);
}
line
}
}
impl Symbol {
fn rendered(&self) -> String {
let mut out = self.name.clone();
for edge in self.callers.iter().chain(&self.callees) {
out.push('\n');
out.push_str(&edge.symbol.name);
out.push(' ');
out.push_str(&edge.symbol.path);
out.push(' ');
out.push_str(&edge.kind);
out.push(' ');
out.push_str(&edge.confidence);
}
out
}
}