Push checkmate through env

crates/compiler/checkmate/src/collector.rs

@@ -24,7 +24,7 @@ impl Collector {
         }
     }
 
-    pub fn unify(&mut self, subs: &s::Subs, to: s::Variable, from: s::Variable) {
+    pub fn unify(&mut self, subs: &s::Subs, from: s::Variable, to: s::Variable) {
         let to = to.as_schema(subs);
         let from = from.as_schema(subs);
         self.add_event(VariableEvent::Unify { to, from });
@@ -61,7 +61,7 @@ impl Collector {
         });
     }
 
-    pub fn start_unification(&mut self, subs: &s::Subs, left: &s::Variable, right: &s::Variable) {
+    pub fn start_unification(&mut self, subs: &s::Subs, left: s::Variable, right: s::Variable) {
         let left = left.as_schema(subs);
         let right = right.as_schema(subs);
         // TODO add mode
@@ -72,11 +72,32 @@ impl Collector {
             right,
             mode,
             subevents,
+            success: None,
         });
     }
 
-    pub fn end_unification(&mut self) {
+    pub fn end_unification(
+        &mut self,
+        subs: &s::Subs,
+        left: s::Variable,
+        right: s::Variable,
+        success: bool,
+    ) {
         let current_event = self.get_path_event();
+        match current_event {
+            Event::Unification {
+                left: l,
+                right: r,
+                success: s,
+                ..
+            } => {
+                assert_eq!(left.as_schema(subs), *l);
+                assert_eq!(right.as_schema(subs), *r);
+                assert!(s.is_none());
+                *s = Some(success);
+            }
+            _ => panic!("end_unification called when not in a unification"),
+        }
         assert!(matches!(current_event, Event::Unification { .. }));
         self.current_event_path.pop();
     }

crates/compiler/checkmate/src/lib.rs

@@ -5,6 +5,39 @@ mod schema;
 pub use collector::Collector;
 
 pub fn is_checkmate_enabled() -> bool {
-    let flag = std::env::var("ROC_CHECKMATE");
-    flag.as_deref() == Ok("1")
+    #[cfg(debug_assertions)]
+    {
+        let flag = std::env::var("ROC_CHECKMATE");
+        flag.as_deref() == Ok("1")
+    }
+    #[cfg(not(debug_assertions))]
+    {
+        false
+    }
+}
+
+#[macro_export]
+macro_rules! debug_checkmate {
+    ($opt_collector:expr, $cm:ident => $expr:expr) => {
+        #[cfg(debug_assertions)]
+        {
+            if let Some($cm) = $opt_collector.as_mut() {
+                $expr
+            }
+        }
+    };
+}
+
+#[macro_export]
+macro_rules! with_checkmate {
+    ($opt_collector:expr, { on => $on:expr, off => $off:expr, }) => {{
+        #[cfg(debug_assertions)]
+        {
+            $on
+        }
+        #[cfg(not(debug_assertions))]
+        {
+            $off
+        }
+    }};
 }

crates/compiler/checkmate/src/schema.rs

@@ -5,7 +5,7 @@ use serde::Serialize;
 #[derive(Serialize)]
 pub enum Constraint {}
 
-#[derive(Serialize)]
+#[derive(Serialize, Debug, PartialEq)]
 pub struct Variable(pub u32);
 
 #[derive(Serialize)]
@@ -197,6 +197,7 @@ pub enum Event {
         left: Variable,
         right: Variable,
         mode: UnificationMode,
+        success: Option<bool>,
         subevents: Vec<Event>,
     },
 }

crates/compiler/derive/Cargo.toml

@@ -9,6 +9,7 @@ version.workspace = true
 
 [dependencies]
 roc_can = { path = "../can" }
+roc_checkmate = { path = "../checkmate" }
 roc_collections = { path = "../collections" }
 roc_derive_key = { path = "../derive_key" }
 roc_error_macros = { path = "../../error_macros" }

crates/compiler/derive/src/util.rs

@@ -1,4 +1,5 @@
 use roc_can::{abilities::SpecializationLambdaSets, module::ExposedByModule};
