hamelin_translation 0.3.10

Lowering and IR for Hamelin query language
Documentation
use std::rc::Rc;

use hamelin_lib::tree::{
    ast::{identifier::SimpleIdentifier, node::Span},
    typed_ast::expression::TypedExpression,
};
use ordermap::OrderMap;

use crate::IRExpression;

/// Tree structure for grouping compound identifier assignments by their root.
///
/// For example, `a.b = 1, a.c = 2` becomes a tree that can be converted to `a = {b: 1, c: 2}`.
#[derive(Default)]
pub enum AssignmentTree {
    #[default]
    Empty,
    /// A leaf value (direct assignment)
    Leaf(Rc<TypedExpression>),
    /// Nested fields to build into a struct literal
    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);
        }
    }

    /// Convert to IRExpression, building struct literals for nested fields.
    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 => {
                // Build empty struct directly without re-typechecking
                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) => {
                // Build struct literal directly from child expressions (already typed)
                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;

                    // Add to struct type
                    struct_type = struct_type.with(
                        name.clone().into(),
                        child_expr.resolved_type.as_ref().clone(),
                    );

                    // Build typed field
                    typed_fields.push((
                        ParsedSimpleIdentifier::Valid(name.clone()),
                        child_expr.clone(),
                    ));

                    // Build AST field
                    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))
            }
        }
    }
}