formalang 0.0.4-beta

//! Lowering for the bare-call form `Path(args)`: struct instantiation,
//! external-type instantiation, closure-binding calls and regular function
//! calls (with default-substitution + let-wrapper for forward references
//! to preceding non-defaulted params).

use crate::ast::Expr;
use crate::ir::lower::expr::helpers::substitute_typeparam_in_resolved;
use crate::ir::lower::IrLowerer;
use crate::ir::{IrBlockStatement, IrExpr, IrFunctionParam, ResolvedType};
use std::collections::{HashMap, HashSet};

/// True if `expr` (or any sub-expression) is an `IrExpr::Reference`
/// with a single-segment path matching one of `names`. Used to
/// decide whether a substituted default needs the let-wrapper that
/// binds preceding non-defaulted params to the call-site values.
fn expr_references_any_name(expr: &IrExpr, names: &HashSet<String>) -> bool {
    match expr {
        IrExpr::Reference { path, .. } => {
            path.first().is_some_and(|seg| names.contains(seg.as_str()))
        }
        IrExpr::LetRef { name, .. } => names.contains(name.as_str()),
        IrExpr::BinaryOp { left, right, .. } => {
            expr_references_any_name(left, names) || expr_references_any_name(right, names)
        }
        IrExpr::UnaryOp { operand, .. } => expr_references_any_name(operand, names),
        IrExpr::FieldAccess { object, .. } => expr_references_any_name(object, names),
        IrExpr::FunctionCall { args, .. } | IrExpr::MethodCall { args, .. } => {
            args.iter().any(|(_, a)| expr_references_any_name(a, names))
        }
        IrExpr::CallClosure { closure, args, .. } => {
            expr_references_any_name(closure, names)
                || args.iter().any(|(_, a)| expr_references_any_name(a, names))
        }
        IrExpr::If {
            condition,
            then_branch,
            else_branch,
            ..
        } => {
            expr_references_any_name(condition, names)
                || expr_references_any_name(then_branch, names)
                || else_branch
                    .as_deref()
                    .is_some_and(|e| expr_references_any_name(e, names))
        }
        IrExpr::Match {
            scrutinee, arms, ..
        } => {
            expr_references_any_name(scrutinee, names)
                || arms
                    .iter()
                    .any(|arm| expr_references_any_name(&arm.body, names))
        }
        IrExpr::Block {
            statements, result, ..
        } => {
            statements.iter().any(|s| match s {
                IrBlockStatement::Let { value, .. } => expr_references_any_name(value, names),
                IrBlockStatement::Assign { target, value, .. } => {
                    expr_references_any_name(target, names)
                        || expr_references_any_name(value, names)
                }
                IrBlockStatement::Expr(e) => expr_references_any_name(e, names),
            }) || expr_references_any_name(result, names)
        }
        IrExpr::Array { elements, .. } => {
            elements.iter().any(|e| expr_references_any_name(e, names))
        }
        IrExpr::Tuple { fields, .. } => fields
            .iter()
            .any(|(_, e)| expr_references_any_name(e, names)),
        IrExpr::StructInst { fields, .. } | IrExpr::EnumInst { fields, .. } => fields
            .iter()
            .any(|(_, _, e)| expr_references_any_name(e, names)),
        IrExpr::DictLiteral { entries, .. } => entries
            .iter()
            .any(|(k, v)| expr_references_any_name(k, names) || expr_references_any_name(v, names)),
        IrExpr::DictAccess { dict, key, .. } => {
            expr_references_any_name(dict, names) || expr_references_any_name(key, names)
        }
        IrExpr::For {
            collection, body, ..
        } => expr_references_any_name(collection, names) || expr_references_any_name(body, names),
        IrExpr::Closure { body, .. } => expr_references_any_name(body, names),
        IrExpr::ClosureRef { env_struct, .. } => expr_references_any_name(env_struct, names),
        IrExpr::Literal { .. } | IrExpr::SelfFieldRef { .. } => false,
    }
}