+use roc_checkmate::with_checkmate;
 use roc_error_macros::internal_error;
 use roc_module::symbol::{IdentIds, Symbol};
 use roc_types::{
@@ -71,10 +72,17 @@ impl Env<'_> {
     }
 
     pub fn unify(&mut self, left: Variable, right: Variable) {
-        use roc_unify::unify::{unify, Env, Mode, Unified};
+        use roc_unify::{
+            unify::{unify, Mode, Unified},
+            Env,
+        };
 
         let unified = unify(
-            &mut Env::new(self.subs),
+            // TODO(checkmate): pass checkmate through
+            &mut with_checkmate!(None, {
+                on => Env::new(self.subs, None),
+                off => Env::new(self.subs),
+            }),
             left,
             right,
             Mode::EQ,
@@ -103,12 +111,19 @@ impl Env<'_> {
         specialization_type: Variable,
         ability_member: Symbol,
     ) -> SpecializationLambdaSets {
-        use roc_unify::unify::{unify_introduced_ability_specialization, Env, Mode, Unified};
+        use roc_unify::{
+            unify::{unify_introduced_ability_specialization, Mode, Unified},
+            Env,
+        };
 
         let member_signature = self.import_builtin_symbol_var(ability_member);
 
         let unified = unify_introduced_ability_specialization(
-            &mut Env::new(self.subs),
+            // TODO(checkmate): pass checkmate through
+            &mut with_checkmate!(None, {
+                on => Env::new(self.subs, None),
+                off => Env::new(self.subs),
+            }),
             member_signature,
             specialization_type,
             Mode::EQ,

crates/compiler/late_solve/Cargo.toml

@@ -9,6 +9,7 @@ version.workspace = true
 
 [dependencies]
 roc_can = { path = "../can" }
+roc_checkmate = { path = "../checkmate" }
 roc_collections = { path = "../collections" }
 roc_derive = { path = "../derive" }
 roc_error_macros = { path = "../../error_macros" }

crates/compiler/late_solve/src/lib.rs

@@ -6,6 +6,7 @@ use std::sync::{Arc, RwLock};
 use bumpalo::Bump;
 use roc_can::abilities::AbilitiesStore;
 use roc_can::module::ExposedByModule;
+use roc_checkmate::with_checkmate;
 use roc_collections::MutMap;
 use roc_derive::SharedDerivedModule;
 use roc_error_macros::internal_error;
@@ -19,7 +20,8 @@ use roc_types::subs::{get_member_lambda_sets_at_region, Content, FlatType, Lambd
 use roc_types::subs::{ExposedTypesStorageSubs, Subs, Variable};
 use roc_types::types::Polarity;
 use roc_unify::unify::MetaCollector;
-use roc_unify::unify::{Env as UEnv, Mode, Unified};
+use roc_unify::unify::{Mode, Unified};
+use roc_unify::Env as UEnv;
 
 pub use roc_solve::ability::{ResolveError, Resolved};
 pub use roc_types::subs::instantiate_rigids;
@@ -360,7 +362,11 @@ pub fn unify(
         "derived module can only unify its subs in its own context!"
     );
     let unified = roc_unify::unify::unify_with_collector::<ChangedVariableCollector>(
-        &mut UEnv::new(subs),
+        // TODO(checkmate): pass checkmate through
+        &mut with_checkmate!(None, {
+            on => UEnv::new(subs, None),
+            off => UEnv::new(subs),
+        }),
         left,
         right,
         Mode::EQ,

crates/compiler/solve/Cargo.toml

@@ -9,6 +9,7 @@ version.workspace = true
 
 [dependencies]
 roc_can = { path = "../can" }
+roc_checkmate = { path = "../checkmate" }
 roc_collections = { path = "../collections" }
 roc_debug_flags = { path = "../debug_flags" }
 roc_derive = { path = "../derive" }

crates/compiler/solve/src/ability.rs

@@ -1,5 +1,6 @@
 use roc_can::abilities::AbilitiesStore;
 use roc_can::expr::PendingDerives;
+use roc_checkmate::with_checkmate;
 use roc_collections::{VecMap, VecSet};
 use roc_debug_flags::dbg_do;
 #[cfg(debug_assertions)]
@@ -18,8 +19,9 @@ use roc_types::subs::{
     TupleElems, Variable,
 };
 use roc_types::types::{AliasKind, Category, MemberImpl, PatternCategory, Polarity, Types};
-use roc_unify::unify::{Env as UEnv, MustImplementConstraints};
+use roc_unify::unify::MustImplementConstraints;
 use roc_unify::unify::{MustImplementAbility, Obligated};
+use roc_unify::Env as UEnv;
 
 use crate::env::InferenceEnv;
 use crate::{aliases::Aliases, to_var::type_to_var};
@@ -1289,9 +1291,12 @@ impl DerivableVisitor for DeriveEq {
 
         // Of the floating-point types,
         // only Dec implements Eq.
-        let mut env = UEnv::new(subs);
+        // TODO(checkmate): pass checkmate through
         let unified = unify(
-            &mut env,
+            &mut with_checkmate!(None, {
+                on => UEnv::new(subs, None),
+                off => UEnv::new(subs),
+            }),
             content_var,
             Variable::DECIMAL,
             Mode::EQ,
@@ -1423,7 +1428,11 @@ pub fn resolve_ability_specialization<R: AbilityResolver>(
 
     instantiate_rigids(subs, signature_var);
     let (_vars, must_implement_ability, _lambda_sets_to_specialize, _meta) = unify(
-        &mut UEnv::new(subs),
+        // TODO(checkmate): pass checkmate through
+        &mut with_checkmate!(None, {
+            on => UEnv::new(subs, None),
+            off => UEnv::new(subs),
+        }),
         specialization_var,
         signature_var,
         Mode::EQ,

crates/compiler/solve/src/env.rs

@@ -1,8 +1,9 @@
 use bumpalo::Bump;
 use roc_can::{constraint::Constraints, module::ExposedByModule};
+use roc_checkmate::with_checkmate;
 use roc_derive::SharedDerivedModule;
 use roc_types::subs::{Content, Descriptor, Mark, OptVariable, Rank, Subs, Variable};
-use roc_unify::unify::Env as UEnv;
+use roc_unify::Env as UEnv;
 
 use crate::{FunctionKind, Pools};
 
@@ -93,7 +94,11 @@ impl<'a> InferenceEnv<'a> {
 
     /// Retrieves an environment for unification.
     pub fn uenv(&mut self) -> UEnv {
-        UEnv::new(self.subs)
+        // TODO(checkmate): pass checkmate through
+        with_checkmate!(None, {
+            on => UEnv::new(self.subs, None),
+            off => UEnv::new(self.subs),
+        })
     }
 
     pub fn as_solve_env(&mut self) -> SolveEnv {

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 => "❔",