From 7fab4e28a9fa96830ab0a3bdfb308af5f2352a8c Mon Sep 17 00:00:00 2001
From: kadmin <julianknodt@gmail.com>
Date: Sat, 15 Aug 2020 06:13:50 +0000
Subject: [PATCH] Add query for layout of variants of types

---
 src/librustc_middle/query/mod.rs |  1 +
 src/librustc_middle/ty/sty.rs    | 20 +++++++++++++++++++-
 2 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/src/librustc_middle/query/mod.rs b/src/librustc_middle/query/mod.rs
index a8f6723a35605..e0d4e18bb250c 100644
--- a/src/librustc_middle/query/mod.rs
+++ b/src/librustc_middle/query/mod.rs
@@ -884,6 +884,7 @@ rustc_queries! {
         query is_sized_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
             desc { "computing whether `{}` is `Sized`", env.value }
         }
+
         /// Query backing `TyS::is_freeze`.
         query is_freeze_raw(env: ty::ParamEnvAnd<'tcx, Ty<'tcx>>) -> bool {
             desc { "computing whether `{}` is freeze", env.value }
diff --git a/src/librustc_middle/ty/sty.rs b/src/librustc_middle/ty/sty.rs
index 05cd1ae456b35..7536899590bb8 100644
--- a/src/librustc_middle/ty/sty.rs
+++ b/src/librustc_middle/ty/sty.rs
@@ -19,7 +19,7 @@ use rustc_hir::def_id::DefId;
 use rustc_index::vec::Idx;
 use rustc_macros::HashStable;
 use rustc_span::symbol::{kw, Ident, Symbol};
-use rustc_target::abi::VariantIdx;
+use rustc_target::abi::{Layout, VariantIdx};
 use rustc_target::spec::abi;
 use std::borrow::Cow;
 use std::cmp::Ordering;
@@ -2251,4 +2251,22 @@ impl<'tcx> TyS<'tcx> {
     pub fn is_zst(&'tcx self, tcx: TyCtxt<'tcx>, did: DefId) -> bool {
         tcx.layout_of(tcx.param_env(did).and(self)).map(|layout| layout.is_zst()).unwrap_or(false)
     }
+
+    /// Returns an iterator over the sized layouts of the variants of a type given context and
+    /// specific definition. If the type is not an ADT, returns None.
+    pub fn layout_of_variants(
+        &'tcx self,
+        tcx: TyCtxt<'tcx>,
+    ) -> Option<impl Iterator<Item = &'tcx Layout> + 'tcx> {
+        let iter = match self.kind {
+            ty::Adt(def, _substs) => def.variants.iter().flat_map(move |var_def| {
+                tcx.layout_of(tcx.param_env(var_def.def_id).and(self))
+                    .ok()
+                    .filter(|ty_and_layout| !ty_and_layout.is_unsized())
+                    .map(|ty_and_layout| ty_and_layout.layout)
+            }),
+            _ => return None,
+        };
+        Some(iter)
+    }
 }