use std::collections::{BTreeMap, BTreeSet};
use harn_lexer::{FixEdit, Span};
use harn_parser::{
visit, BindingPattern, DiagnosticCode as Code, Node, Repair, SNode, TypeExpr, TypedParam,
};
use super::capability_migrations::collect_callable_node_calls;
use super::CallableInfo;
pub(super) fn collect_callable_infos(
program: &[SNode],
source: &str,
exported_names: &BTreeSet<String>,
) -> Vec<CallableInfo> {
let mut infos = Vec::new();
for node in program {
let inner = match &node.node {
Node::AttributedDecl { inner, .. } => inner.as_ref(),
_ => node,
};
match &inner.node {
Node::FnDecl {
name,
params,
body,
is_pub,
..
}
| Node::ToolDecl {
name,
params,
body,
is_pub,
..
} => {
let mut calls = Vec::new();
let mut ambient_capability_calls = Vec::new();
visit_callable_body(inner, &mut |child| {
collect_callable_node_calls(
child,
source,
&mut calls,
&mut ambient_capability_calls,
);
});
let Some((insert_offset, has_params)) = callable_param_insert(source, inner.span)
else {
continue;
};
let bound_names = callable_bound_names(params, body);
infos.push(CallableInfo {
name: name.clone(),
span: inner.span,
is_exported: *is_pub || exported_names.contains(name),
insert_offset,
has_params: has_params || !params.is_empty(),
bound_names,
harness_binding: harness_param_name(params).map(str::to_string),
can_add_harness_param: true,
calls,
ambient_capability_calls,
});
}
Node::Pipeline {
name,
params,
body,
is_pub,
..
} => {
let mut calls = Vec::new();
let mut ambient_capability_calls = Vec::new();
visit_callable_body(inner, &mut |child| {
collect_callable_node_calls(
child,
source,
&mut calls,
&mut ambient_capability_calls,
);
});
let Some((insert_offset, has_params)) = callable_param_insert(source, inner.span)
else {
continue;
};
let bound_names = callable_bound_names(params, body);
infos.push(CallableInfo {
name: name.clone(),
span: inner.span,
is_exported: *is_pub || exported_names.contains(name),
insert_offset,
has_params: has_params || !params.is_empty(),
bound_names,
harness_binding: harness_param_name(params).map(str::to_string),
can_add_harness_param: true,
calls,
ambient_capability_calls,
});
}
_ => {}
}
}
infos
}
fn callable_bound_names(params: &[TypedParam], body: &[SNode]) -> BTreeSet<String> {
let mut names = params
.iter()
.map(|param| param.name.clone())
.collect::<BTreeSet<_>>();
collect_binding_names(body, &mut names);
names
}
fn collect_binding_names(nodes: &[SNode], names: &mut BTreeSet<String>) {
for node in nodes {
visit::walk_node(node, &mut |child| match &child.node {
Node::LetBinding { pattern, .. } | Node::ConstBinding { pattern, .. } => {
collect_pattern_names(pattern, names);
}
Node::ForIn { pattern, .. } => {
collect_pattern_names(pattern, names);
}
Node::Parallel {
variable: Some(name),
..
} => {
names.insert(name.clone());
}
Node::TryCatch {
error_var: Some(name),
..
} => {
names.insert(name.clone());
}
Node::Closure { params, .. } => {
names.extend(params.iter().map(|param| param.name.clone()));
}
_ => {}
});
}
}
fn collect_pattern_names(pattern: &BindingPattern, names: &mut BTreeSet<String>) {
match pattern {
BindingPattern::Identifier(name) => {
names.insert(name.clone());
}
BindingPattern::Dict(fields) => {
names.extend(
fields
.iter()
.map(|field| field.alias.as_ref().unwrap_or(&field.key).clone()),
);
}
BindingPattern::List(elements) => {
names.extend(elements.iter().map(|element| element.name.clone()));
}
BindingPattern::Pair(left, right) => {
names.insert(left.clone());
names.insert(right.clone());
}
}
}
fn visit_callable_body(node: &SNode, visitor: &mut impl FnMut(&SNode)) {
let (params, body) = match &node.node {
Node::FnDecl { params, body, .. }
| Node::ToolDecl { params, body, .. }
| Node::Pipeline { params, body, .. } => (params, body),
_ => return,
};
for param in params {
if let Some(default) = ¶m.default_value {
visit::walk_node(default, visitor);
}
}
for stmt in body {
visit::walk_node(stmt, visitor);
}
}
pub(super) fn callable_param_insert(source: &str, span: Span) -> Option<(usize, bool)> {
let region = source.get(span.start..span.end)?;
let open_paren = region.find('(')?;
let mut depth = 0usize;
let mut quote = None;
let mut escaped = false;
let mut close_paren = None;
for (offset, ch) in region[open_paren..].char_indices() {
if let Some(delimiter) = quote {
if escaped {
escaped = false;
} else if ch == '\\' {
escaped = true;
} else if ch == delimiter {
quote = None;
}
continue;
}
if matches!(ch, '"' | '\'') {
quote = Some(ch);
continue;
}
match ch {
'(' => depth += 1,
')' => {
depth = depth.checked_sub(1)?;
if depth == 0 {
close_paren = Some(open_paren + offset);
break;
}
}
_ => {}
}
}
let close_paren = close_paren?;
let has_params = !region[open_paren + 1..close_paren].trim().is_empty();
Some((span.start + open_paren + 1, has_params))
}
fn harness_param_name(params: &[TypedParam]) -> Option<&str> {
params.iter().find_map(|param| {
let TypeExpr::Named(name) = param.type_expr.as_ref()? else {
return None;
};
if name == "Harness" && matches!(param.name.as_str(), "harness" | "_harness") {
Some(param.name.as_str())
} else {
None
}
})
}
pub(super) fn build_reverse_callers(infos: &[CallableInfo]) -> Vec<Vec<(usize, usize)>> {
let by_name = infos
.iter()
.enumerate()
.map(|(idx, info)| (info.name.as_str(), idx))
.collect::<BTreeMap<_, _>>();
let mut reverse = vec![Vec::new(); infos.len()];
for (caller_idx, info) in infos.iter().enumerate() {
for (call_idx, call) in info.calls.iter().enumerate() {
let Some(&callee_idx) = by_name.get(call.callee.as_str()) else {
continue;
};
reverse[callee_idx].push((caller_idx, call_idx));
}
}
reverse
}
pub(super) fn propagate_harness_requirements(
infos: &[CallableInfo],
reverse_callers: &[Vec<(usize, usize)>],
owner_idx: usize,
) -> BTreeSet<usize> {
let mut needed = BTreeSet::from([owner_idx]);
let mut changed = true;
while changed {
changed = false;
let snapshot = needed.iter().copied().collect::<Vec<_>>();
for callee_idx in snapshot {
for &(caller_idx, _) in &reverse_callers[callee_idx] {
if infos[caller_idx].harness_binding.is_none()
&& infos[caller_idx].can_add_harness_param
&& needed.insert(caller_idx)
{
changed = true;
}
}
}
}
needed
}
pub(super) fn repair_for_ambient_capability_plan(
code: Code,
infos: &[CallableInfo],
reverse_callers: &[Vec<(usize, usize)>],
needed: &BTreeSet<usize>,
) -> Option<Repair> {
let surface_changing = needed.iter().any(|&idx| {
let info = &infos[idx];
info.is_exported || info.name == "main" || reverse_callers[idx].is_empty()
});
if surface_changing {
Some(Repair::from_template(
Code::InvalidMainSignature.repair_template()?,
))
} else {
Some(Repair::from_template(code.repair_template()?))
}
}
pub(super) fn add_harness_param_edit(info: &CallableInfo) -> Option<FixEdit> {
let name = harness_param_name_for_insert(info)?;
Some(FixEdit {
span: Span::with_offsets(
info.insert_offset,
info.insert_offset,
info.span.line,
info.span.column,
),
replacement: if info.has_params {
format!("{name}: Harness, ")
} else {
format!("{name}: Harness")
},
})
}
pub(super) fn harness_param_name_for_insert(info: &CallableInfo) -> Option<&'static str> {
if !info.bound_names.contains("harness") {
return Some("harness");
}
if !info.bound_names.contains("_harness") {
return Some("_harness");
}
None
}
pub(super) fn add_call_argument_edit(source: &str, span: &Span, arg_name: &str) -> Option<FixEdit> {
let region = source.get(span.start..span.end)?;
let open_paren = region.find('(')?;
let close_paren = region[open_paren + 1..].find(')')? + open_paren + 1;
let has_args = !region[open_paren + 1..close_paren].trim().is_empty();
let insert_at = span.start + open_paren + 1;
Some(FixEdit {
span: Span::with_offsets(insert_at, insert_at, span.line, span.column),
replacement: if has_args {
format!("{arg_name}, ")
} else {
arg_name.to_string()
},
})
}