Push checkmate through env

crates/compiler/unify/src/env.rs

@@ -1,8 +1,13 @@
+use roc_checkmate::debug_checkmate;
 use roc_collections::VecSet;
-use roc_types::subs::{Subs, Variable};
+use roc_types::subs::{Descriptor, Subs, Variable};
+
+use crate::unify::Mode;
 
 pub struct Env<'a> {
     subs: &'a mut Subs,
+    #[cfg(debug_assertions)]
+    cm: Option<&'a mut roc_checkmate::Collector>,
     seen_recursion: VecSet<(Variable, Variable)>,
     fixed_variables: VecSet<Variable>,
 }
@@ -22,6 +27,17 @@ impl std::ops::DerefMut for Env<'_> {
 }
 
 impl<'a> Env<'a> {
+    #[cfg(debug_assertions)]
+    pub fn new(subs: &'a mut Subs, cm: Option<&'a mut roc_checkmate::Collector>) -> Self {
+        Self {
+            subs,
+            cm,
+            seen_recursion: Default::default(),
+            fixed_variables: Default::default(),
+        }
+    }
+
+    #[cfg(not(debug_assertions))]
     pub fn new(subs: &'a mut Subs) -> Self {
         Self {
             subs,
@@ -30,7 +46,7 @@ impl<'a> Env<'a> {
         }
     }
 
-    pub fn add_recursion_pair(&mut self, var1: Variable, var2: Variable) {
+    pub(crate) fn add_recursion_pair(&mut self, var1: Variable, var2: Variable) {
         let pair = (
             self.subs.get_root_key_without_compacting(var1),
             self.subs.get_root_key_without_compacting(var2),
@@ -40,7 +56,7 @@ impl<'a> Env<'a> {
         debug_assert!(!already_seen);
     }
 
-    pub fn remove_recursion_pair(&mut self, var1: Variable, var2: Variable) {
+    pub(crate) fn remove_recursion_pair(&mut self, var1: Variable, var2: Variable) {
         #[cfg(debug_assertions)]
         let size_before = self.seen_recursion.len();
 
@@ -57,7 +73,7 @@ impl<'a> Env<'a> {
         debug_assert!(size_after < size_before, "nothing was removed");
     }
 
-    pub fn seen_recursion_pair(&self, var1: Variable, var2: Variable) -> bool {
+    pub(crate) fn seen_recursion_pair(&self, var1: Variable, var2: Variable) -> bool {
         let (var1, var2) = (
             self.subs.get_root_key_without_compacting(var1),
             self.subs.get_root_key_without_compacting(var2),
@@ -66,13 +82,38 @@ impl<'a> Env<'a> {
         self.seen_recursion.contains(&(var1, var2))
     }
 
-    pub fn was_fixed(&self, var: Variable) -> bool {
+    pub(crate) fn was_fixed(&self, var: Variable) -> bool {
         self.fixed_variables
             .iter()
             .any(|fixed_var| self.subs.equivalent_without_compacting(*fixed_var, var))
     }
 
-    pub fn extend_fixed_variables(&mut self, vars: impl IntoIterator<Item = Variable>) {
+    pub(crate) fn extend_fixed_variables(&mut self, vars: impl IntoIterator<Item = Variable>) {
         self.fixed_variables.extend(vars);
     }
+
+    pub(crate) fn union(&mut self, left: Variable, right: Variable, desc: Descriptor) {
+        self.subs.union(left, right, desc);
+
+        debug_checkmate!(self.cm, cm => {
+            let new_root = self.subs.get_root_key_without_compacting(left);
+            cm.set_descriptor(self.subs, new_root, desc);
+            cm.unify(self.subs, left, new_root);
+            cm.unify(self.subs, right, new_root);
+        });
+    }
+
+    #[cfg(debug_assertions)]
+    pub(crate) fn debug_start_unification(&mut self, left: Variable, right: Variable, mode: Mode) {
+        debug_checkmate!(self.cm, cm => {
+            cm.start_unification(self.subs, left, right);
+        });
+    }
+
+    #[cfg(debug_assertions)]
+    pub(crate) fn debug_end_unification(&mut self, left: Variable, right: Variable, success: bool) {
+        debug_checkmate!(self.cm, cm => {
+            cm.end_unification(self.subs, left, right, success);
+        });
+    }
 }

crates/compiler/unify/src/lib.rs

@@ -7,3 +7,4 @@
 mod env;
 mod fix;
 pub mod unify;
+pub use env::Env;

crates/compiler/unify/src/unify.rs

@@ -468,6 +468,14 @@ fn debug_print_unified_types<M: MetaCollector>(
 
     static mut UNIFICATION_DEPTH: usize = 0;
 
+    match opt_outcome {
+        None => env.debug_start_unification(ctx.first, ctx.second, ctx.mode),
+        Some(outcome) => {
+            let success = outcome.mismatches.is_empty();
+            env.debug_end_unification(ctx.first, ctx.second, success);
+        }
+    }
+
     dbg_do!(ROC_PRINT_UNIFICATIONS, {
         let prefix = match opt_outcome {
             None => "❔",