wip: more salsa

Author: tchupp

compiler/hir/hir_def/src/lib.rs

@@ -5,10 +5,12 @@ mod dependency_graph;
 mod diagnostics;
 mod hir;
 mod hir_module;
+mod value_def;
 
 pub use diagnostics::*;
 pub use hir::*;
 pub use hir_module::HirModule;
+pub use value_def::{module_value_def, ValueDef};
 
 mod fqn;
 pub use fqn::{Fqn, FqnResolutionError};

compiler/hir/hir_def/src/value_def.rs

@@ -0,0 +1,32 @@
+use crate::{hir::lower_file, ExpressionIdx, HirDefDatabase, Name, TypeIdx};
+use alloy_workspace::ModuleId;
+
+/// A Salsa-tracked value definition.
+///
+/// Each field is independently tracked by Salsa, so downstream queries
+/// that only read `type_annotation` won't re-run when `expression_idx` changes.
+#[salsa::tracked]
+pub struct ValueDef<'db> {
+    pub module_id: ModuleId,
+    #[returns(ref)]
+    pub name: Name,
+    pub type_annotation: Option<TypeIdx>,
+    pub expression_idx: ExpressionIdx,
+}
+
+/// Returns one `ValueDef` per top-level value definition in the module.
+///
+/// These are created from the same data as `lower_file`, but as Salsa-tracked
+/// structs so that downstream queries can depend on individual fields.
+#[salsa::tracked]
+pub fn module_value_def<'db>(
+    db: &'db dyn HirDefDatabase,
+    module_id: ModuleId,
+    expr_idx: ExpressionIdx,
+) -> Option<ValueDef<'db>> {
+    let (hir_module, _) = lower_file(db, module_id);
+
+    hir_module
+        .get_value_by_id(&expr_idx)
+        .map(|v| ValueDef::new(db, module_id, v.name.clone(), v.type_annotation, v.value))
+}

compiler/hir/hir_infer/src/hir_ty/hm/constraint_gen/expr/mod.rs

@@ -1,9 +1,11 @@
 use super::{HMInferenceContext, MonoType};
-use crate::hir_ty::hm::inference::annotated_to_mono;
 use crate::hir_ty::hm::PolyType;
+use crate::hir_ty::InferredType;
+use crate::value_signature::infer_value_signature;
 use alloy_hir_def as hir;
 use alloy_hir_resolved as res;
-use alloy_hir_resolved::{resolve_annotated_type, Fql};
+use alloy_hir_resolved::Fql;
+use rustc_hash::FxHashMap;
 
 mod binary;
 mod function_call;
@@ -23,35 +25,35 @@ pub(crate) fn infer_expr_hm(
         return existing_ty;
     }
 
