Fix generic inference on function calls

crates/rue-compiler/src/compile/expr/function_call.rs

@@ -5,7 +5,7 @@ use rue_ast::{AstFunctionCallExpr, AstNode};
 use rue_diagnostic::DiagnosticKind;
 use rue_hir::{BinaryOp, Builtin, FunctionCall, Hir, Symbol, UnaryOp, Value};
 use rue_lir::ClvmOp;
-use rue_types::{Pair, Type, Union, substitute_with_mappings};
+use rue_types::{Pair, Type, TypeId, Union, substitute_with_mappings};
 
 use crate::{Compiler, compile_expr};
 
@@ -69,14 +69,32 @@ pub fn compile_function_call_expr(ctx: &mut Compiler, call: &AstFunctionCallExpr
             );
         }
 
-        let mut mappings = HashMap::new();
+        let mut mappings: HashMap<TypeId, Vec<TypeId>> = HashMap::new();
         let mut results = Vec::new();
 
         for (i, (_, param)) in function.params.iter().enumerate() {
+            let substitute_mappings = mappings
+                .iter()
+                .map(|(k, v)| {
+                    (
+                        *k,
+                        if v.is_empty() {
+                            ctx.builtins().types.never
+                        } else if v.len() == 1 {
+                            v[0]
+                        } else {
+                            ctx.alloc_type(Type::Union(Union::new(v.clone())))
+                        },
+                    )
+                })
+                .collect();
+
+            let param = substitute_with_mappings(ctx.types_mut(), *param, &substitute_mappings);
+
             if let Some(expr) = args.get(i) {
-                let value = compile_expr(ctx, expr, Some(*param));
+                let value = compile_expr(ctx, expr, Some(param));
                 results.push(value.hir);
-                ctx.infer_type(expr.syntax(), value.ty, *param, &mut mappings);
+                ctx.infer_type(expr.syntax(), value.ty, param, &mut mappings);
             } else {
                 debug!("Unresolved function call argument");
                 results.push(ctx.builtins().unresolved.hir);
@@ -102,6 +120,22 @@ pub fn compile_function_call_expr(ctx: &mut Compiler, call: &AstFunctionCallExpr
         )))
     };
 
+    let mappings = mappings
+        .into_iter()
+        .map(|(k, v)| {
+            (
+                k,
+                if v.is_empty() {
+                    ctx.builtins().types.never
+                } else if v.len() == 1 {
+                    v[0]
+                } else {
+                    ctx.alloc_type(Type::Union(Union::new(v)))
+                },
+            )
+        })
+        .collect();
+
     let ty = substitute_with_mappings(ctx.types_mut(), ty, &mappings);
 
     let hir = ctx.alloc_hir(Hir::FunctionCall(FunctionCall {

crates/rue-compiler/src/compiler.rs

@@ -399,7 +399,7 @@ impl Compiler {
         node: &impl GetTextRange,
         from: TypeId,
         to: TypeId,
-        infer: &mut HashMap<TypeId, TypeId>,
+        infer: &mut HashMap<TypeId, Vec<TypeId>>,
     ) {
         self.compare_type(node, from, to, false, Some(infer));
     }
@@ -410,7 +410,7 @@ impl Compiler {
         from: TypeId,
         to: TypeId,
         cast: bool,
-        infer: Option<&mut HashMap<TypeId, TypeId>>,
+        infer: Option<&mut HashMap<TypeId, Vec<TypeId>>>,
     ) {
         let comparison = rue_types::compare_with_inference(
             self.db.types_mut(),

crates/rue-types/src/compare.rs

@@ -20,7 +20,7 @@ pub enum Comparison {
 
 #[derive(Debug, Default)]
 pub(crate) struct ComparisonContext<'a> {
-    infer: Option<&'a mut HashMap<TypeId, TypeId>>,
+    infer: Option<&'a mut HashMap<TypeId, Vec<TypeId>>>,
     stack: IndexSet<(TypeId, TypeId)>,
 }
 
@@ -29,7 +29,7 @@ pub fn compare_with_inference(
     builtins: &BuiltinTypes,
     lhs: TypeId,
     rhs: TypeId,
-    infer: Option<&mut HashMap<TypeId, TypeId>>,
+    infer: Option<&mut HashMap<TypeId, Vec<TypeId>>>,
 ) -> Comparison {
     let lhs = substitute(arena, lhs);
     let rhs = substitute(arena, rhs);
@@ -156,17 +156,13 @@ pub(crate) fn compare_impl(
             if lhs == rhs {
                 Comparison::Assign
             } else if let Some(infer) = &mut ctx.infer {
-                if let Some(rhs) = infer.get(&rhs).copied() {
-                    compare_impl(arena, builtins, ctx, lhs, rhs, lhs_semantic, rhs_semantic)
-                } else {
-                    debug!(
-                        "Inferring {} is {}",
-                        stringify_impl(arena, rhs, &mut IndexMap::new()),
-                        stringify_impl(arena, lhs, &mut IndexMap::new())
-                    );
-                    infer.insert(rhs, lhs);
-                    Comparison::Assign
-                }
+                debug!(
+                    "Inferring {} could include {}",
+                    stringify_impl(arena, rhs, &mut IndexMap::new()),
+                    stringify_impl(arena, lhs, &mut IndexMap::new())
+                );
+                infer.entry(rhs).or_default().push(lhs);
+                Comparison::Assign
             } else if let Type::Union(lhs) = arena[lhs].clone() {
                 let mut result = Comparison::Assign;