use std::rc::Rc;
use hamelin_lib::tree::{
ast::{identifier::SimpleIdentifier, node::Span},
typed_ast::expression::TypedExpression,
};
use ordermap::OrderMap;
use crate::IRExpression;
#[derive(Default)]
pub enum AssignmentTree {
#[default]
Empty,
Leaf(Rc<TypedExpression>),
Nested(OrderMap<SimpleIdentifier, AssignmentTree>),
}
impl AssignmentTree {
pub fn insert_leaf(&mut self, expr: Rc<TypedExpression>) {
*self = AssignmentTree::Leaf(expr);
}
pub fn insert_at_path(&mut self, path: &[SimpleIdentifier], expr: Rc<TypedExpression>) {
if path.is_empty() {
*self = AssignmentTree::Leaf(expr);
return;
}
if !matches!(self, AssignmentTree::Nested(_)) {
*self = AssignmentTree::Nested(OrderMap::new());
}
if let AssignmentTree::Nested(fields) = self {
fields
.entry(path[0].clone())
.or_default()
.insert_at_path(&path[1..], expr);
}
}
pub fn into_ir_expression(self) -> IRExpression {
use hamelin_lib::tree::ast::expression::{Expression, StructLiteral};
use hamelin_lib::tree::ast::identifier::ParsedSimpleIdentifier;
use hamelin_lib::tree::typed_ast::expression::TypedStructLiteral;
use hamelin_lib::types::struct_type::Struct;
match self {
AssignmentTree::Empty => {
let typed = TypedExpression {
ast: Rc::new(Expression {
span: Span::default(),
kind: StructLiteral { fields: vec![] }.into(),
}),
resolved_type: Rc::new(Struct::default().into()),
kind: TypedStructLiteral { fields: vec![] }.into(),
};
IRExpression::new(Rc::new(typed))
}
AssignmentTree::Leaf(expr) => IRExpression::new(expr),
AssignmentTree::Nested(fields) => {
let mut struct_type = Struct::default();
let mut typed_fields = Vec::with_capacity(fields.len());
let mut ast_fields = Vec::with_capacity(fields.len());
for (name, child) in fields {
let child_ir = child.into_ir_expression();
let child_expr = child_ir.0;
struct_type = struct_type.with(
name.clone().into(),
child_expr.resolved_type.as_ref().clone(),
);
typed_fields.push((
ParsedSimpleIdentifier::Valid(name.clone()),
child_expr.clone(),
));
ast_fields.push((ParsedSimpleIdentifier::Valid(name), child_expr.ast.clone()));
}
let typed = TypedExpression {
ast: Rc::new(Expression {
span: Span::default(),
kind: StructLiteral { fields: ast_fields }.into(),
}),
resolved_type: Rc::new(struct_type.into()),
kind: TypedStructLiteral {
fields: typed_fields,
}
.into(),
};
IRExpression::new(Rc::new(typed))
}
}
}
}