use std::collections::HashMap;
use std::hash::BuildHasher;
use crate::ast::function::Function;
use crate::ast::statement::Statement;
use crate::meta::origin::Origin;
use crate::module::Module;
pub fn remap_origins_with_anchors<S: BuildHasher>(
target: &mut Module,
anchors: &HashMap<String, Origin, S>,
) {
if anchors.is_empty() {
return;
}
remap_block_with_anchors(&mut target.body.statements, anchors);
for symbol in &mut target.symbols {
if let Statement::FunctionDecl(function) = &mut symbol.statement {
remap_function_with_anchor(function, anchors);
}
}
}
fn remap_block_with_anchors<S: BuildHasher>(
statements: &mut [Statement],
anchors: &HashMap<String, Origin, S>,
) {
for statement in statements {
match statement {
Statement::FunctionDecl(function) => remap_function_with_anchor(function, anchors),
Statement::StructDecl(struct_decl) => {
for method in &mut struct_decl.methods {
remap_function_with_anchor(method, anchors);
}
}
Statement::EnumDecl(enum_decl) => {
for method in &mut enum_decl.methods {
remap_function_with_anchor(method, anchors);
}
}
_ => {}
}
}
}
fn remap_function_with_anchor<S: BuildHasher>(
function: &mut Function,
anchors: &HashMap<String, Origin, S>,
) {
let Some(anchor) = anchors.get(&function.name) else {
return;
};
let file = anchor.file.clone();
let line = anchor.line;
for origin in collect_origins_mut(&mut function.body.statements) {
*origin = Origin::new(file.clone(), line, origin.column);
}
}
fn collect_origins_mut(statements: &mut [Statement]) -> Vec<&mut Origin> {
let mut out = Vec::new();
walk_origins_mut(statements, &mut out);
out
}
fn walk_origins_mut<'a>(statements: &'a mut [Statement], out: &mut Vec<&'a mut Origin>) {
for statement in statements {
match statement {
Statement::SourceOrigin { origin } => out.push(origin),
Statement::FunctionDecl(function) => {
walk_origins_mut(&mut function.body.statements, out);
}
Statement::StructDecl(struct_decl) => {
for method in &mut struct_decl.methods {
walk_origins_mut(&mut method.body.statements, out);
}
}
Statement::EnumDecl(enum_decl) => {
for method in &mut enum_decl.methods {
walk_origins_mut(&mut method.body.statements, out);
}
}
Statement::Conditional {
then_block,
else_block,
..
} => {
walk_origins_mut(then_block, out);
walk_origins_mut(else_block, out);
}
Statement::ForIn { body, .. }
| Statement::ForNumeric { body, .. }
| Statement::While { body, .. } => walk_origins_mut(body, out),
Statement::VariableDecl { value, .. }
| Statement::Assignment { value, .. }
| Statement::Return(Some(value))
| Statement::Expr(value) => walk_origins_in_expr_mut(value, out),
_ => {}
}
}
}
fn walk_origins_in_expr_mut<'a>(
expr: &'a mut crate::ast::expression::Expression,
out: &mut Vec<&'a mut Origin>,
) {
use crate::ast::expression::Expression;
match expr {
Expression::Closure { body, .. } => walk_origins_mut(&mut body.statements, out),
Expression::Call { func, args, .. } => {
walk_origins_in_expr_mut(func, out);
for arg in args {
walk_origins_in_expr_mut(arg, out);
}
}
Expression::MethodCall { receiver, args, .. }
| Expression::DynMethodCall { receiver, args, .. } => {
walk_origins_in_expr_mut(receiver, out);
for arg in args {
walk_origins_in_expr_mut(arg, out);
}
}
Expression::FieldAccess { base, .. } | Expression::Not(base) | Expression::Len(base) => {
walk_origins_in_expr_mut(base, out);
}
Expression::BinaryOp { lhs, rhs, .. } => {
walk_origins_in_expr_mut(lhs, out);
walk_origins_in_expr_mut(rhs, out);
}
Expression::FormatConcat { parts } | Expression::Array { elements: parts } => {
for part in parts {
walk_origins_in_expr_mut(part, out);
}
}
Expression::Index { base, key } => {
walk_origins_in_expr_mut(base, out);
walk_origins_in_expr_mut(key, out);
}
Expression::If {
condition,
then_expr,
else_expr,
} => {
walk_origins_in_expr_mut(condition, out);
walk_origins_in_expr_mut(then_expr, out);
walk_origins_in_expr_mut(else_expr, out);
}
Expression::StructLiteral { fields, .. } | Expression::EnumLiteral { fields, .. } => {
for (_, value) in fields {
walk_origins_in_expr_mut(value, out);
}
}
Expression::FatPointer { data, .. } => walk_origins_in_expr_mut(data, out),
_ => {}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used)]
use super::*;
use crate::ast::block::Block;
use crate::ast::expression::Expression;
use crate::ast::function::Function;
use crate::ast::literal::Literal;
use crate::ast::scope::Scope;
use crate::module::{Module, Symbol, stage::Stage};
use std::sync::Arc;
fn origin(file: &str, line: u32) -> Origin {
Origin::new(Arc::<str>::from(file), line, 1)
}
fn fn_with_origins(name: &str, file: &str, lines: &[u32]) -> Function {
let mut statements = Vec::new();
for &line in lines {
statements.push(Statement::SourceOrigin {
origin: origin(file, line),
});
statements.push(Statement::Expr(Expression::Literal(Literal::Nil)));
}
Function {
name: name.into(),
params: vec![],
body: Block { statements },
doc: None,
debug: None,
event: None,
event_filter: None,
export: None,
inline: false,
}
}
#[test]
fn shifts_function_origins_to_anchor() {
let mut target = Module {
name: "control".into(),
stage: Stage::Control,
body: Block { statements: vec![] },
imports: vec![],
submodules: vec![],
locales: vec![],
pending_locales: vec![],
vtables: vec![],
symbols: vec![Symbol {
scope: Scope::Public,
statement: Statement::FunctionDecl(fn_with_origins(
"on_tick",
"src/control.rs",
&[126, 127],
)),
}],
};
let mut anchors = HashMap::new();
anchors.insert("on_tick".into(), origin("src/lib.rs", 139));
remap_origins_with_anchors(&mut target, &anchors);
let Statement::FunctionDecl(function) = &target.symbols[0].statement else {
panic!("expected function");
};
let mut lines = Vec::new();
let mut file = String::new();
for statement in &function.body.statements {
if let Statement::SourceOrigin { origin } = statement {
file = origin.file.to_string();
lines.push(origin.line);
}
}
assert_eq!(file, "src/lib.rs");
assert_eq!(lines, vec![139, 139]);
}
}