hamelin_translation 0.4.4

use std::sync::Arc;

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}`.
///
/// Note: overlapping assignments (e.g., `x = 1` and `x.a = 2`) are not detected here —
/// the type checker rejects these as "conflict in field definitions" before IR conversion,
/// so they cannot reach this code through normal compilation.
///
/// Tracks a representative span from the first assignment that contributed to this tree,
/// which is propagated to synthesized struct literals for better error diagnostics.
#[derive(Default)]
pub struct AssignmentTree {
    /// Representative span from the first assignment that contributed to this tree
    span: Span,
    /// The tree structure
    kind: AssignmentTreeKind,
}

#[derive(Default)]
enum AssignmentTreeKind {
    #[default]
    Empty,
    /// A leaf value (direct assignment)
    Leaf(Arc<TypedExpression>),
    /// Nested fields to build into a struct literal
    Nested(OrderMap<SimpleIdentifier, AssignmentTree>),
}

impl AssignmentTree {
    pub fn insert_leaf(&mut self, expr: Arc<TypedExpression>) {
        // Capture span from the first assignment if we don't have one yet
        if self.span.is_none() {
            self.span = expr.ast.span;
        }
        self.kind = AssignmentTreeKind::Leaf(expr);
    }

    pub fn insert_at_path(&mut self, path: &[SimpleIdentifier], expr: Arc<TypedExpression>) {
        // Capture span from the first assignment if we don't have one yet
        if self.span.is_none() {
            self.span = expr.ast.span;
        }

        if path.is_empty() {
            self.kind = AssignmentTreeKind::Leaf(expr);
            return;
        }

        if !matches!(self.kind, AssignmentTreeKind::Nested(_)) {
            self.kind = AssignmentTreeKind::Nested(OrderMap::new());
        }

        if let AssignmentTreeKind::Nested(fields) = &mut self.kind {
            fields
                .entry(path[0].clone())
                .or_default()
                .insert_at_path(&path[1..], expr);
        }
    }

    /// Convert to IRExpression, building struct literals for nested fields.
    ///
    /// The synthesized struct literals use the representative span from the first
    /// assignment that contributed to this tree, preserving source location context
    /// for error diagnostics.
    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.kind {
            AssignmentTreeKind::Empty => {
                // Build empty struct directly without re-typechecking
                let typed = TypedExpression {
                    ast: Arc::new(Expression {
                        span: self.span,
                        kind: StructLiteral { fields: vec![] }.into(),
                    }),
                    resolved_type: Arc::new(Struct::default().into()),
                    kind: TypedStructLiteral { fields: vec![] }.into(),
                };
                IRExpression::new(Arc::new(typed))
            }
            AssignmentTreeKind::Leaf(expr) => IRExpression::new(expr),
            AssignmentTreeKind::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: Arc::new(Expression {
                        span: self.span,
                        kind: StructLiteral { fields: ast_fields }.into(),
                    }),
                    resolved_type: Arc::new(struct_type.into()),
                    kind: TypedStructLiteral {
                        fields: typed_fields,
                    }
                    .into(),
                };
                IRExpression::new(Arc::new(typed))
            }
        }
    }
}