From 41aaf7bc468759b4775b28fc039ff07c538d4ccb Mon Sep 17 00:00:00 2001
From: David Wood <david@davidtw.co>
Date: Tue, 14 May 2019 21:34:43 +0100
Subject: [PATCH] Fix ICE with struct ctors and const generics.

This commit fixes a ICE where struct constructors were resulting in an
ICE with const generics. Previously, a `match` in `type_of` did not have
an arm for the `DefKind::Ctor` resolutions and therefore would assume
that the type did not have generics.
---
 src/librustc/hir/mod.rs                       |  7 ++
 src/librustc_typeck/collect.rs                | 98 +++++++++----------
 .../cannot-infer-type-for-const-param.rs      |  5 +-
 .../cannot-infer-type-for-const-param.stderr  | 21 +---
 .../issue-60818-struct-constructors.rs        | 10 ++
 .../issue-60818-struct-constructors.stderr    |  6 ++
 6 files changed, 74 insertions(+), 73 deletions(-)
 create mode 100644 src/test/ui/const-generics/issue-60818-struct-constructors.rs
 create mode 100644 src/test/ui/const-generics/issue-60818-struct-constructors.stderr

diff --git a/src/librustc/hir/mod.rs b/src/librustc/hir/mod.rs
index 1a6f5d3733e7a..f03a8ddc90825 100644
--- a/src/librustc/hir/mod.rs
+++ b/src/librustc/hir/mod.rs
@@ -425,6 +425,13 @@ impl GenericArg {
             GenericArg::Const(c) => c.value.hir_id,
         }
     }
+
+    pub fn is_const(&self) -> bool {
+        match self {
+            GenericArg::Const(_) => true,
+            _ => false,
+        }
+    }
 }
 
 #[derive(Clone, RustcEncodable, RustcDecodable, Debug, HashStable)]
diff --git a/src/librustc_typeck/collect.rs b/src/librustc_typeck/collect.rs
index 00990a5c5b579..3806fd0998b5e 100644
--- a/src/librustc_typeck/collect.rs
+++ b/src/librustc_typeck/collect.rs
@@ -27,7 +27,7 @@ use rustc::ty::subst::{Subst, InternalSubsts};
 use rustc::ty::util::Discr;
 use rustc::ty::util::IntTypeExt;
 use rustc::ty::subst::UnpackedKind;
-use rustc::ty::{self, AdtKind, ToPolyTraitRef, Ty, TyCtxt};
+use rustc::ty::{self, AdtKind, DefIdTree, ToPolyTraitRef, Ty, TyCtxt};
 use rustc::ty::{ReprOptions, ToPredicate};
 use rustc::util::captures::Captures;
 use rustc::util::nodemap::FxHashMap;