-    // Lazy constraint generation: if this expression is in a later group,
-    // don't infer it now - just return a fresh type variable.
-    // Only applies to expressions in the current module — cross-module expressions
-    // have their own independent ordering and must always be inferred immediately.
-    // DON'T assign to type_env to avoid polluting env_type_vars for generalization
-    if source_fql.module_id == ctx.module_id && ctx.is_in_later_group(source_fql.local_id) {
-        return ctx.fresh_type_var();
-    }
+    // Use `infer_value_signature` for cross-module value definitions and
+    // same-module definitions with polymorphic annotations (when referenced
+    // from other definitions, not when being directly inferred).
+    //
+    // Skip the shortcut for:
+    // - The expression currently being inferred (inferring_expr) to prevent cycles
+    //   and ensure sub-expression types are collected
+    // - Same-module unannotated definitions, so call-site constraints flow back
+    let is_self = ctx
+        .inferring_expr
+        .as_ref()
+        .is_some_and(|e| *e == source_fql);
+    if !is_self {
+        if let Some(value_def) =
+            hir::module_value_def(ctx.db, source_fql.module_id, source_fql.local_id)
+        {
+            let is_cross_module = source_fql.module_id != ctx.module_id;
+            let has_poly_annotation = value_def
+                .type_annotation(ctx.db)
+                .map(|ta| {
+                    alloy_hir_resolved::resolve_annotated_type(ctx.db, source_fql.module_id, ta)
+                        .is_polymorphic()
+                })
+                .unwrap_or(false);
 
-    // For cross-module expressions with polymorphic type annotations, use the annotation
-    // directly rather than re-inferring from the body. This ensures each use site gets
-    // fresh type variables via poly_env instantiation, avoiding stale specialization
-    // when a polymorphic cross-module function (e.g., <|) is used multiple times.
-    if source_fql.module_id != ctx.module_id {
-        let (other_hir_module, _) = hir::lower_file(ctx.db, source_fql.module_id);
-        if let Some(value) = other_hir_module.get_value_by_id(&source_fql.local_id) {
-            if let Some(type_annotation) = value.type_annotation {
-                let annotated =
-                    resolve_annotated_type(ctx.db, source_fql.module_id, type_annotation);
-                if annotated.is_polymorphic() {
-                    if let Some(mono_ty) = annotated_to_mono(&annotated, ctx) {
-                        let poly_ty = PolyType::generalize_all(mono_ty);
-                        ctx.poly_env.insert(source_fql.clone().into(), poly_ty);
-                        // Return a fresh instantiation
-                        if let Some(instantiated) = ctx.maybe_find_type(&source_fql) {
-                            return instantiated;
-                        }
-                    }
-                }
+            if is_cross_module || has_poly_annotation {
+                let sig = infer_value_signature(ctx.db, value_def);
+                let mono_ty = inferred_to_mono(&sig, ctx);
+                return ctx.generalize_to_poly(mono_ty, &source_fql);
             }
         }
     }
@@ -108,6 +110,80 @@ pub(crate) fn infer_expr_hm(
     }
 }
 
+/// Convert an InferredType (from `infer_value_signature`) to a MonoType for use
+/// in constraint generation. Generic IDs are mapped to fresh type variables,
+/// with consistent mapping so the same Generic(id) produces the same TypeVarId.
+fn inferred_to_mono(inferred: &InferredType, ctx: &mut HMInferenceContext) -> MonoType {
+    let mut generic_map: FxHashMap<usize, super::super::TypeVarId> = FxHashMap::default();
+    inferred_to_mono_inner(inferred, ctx, &mut generic_map)
+}
+
+fn inferred_to_mono_inner(
+    inferred: &InferredType,
+    ctx: &mut HMInferenceContext,
+    generic_map: &mut FxHashMap<usize, super::super::TypeVarId>,
+) -> MonoType {
+    match inferred {
+        InferredType::Unconstrained => MonoType::Unconstrained,
+        InferredType::Missing => ctx.fresh_type_var(),
+        InferredType::Unit => MonoType::Unit,
+        InferredType::BuiltIn(b) => MonoType::Concrete(*b),
+        InferredType::TypeDef(fql, name) => MonoType::TypeDef {
+            fql: fql.clone(),
+            type_args: vec![],
+            type_def_name: name.clone(),
+        },
+        InferredType::Lambda {
+            arg_type,
+            return_type,
+        } => MonoType::Function(
+            Box::new(inferred_to_mono_inner(arg_type, ctx, generic_map)),
+            Box::new(inferred_to_mono_inner(return_type, ctx, generic_map)),
+        ),
+        InferredType::Tuple(elements) => MonoType::Tuple(
+            elements
+                .iter()
+                .map(|e| inferred_to_mono_inner(e, ctx, generic_map))
+                .collect(),
+        ),
+        InferredType::Bounded { base, args } => {
+            // When the base is a TypeDef, populate its type_args with fresh TypeVarIds
+            // to preserve arity information (used in error messages and display)
+            let constructor = match base.as_ref() {
+                InferredType::TypeDef(fql, name) => {
+                    let type_args = args.iter().map(|_| ctx.type_var_gen.fresh()).collect();
+                    MonoType::TypeDef {
+                        fql: fql.clone(),
+                        type_args,
+                        type_def_name: name.clone(),
+                    }
+                }
+                other => inferred_to_mono_inner(other, ctx, generic_map),
+            };
+            MonoType::App {
+                constructor: Box::new(constructor),
+                args: args
+                    .iter()
+                    .map(|a| inferred_to_mono_inner(a, ctx, generic_map))
+                    .collect(),
+            }
+        }
+        InferredType::Generic(id) => {
+            let var_id = *generic_map
+                .entry(*id)
+                .or_insert_with(|| ctx.type_var_gen.fresh());
+            MonoType::Var(var_id)
+        }
+        InferredType::ConstrainedGeneric { id, .. } => {
+            // TODO: Track constraints during solving
+            let var_id = *generic_map
+                .entry(*id)
+                .or_insert_with(|| ctx.type_var_gen.fresh());
+            MonoType::Var(var_id)
+        }
+    }
+}
+
 fn infer_missing_expr(ctx: &mut HMInferenceContext, source_fql: Fql<hir::Expression>) -> MonoType {
     let ty = ctx.fresh_type_var();
     ctx.assign_type(source_fql, ty)

compiler/hir/hir_infer/src/hir_ty/hm/constraint_gen/expr/type_def.rs

@@ -70,22 +70,8 @@ pub(super) fn infer_variant_constructor(
     // Build the constructor type using the shared helper
     let constructor_ty = build_constructor_type(ctx, &type_def_fql, &type_def, &member);
 
-    // Check if this is a polymorphic constructor
-    let is_polymorphic = has_type_variables(&constructor_ty);
-
-    if is_polymorphic {
-        // Generalize and store in poly_env for proper instantiation
-        // For variant constructors, quantify over ALL free variables (not filtered by env_type_vars)
-        // since constructors are top-level polymorphic values
-        let poly_ty = PolyType::generalize_all(constructor_ty.clone());
-        ctx.poly_env.insert(variant_fql.clone(), poly_ty);
-
-        if let Some(tracked_ty) = ctx.maybe_find_type(variant_fql) {
-            return ctx.assign_type(source_fql, tracked_ty);
-        }
-    }
-
-    ctx.assign_type(source_fql, constructor_ty)
+    let tracked_ty = ctx.generalize_to_poly(constructor_ty, variant_fql);
+    ctx.assign_type(source_fql, tracked_ty)
 }
 
 pub(super) fn infer_type_definition(
@@ -115,18 +101,7 @@ pub(super) fn infer_type_definition(
             // When you call Identity(...), it's the same as Id(...)
             let constructor_ty = build_constructor_type(ctx, &td_fql, &type_def, member);
 
-            // Check if this is a polymorphic constructor (has App with type variables)
-            let is_polymorphic = has_type_variables(&constructor_ty);
-
-            if is_polymorphic {
-                // Generalize and store in poly_env for proper instantiation
-                // For type definition constructors, quantify over ALL free variables
-                // since constructors are top-level polymorphic values
-                let poly_ty = PolyType::generalize_all(constructor_ty.clone());
-                ctx.poly_env.insert(td_fql.clone().into(), poly_ty);
-            }
-
-            constructor_ty
+            ctx.generalize_to_poly(constructor_ty, &td_fql)
         }
         TypeDefinitionKind::Union(_members) => {
             // Multi-variant type - cannot be called as a function directly
@@ -209,19 +184,3 @@ fn build_constructor_type(
             MonoType::Function(Box::new(param_ty), Box::new(acc))
         })
 }
-
-/// Check if a MonoType contains type variables (is polymorphic)
-fn has_type_variables(ty: &MonoType) -> bool {
-    match ty {
-        MonoType::Var(_) => true,
-        MonoType::Function(arg, ret) => has_type_variables(arg) || has_type_variables(ret),
-        MonoType::Tuple(elements) => elements.iter().any(has_type_variables),
-        MonoType::App { constructor, args } => {
-            has_type_variables(constructor) || args.iter().any(has_type_variables)
-        }
-        MonoType::Unconstrained
-        | MonoType::Concrete(_)
-        | MonoType::TypeDef { .. }
-        | MonoType::Unit => false,
-    }
-}