Auto merge of #77049 - lcnr:const-eval-function-signature, r=oli-obk

const_evaluatable_checked: extend predicate collection

We now walk the hir instead of using `ty` so that we get better spans here, While I am still not completely sure if that's
what we want in the end, it does seem a lot closer to the final goal than the previous version.

We also look into type aliases (and use a `TypeVisitor` here), about which I am not completely sure, but we will see how well this works.

We also look into fn decls, so the following should work now.
```rust
fn test<T>() -> [u8; std::mem::size_of::<T>()] {
    [0; std::mem::size_of::<T>()]
}
```
Additionally, we visit the optional trait and self type of impls.

r? `@oli-obk`

9 files changed, 206 insertions, 81 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/wf.rs b/compiler/rustc_trait_selection/src/traits/wf.rs
index 998990f374c..909cd2aa155 100644
--- a/compiler/rustc_trait_selection/src/traits/wf.rs
+++ b/compiler/rustc_trait_selection/src/traits/wf.rs
@@ -7,8 +7,9 @@ use rustc_hir::lang_items::LangItem;
 use rustc_middle::ty::subst::{GenericArg, GenericArgKind, SubstsRef};
 use rustc_middle::ty::{self, ToPredicate, Ty, TyCtxt, TypeFoldable, WithConstness};
 use rustc_span::Span;
-use std::rc::Rc;
 
+use std::iter;
+use std::rc::Rc;
 /// Returns the set of obligations needed to make `arg` well-formed.
 /// If `arg` contains unresolved inference variables, this may include
 /// further WF obligations. However, if `arg` IS an unresolved
@@ -616,13 +617,24 @@ impl<'a, 'tcx> WfPredicates<'a, 'tcx> {
         def_id: DefId,
         substs: SubstsRef<'tcx>,
     ) -> Vec<traits::PredicateObligation<'tcx>> {
-        let predicates = self.infcx.tcx.predicates_of(def_id).instantiate(self.infcx.tcx, substs);
+        let predicates = self.infcx.tcx.predicates_of(def_id);
+        let mut origins = vec![def_id; predicates.predicates.len()];
+        let mut head = predicates;
+        while let Some(parent) = head.parent {
+            head = self.infcx.tcx.predicates_of(parent);
+            origins.extend(iter::repeat(parent).take(head.predicates.len()));
+        }
+
+        let predicates = predicates.instantiate(self.infcx.tcx, substs);
+        debug_assert_eq!(predicates.predicates.len(), origins.len());
+
         predicates
             .predicates
             .into_iter()
             .zip(predicates.spans.into_iter())
-            .map(|(pred, span)| {
-                let cause = self.cause(traits::BindingObligation(def_id, span));
+            .zip(origins.into_iter().rev())
+            .map(|((pred, span), origin_def_id)| {
+                let cause = self.cause(traits::BindingObligation(origin_def_id, span));
                 traits::Obligation::new(cause, self.param_env, pred)
             })
             .filter(|pred| !pred.has_escaping_bound_vars())
diff --git a/compiler/rustc_typeck/src/collect.rs b/compiler/rustc_typeck/src/collect.rs
index 458622dd65a..a571bd58abc 100644
--- a/compiler/rustc_typeck/src/collect.rs
+++ b/compiler/rustc_typeck/src/collect.rs
@@ -37,12 +37,11 @@ use rustc_middle::hir::map::Map;
 use rustc_middle::middle::codegen_fn_attrs::{CodegenFnAttrFlags, CodegenFnAttrs};
 use rustc_middle::mir::mono::Linkage;
 use rustc_middle::ty::query::Providers;
-use rustc_middle::ty::subst::{InternalSubsts, SubstsRef};
+use rustc_middle::ty::subst::InternalSubsts;
 use rustc_middle::ty::util::Discr;
 use rustc_middle::ty::util::IntTypeExt;
 use rustc_middle::ty::{self, AdtKind, Const, ToPolyTraitRef, Ty, TyCtxt};
 use rustc_middle::ty::{ReprOptions, ToPredicate, WithConstness};
-use rustc_middle::ty::{TypeFoldable, TypeVisitor};
 use rustc_session::config::SanitizerSet;
 use rustc_session::lint;
 use rustc_session::parse::feature_err;
@@ -51,8 +50,6 @@ use rustc_span::{Span, DUMMY_SP};
 use rustc_target::spec::abi;
 use rustc_trait_selection::traits::error_reporting::suggestions::NextTypeParamName;
 
-use smallvec::SmallVec;
-
 mod type_of;
 
 struct OnlySelfBounds(bool);
@@ -1676,47 +1673,10 @@ fn predicates_defined_on(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicate
         }
     }
 
-    if tcx.features().const_evaluatable_checked {
-        let const_evaluatable = const_evaluatable_predicates_of(tcx, def_id, &result);
-        result.predicates =
-            tcx.arena.alloc_from_iter(result.predicates.iter().copied().chain(const_evaluatable));
-    }
-
     debug!("predicates_defined_on({:?}) = {:?}", def_id, result);
     result
 }
 
