From 280f05856071ff33c472db10631f4e228ce0062d Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 14 Sep 2023 17:24:18 +0000
Subject: [PATCH] Canonicalize effect vars in new solver

---
 compiler/rustc_infer/src/infer/mod.rs         |  4 ++++
 .../src/solve/canonicalize.rs                 | 12 ++++++++--
 .../src/solve/eval_ctxt/canonical.rs          | 11 ++++++++++
 .../new-solver/canonicalize-effect-var.rs     | 22 +++++++++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 tests/ui/traits/new-solver/canonicalize-effect-var.rs

diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index abb59a789da..3706c9646af 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -1335,6 +1335,10 @@ pub fn root_const_var(&self, var: ty::ConstVid<'tcx>) -> ty::ConstVid<'tcx> {
         self.inner.borrow_mut().const_unification_table().find(var)
     }
 
+    pub fn root_effect_var(&self, var: ty::EffectVid<'tcx>) -> ty::EffectVid<'tcx> {
+        self.inner.borrow_mut().effect_unification_table().find(var)
+    }
+
     /// Resolves an int var to a rigid int type, if it was constrained to one,
     /// or else the root int var in the unification table.
     pub fn opportunistic_resolve_int_var(&self, vid: ty::IntVid) -> Ty<'tcx> {
diff --git a/compiler/rustc_trait_selection/src/solve/canonicalize.rs b/compiler/rustc_trait_selection/src/solve/canonicalize.rs
index 77c0f9f090c..f5901057a9d 100644
--- a/compiler/rustc_trait_selection/src/solve/canonicalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/canonicalize.rs
@@ -365,8 +365,16 @@ fn fold_const(&mut self, c: ty::Const<'tcx>) -> ty::Const<'tcx> {
                 // FIXME: we should fold this ty eventually
                 CanonicalVarKind::Const(ui, c.ty())
             }
-            ty::ConstKind::Infer(ty::InferConst::EffectVar(_)) => {
-                bug!("effect var has no universe")
+            ty::ConstKind::Infer(ty::InferConst::EffectVar(vid)) => {
+                assert_eq!(
+                    self.infcx.root_effect_var(vid),
+                    vid,
+                    "effect var should have been resolved"
+                );
+                let None = self.infcx.probe_effect_var(vid) else {
+                    bug!("effect var should have been resolved");
+                };
+                CanonicalVarKind::Effect
             }
             ty::ConstKind::Infer(ty::InferConst::Fresh(_)) => {
                 bug!("fresh var during canonicalization: {c:?}")
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
index 523841951b0..790f9b840f2 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
@@ -382,6 +382,17 @@ fn fold_const(&mut self, c: ty::Const<'tcx>) -> ty::Const<'tcx> {
                     }
                 }
             }
+            ty::ConstKind::Infer(ty::InferConst::EffectVar(vid)) => {
+                debug_assert_eq!(c.ty(), self.infcx.tcx.types.bool);
+                match self.infcx.probe_effect_var(vid) {
+                    Some(c) => c.as_const(self.infcx.tcx),
+                    None => ty::Const::new_infer(
+                        self.infcx.tcx,
+                        ty::InferConst::EffectVar(self.infcx.root_effect_var(vid)),
+                        self.infcx.tcx.types.bool,
+                    ),
+                }
+            }
             _ => {
                 if c.has_infer() {
                     c.super_fold_with(self)
diff --git a/tests/ui/traits/new-solver/canonicalize-effect-var.rs b/tests/ui/traits/new-solver/canonicalize-effect-var.rs
new file mode 100644
index 00000000000..35b69ed1a6b
--- /dev/null
+++ b/tests/ui/traits/new-solver/canonicalize-effect-var.rs
@@ -0,0 +1,22 @@
+// compile-flags: -Ztrait-solver=next
+// check-pass
+
+#![feature(effects)]
+#![feature(const_trait_impl)]
+
+#[const_trait]
+trait Foo {
+    fn foo();
+}
+
+trait Bar {}
+
+impl const Foo for i32 {
+    fn foo() {}
+}
+
+impl<T> const Foo for T where T: Bar {
+    fn foo() {}
+}
+
+fn main() {}