/// Walk `pattern` (a struct field's declared type containing `TypeParam`s)
/// alongside `concrete` (a lowered argument's type) and record each
/// `TypeParam(name) -> concrete` binding into `out`. Conflicts (the same
/// name bound to two different types) and shape mismatches are silently
/// skipped — caller checks coverage afterwards.
fn unify_type_args(
    pattern: &ResolvedType,
    concrete: &ResolvedType,
    out: &mut HashMap<String, ResolvedType>,
) {
    match (pattern, concrete) {
        (ResolvedType::TypeParam(name), other) => {
            out.entry(name.clone()).or_insert_with(|| other.clone());
        }
        (ResolvedType::Tuple(p_fields), ResolvedType::Tuple(c_fields)) => {
            for ((_, pt), (_, ct)) in p_fields.iter().zip(c_fields.iter()) {
                unify_type_args(pt, ct, out);
            }
        }
        (
            ResolvedType::Generic {
                args: p_args,
                base: p_base,
            },
            ResolvedType::Generic {
                args: c_args,
                base: c_base,
            },
        ) if p_base == c_base => {
            for (pa, ca) in p_args.iter().zip(c_args.iter()) {
                unify_type_args(pa, ca, out);
            }
        }
        (
            ResolvedType::Closure {
                param_tys: p_params,
                return_ty: p_ret,
            },
            ResolvedType::Closure {
                param_tys: c_params,
                return_ty: c_ret,
            },
        ) => {
            for ((_, pt), (_, ct)) in p_params.iter().zip(c_params.iter()) {
                unify_type_args(pt, ct, out);
            }
            unify_type_args(p_ret, c_ret, out);
        }
        _ => {}
    }
}