-pub fn const_evaluatable_predicates_of<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    def_id: DefId,
-    predicates: &ty::GenericPredicates<'tcx>,
-) -> impl Iterator<Item = (ty::Predicate<'tcx>, Span)> {
-    #[derive(Default)]
-    struct ConstCollector<'tcx> {
-        ct: SmallVec<[(ty::WithOptConstParam<DefId>, SubstsRef<'tcx>, Span); 4]>,
-        curr_span: Span,
-    }
-
-    impl<'tcx> TypeVisitor<'tcx> for ConstCollector<'tcx> {
-        fn visit_const(&mut self, ct: &'tcx Const<'tcx>) -> bool {
-            if let ty::ConstKind::Unevaluated(def, substs, None) = ct.val {
-                self.ct.push((def, substs, self.curr_span));
-            }
-            false
-        }
-    }
-
-    let mut collector = ConstCollector::default();
-    for &(pred, span) in predicates.predicates.iter() {
-        collector.curr_span = span;
-        pred.visit_with(&mut collector);
-    }
-    warn!("const_evaluatable_predicates_of({:?}) = {:?}", def_id, collector.ct);
-    collector.ct.into_iter().map(move |(def_id, subst, span)| {
-        (ty::PredicateAtom::ConstEvaluatable(def_id, subst).to_predicate(tcx), span)
-    })
-}
-
 /// Returns a list of all type predicates (explicit and implicit) for the definition with
 /// ID `def_id`. This includes all predicates returned by `predicates_defined_on`, plus
 /// `Self: Trait` predicates for traits.
@@ -1754,29 +1714,6 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
 
     debug!("explicit_predicates_of(def_id={:?})", def_id);
 