@@ -1349,65 +1349,61 @@ pub fn checked_type_of<'a, 'tcx>(
 
                     match path {
                         QPath::Resolved(_, ref path) => {
-                            let mut arg_index = 0;
-                            let mut found_const = false;
-                            for seg in &path.segments {
-                                if let Some(generic_args) = &seg.args {
-                                    let args = &generic_args.args;
-                                    for arg in args {
-                                        if let GenericArg::Const(ct) = arg {
-                                            if ct.value.hir_id == hir_id {
-                                                found_const = true;
-                                                break;
-                                            }
-                                            arg_index += 1;
-                                        }
-                                    }
-                                }
-                            }
-                            // Sanity check to make sure everything is as expected.
-                            if !found_const {
-                                if !fail {
-                                    return None;
-                                }
-                                bug!("no arg matching AnonConst in path")
-                            }
-                            match path.res {
-                                // We've encountered an `AnonConst` in some path, so we need to
-                                // figure out which generic parameter it corresponds to and return
-                                // the relevant type.
-                                Res::Def(DefKind::Struct, def_id)
-                                | Res::Def(DefKind::Union, def_id)
-                                | Res::Def(DefKind::Enum, def_id)
-                                | Res::Def(DefKind::Fn, def_id) => {
-                                    let generics = tcx.generics_of(def_id);
-                                    let mut param_index = 0;
-                                    for param in &generics.params {
-                                        if let ty::GenericParamDefKind::Const = param.kind {
-                                            if param_index == arg_index {
-                                                return Some(tcx.type_of(param.def_id));
-                                            }
-                                            param_index += 1;
-                                        }
-                                    }
-                                    // This is no generic parameter associated with the arg. This is
-                                    // probably from an extra arg where one is not needed.
-                                    return Some(tcx.types.err);
-                                }
-                                Res::Err => tcx.types.err,
-                                x => {
+                            let arg_index = path.segments.iter()
+                                .filter_map(|seg| seg.args.as_ref())
+                                .map(|generic_args| generic_args.args.as_ref())
+                                .find_map(|args| {
+                                    args.iter()
+                                        .filter(|arg| arg.is_const())
+                                        .enumerate()
+                                        .filter(|(_, arg)| arg.id() == hir_id)
+                                        .map(|(index, _)| index)
+                                        .next()
+                                })
+                                .or_else(|| {
                                     if !fail {
-                                        return None;
+                                        None
+                                    } else {
+                                        bug!("no arg matching AnonConst in path")
                                     }
+                                })?;
+
+                            // We've encountered an `AnonConst` in some path, so we need to
+                            // figure out which generic parameter it corresponds to and return
+                            // the relevant type.
+                            let generics = match path.res {
+                                Res::Def(DefKind::Ctor(..), def_id) =>
+                                    tcx.generics_of(tcx.parent(def_id).unwrap()),
+                                Res::Def(_, def_id) =>
+                                    tcx.generics_of(def_id),
+                                Res::Err =>
+                                    return Some(tcx.types.err),
+                                _ if !fail =>
+                                    return None,
+                                x => {
                                     tcx.sess.delay_span_bug(
                                         DUMMY_SP,
                                         &format!(
                                             "unexpected const parent path def {:?}", x
                                         ),
                                     );
-                                    tcx.types.err
+                                    return Some(tcx.types.err);
                                 }
-                            }
+                            };
+
+                            generics.params.iter()
+                                .filter(|param| {
+                                    if let ty::GenericParamDefKind::Const = param.kind {
+                                        true
+                                    } else {
+                                        false
+                                    }
+                                })
+                                .nth(arg_index)
+                                .map(|param| tcx.type_of(param.def_id))
+                                // This is no generic parameter associated with the arg. This is
+                                // probably from an extra arg where one is not needed.
+                                .unwrap_or(tcx.types.err)
                         }
                         x => {
                             if !fail {
diff --git a/src/test/ui/const-generics/cannot-infer-type-for-const-param.rs b/src/test/ui/const-generics/cannot-infer-type-for-const-param.rs
index 26496ec4a90b9..f592e486be951 100644
--- a/src/test/ui/const-generics/cannot-infer-type-for-const-param.rs
+++ b/src/test/ui/const-generics/cannot-infer-type-for-const-param.rs
@@ -1,8 +1,9 @@
+// compile-pass
 #![feature(const_generics)]
 //~^ WARN the feature `const_generics` is incomplete and may cause the compiler to crash
 
-// We should probably be able to infer the types here. However, this test is checking that we don't
-// get an ICE in this case. It may be modified later to not be an error.
+// This test confirms that the types can be inferred correctly for this example with const
+// generics. Previously this would ICE, and more recently error.
 
 struct Foo<const NUM_BYTES: usize>(pub [u8; NUM_BYTES]);
 
diff --git a/src/test/ui/const-generics/cannot-infer-type-for-const-param.stderr b/src/test/ui/const-generics/cannot-infer-type-for-const-param.stderr
index fb151648f2f9b..52907bbb67720 100644
--- a/src/test/ui/const-generics/cannot-infer-type-for-const-param.stderr
+++ b/src/test/ui/const-generics/cannot-infer-type-for-const-param.stderr
@@ -1,25 +1,6 @@
 warning: the feature `const_generics` is incomplete and may cause the compiler to crash
-  --> $DIR/cannot-infer-type-for-const-param.rs:1:12
+  --> $DIR/cannot-infer-type-for-const-param.rs:2:12
    |
 LL | #![feature(const_generics)]
    |            ^^^^^^^^^^^^^^
 
-error[E0282]: type annotations needed
-  --> $DIR/cannot-infer-type-for-const-param.rs:10:19
-   |
-LL |     let _ = Foo::<3>([1, 2, 3]);
-   |                   ^ cannot infer type for `{integer}`
-
-error[E0308]: mismatched types
-  --> $DIR/cannot-infer-type-for-const-param.rs:10:22
-   |
-LL |     let _ = Foo::<3>([1, 2, 3]);
-   |                      ^^^^^^^^^ expected `3`, found `3usize`
-   |
-   = note: expected type `[u8; _]`
-              found type `[u8; 3]`
-
-error: aborting due to 2 previous errors
-
-Some errors have detailed explanations: E0282, E0308.
-For more information about an error, try `rustc --explain E0282`.
diff --git a/src/test/ui/const-generics/issue-60818-struct-constructors.rs b/src/test/ui/const-generics/issue-60818-struct-constructors.rs
new file mode 100644
index 0000000000000..0b4aeae7a4a39
--- /dev/null
+++ b/src/test/ui/const-generics/issue-60818-struct-constructors.rs
@@ -0,0 +1,10 @@
+// compile-pass
+
+#![feature(const_generics)]
+//~^ WARN the feature `const_generics` is incomplete and may cause the compiler to crash
+
+struct Generic<const V: usize>;
+
+fn main() {
+    let _ = Generic::<0>;
+}
diff --git a/src/test/ui/const-generics/issue-60818-struct-constructors.stderr b/src/test/ui/const-generics/issue-60818-struct-constructors.stderr
new file mode 100644
index 0000000000000..4b8f50b9b0219
--- /dev/null
+++ b/src/test/ui/const-generics/issue-60818-struct-constructors.stderr
@@ -0,0 +1,6 @@
+warning: the feature `const_generics` is incomplete and may cause the compiler to crash
+  --> $DIR/issue-60818-struct-constructors.rs:3:12
+   |
+LL | #![feature(const_generics)]
+   |            ^^^^^^^^^^^^^^
+