use std::collections::HashMap;
use std::path::Path;
use crate::builtins::BuiltinFunction;
use super::{
ParseError, SourceError, SourcePathError,
ir::{Expr, FrontendIr, FunctionDecl, FunctionImpl, LocalSlot, Stmt, StructDecl},
};
pub(super) struct ParsedUnit {
pub(super) parsed: FrontendIr,
pub(super) scope_prefix: Option<String>,
pub(super) source_name: String,
}
pub(super) fn sanitize_scope_prefix(path: &Path) -> String {
path.file_stem()
.and_then(|value| value.to_str())
.unwrap_or("module")
.chars()
.map(|ch| {
if ch.is_ascii_alphanumeric() || ch == '_' {
ch
} else {
'_'
}
})
.collect()
}
pub(super) fn merge_units(units: Vec<ParsedUnit>) -> Result<FrontendIr, SourcePathError> {
let mut merged_stmts = Vec::new();
let mut merged_stmt_sources = Vec::new();
let mut merged_local_bindings = Vec::new();
let mut merged_struct_schemas = HashMap::<String, StructDecl>::new();
let mut merged_unknown_type_spans = Vec::new();
let mut merged_functions = Vec::new();
let mut merged_function_impls = HashMap::<u16, FunctionImpl>::new();
let mut merged_function_sources = HashMap::<u16, String>::new();
let mut function_index_by_name = HashMap::<String, u16>::new();
let mut local_base = 0usize;
for unit in units {
let source_name = unit.source_name.clone();
let function_map = remap_functions(
&unit.parsed.functions,
&mut merged_functions,
&mut function_index_by_name,
)?;
let unit_local_base = local_base;
let unit_local_count = unit.parsed.locals;
let mut remapped_stmts = unit.parsed.stmts;
for stmt in &mut remapped_stmts {
remap_stmt_indices(stmt, unit_local_base, &function_map)?;
}
merged_stmt_sources.extend(std::iter::repeat_n(
Some(source_name.clone()),
remapped_stmts.len(),
));
merged_stmts.extend(remapped_stmts);
for (name, index) in unit.parsed.local_bindings {
let remapped_index = remap_local_index(index, unit_local_base)?;
let scoped_name = if let Some(prefix) = &unit.scope_prefix {
format!("{prefix}::{name}")
} else {
name
};
merged_local_bindings.push((scoped_name, remapped_index));
}
for (name, schema) in unit.parsed.struct_schemas {
if let Some(existing) = merged_struct_schemas.get(&name) {
if existing != &schema {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"struct schema '{}' declared with conflicting definitions across imported modules",
name
),
})));
}
continue;
}
merged_struct_schemas.insert(name, schema);
}
merged_unknown_type_spans.extend(unit.parsed.unknown_type_spans);
for (unit_index, mut function_impl) in unit.parsed.function_impls {
let merged_index = function_map.get(&unit_index).copied().ok_or_else(|| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "function implementation remap failed while merging imported modules"
.to_string(),
}))
})?;
for param_slot in &mut function_impl.param_slots {
*param_slot = remap_local_index(*param_slot, unit_local_base)?;
}
for (source_slot, captured_slot) in &mut function_impl.capture_copies {
*source_slot = remap_local_index(*source_slot, unit_local_base)?;
*captured_slot = remap_local_index(*captured_slot, unit_local_base)?;
}
for stmt in &mut function_impl.body_stmts {
remap_stmt_indices(stmt, unit_local_base, &function_map)?;
}
remap_expr_indices(&mut function_impl.body_expr, unit_local_base, &function_map)?;
if merged_function_impls
.insert(merged_index, function_impl)
.is_some()
{
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "duplicate RSS function implementation while merging imported modules"
.to_string(),
})));
}
merged_function_sources.insert(merged_index, source_name.clone());
}
local_base = local_base.checked_add(unit_local_count).ok_or_else(|| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "local count overflow while merging imported modules".to_string(),
}))
})?;
if local_base > (LocalSlot::MAX as usize + 1) {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "too many locals across imported modules".to_string(),
})));
}
}
Ok(FrontendIr {
stmts: merged_stmts,
locals: local_base,
local_bindings: merged_local_bindings,
struct_schemas: merged_struct_schemas,
unknown_type_spans: merged_unknown_type_spans,
functions: merged_functions,
function_impls: merged_function_impls,
stmt_sources: merged_stmt_sources,
function_sources: merged_function_sources,
})
}
fn remap_functions(
unit_functions: &[FunctionDecl],
merged_functions: &mut Vec<FunctionDecl>,
function_index_by_name: &mut HashMap<String, u16>,
) -> Result<HashMap<u16, u16>, SourcePathError> {
let mut map = HashMap::new();
for func in unit_functions {
let merged_index = if let Some(existing_index) = function_index_by_name.get(&func.name) {
let existing = &mut merged_functions[*existing_index as usize];
if existing.arity != func.arity {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"function '{}' declared with conflicting arity {} vs {}",
func.name, existing.arity, func.arity
),
})));
}
if existing.return_type != func.return_type {
match (existing.return_type, func.return_type) {
(crate::ValueType::Unknown, known) => existing.return_type = known,
(known, crate::ValueType::Unknown) => existing.return_type = known,
(lhs, rhs) => {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"function '{}' declared with conflicting return type {} vs {}",
func.name,
value_type_name(lhs),
value_type_name(rhs)
),
})));
}
}
}
if existing.return_schema != func.return_schema {
match (&existing.return_schema, &func.return_schema) {
(None, Some(schema)) => existing.return_schema = Some(schema.clone()),
(Some(_), None) => {}
(Some(lhs), Some(rhs)) if lhs == rhs => {}
_ => {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"function '{}' declared with conflicting return schemas across imported modules",
func.name
),
})));
}
}
}
if existing.type_params != func.type_params {
if existing.type_params.is_empty() {
existing.type_params = func.type_params.clone();
} else if !func.type_params.is_empty() {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"function '{}' declared with conflicting type parameters across imported modules",
func.name
),
})));
}
}
if existing.arg_schemas != func.arg_schemas {
if existing.arg_schemas.iter().all(Option::is_none) {
existing.arg_schemas = func.arg_schemas.clone();
} else if !func.arg_schemas.iter().all(Option::is_none) {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: format!(
"function '{}' declared with conflicting parameter schemas across imported modules",
func.name
),
})));
}
}
if function_args_are_placeholders(&existing.args)
&& !function_args_are_placeholders(&func.args)
{
existing.args = func.args.clone();
}
existing.exported = existing.exported || func.exported;
*existing_index
} else {
let next_index = u16::try_from(merged_functions.len()).map_err(|_| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "too many functions across imported modules".to_string(),
}))
})?;
merged_functions.push(FunctionDecl {
name: func.name.clone(),
arity: func.arity,
index: next_index,
args: func.args.clone(),
arg_schemas: func.arg_schemas.clone(),
return_schema: func.return_schema.clone(),
type_params: func.type_params.clone(),
exported: func.exported,
return_type: func.return_type,
});
function_index_by_name.insert(func.name.clone(), next_index);
next_index
};
map.insert(func.index, merged_index);
}
Ok(map)
}
fn function_args_are_placeholders(args: &[String]) -> bool {
args.iter()
.enumerate()
.all(|(index, arg)| arg == &format!("arg{index}"))
}
fn value_type_name(ty: crate::ValueType) -> &'static str {
match ty {
crate::ValueType::Unknown => "unknown",
crate::ValueType::Null => "null",
crate::ValueType::Int => "int",
crate::ValueType::Float => "float",
crate::ValueType::Bool => "bool",
crate::ValueType::String => "string",
crate::ValueType::Bytes => "bytes",
crate::ValueType::Array => "array",
crate::ValueType::Map => "map",
}
}
fn remap_local_index(index: LocalSlot, local_base: usize) -> Result<LocalSlot, SourcePathError> {
let remapped = (index as usize).checked_add(local_base).ok_or_else(|| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "local index overflow while merging imported modules".to_string(),
}))
})?;
LocalSlot::try_from(remapped).map_err(|_| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "local index overflow while merging imported modules".to_string(),
}))
})
}
fn remap_stmt_indices(
stmt: &mut Stmt,
local_base: usize,
function_map: &HashMap<u16, u16>,
) -> Result<(), SourcePathError> {
match stmt {
Stmt::Noop { .. } => {}
Stmt::Let { index, expr, .. } => {
*index = remap_local_index(*index, local_base)?;
remap_expr_indices(expr, local_base, function_map)?;
}
Stmt::Assign { index, expr, .. } => {
*index = remap_local_index(*index, local_base)?;
remap_expr_indices(expr, local_base, function_map)?;
}
Stmt::ClosureLet { closure, .. } => {
for (source_index, captured_slot) in &mut closure.capture_copies {
*source_index = remap_local_index(*source_index, local_base)?;
*captured_slot = remap_local_index(*captured_slot, local_base)?;
}
remap_expr_indices(&mut closure.body, local_base, function_map)?;
}
Stmt::FuncDecl { index, .. } => {
*index = function_map.get(index).copied().ok_or_else(|| {
SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "function index remap failed while merging imported modules"
.to_string(),
}))
})?;
}
Stmt::Expr { expr, .. } => {
remap_expr_indices(expr, local_base, function_map)?;
}
Stmt::IfElse {
condition,
then_branch,
else_branch,
..
} => {
remap_expr_indices(condition, local_base, function_map)?;
for stmt in then_branch {
remap_stmt_indices(stmt, local_base, function_map)?;
}
for stmt in else_branch {
remap_stmt_indices(stmt, local_base, function_map)?;
}
}
Stmt::For {
init,
condition,
post,
body,
..
} => {
remap_stmt_indices(init, local_base, function_map)?;
remap_expr_indices(condition, local_base, function_map)?;
remap_stmt_indices(post, local_base, function_map)?;
for stmt in body {
remap_stmt_indices(stmt, local_base, function_map)?;
}
}
Stmt::While {
condition, body, ..
} => {
remap_expr_indices(condition, local_base, function_map)?;
for stmt in body {
remap_stmt_indices(stmt, local_base, function_map)?;
}
}
Stmt::Break { .. } | Stmt::Continue { .. } => {}
Stmt::Drop { index, .. } => {
*index = remap_local_index(*index, local_base)?;
}
}
Ok(())
}
fn remap_expr_indices(
expr: &mut Expr,
local_base: usize,
function_map: &HashMap<u16, u16>,
) -> Result<(), SourcePathError> {
match expr {
Expr::Null
| Expr::Int(_)
| Expr::Float(_)
| Expr::Bool(_)
| Expr::Bytes(_)
| Expr::String(_) => {}
Expr::FunctionRef(index) => {
if let Some(remapped_index) = function_map.get(index).copied() {
*index = remapped_index;
} else if BuiltinFunction::from_call_index(*index).is_none() {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "function index remap failed while merging imported modules"
.to_string(),
})));
}
}
Expr::Call(index, _, args) => {
if let Some(remapped_index) = function_map.get(index).copied() {
*index = remapped_index;
} else if BuiltinFunction::from_call_index(*index).is_none() {
return Err(SourcePathError::Source(SourceError::Parse(ParseError {
span: None,
code: None,
line: 1,
message: "function index remap failed while merging imported modules"
.to_string(),
})));
}
for arg in args {
remap_expr_indices(arg, local_base, function_map)?;
}
}
Expr::OptionalGet {
container,
key,
container_slot,
key_slot,
} => {
*container_slot = remap_local_index(*container_slot, local_base)?;
*key_slot = remap_local_index(*key_slot, local_base)?;
remap_expr_indices(container, local_base, function_map)?;
remap_expr_indices(key, local_base, function_map)?;
}
Expr::OptionUnwrapOr {
value,
value_slot,
fallback,
} => {
*value_slot = remap_local_index(*value_slot, local_base)?;
remap_expr_indices(value, local_base, function_map)?;
remap_expr_indices(fallback, local_base, function_map)?;
}
Expr::LocalCall(index, _, args) => {
*index = remap_local_index(*index, local_base)?;
for arg in args {
remap_expr_indices(arg, local_base, function_map)?;
}
}
Expr::Closure(closure) => {
for param_slot in &mut closure.param_slots {
*param_slot = remap_local_index(*param_slot, local_base)?;
}
for (source_index, captured_slot) in &mut closure.capture_copies {
*source_index = remap_local_index(*source_index, local_base)?;
*captured_slot = remap_local_index(*captured_slot, local_base)?;
}
remap_expr_indices(&mut closure.body, local_base, function_map)?;
}
Expr::ClosureCall(closure, args) => {
for param_slot in &mut closure.param_slots {
*param_slot = remap_local_index(*param_slot, local_base)?;
}
for (source_index, captured_slot) in &mut closure.capture_copies {
*source_index = remap_local_index(*source_index, local_base)?;
*captured_slot = remap_local_index(*captured_slot, local_base)?;
}
remap_expr_indices(&mut closure.body, local_base, function_map)?;
for arg in args {
remap_expr_indices(arg, local_base, function_map)?;
}
}
Expr::Add(lhs, rhs)
| Expr::Sub(lhs, rhs)
| Expr::Mul(lhs, rhs)
| Expr::Div(lhs, rhs)
| Expr::Mod(lhs, rhs)
| Expr::And(lhs, rhs)
| Expr::Or(lhs, rhs)
| Expr::Eq(lhs, rhs)
| Expr::Lt(lhs, rhs)
| Expr::Gt(lhs, rhs) => {
remap_expr_indices(lhs, local_base, function_map)?;
remap_expr_indices(rhs, local_base, function_map)?;
}
Expr::Neg(inner)
| Expr::Not(inner)
| Expr::ToOwned(inner)
| Expr::Borrow(inner)
| Expr::BorrowMut(inner) => {
remap_expr_indices(inner, local_base, function_map)?;
}
Expr::Var(index) | Expr::MoveVar(index) => {
*index = remap_local_index(*index, local_base)?;
}
Expr::MoveField { root, .. } | Expr::MoveIndex { root, .. } => {
*root = remap_local_index(*root, local_base)?;
}
Expr::IfElse {
condition,
then_expr,
else_expr,
} => {
remap_expr_indices(condition, local_base, function_map)?;
remap_expr_indices(then_expr, local_base, function_map)?;
remap_expr_indices(else_expr, local_base, function_map)?;
}
Expr::Match {
value_slot,
result_slot,
value,
arms,
default,
} => {
*value_slot = remap_local_index(*value_slot, local_base)?;
*result_slot = remap_local_index(*result_slot, local_base)?;
remap_expr_indices(value, local_base, function_map)?;
for (pattern, arm_expr) in arms {
if let crate::compiler::ir::MatchPattern::SomeBinding(binding_slot) = pattern {
*binding_slot = remap_local_index(*binding_slot, local_base)?;
}
remap_expr_indices(arm_expr, local_base, function_map)?;
}
remap_expr_indices(default, local_base, function_map)?;
}
Expr::Block { stmts, expr } => {
for stmt in stmts {
remap_stmt_indices(stmt, local_base, function_map)?;
}
remap_expr_indices(expr, local_base, function_map)?;
}
}
Ok(())
}