Move abstract const to rustc_middle::ty

8 files changed, 309 insertions, 69 deletions

diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index cbc45526e89..bdae7e5fcd6 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -351,7 +351,7 @@ rustc_queries! {
     /// Try to build an abstract representation of the given constant.
     query thir_abstract_const(
         key: DefId
-    ) -> Result<Option<&'tcx [thir::abstract_const::Node<'tcx>]>, ErrorGuaranteed> {
+    ) -> Result<Option<&'tcx [ty::abstract_const::Node<'tcx>]>, ErrorGuaranteed> {
         desc {
             |tcx| "building an abstract representation for {}", tcx.def_path_str(key),
         }
@@ -360,7 +360,7 @@ rustc_queries! {
     /// Try to build an abstract representation of the given constant.
     query thir_abstract_const_of_const_arg(
         key: (LocalDefId, DefId)
-    ) -> Result<Option<&'tcx [thir::abstract_const::Node<'tcx>]>, ErrorGuaranteed> {
+    ) -> Result<Option<&'tcx [ty::abstract_const::Node<'tcx>]>, ErrorGuaranteed> {
         desc {
             |tcx|
             "building an abstract representation for the const argument {}",
diff --git a/compiler/rustc_middle/src/thir.rs b/compiler/rustc_middle/src/thir.rs
index 3fe6394ad7e..d5700692b9d 100644
--- a/compiler/rustc_middle/src/thir.rs
+++ b/compiler/rustc_middle/src/thir.rs
@@ -30,7 +30,6 @@ use rustc_target::asm::InlineAsmRegOrRegClass;
 use std::fmt;
 use std::ops::Index;
 
-pub mod abstract_const;
 pub mod visit;
 
 newtype_index! {
diff --git a/compiler/rustc_middle/src/thir/abstract_const.rs b/compiler/rustc_middle/src/thir/abstract_const.rs
deleted file mode 100644
index 527dbd1cd09..00000000000
--- a/compiler/rustc_middle/src/thir/abstract_const.rs
+++ /dev/null
@@ -1,61 +0,0 @@
-//! A subset of a mir body used for const evaluatability checking.
-use crate::mir;
-use crate::ty::{self, Ty, TyCtxt};
-use rustc_errors::ErrorGuaranteed;
-
-rustc_index::newtype_index! {
-    /// An index into an `AbstractConst`.
-    pub struct NodeId {
-        derive [HashStable]
-        DEBUG_FORMAT = "n{}",
-    }
-}
-
-#[derive(Debug, Clone, Copy, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
-pub enum CastKind {
-    /// thir::ExprKind::As
-    As,
-    /// thir::ExprKind::Use
-    Use,
-}
-
-/// A node of an `AbstractConst`.
-#[derive(Debug, Clone, Copy, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
-pub enum Node<'tcx> {
-    Leaf(ty::Const<'tcx>),
-    Binop(mir::BinOp, NodeId, NodeId),
-    UnaryOp(mir::UnOp, NodeId),
-    FunctionCall(NodeId, &'tcx [NodeId]),
-    Cast(CastKind, NodeId, Ty<'tcx>),
-}
-
-#[derive(Debug, Copy, Clone, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
-pub enum NotConstEvaluatable {
-    Error(ErrorGuaranteed),
-    MentionsInfer,
-    MentionsParam,
-}
-
-impl From<ErrorGuaranteed> for NotConstEvaluatable {
-    fn from(e: ErrorGuaranteed) -> NotConstEvaluatable {
-        NotConstEvaluatable::Error(e)
-    }
-}
-
-TrivialTypeTraversalAndLiftImpls! {
-    NotConstEvaluatable,
-}
-
-impl<'tcx> TyCtxt<'tcx> {
-    #[inline]
-    pub fn thir_abstract_const_opt_const_arg(
-        self,
-        def: ty::WithOptConstParam<rustc_hir::def_id::DefId>,
-    ) -> Result<Option<&'tcx [Node<'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)
-        }
-    }
-}
diff --git a/compiler/rustc_middle/src/traits/mod.rs b/compiler/rustc_middle/src/traits/mod.rs
index fe7f72024d3..955f2bdfa1d 100644
--- a/compiler/rustc_middle/src/traits/mod.rs
+++ b/compiler/rustc_middle/src/traits/mod.rs
@@ -10,7 +10,7 @@ mod structural_impls;
 pub mod util;
 
 use crate::infer::canonical::Canonical;
-use crate::thir::abstract_const::NotConstEvaluatable;
+use crate::ty::abstract_const::NotConstEvaluatable;
 use crate::ty::subst::SubstsRef;
 use crate::ty::{self, AdtKind, Ty, TyCtxt};
 
diff --git a/compiler/rustc_middle/src/ty/abstract_const.rs b/compiler/rustc_middle/src/ty/abstract_const.rs
new file mode 100644
index 00000000000..3e376d54000
--- /dev/null
+++ b/compiler/rustc_middle/src/ty/abstract_const.rs
@@ -0,0 +1,302 @@
+//! A subset of a mir body used for const evaluatability checking.
+use crate::mir;
+use crate::ty::{self, subst::Subst, DelaySpanBugEmitted, EarlyBinder, SubstsRef, Ty, TyCtxt};
+use rustc_errors::ErrorGuaranteed;
+use std::iter;
+use std::ops::ControlFlow;
+
+rustc_index::newtype_index! {
+    /// An index into an `AbstractConst`.
+    pub struct NodeId {
+        derive [HashStable]
+        DEBUG_FORMAT = "n{}",
+    }
+}
+
+/// A tree representing an anonymous constant.
+///
+/// This is only able to represent a subset of `MIR`,
+/// and should not leak any information about desugarings.
+#[derive(Debug, Clone, Copy)]
+pub struct AbstractConst<'tcx> {
+    // FIXME: Consider adding something like `IndexSlice`
+    // and use this here.
+    inner: &'tcx [Node<'tcx>],
+    substs: SubstsRef<'tcx>,
+}
+
+impl<'tcx> AbstractConst<'tcx> {
+    pub fn new(
+        tcx: TyCtxt<'tcx>,
+        uv: ty::Unevaluated<'tcx, ()>,
+    ) -> Result<Option<AbstractConst<'tcx>>, ErrorGuaranteed> {
+        let inner = tcx.thir_abstract_const_opt_const_arg(uv.def)?;
+        debug!("AbstractConst::new({:?}) = {:?}", uv, inner);
+        Ok(inner.map(|inner| AbstractConst { inner, substs: tcx.erase_regions(uv.substs) }))
+    }
+
+    pub fn from_const(
+        tcx: TyCtxt<'tcx>,
+        ct: ty::Const<'tcx>,
+    ) -> Result<Option<AbstractConst<'tcx>>, ErrorGuaranteed> {
+        match ct.kind() {
+            ty::ConstKind::Unevaluated(uv) => AbstractConst::new(tcx, uv.shrink()),
+            ty::ConstKind::Error(DelaySpanBugEmitted { reported, .. }) => Err(reported),
+            _ => Ok(None),
+        }
+    }
+
+    #[inline]
+    pub fn subtree(self, node: NodeId) -> AbstractConst<'tcx> {
+        AbstractConst { inner: &self.inner[..=node.index()], substs: self.substs }
+    }
+
+    #[inline]
+    pub fn root(self, tcx: TyCtxt<'tcx>) -> Node<'tcx> {
+        let node = self.inner.last().copied().unwrap();
+        match node {
+            Node::Leaf(leaf) => Node::Leaf(EarlyBinder(leaf).subst(tcx, self.substs)),
+            Node::Cast(kind, operand, ty) => {
+                Node::Cast(kind, operand, EarlyBinder(ty).subst(tcx, self.substs))
+            }
+            // Don't perform substitution on the following as they can't directly contain generic params
+            Node::Binop(_, _, _) | Node::UnaryOp(_, _) | Node::FunctionCall(_, _) => node,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
+pub enum CastKind {
+    /// thir::ExprKind::As
+    As,
+    /// thir::ExprKind::Use
+    Use,
+}
+
+/// A node of an `AbstractConst`.
+#[derive(Debug, Clone, Copy, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
+pub enum Node<'tcx> {
+    Leaf(ty::Const<'tcx>),
+    Binop(mir::BinOp, NodeId, NodeId),
+    UnaryOp(mir::UnOp, NodeId),
+    FunctionCall(NodeId, &'tcx [NodeId]),
+    Cast(CastKind, NodeId, Ty<'tcx>),
+}
+
+#[derive(Debug, Copy, Clone, PartialEq, Eq, HashStable, TyEncodable, TyDecodable)]
+pub enum NotConstEvaluatable {
+    Error(ErrorGuaranteed),
+    MentionsInfer,
+    MentionsParam,
+}
+
+impl From<ErrorGuaranteed> for NotConstEvaluatable {
+    fn from(e: ErrorGuaranteed) -> NotConstEvaluatable {
+        NotConstEvaluatable::Error(e)
+    }
+}
+
+TrivialTypeTraversalAndLiftImpls! {
+    NotConstEvaluatable,
+}
+
+impl<'tcx> TyCtxt<'tcx> {
+    #[inline]
+    pub fn thir_abstract_const_opt_const_arg(
+        self,
+        def: ty::WithOptConstParam<rustc_hir::def_id::DefId>,
+    ) -> Result<Option<&'tcx [Node<'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)
+        }
+    }
+}
+
+#[instrument(skip(tcx), level = "debug")]
+pub fn try_unify_abstract_consts<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    (a, b): (ty::Unevaluated<'tcx, ()>, ty::Unevaluated<'tcx, ()>),
+    param_env: ty::ParamEnv<'tcx>,
+) -> bool {
+    (|| {
+        if let Some(a) = AbstractConst::new(tcx, a)? {
+            if let Some(b) = AbstractConst::new(tcx, b)? {
+                let const_unify_ctxt = ConstUnifyCtxt { tcx, param_env };
+                return Ok(const_unify_ctxt.try_unify(a, b));
+            }
+        }
+
+        Ok(false)
+    })()
+    .unwrap_or_else(|_: ErrorGuaranteed| true)
+    // FIXME(generic_const_exprs): We should instead have this
+    // method return the resulting `ty::Const` and return `ConstKind::Error`
+    // on `ErrorGuaranteed`.
+}
+
+#[instrument(skip(tcx, f), level = "debug")]
+pub fn walk_abstract_const<'tcx, R, F>(
+    tcx: TyCtxt<'tcx>,
+    ct: AbstractConst<'tcx>,
+    mut f: F,
+) -> ControlFlow<R>
+where
+    F: FnMut(AbstractConst<'tcx>) -> ControlFlow<R>,
+{
+    #[instrument(skip(tcx, f), level = "debug")]
+    fn recurse<'tcx, R>(
+        tcx: TyCtxt<'tcx>,
+        ct: AbstractConst<'tcx>,
+        f: &mut dyn FnMut(AbstractConst<'tcx>) -> ControlFlow<R>,
+    ) -> ControlFlow<R> {
+        f(ct)?;
+        let root = ct.root(tcx);
+        debug!(?root);
+        match root {
+            Node::Leaf(_) => ControlFlow::CONTINUE,
+            Node::Binop(_, l, r) => {
+                recurse(tcx, ct.subtree(l), f)?;
+                recurse(tcx, ct.subtree(r), f)
+            }
+            Node::UnaryOp(_, v) => recurse(tcx, ct.subtree(v), f),
+            Node::FunctionCall(func, args) => {
+                recurse(tcx, ct.subtree(func), f)?;
+                args.iter().try_for_each(|&arg| recurse(tcx, ct.subtree(arg), f))
+            }
+            Node::Cast(_, operand, _) => recurse(tcx, ct.subtree(operand), f),
+        }
+    }
+
+    recurse(tcx, ct, &mut f)
+}
+
+pub struct ConstUnifyCtxt<'tcx> {
+    pub tcx: TyCtxt<'tcx>,
+    pub param_env: ty::ParamEnv<'tcx>,
+}
+
+impl<'tcx> ConstUnifyCtxt<'tcx> {
+    // Substitutes generics repeatedly to allow AbstractConsts to unify where a
+    // ConstKind::Unevaluated could be turned into an AbstractConst that would unify e.g.
+    // Param(N) should unify with Param(T), substs: [Unevaluated("T2", [Unevaluated("T3", [Param(N)])])]
+    #[inline]
+    #[instrument(skip(self), level = "debug")]
+    fn try_replace_substs_in_root(
+        &self,
+        mut abstr_const: AbstractConst<'tcx>,
+    ) -> Option<AbstractConst<'tcx>> {
+        while let Node::Leaf(ct) = abstr_const.root(self.tcx) {
+            match AbstractConst::from_const(self.tcx, ct) {
+                Ok(Some(act)) => abstr_const = act,
+                Ok(None) => break,
+                Err(_) => return None,
+            }
+        }
+
+        Some(abstr_const)
+    }
+
+    /// Tries to unify two abstract constants using structural equality.
+    #[instrument(skip(self), level = "debug")]
+    pub fn try_unify(&self, a: AbstractConst<'tcx>, b: AbstractConst<'tcx>) -> bool {
+        let a = if let Some(a) = self.try_replace_substs_in_root(a) {
+            a
+        } else {
+            return true;
+        };
+
+        let b = if let Some(b) = self.try_replace_substs_in_root(b) {
+            b
+        } else {
+            return true;
+        };
+
+        let a_root = a.root(self.tcx);
+        let b_root = b.root(self.tcx);
+        debug!(?a_root, ?b_root);
+
+        match (a_root, b_root) {
+            (Node::Leaf(a_ct), Node::Leaf(b_ct)) => {
+                let a_ct = a_ct.eval(self.tcx, self.param_env);
+                debug!("a_ct evaluated: {:?}", a_ct);
+                let b_ct = b_ct.eval(self.tcx, self.param_env);
+                debug!("b_ct evaluated: {:?}", b_ct);
+
+                if a_ct.ty() != b_ct.ty() {
+                    return false;
+                }
+
+                match (a_ct.kind(), b_ct.kind()) {
+                    // We can just unify errors with everything to reduce the amount of
+                    // emitted errors here.
+                    (ty::ConstKind::Error(_), _) | (_, ty::ConstKind::Error(_)) => true,
+                    (ty::ConstKind::Param(a_param), ty::ConstKind::Param(b_param)) => {
+                        a_param == b_param
+                    }
+                    (ty::ConstKind::Value(a_val), ty::ConstKind::Value(b_val)) => a_val == b_val,
+                    // If we have `fn a<const N: usize>() -> [u8; N + 1]` and `fn b<const M: usize>() -> [u8; 1 + M]`
+                    // we do not want to use `assert_eq!(a(), b())` to infer that `N` and `M` have to be `1`. This
+                    // means that we only allow inference variables if they are equal.
+                    (ty::ConstKind::Infer(a_val), ty::ConstKind::Infer(b_val)) => a_val == b_val,
+                    // We expand generic anonymous constants at the start of this function, so this
+                    // branch should only be taking when dealing with associated constants, at
+                    // which point directly comparing them seems like the desired behavior.
+                    //
+                    // FIXME(generic_const_exprs): This isn't actually the case.
+                    // We also take this branch for concrete anonymous constants and
+                    // expand generic anonymous constants with concrete substs.
+                    (ty::ConstKind::Unevaluated(a_uv), ty::ConstKind::Unevaluated(b_uv)) => {
+                        a_uv == b_uv
+                    }
+                    // FIXME(generic_const_exprs): We may want to either actually try
+                    // to evaluate `a_ct` and `b_ct` if they are are fully concrete or something like
+                    // this, for now we just return false here.
+                    _ => false,
+                }
+            }
+            (Node::Binop(a_op, al, ar), Node::Binop(b_op, bl, br)) if a_op == b_op => {
+                self.try_unify(a.subtree(al), b.subtree(bl))
+                    && self.try_unify(a.subtree(ar), b.subtree(br))
+            }
+            (Node::UnaryOp(a_op, av), Node::UnaryOp(b_op, bv)) if a_op == b_op => {
+                self.try_unify(a.subtree(av), b.subtree(bv))
+            }
+            (Node::FunctionCall(a_f, a_args), Node::FunctionCall(b_f, b_args))
+                if a_args.len() == b_args.len() =>
+            {
+                self.try_unify(a.subtree(a_f), b.subtree(b_f))
+                    && iter::zip(a_args, b_args)
+                        .all(|(&an, &bn)| self.try_unify(a.subtree(an), b.subtree(bn)))
+            }
+            (Node::Cast(a_kind, a_operand, a_ty), Node::Cast(b_kind, b_operand, b_ty))
+                if (a_ty == b_ty) && (a_kind == b_kind) =>
+            {
+                self.try_unify(a.subtree(a_operand), b.subtree(b_operand))
+            }
+            // use this over `_ => false` to make adding variants to `Node` less error prone
+            (Node::Cast(..), _)
+            | (Node::FunctionCall(..), _)
+            | (Node::UnaryOp(..), _)
+            | (Node::Binop(..), _)
+            | (Node::Leaf(..), _) => false,
+        }
+    }
+}
+
+// We were unable to unify the abstract constant with
+// a constant found in the caller bounds, there are
+// now three possible cases here.
+#[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub enum FailureKind {
+    /// The abstract const still references an inference
+    /// variable, in this case we return `TooGeneric`.
+    MentionsInfer,
+    /// The abstract const references a generic parameter,
+    /// this means that we emit an error here.
+    MentionsParam,
+    /// The substs are concrete enough that we can simply
+    /// try and evaluate the given constant.
+    Concrete,
+}
diff --git a/compiler/rustc_middle/src/ty/codec.rs b/compiler/rustc_middle/src/ty/codec.rs
index 9a363914dc3..e6ea3d88853 100644
--- a/compiler/rustc_middle/src/ty/codec.rs
+++ b/compiler/rustc_middle/src/ty/codec.rs
@@ -12,7 +12,6 @@ use crate::mir::{
     self,
     interpret::{AllocId, ConstAllocation},
 };
