Add expand_abstract_const

Adds the ability to directly expand a const to an expr without having to deal with intermediate
steps.

2 files changed, 64 insertions, 53 deletions

diff --git a/compiler/rustc_middle/src/ty/abstract_const.rs b/compiler/rustc_middle/src/ty/abstract_const.rs
index 8a8c46d6f21..7c1029e6604 100644
--- a/compiler/rustc_middle/src/ty/abstract_const.rs
+++ b/compiler/rustc_middle/src/ty/abstract_const.rs
@@ -1,5 +1,8 @@
 //! A subset of a mir body used for const evaluatability checking.
-use crate::ty::{self, Const, EarlyBinder, FallibleTypeFolder, GenericArg, TyCtxt, TypeFoldable};
+use crate::ty::{
+    self, subst::SubstsRef, Const, EarlyBinder, FallibleTypeFolder, Ty, TyCtxt, TypeFoldable,
+    TypeSuperFoldable, TypeVisitable,
+};
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir::def_id::DefId;
 
@@ -33,71 +36,79 @@ pub type BoundAbstractConst<'tcx> = Result<Option<EarlyBinder<ty::Const<'tcx>>>,
 
 impl<'tcx> TyCtxt<'tcx> {
     /// Returns a const with substs applied by
-    pub fn bound_abstract_const(
+    fn bound_abstract_const(self, uv: ty::WithOptConstParam<DefId>) -> BoundAbstractConst<'tcx> {
+        let ac = if let Some((did, param_did)) = uv.as_const_arg() {
+            self.thir_abstract_const_of_const_arg((did, param_did))
+        } else {
+            self.thir_abstract_const(uv.did)
+        };
+        Ok(ac?.map(|ac| EarlyBinder(ac)))
+    }
+
+    pub fn expand_abstract_consts<T: TypeFoldable<'tcx>>(
         self,
-        uv: ty::WithOptConstParam<DefId>,
-    ) -> BoundAbstractConst<'tcx> {
-        self.thir_abstract_const_opt_const_arg(uv).map(|ac| ac.map(|ac| EarlyBinder(ac)))
+        ac: T,
+    ) -> Result<Option<T>, ErrorGuaranteed> {
+        self._expand_abstract_consts(ac, true)
     }
-    #[inline]
-    pub fn thir_abstract_const_opt_const_arg(
+
+    pub fn expand_unevaluated_abstract_const(
         self,
-        def: ty::WithOptConstParam<DefId>,
+        did: ty::WithOptConstParam<DefId>,
+        substs: SubstsRef<'tcx>,
     ) -> Result<Option<ty::Const<'tcx>>, ErrorGuaranteed> {
-        if let Some((did, param_did)) = def.as_const_arg() {
-            self.thir_abstract_const_of_const_arg((did, param_did))
-        } else {
-            self.thir_abstract_const(def.did)
-        }
+        let Some(ac) = self.bound_abstract_const(did)? else {
+            return Ok(None);
+        };
+        let substs = self.erase_regions(substs);
+        let ac = ac.subst(self, substs);
+        self._expand_abstract_consts(ac, false)
     }
 
-    pub fn expand_bound_abstract_const(
+    fn _expand_abstract_consts<T: TypeFoldable<'tcx>>(
         self,
-        ct: BoundAbstractConst<'tcx>,
-        substs: &[GenericArg<'tcx>],
-    ) -> Result<Option<Const<'tcx>>, ErrorGuaranteed> {
+        ac: T,
+        first: bool,
+    ) -> Result<Option<T>, ErrorGuaranteed> {
         struct Expander<'tcx> {
             tcx: TyCtxt<'tcx>,
+            first: bool,
         }
+
         impl<'tcx> FallibleTypeFolder<'tcx> for Expander<'tcx> {
-            type Error = ErrorGuaranteed;
+            type Error = Option<ErrorGuaranteed>;
             fn tcx(&self) -> TyCtxt<'tcx> {
                 self.tcx
             }