-    /// A data structure with unique elements, which preserves order of insertion.
-    /// Preserving the order of insertion is important here so as not to break
-    /// compile-fail UI tests.
-    struct UniquePredicates<'tcx> {
-        predicates: FxIndexSet<(ty::Predicate<'tcx>, Span)>,
-    }
-
-    impl<'tcx> UniquePredicates<'tcx> {
-        fn new() -> Self {
-            UniquePredicates { predicates: FxIndexSet::default() }
-        }
-
-        fn push(&mut self, value: (ty::Predicate<'tcx>, Span)) {
-            self.predicates.insert(value);
-        }
-
-        fn extend<I: IntoIterator<Item = (ty::Predicate<'tcx>, Span)>>(&mut self, iter: I) {
-            for value in iter {
-                self.push(value);
-            }
-        }
-    }
-
     let hir_id = tcx.hir().local_def_id_to_hir_id(def_id.expect_local());
     let node = tcx.hir().get(hir_id);
 
@@ -1789,7 +1726,10 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
 
     const NO_GENERICS: &hir::Generics<'_> = &hir::Generics::empty();
 
-    let mut predicates = UniquePredicates::new();
+    // We use an `IndexSet` to preserves order of insertion.
+    // Preserving the order of insertion is important here so as not to break
+    // compile-fail UI tests.
+    let mut predicates: FxIndexSet<(ty::Predicate<'_>, Span)> = FxIndexSet::default();
 
     let ast_generics = match node {
         Node::TraitItem(item) => {
@@ -1891,7 +1831,7 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
     // (see below). Recall that a default impl is not itself an impl, but rather a
     // set of defaults that can be incorporated into another impl.
     if let Some(trait_ref) = is_default_impl_trait {
-        predicates.push((
+        predicates.insert((
             trait_ref.to_poly_trait_ref().without_const().to_predicate(tcx),
             tcx.def_span(def_id),
         ));
@@ -1915,7 +1855,7 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
                     hir::GenericBound::Outlives(lt) => {
                         let bound = AstConv::ast_region_to_region(&icx, &lt, None);
                         let outlives = ty::Binder::bind(ty::OutlivesPredicate(region, bound));
-                        predicates.push((outlives.to_predicate(tcx), lt.span));
+                        predicates.insert((outlives.to_predicate(tcx), lt.span));
                     }
                     _ => bug!(),
                 });
@@ -1970,7 +1910,7 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
                         let span = bound_pred.bounded_ty.span;
                         let re_root_empty = tcx.lifetimes.re_root_empty;
                         let predicate = ty::OutlivesPredicate(ty, re_root_empty);
-                        predicates.push((
+                        predicates.insert((
                             ty::PredicateAtom::TypeOutlives(predicate)
                                 .potentially_quantified(tcx, ty::PredicateKind::ForAll),
                             span,
@@ -2014,11 +1954,11 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
 
                         &hir::GenericBound::Outlives(ref lifetime) => {
                             let region = AstConv::ast_region_to_region(&icx, lifetime, None);
-                            predicates.push((
+                            predicates.insert((
                                 ty::PredicateAtom::TypeOutlives(ty::OutlivesPredicate(ty, region))
                                     .potentially_quantified(tcx, ty::PredicateKind::ForAll),
                                 lifetime.span,
-                            ))
+                            ));
                         }
                     }
                 }
@@ -2063,7 +2003,11 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
         }))
     }
 
-    let mut predicates: Vec<_> = predicates.predicates.into_iter().collect();
+    if tcx.features().const_evaluatable_checked {
+        predicates.extend(const_evaluatable_predicates_of(tcx, def_id.expect_local()));
+    }
+
+    let mut predicates: Vec<_> = predicates.into_iter().collect();
 
     // Subtle: before we store the predicates into the tcx, we
     // sort them so that predicates like `T: Foo<Item=U>` come
@@ -2089,6 +2033,97 @@ fn explicit_predicates_of(tcx: TyCtxt<'_>, def_id: DefId) -> ty::GenericPredicat
     result
 }
 
+fn const_evaluatable_predicates_of<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: LocalDefId,
+) -> FxIndexSet<(ty::Predicate<'tcx>, Span)> {
+    struct ConstCollector<'tcx> {
+        tcx: TyCtxt<'tcx>,
+        preds: FxIndexSet<(ty::Predicate<'tcx>, Span)>,
+    }
+
+    impl<'tcx> intravisit::Visitor<'tcx> for ConstCollector<'tcx> {
+        type Map = Map<'tcx>;
+
+        fn nested_visit_map(&mut self) -> intravisit::NestedVisitorMap<Self::Map> {
+            intravisit::NestedVisitorMap::None
+        }
+
+        fn visit_anon_const(&mut self, c: &'tcx hir::AnonConst) {
+            let def_id = self.tcx.hir().local_def_id(c.hir_id);
+            let ct = ty::Const::from_anon_const(self.tcx, def_id);
+            if let ty::ConstKind::Unevaluated(def, substs, None) = ct.val {
+                let span = self.tcx.hir().span(c.hir_id);
+                self.preds.insert((
+                    ty::PredicateAtom::ConstEvaluatable(def, substs).to_predicate(self.tcx),
+                    span,
+                ));
+            }
+        }
+
+        // Look into `TyAlias`.
+        fn visit_ty(&mut self, ty: &'tcx hir::Ty<'tcx>) {
+            use ty::fold::{TypeFoldable, TypeVisitor};
+            struct TyAliasVisitor<'a, 'tcx> {
+                tcx: TyCtxt<'tcx>,
+                preds: &'a mut FxIndexSet<(ty::Predicate<'tcx>, Span)>,
+                span: Span,
+            }
+
+            impl<'a, 'tcx> TypeVisitor<'tcx> for TyAliasVisitor<'a, 'tcx> {
+                fn visit_const(&mut self, ct: &'tcx Const<'tcx>) -> bool {
+                    if let ty::ConstKind::Unevaluated(def, substs, None) = ct.val {
+                        self.preds.insert((
+                            ty::PredicateAtom::ConstEvaluatable(def, substs).to_predicate(self.tcx),
+                            self.span,
+                        ));
+                    }
+                    false
+                }
+            }
+
+            if let hir::TyKind::Path(hir::QPath::Resolved(None, path)) = ty.kind {
+                if let Res::Def(DefKind::TyAlias, def_id) = path.res {
+                    let mut visitor =
+                        TyAliasVisitor { tcx: self.tcx, preds: &mut self.preds, span: path.span };
+                    self.tcx.type_of(def_id).visit_with(&mut visitor);
+                }
+            }
+
+            intravisit::walk_ty(self, ty)
+        }
+    }
+
+    let hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
+    let node = tcx.hir().get(hir_id);
+
+    let mut collector = ConstCollector { tcx, preds: FxIndexSet::default() };
+    if let hir::Node::Item(item) = node {
+        if let hir::ItemKind::Impl { ref of_trait, ref self_ty, .. } = item.kind {
+            if let Some(of_trait) = of_trait {
+                warn!("const_evaluatable_predicates_of({:?}): visit impl trait_ref", def_id);
+                collector.visit_trait_ref(of_trait);
+            }
+
+            warn!("const_evaluatable_predicates_of({:?}): visit_self_ty", def_id);
+            collector.visit_ty(self_ty);
+        }
+    }
+
+    if let Some(generics) = node.generics() {
+        warn!("const_evaluatable_predicates_of({:?}): visit_generics", def_id);
+        collector.visit_generics(generics);
+    }
+
+    if let Some(fn_sig) = tcx.hir().fn_sig_by_hir_id(hir_id) {
+        warn!("const_evaluatable_predicates_of({:?}): visit_fn_decl", def_id);
+        collector.visit_fn_decl(fn_sig.decl);
+    }
+    warn!("const_evaluatable_predicates_of({:?}) = {:?}", def_id, collector.preds);
+
+    collector.preds
+}
+
 fn projection_ty_from_predicates(
     tcx: TyCtxt<'tcx>,
     key: (
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.rs b/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.rs
index 22369923329..52b89cfa045 100644
--- a/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.rs
+++ b/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.rs
@@ -8,6 +8,7 @@ fn user<T>() {
     //~^ ERROR constant expression depends
     //~| ERROR constant expression depends
     //~| ERROR constant expression depends
+    //~| ERROR constant expression depends
 }
 
 fn main() {}
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.stderr b/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.stderr
index 63abb782b93..4af68118be3 100644
--- a/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.stderr
+++ b/src/test/ui/const-generics/const_evaluatable_checked/cross_crate_predicate.stderr
@@ -4,10 +4,23 @@ error: constant expression depends on a generic parameter
 LL |     let _ = const_evaluatable_lib::test1::<T>();
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    | 
-  ::: $DIR/auxiliary/const_evaluatable_lib.rs:6:41
+  ::: $DIR/auxiliary/const_evaluatable_lib.rs:6:10
    |
 LL |     [u8; std::mem::size_of::<T>() - 1]: Sized,
-   |                                         ----- required by this bound in `test1`
+   |          ---------------------------- required by this bound in `test1`
+   |
+   = note: this may fail depending on what value the parameter takes
+
+error: constant expression depends on a generic parameter
+  --> $DIR/cross_crate_predicate.rs:7:13
+   |
+LL |     let _ = const_evaluatable_lib::test1::<T>();
+   |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   | 
+  ::: $DIR/auxiliary/const_evaluatable_lib.rs:4:27
+   |
+LL | pub fn test1<T>() -> [u8; std::mem::size_of::<T>() - 1]
+   |                           ---------------------------- required by this bound in `test1`
    |
    = note: this may fail depending on what value the parameter takes
 
@@ -17,10 +30,10 @@ error: constant expression depends on a generic parameter
 LL |     let _ = const_evaluatable_lib::test1::<T>();
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    | 
-  ::: $DIR/auxiliary/const_evaluatable_lib.rs:6:41
+  ::: $DIR/auxiliary/const_evaluatable_lib.rs:6:10
    |
 LL |     [u8; std::mem::size_of::<T>() - 1]: Sized,
-   |                                         ----- required by this bound in `test1::{{constant}}#1`
+   |          ---------------------------- required by this bound in `test1`
    |
    = note: this may fail depending on what value the parameter takes
 
@@ -29,8 +42,13 @@ error: constant expression depends on a generic parameter
    |
 LL |     let _ = const_evaluatable_lib::test1::<T>();
    |             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+   | 
+  ::: $DIR/auxiliary/const_evaluatable_lib.rs:4:27
+   |
+LL | pub fn test1<T>() -> [u8; std::mem::size_of::<T>() - 1]
+   |                           ---------------------------- required by this bound in `test1`
    |
    = note: this may fail depending on what value the parameter takes
 
-error: aborting due to 3 previous errors
+error: aborting due to 4 previous errors
 
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.rs b/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.rs
new file mode 100644
index 00000000000..3da4688702c
--- /dev/null
+++ b/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.rs
@@ -0,0 +1,11 @@
+#![feature(const_generics, const_evaluatable_checked)]
+#![allow(incomplete_features)]
+
+fn test<const N: usize>() -> [u8; N - 1] {
+    //~^ ERROR evaluation of constant
+    todo!()
+}
+
+fn main() {
+    test::<0>();
+}
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.stderr b/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.stderr
new file mode 100644
index 00000000000..a5acfec34aa
--- /dev/null
+++ b/src/test/ui/const-generics/const_evaluatable_checked/from-sig-fail.stderr
@@ -0,0 +1,9 @@
+error[E0080]: evaluation of constant value failed
+  --> $DIR/from-sig-fail.rs:4:35
+   |
+LL | fn test<const N: usize>() -> [u8; N - 1] {
+   |                                   ^^^^^ attempt to compute `0_usize - 1_usize` which would overflow
+
+error: aborting due to previous error
+
+For more information about this error, try `rustc --explain E0080`.
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/from-sig.rs b/src/test/ui/const-generics/const_evaluatable_checked/from-sig.rs
new file mode 100644
index 00000000000..5c05a5acfe9
--- /dev/null
+++ b/src/test/ui/const-generics/const_evaluatable_checked/from-sig.rs
@@ -0,0 +1,14 @@
+// run-pass
+#![feature(const_generics, const_evaluatable_checked)]
+#![allow(incomplete_features)]
+
+struct Foo<const B: bool>;
+
+fn test<const N: usize>() -> Foo<{ N > 10 }> {
+    Foo
+}
+
+fn main() {
+    let _: Foo<true> = test::<12>();
+    let _: Foo<false> = test::<9>();
+}
diff --git a/src/test/ui/const-generics/const_evaluatable_checked/impl-bounds.rs b/src/test/ui/const-generics/const_evaluatable_checked/impl-bounds.rs
new file mode 100644
index 00000000000..193a365f9b6
--- /dev/null
+++ b/src/test/ui/const-generics/const_evaluatable_checked/impl-bounds.rs
@@ -0,0 +1,25 @@
+// check-pass
+#![feature(const_generics, const_evaluatable_checked)]
+#![allow(incomplete_features)]
+
+use std::mem::size_of;
+
+struct Foo<T, const N: usize>(T);
+
+impl<T> Foo<T, { size_of::<T>() }> {
+    fn test() {
+        let _: [u8; std::mem::size_of::<T>()];
+    }
+}
+
+trait Bar<const N: usize> {
+    fn test_me();
+}
+
+impl<T> Bar<{ size_of::<T>() }> for Foo<T, 3> {
+    fn test_me() {
+        let _: [u8; std::mem::size_of::<T>()];
+    }
+}
+
+fn main() {}
diff --git a/src/test/ui/const-generics/issues/issue-76595.stderr b/src/test/ui/const-generics/issues/issue-76595.stderr
index 2e457595393..bbc81693fc0 100644
--- a/src/test/ui/const-generics/issues/issue-76595.stderr
+++ b/src/test/ui/const-generics/issues/issue-76595.stderr
@@ -8,7 +8,7 @@ error: constant expression depends on a generic parameter
   --> $DIR/issue-76595.rs:15:5
    |
 LL | fn test<T, const P: usize>() where Bool<{core::mem::size_of::<T>() > 4}>: True {
-   |                                                                           ---- required by this bound in `test`
+   |                                         ------------------------------- required by this bound in `test`
 ...
 LL |     test::<2>();
    |     ^^^^^^^^^