Improve constraint generalisation algorithm

Author: purefunctor

compiler-core/checking/src/algorithm/quantify.rs

@@ -9,12 +9,12 @@ use petgraph::visit::{DfsPostOrder, Reversed};
 use rustc_hash::FxHashSet;
 use smol_str::SmolStrBuilder;
 
-use crate::ExternalQueries;
 use crate::algorithm::constraint::{self, ConstraintApplication};
 use crate::algorithm::fold::Zonk;
 use crate::algorithm::state::{CheckContext, CheckState};
 use crate::algorithm::substitute::{ShiftBound, SubstituteUnification, UniToLevel};
 use crate::core::{Class, ForallBinder, Instance, RowType, Type, TypeId, Variable, debruijn};
+use crate::{ExternalQueries, safe_loop};
 
 pub fn quantify(state: &mut CheckState, id: TypeId) -> Option<(TypeId, debruijn::Size)> {
     let graph = collect_unification(state, id);
@@ -96,28 +96,25 @@ where
         return Ok(Some(quantified_with_constraints));
     }
 
-    let unsolved_graph = collect_unification(state, type_id);
-    let unsolved_nodes: FxHashSet<u32> = unsolved_graph.nodes().collect();
-
-    let mut valid: FxHashSet<TypeId> = FxHashSet::default();
-    let mut ambiguous = vec![];
+    let mut pending = vec![];
     let mut unsatisfied = vec![];
 
     for constraint in constraints {
         let constraint = Zonk::on(state, constraint);
-        let unsolved_graph = collect_unification(state, constraint);
-        if unsolved_graph.node_count() == 0 {
+        let unification: FxHashSet<u32> = collect_unification(state, constraint).nodes().collect();
+        if unification.is_empty() {
             unsatisfied.push(constraint);
-        } else if unsolved_graph.nodes().all(|unification| unsolved_nodes.contains(&unification)) {
-            valid.insert(constraint);
         } else {
-            ambiguous.push(constraint);
+            pending.push((constraint, unification));
         }
     }
 
+    let in_signature = collect_unification(state, type_id).nodes().collect();
+    let (generalised, ambiguous) = classify_constraints_by_reachability(pending, in_signature);
+
     // Subtle: stable ordering for consistent output
-    let valid = valid.into_iter().sorted().collect_vec();
-    let minimized = minimize_by_superclasses(state, context, valid)?;
+    let generalised = generalised.into_iter().sorted().collect_vec();
+    let minimized = minimize_by_superclasses(state, context, generalised)?;
 
     let constrained_type = minimized.into_iter().rfold(type_id, |constrained, constraint| {
         state.storage.intern(Type::Constrained(constraint, constrained))
@@ -186,6 +183,40 @@ where
         .collect())
 }
 
+/// Classifies constraints as valid or ambiguous based on variable reachability.
+///
+/// A constraint is valid if its variables are transitively reachable from the
+/// signature variables. The algorithm uses fixed-point iteration. As long as
+/// a constraint shares any variable with the reachable set, all its variables
+/// become reachable.
+fn classify_constraints_by_reachability(
+    pending: Vec<(TypeId, FxHashSet<u32>)>,
+    in_signature: FxHashSet<u32>,
+) -> (FxHashSet<TypeId>, Vec<TypeId>) {
+    let mut reachable = in_signature;
+    let mut valid = FxHashSet::default();
+    let mut remaining = pending;
+
+    safe_loop! {
+        let (connected, disconnected): (Vec<_>, Vec<_>) =
+            remaining.into_iter().partition(|(_, unification)| {
+                unification.iter().any(|variable| reachable.contains(variable))
+            });
+
+        if connected.is_empty() {
+            let ambiguous = disconnected.into_iter().map(|(id, _)| id).collect();
+            return (valid, ambiguous);
+        }
+
+        for (constraint, unification) in connected {
+            valid.insert(constraint);
+            reachable.extend(unification);
+        }
+
+        remaining = disconnected;
+    }
+}
+
 fn generate_type_name(id: u32) -> smol_str::SmolStr {
     let mut builder = SmolStrBuilder::default();
     write!(builder, "t{id}").unwrap();

tests-integration/fixtures/checking/210_row_constraint_generalization/Main.purs

@@ -0,0 +1,21 @@
+module Main where
+
+import Prim.Row as Prim.Row
+
+foreign import unsafeCoerce :: forall a b. a -> b
+
+merge
+  :: forall r1 r2 r3 r4
+   . Prim.Row.Union r1 r2 r3
+  => Prim.Row.Nub r3 r4
+  => Record r1
+  -> Record r2
+  -> Record r4
+merge _ _ = unsafeCoerce {}
+
+-- This should generalise the unsolved Union and Nub constraints
+-- rather than emitting AmbiguousConstraint errors
+a = merge { a: 123 }
+
+-- Fully applied, should resolve to a concrete record type
+b = a { b: 123 }

tests-integration/fixtures/checking/210_row_constraint_generalization/Main.snap

@@ -0,0 +1,14 @@
+module Main where
+
+import Prim.Row as Prim.Row
+
+foreign import unsafeCoerce :: forall a b. a -> b
+
+fromUnion
+  :: forall r1 r2 r3
+   . Prim.Row.Union r1 r2 r3
+  => Record r3
+  -> Int
+fromUnion _ = unsafeCoerce 0
+
+test = fromUnion

tests-integration/fixtures/checking/211_row_constraint_result/Main.purs

@@ -0,0 +1,15 @@
+module Main where
+
+import Prim.Row as Prim.Row
+
+foreign import unsafeCoerce :: forall a b. a -> b
+
+chainedUnion
+  :: forall r1 r2 r3 r4 r5
+   . Prim.Row.Union r1 r2 r3
+  => Prim.Row.Union r3 r4 r5
+  => Record r1
+  -> Record r5
+chainedUnion _ = unsafeCoerce {}
+
+test = chainedUnion { x: 1 }

tests-integration/fixtures/checking/211_row_constraint_result/Main.snap

@@ -0,0 +1,19 @@
+module Main where
+
+import Prim.Row as Prim.Row
+
+foreign import unsafeCoerce :: forall a b. a -> b
+
+multiMerge
+  :: forall r1 r2 r3 r4 r5 r6
+   . Prim.Row.Union r1 r2 r3
+  => Prim.Row.Nub r3 r4
+  => Prim.Row.Union r4 r5 r6
+  => Record r1
+  -> Record r5
+  -> Record r6
+multiMerge _ _ = unsafeCoerce {}
+
+test1 = multiMerge { a: 1 }
+
+test2 = multiMerge { a: 1 } { b: 2 }

tests-integration/fixtures/checking/212_row_constraint_chained/Main.purs

@@ -445,3 +445,11 @@ fn run_test(folder: &str, file: &str) {
 #[rustfmt::skip] #[test] fn test_208_int_add_constraint_main() { run_test("208_int_add_constraint", "Main"); }
 
 #[rustfmt::skip] #[test] fn test_209_int_cons_constraint_main() { run_test("209_int_cons_constraint", "Main"); }
+
+#[rustfmt::skip] #[test] fn test_210_row_constraint_generalization_main() { run_test("210_row_constraint_generalization", "Main"); }
+
+#[rustfmt::skip] #[test] fn test_211_row_constraint_result_main() { run_test("211_row_constraint_result", "Main"); }
+
+#[rustfmt::skip] #[test] fn test_212_row_constraint_chained_main() { run_test("212_row_constraint_chained", "Main"); }
+
+#[rustfmt::skip] #[test] fn test_213_row_constraint_multiple_main() { run_test("213_row_constraint_multiple", "Main"); }