-            fn try_fold_const(&mut self, c: Const<'tcx>) -> Result<Const<'tcx>, ErrorGuaranteed> {
-                use ty::ConstKind::*;
-                let uv = match c.kind() {
-                    Unevaluated(uv) => uv,
-                    Param(..) | Infer(..) | Bound(..) | Placeholder(..) | Value(..) | Error(..) => {
-                        return Ok(c);
-                    }
-                    Expr(e) => {
-                        let new_expr = match e {
-                            ty::Expr::Binop(op, l, r) => {
-                                ty::Expr::Binop(op, l.try_fold_with(self)?, r.try_fold_with(self)?)
-                            }
-                            ty::Expr::UnOp(op, v) => ty::Expr::UnOp(op, v.try_fold_with(self)?),
-                            ty::Expr::Cast(k, c, t) => {
-                                ty::Expr::Cast(k, c.try_fold_with(self)?, t.try_fold_with(self)?)
-                            }
-                            ty::Expr::FunctionCall(func, args) => ty::Expr::FunctionCall(
-                                func.try_fold_with(self)?,
-                                args.try_fold_with(self)?,
-                            ),
-                        };
-                        return Ok(self.tcx().mk_const(ty::ConstKind::Expr(new_expr), c.ty()));
+            fn try_fold_ty(&mut self, ty: Ty<'tcx>) -> Result<Ty<'tcx>, Self::Error> {
+                if ty.has_type_flags(ty::TypeFlags::HAS_CT_PROJECTION) {
+                    ty.try_super_fold_with(self)
+                } else {
+                    Ok(ty)
+                }
+            }
+            fn try_fold_const(&mut self, c: Const<'tcx>) -> Result<Const<'tcx>, Self::Error> {
+                let ct = match c.kind() {
+                    ty::ConstKind::Unevaluated(uv) => {
+                        if let Some(bac) = self.tcx.bound_abstract_const(uv.def)? {
+                            let substs = self.tcx.erase_regions(uv.substs);
+                            bac.subst(self.tcx, substs)
+                        } else if self.first {
+                            return Err(None);
+                        } else {
+                            c
+                        }
                     }
+                    _ => c,
                 };
-                let bac = self.tcx.bound_abstract_const(uv.def);
-                let ac = self.tcx.expand_bound_abstract_const(bac, uv.substs);
-                if let Ok(Some(ac)) = ac { ac.try_fold_with(self) } else { Ok(c) }
+                self.first = false;
+                ct.try_super_fold_with(self)
             }
         }
-
-        let Some(ac) = ct? else {
-            return Ok(None);
-        };
-        let ac = ac.subst(self, substs);
-        Ok(Some(ac.try_fold_with(&mut Expander { tcx: self })?))
+        match ac.try_fold_with(&mut Expander { tcx: self, first }) {
+            Ok(c) => Ok(Some(c)),
+            Err(None) => Ok(None),
+            Err(Some(e)) => Err(e),
+        }
     }
 }
diff --git a/compiler/rustc_middle/src/ty/relate.rs b/compiler/rustc_middle/src/ty/relate.rs
index 4249decc88f..84eeb81f1db 100644
--- a/compiler/rustc_middle/src/ty/relate.rs
+++ b/compiler/rustc_middle/src/ty/relate.rs
@@ -626,7 +626,7 @@ pub fn super_relate_consts<'tcx, R: TypeRelation<'tcx>>(
     // an unnormalized (i.e. unevaluated) const in the param-env.
     // FIXME(generic_const_exprs): Once we always lazily unify unevaluated constants
     // these `eval` calls can be removed.
-    if !relation.tcx().features().generic_const_exprs {
+    if !tcx.features().generic_const_exprs {
         a = a.eval(tcx, relation.param_env());
         b = b.eval(tcx, relation.param_env());
     }
@@ -647,12 +647,12 @@ pub fn super_relate_consts<'tcx, R: TypeRelation<'tcx>>(
         (ty::ConstKind::Placeholder(p1), ty::ConstKind::Placeholder(p2)) => p1 == p2,
         (ty::ConstKind::Value(a_val), ty::ConstKind::Value(b_val)) => a_val == b_val,
 
-        (ty::ConstKind::Unevaluated(au), ty::ConstKind::Unevaluated(bu))
+        (ty::ConstKind::Unevaluated(_au), ty::ConstKind::Unevaluated(_bu))
             if tcx.features().generic_const_exprs =>
         {
             if let (Ok(Some(a)), Ok(Some(b))) = (
-                tcx.expand_bound_abstract_const(tcx.bound_abstract_const(au.def), au.substs),
-                tcx.expand_bound_abstract_const(tcx.bound_abstract_const(bu.def), bu.substs),
+                tcx.expand_abstract_consts(a),
+                tcx.expand_abstract_consts(b),
             ) && a.ty() == b.ty() {
                 return relation.consts(a, b);
             } else {