impl IrLowerer<'_> {
    /// Infer the type arguments for a generic struct constructor invoked
    /// without explicit `<...>`. Walks each generic parameter, finds a
    /// field whose declared type mentions the param, and unifies it
    /// against the corresponding lowered argument's type. Returns the
    /// inferred argument vector when every parameter is bound, or an
    /// empty vector when inference can't cover all of them (in which
    /// case the caller falls back to the bare struct type).
    fn infer_struct_type_args(
        &self,
        struct_id: crate::ir::StructId,
        field_target: &HashMap<String, ResolvedType>,
        named_fields: &[(String, crate::ir::FieldIdx, IrExpr)],
    ) -> Vec<ResolvedType> {
        let Some(struct_def) = self.module.get_struct(struct_id) else {
            return Vec::new();
        };
        if struct_def.generic_params.is_empty() {
            return Vec::new();
        }
        let param_names: Vec<String> = struct_def
            .generic_params
            .iter()
            .map(|p| p.name.clone())
            .collect();
        let mut bindings: HashMap<String, ResolvedType> = HashMap::new();
        for (arg_name, _, lowered) in named_fields {
            let Some(declared) = field_target.get(arg_name) else {
                continue;
            };
            unify_type_args(declared, lowered.ty(), &mut bindings);
        }
        let mut resolved = Vec::with_capacity(param_names.len());
        for name in &param_names {
            match bindings.get(name) {
                Some(ty) => resolved.push(ty.clone()),
                None => return Vec::new(),
            }
        }
        resolved
    }

    /// resolve a `ResolvedType` to its enum
    /// type-name (used as the inferred-enum target for a struct-arg
    /// expression). Returns the empty string for non-enum, non-optional-
    /// of-enum types, which the caller filters out.
    fn enum_name_of(module: &crate::ir::IrModule, ty: &ResolvedType) -> String {
        match ty {
            ResolvedType::Enum(eid) => module
                .get_enum(*eid)
                .map_or_else(String::new, |e| e.name.clone()),
            ResolvedType::Generic {
                base: crate::ir::GenericBase::Enum(eid),
                args,
            } => {
                // Optional<T>: peel to its T so an inferred-enum target on
                // a `String?` field reaches the inner enum's variants.
                if Some(*eid) == module.prelude_optional_id() {
                    if let [t] = args.as_slice() {
                        return Self::enum_name_of(module, t);
                    }
                }
                module
                    .get_enum(*eid)
                    .map_or_else(String::new, |e| e.name.clone())
            }
            ResolvedType::Primitive(_)
            | ResolvedType::Struct(_)
            | ResolvedType::Trait(_)
            | ResolvedType::Tuple(_)
            | ResolvedType::Generic { .. }
            | ResolvedType::TypeParam(_)
            | ResolvedType::External { .. }
            | ResolvedType::Closure { .. }
            | ResolvedType::Error => String::new(),
        }
    }

    #[expect(
        clippy::too_many_lines,
        reason = "three branches (struct / external / function) each with their own arg-lowering plumbing; splitting hides the contract"
    )]
    pub(in crate::ir::lower::expr) fn lower_invocation(
        &mut self,
        path: &[crate::ast::Ident],
        type_args: &[crate::ast::Type],
        args: &[(Option<crate::ast::Ident>, Expr)],
    ) -> IrExpr {
        let name = path
            .iter()
            .map(|id| id.name.as_str())
            .collect::<Vec<_>>()
            .join("::");
        let type_args_resolved: Vec<ResolvedType> =
            type_args.iter().map(|t| self.lower_type(t)).collect();

        if let Some(id) = self.module.struct_id(&name) {
            // build a name->type-name map of the
            // struct's fields so each named-arg lowers with the field's
            // declared type as the inferred-enum target. Without this,
            // `Size(width: .auto)` inherits whatever outer
            // `current_function_return_type` was set to and `.auto` can't
            // resolve.
            let field_target: HashMap<String, ResolvedType> = self
                .module
                .get_struct(id)
                .map(|s| {
                    s.fields
                        .iter()
                        .map(|f| (f.name.clone(), f.ty.clone()))
                        .collect()
                })
                .unwrap_or_default();
            let named_fields: Vec<(String, crate::ir::FieldIdx, IrExpr)> = args
                .iter()
                .filter_map(|(name_opt, expr)| {
                    name_opt.as_ref().map(|n| {
                        let saved = self.current_function_return_type.take();
                        let saved_closure = self.expected_closure_type.take();
                        self.current_function_return_type = field_target
                            .get(&n.name)
                            .map(|t| Self::enum_name_of(&self.module, t))
                            .filter(|s| !s.is_empty());
                        // thread closure-typed field annotations
                        // into the closure-literal lowering so untyped params
                        // pick up the field's expected param types.
                        if let Some(t) = field_target.get(&n.name) {
                            if matches!(t, ResolvedType::Closure { .. }) {
                                self.expected_closure_type = Some(t.clone());
                            }
                        }
                        let lowered = self.lower_expr(expr);
                        self.expected_closure_type = saved_closure;
                        self.current_function_return_type = saved;
                        (n.name.clone(), crate::ir::FieldIdx(0), lowered)
                    })
                })
                .collect();
            // Infer type args when the call site omits them. Walks each
            // generic parameter, finds the first struct field whose
            // declared type mentions the param, and unifies it against
            // the corresponding lowered arg's type to recover a concrete
            // binding (e.g. `Box(value: 7)` infers `Box<I32>`). Falls
            // back to the bare struct type when inference can't fill
            // every param.
            let inferred_type_args: Vec<ResolvedType> = if type_args_resolved.is_empty() {
                self.infer_struct_type_args(id, &field_target, &named_fields)
            } else {
                type_args_resolved
            };
            let ty = if inferred_type_args.is_empty() {
                ResolvedType::Struct(id)
            } else {
                ResolvedType::Generic {
                    base: crate::ir::GenericBase::Struct(id),
                    args: inferred_type_args.clone(),
                }
            };
            IrExpr::StructInst {
                struct_id: Some(id),
                type_args: inferred_type_args,
                fields: named_fields,
                ty,
                span: self.current_ir_span(),
            }
        } else if let Some(external_ty) = self.try_external_type(&name, type_args_resolved.clone())
        {
            let named_fields: Vec<(String, crate::ir::FieldIdx, IrExpr)> = args
                .iter()
                .filter_map(|(name_opt, expr)| {
                    name_opt.as_ref().map(|n| {
                        (
                            n.name.clone(),
                            crate::ir::FieldIdx(0),
                            self.lower_expr(expr),
                        )
                    })
                })
                .collect();
            IrExpr::StructInst {
                struct_id: None,
                type_args: type_args_resolved,
                fields: named_fields,
                ty: external_ty,
                span: self.current_ir_span(),
            }
        } else if let Some(call) = self.try_lower_closure_invocation(path, args) {
            call
        } else {
            let path_strs: Vec<String> = path.iter().map(|i| i.name.clone()).collect();
            let fn_name = path_strs.last().map_or("", std::string::String::as_str);
            // Resolve the call to a `FunctionId` first — module-aware
            // for single-segment calls (try the current `mod`'s
            // qualified form, fall back to bare); joined-name lookup
            // for multi-segment. Cross-module / forward-reference
            // cases stay `None` and `ResolveReferencesPass` finishes
            // the job.
            let function_id = if path_strs.len() == 1 {
                self.find_function_in_scope(fn_name)
            } else {
                self.module
                    .function_id(&path_strs.join("::"))
                    .or_else(|| self.find_function_in_scope(fn_name))
            };
            // Derive expected param types from the resolved id
            // (covers cross-module qualified calls correctly), or
            // fall back to scanning by bare name for forward
            // references.
            let expected_param_tys: Vec<(String, ResolvedType)> = function_id
                .and_then(|id| self.module.functions.get(id.0 as usize))
                .map_or_else(
                    || self.lookup_function_param_types(fn_name),
                    |f| {
                        f.params
                            .iter()
                            .filter_map(|p| p.ty.as_ref().map(|t| (p.name.clone(), t.clone())))
                            .collect()
                    },
                );
            let mut lowered_args: Vec<(Option<String>, IrExpr)> = args
                .iter()
                .enumerate()
                .map(|(i, (name_opt, expr))| {
                    let saved_closure = self.expected_closure_type.take();
                    self.expected_closure_type =
                        Self::expected_arg_closure_ty(&expected_param_tys, i, name_opt.as_ref());
                    let lowered = self.lower_expr(expr);
                    self.expected_closure_type = saved_closure;
                    (name_opt.as_ref().map(|n| n.name.clone()), lowered)
                })
                .collect();
            // DP-2 / DP-4 / DP-7: substitute defaults for missing args.
            // For all-positional calls, append trailing defaults. For
            // labeled calls (mode A), walk callee params in order and
            // fill any whose label is missing from the call. If any
            // substituted default references a preceding non-defaulted
            // param by name, wrap the entire FunctionCall in a Block
            // whose Let statements bind those param names to the
            // explicit args (so the default's Reference resolves via
            // path lookup to the new binding, not to the callee's
            // stale binding-id, and side-effects don't duplicate).
            let mut needs_let_wrapper = false;
            let mut wrapper_param_names: Vec<String> = Vec::new();
            let mut wrapper_param_types: Vec<Option<ResolvedType>> = Vec::new();
            if let Some(func_id) = function_id {
                if let Some(func) = self.module.functions.get(func_id.0 as usize) {
                    let non_self_params: Vec<IrFunctionParam> = func
                        .params
                        .iter()
                        .filter(|p| p.name != "self")
                        .cloned()
                        .collect();
                    let want = non_self_params.len();
                    if lowered_args.len() < want {
                        let any_labeled = lowered_args.iter().any(|(l, _)| l.is_some());
                        // Names of params that already have a value in
                        // the call (used to detect earlier-param refs
                        // that need the let-wrapper).
                        let already_provided_names: HashSet<String> = if any_labeled {
                            lowered_args.iter().filter_map(|(l, _)| l.clone()).collect()
                        } else {
                            non_self_params
                                .iter()
                                .take(lowered_args.len())
                                .map(|p| p.name.clone())
                                .collect()
                        };
                        if any_labeled {
                            // Mode A — labeled call. Build a new
                            // ordered args list that walks callee
                            // params in order, picking up the
                            // explicit-call value when its label is
                            // present and substituting the default
                            // otherwise. Mid-list omissions get
                            // filled at the right position.
                            let mut new_args: Vec<(Option<String>, IrExpr)> =
                                Vec::with_capacity(want);
                            for param in &non_self_params {
                                if let Some(pos) = lowered_args.iter().position(|(l, _)| {
                                    l.as_ref().is_some_and(|name| name == &param.name)
                                }) {
                                    let (label, value) = lowered_args.remove(pos);
                                    new_args.push((label, value));
                                } else if let Some(default) = &param.default {
                                    if expr_references_any_name(default, &already_provided_names) {
                                        needs_let_wrapper = true;
                                    }
                                    new_args.push((Some(param.name.clone()), default.clone()));
                                } else {
                                    // Required label missing: validator
                                    // should have rejected. Stop on the
                                    // first gap to preserve some signal.
                                    break;
                                }
                            }
                            lowered_args = new_args;
                        } else {
                            // Positional — append trailing defaults.
                            for param in non_self_params.iter().skip(lowered_args.len()) {
                                if let Some(default) = &param.default {
                                    if expr_references_any_name(default, &already_provided_names) {
                                        needs_let_wrapper = true;
                                    }
                                    lowered_args.push((None, default.clone()));
                                } else {
                                    break;
                                }
                            }
                        }
                        if needs_let_wrapper {
                            for param in non_self_params
                                .iter()
                                .filter(|p| already_provided_names.contains(&p.name))
                            {
                                wrapper_param_names.push(param.name.clone());
                                wrapper_param_types.push(param.ty.clone());
                            }
                        }
                    }
                }
            }
            // Return type lookup uses the same id when available; the
            // bare-name lookup is the fallback for forward refs.
            let mut ty = function_id
                .and_then(|id| self.module.functions.get(id.0 as usize))
                .and_then(|f| f.return_type.clone())
                .unwrap_or_else(|| self.resolve_function_return_type(fn_name, &lowered_args));
            // Substitute the callee's generic-parameter slots in the
            // returned `ty` using the explicit `<...>` type arguments
            // at the call site. Without this, a call like
            // `pair_of<I32>(x: 3, y: 4)` keeps `Pair<T, T>` as its IR
            // type slot until `MonomorphisePass`, and downstream uses
            // (a `let p = …; p.first` field access typed at lowering
            // time) carry `TypeParam(T)` until the leftover scanner
            // surfaces them.
            if !type_args_resolved.is_empty() {
                if let Some(func) = function_id
                    .and_then(|id| self.module.functions.get(id.0 as usize))
                    .filter(|f| !f.generic_params.is_empty())
                {
                    if func.generic_params.len() == type_args_resolved.len() {
                        let subs: HashMap<String, ResolvedType> = func
                            .generic_params
                            .iter()
                            .zip(type_args_resolved.iter())
                            .map(|(p, a)| (p.name.clone(), a.clone()))
                            .collect();
                        substitute_typeparam_in_resolved(&mut ty, &subs);
                    }
                }
            }

            if needs_let_wrapper {
                // Build let bindings for each preceding non-defaulted
                // param. Move the explicit lowered_args[i] into the
                // let value; replace the call-site arg with a
                // Reference to the binding name. The default's
                // Reference{path:[name]} resolves to the let-binding
                // post-ResolveReferencesPass.
                let mut statements = Vec::with_capacity(wrapper_param_names.len());
                for ((name, ty), arg) in wrapper_param_names
                    .iter()
                    .zip(wrapper_param_types.iter())
                    .zip(lowered_args.iter_mut())
                {
                    let value = std::mem::replace(
                        &mut arg.1,
                        IrExpr::Reference {
                            path: vec![name.clone()],
                            target: crate::ir::ReferenceTarget::Unresolved,
                            ty: ty.clone().unwrap_or(ResolvedType::Error),
                            span: self.current_ir_span(),
                        },
                    );
                    statements.push(IrBlockStatement::Let {
                        binding_id: crate::ir::BindingId(0),
                        name: name.clone(),
                        mutable: false,
                        ty: ty.clone(),
                        value,
                        span: self.current_ir_span(),
                    });
                }
                let call = IrExpr::FunctionCall {
                    path: path_strs,
                    function_id,
                    args: lowered_args,
                    ty: ty.clone(),
                    span: self.current_ir_span(),
                };
                IrExpr::Block {
                    statements,
                    result: Box::new(call),
                    ty,
                    span: self.current_ir_span(),
                }
            } else {
                IrExpr::FunctionCall {
                    path: path_strs,
                    function_id,
                    args: lowered_args,
                    ty,
                    span: self.current_ir_span(),
                }
            }
        }
    }

    /// Detect when an `Invocation` whose path is a single segment
    /// resolves to a closure-typed local binding rather than a top-
    /// level function, and lower it as [`IrExpr::CallClosure`]
    /// targeting that binding.
    ///
    /// Returns `None` when the path doesn't refer to a closure-typed
    /// local — the caller falls through to the regular
    /// [`IrExpr::FunctionCall`] path.
    fn try_lower_closure_invocation(
        &mut self,
        path: &[crate::ast::Ident],
        args: &[(Option<crate::ast::Ident>, Expr)],
    ) -> Option<IrExpr> {
        let [ident] = path else {
            return None;
        };
        let name = &ident.name;
        let local_ty = self.lookup_local_binding(name)?.clone();
        let ResolvedType::Closure {
            param_tys,
            return_ty,
        } = &local_ty
        else {
            return None;
        };
        let return_ty = (**return_ty).clone();
        let expected_param_tys: Vec<(String, ResolvedType)> = param_tys
            .iter()
            .enumerate()
            .map(|(i, (_, ty))| (format!("__closure_arg_{i}"), ty.clone()))
            .collect();
        let lowered_args: Vec<(Option<String>, IrExpr)> = args
            .iter()
            .enumerate()
            .map(|(i, (arg_name, expr))| {
                let saved_closure = self.expected_closure_type.take();
                self.expected_closure_type =
                    Self::expected_arg_closure_ty(&expected_param_tys, i, arg_name.as_ref());
                let lowered = self.lower_expr(expr);
                self.expected_closure_type = saved_closure;
                (arg_name.as_ref().map(|n| n.name.clone()), lowered)
            })
            .collect();
        Some(IrExpr::CallClosure {
            closure: Box::new(IrExpr::LetRef {
                name: name.clone(),
                binding_id: crate::ir::BindingId(0),
                ty: local_ty,
                span: self.current_ir_span(),
            }),
            args: lowered_args,
            ty: return_ty,
            span: self.current_ir_span(),
        })
    }
}