-use crate::thir;
 use crate::traits;
 use crate::ty::subst::SubstsRef;
 use crate::ty::{self, AdtDef, Ty};
@@ -346,7 +345,7 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
 }
 
 impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
-    for [thir::abstract_const::Node<'tcx>]
+    for [ty::abstract_const::Node<'tcx>]
 {
     fn decode(decoder: &mut D) -> &'tcx Self {
         decoder.interner().arena.alloc_from_iter(
@@ -356,7 +355,7 @@ impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
 }
 
 impl<'tcx, D: TyDecoder<I = TyCtxt<'tcx>>> RefDecodable<'tcx, D>
-    for [thir::abstract_const::NodeId]
+    for [ty::abstract_const::NodeId]
 {
     fn decode(decoder: &mut D) -> &'tcx Self {
         decoder.interner().arena.alloc_from_iter(
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index f15108fb750..d23019e1521 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -92,6 +92,7 @@ pub use self::sty::{
 pub use self::trait_def::TraitDef;
 
 pub mod _match;
+pub mod abstract_const;
 pub mod adjustment;
 pub mod binding;
 pub mod cast;
diff --git a/compiler/rustc_middle/src/ty/parameterized.rs b/compiler/rustc_middle/src/ty/parameterized.rs
index 54ba9e84fdb..e189ee2fc4d 100644
--- a/compiler/rustc_middle/src/ty/parameterized.rs
+++ b/compiler/rustc_middle/src/ty/parameterized.rs
@@ -3,7 +3,7 @@ use rustc_index::vec::{Idx, IndexVec};
 
 use crate::middle::exported_symbols::ExportedSymbol;
 use crate::mir::Body;
-use crate::thir::abstract_const::Node;
+use crate::ty::abstract_const::Node;
 use crate::ty::{
     self, Const, FnSig, GeneratorDiagnosticData, GenericPredicates, Predicate, TraitRef, Ty,
 };