Generator go brrr

Author: nbdd0121

Cargo.lock

@@ -2949,9 +2949,8 @@ dependencies = [
 
 [[package]]
 name = "psm"
-version = "0.1.16"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "cd136ff4382c4753fc061cb9e4712ab2af263376b95bbd5bd8cd50c020b78e69"
+version = "0.1.18"
+source = "git+https://github.com/nbdd0121/stacker.git#eae7e163f2286071b4abb3c2ba85ad6632df848d"
 dependencies = [
  "cc",
 ]
@@ -4496,6 +4495,7 @@ dependencies = [
  "rustc_trait_selection",
  "rustc_ty_utils",
  "smallvec",
+ "stackful",
  "tracing",
 ]
 
@@ -4896,8 +4896,7 @@ checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3"
 [[package]]
 name = "stacker"
 version = "0.1.14"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "90939d5171a4420b3ff5fbc8954d641e7377335454c259dcb80786f3f21dc9b4"
+source = "git+https://github.com/nbdd0121/stacker.git#eae7e163f2286071b4abb3c2ba85ad6632df848d"
 dependencies = [
  "cc",
  "cfg-if 1.0.0",
@@ -4906,6 +4905,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "stackful"
+version = "0.1.3"
+source = "git+https://github.com/nbdd0121/stackful.git?branch=dev#5fedf72a244f7581bb10d063ae89919636b86ea1"
+dependencies = [
+ "cc",
+ "libc",
+ "stacker",
+]
+
 [[package]]
 name = "static_assertions"
 version = "1.1.0"

compiler/rustc_data_structures/Cargo.toml

@@ -25,7 +25,7 @@ rustc_index = { path = "../rustc_index", package = "rustc_index" }
 bitflags = "1.2.1"
 measureme = "10.0.0"
 libc = "0.2"
-stacker = "0.1.14"
+stacker = { git = "https://github.com/nbdd0121/stacker.git" }
 tempfile = "3.2"
 
 [dependencies.parking_lot]

compiler/rustc_middle/src/arena.rs

@@ -102,6 +102,8 @@ macro_rules! arena_types {
             [decode] impl_source: rustc_middle::traits::ImplSource<'tcx, ()>,
 
             [] dep_kind: rustc_middle::dep_graph::DepKindStruct,
+
+            [] steal_typeck_generator: rustc_data_structures::steal::Steal<rustc_middle::ty::TypeckResultGenerator<'tcx>>,
         ]);
     )
 }

compiler/rustc_middle/src/lib.rs

@@ -59,6 +59,7 @@
 #![feature(decl_macro)]
 #![feature(drain_filter)]
 #![feature(intra_doc_pointers)]
+#![feature(generator_trait)]
 #![recursion_limit = "512"]
 #![allow(rustc::potential_query_instability)]
 

compiler/rustc_middle/src/query/mod.rs

@@ -808,6 +808,14 @@ rustc_queries! {
         desc { "type-checking all item bodies" }
     }
 
+    query typeck_generator(key: (LocalDefId, u32)) -> (
+        &'tcx Steal<ty::TypeckResultGenerator<'tcx>>,
+        std::ops::GeneratorState<(LocalDefId, DefId), &'tcx ty::TypeckResults<'tcx>>,
+    ) {
+        no_hash
+        desc { |tcx| "type-checking `{}`, step {}", tcx.def_path_str(key.0.to_def_id()), key.1 }
+    }
+
     query typeck(key: LocalDefId) -> &'tcx ty::TypeckResults<'tcx> {
         desc { |tcx| "type-checking `{}`", tcx.def_path_str(key.to_def_id()) }
         cache_on_disk_if { true }

compiler/rustc_middle/src/ty/context.rs

@@ -377,6 +377,25 @@ pub struct GeneratorDiagnosticData<'tcx> {
     pub adjustments: ItemLocalMap<Vec<ty::adjustment::Adjustment<'tcx>>>,
 }
 
+pub struct TypeckResultGenerator<'tcx>(
+    pub  std::mem::ManuallyDrop<
+        std::pin::Pin<
+            Box<
+                dyn std::ops::Generator<
+                        Yield = (LocalDefId, DefId),
+                        Return = &'tcx TypeckResults<'tcx>,
+                    > + 'tcx,
+            >,
+        >,
+    >,
+);
+
+impl fmt::Debug for TypeckResultGenerator<'_> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple("TypeckResultGenerator").finish()
+    }
+}
+
 #[derive(TyEncodable, TyDecodable, Debug, HashStable)]
 pub struct TypeckResults<'tcx> {
     /// The `HirId::owner` all `ItemLocalId`s in this table are relative to.

compiler/rustc_middle/src/ty/mod.rs

@@ -68,8 +68,8 @@ pub use self::consts::{
 pub use self::context::{
     tls, CanonicalUserType, CanonicalUserTypeAnnotation, CanonicalUserTypeAnnotations,
     CtxtInterners, DelaySpanBugEmitted, FreeRegionInfo, GeneratorDiagnosticData,
-    GeneratorInteriorTypeCause, GlobalCtxt, Lift, OnDiskCache, TyCtxt, TypeckResults, UserType,
-    UserTypeAnnotationIndex,
+    GeneratorInteriorTypeCause, GlobalCtxt, Lift, OnDiskCache, TyCtxt, TypeckResultGenerator,
+    TypeckResults, UserType, UserTypeAnnotationIndex,
 };
 pub use self::instance::{Instance, InstanceDef};
 pub use self::list::List;

compiler/rustc_query_impl/src/keys.rs

@@ -513,3 +513,16 @@ impl<'tcx> Key for (Ty<'tcx>, ty::ValTree<'tcx>) {
         DUMMY_SP
     }
 }
+
+impl Key for (LocalDefId, u32) {
+    #[inline(always)]
+    fn query_crate_is_local(&self) -> bool {
+        true
+    }
+    fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
+        self.0.to_def_id().default_span(tcx)
+    }
+    fn key_as_def_id(&self) -> Option<DefId> {
+        Some(self.0.to_def_id())
+    }
+}

compiler/rustc_typeck/Cargo.toml

@@ -29,3 +29,4 @@ rustc_trait_selection = { path = "../rustc_trait_selection" }
 rustc_ty_utils = { path = "../rustc_ty_utils" }
 rustc_lint = { path = "../rustc_lint" }
 rustc_serialize = { path = "../rustc_serialize" }
+stackful = { git = "https://github.com/nbdd0121/stackful.git", branch = "dev", default-features = false, features = ["nightly"] }

compiler/rustc_typeck/src/check/fn_ctxt/_impl.rs

@@ -537,6 +537,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             did: self.tcx.hir().local_def_id(ast_c.hir_id),
             const_param_did: Some(param_def_id),
         };
+
+        if let Some(y) = self.yield_handle {
+            y.yeet((const_def.did, param_def_id));
+        }
+
         let c = ty::Const::from_opt_const_arg_anon_const(self.tcx, const_def);
         self.register_wf_obligation(
             c.into(),

compiler/rustc_typeck/src/check/inherited.rs

@@ -3,7 +3,7 @@ use super::MaybeInProgressTables;
 
 use rustc_data_structures::fx::FxHashSet;
 use rustc_hir as hir;
-use rustc_hir::def_id::{DefIdMap, LocalDefId};
+use rustc_hir::def_id::{DefId, DefIdMap, LocalDefId};
 use rustc_hir::HirIdMap;
 use rustc_infer::infer;
 use rustc_infer::infer::{InferCtxt, InferOk, TyCtxtInferExt};
@@ -59,6 +59,8 @@ pub struct Inherited<'a, 'tcx> {
     /// we record that type variable here. This is later used to inform
     /// fallback. See the `fallback` module for details.
     pub(super) diverging_type_vars: RefCell<FxHashSet<Ty<'tcx>>>,
+
+    pub(super) yield_handle: Option<&'a stackful::generator::YieldHandle<(LocalDefId, DefId)>>,
 }
 
 impl<'a, 'tcx> Deref for Inherited<'a, 'tcx> {
@@ -95,10 +97,26 @@ impl<'tcx> InheritedBuilder<'tcx> {
         let def_id = self.def_id;
         self.infcx.enter(|infcx| f(Inherited::new(infcx, def_id)))
     }
+
+    pub fn enter_with_yield_handle<F, R>(
+        &mut self,
+        yield_handle: Option<&stackful::generator::YieldHandle<(LocalDefId, DefId)>>,
+        f: F,
+    ) -> R
+    where
+        F: for<'a> FnOnce(Inherited<'a, 'tcx>) -> R,
+    {
+        let def_id = self.def_id;
+        self.infcx.enter(|infcx| f(Inherited::new_with_yield_handle(infcx, def_id, yield_handle)))
+    }
 }
 
 impl<'a, 'tcx> Inherited<'a, 'tcx> {
-    pub(super) fn new(infcx: InferCtxt<'a, 'tcx>, def_id: LocalDefId) -> Self {
+    pub(super) fn new_with_yield_handle(
+        infcx: InferCtxt<'a, 'tcx>,
+        def_id: LocalDefId,
+        yield_handle: Option<&'a stackful::generator::YieldHandle<(LocalDefId, DefId)>>,
+    ) -> Self {
         let tcx = infcx.tcx;
         let item_id = tcx.hir().local_def_id_to_hir_id(def_id);
         let body_id = tcx.hir().maybe_body_owned_by(item_id);
@@ -116,9 +134,14 @@ impl<'a, 'tcx> Inherited<'a, 'tcx> {
             deferred_generator_interiors: RefCell::new(Vec::new()),
             diverging_type_vars: RefCell::new(Default::default()),
             body_id,
+            yield_handle,
         }
     }
 
+    pub(super) fn new(infcx: InferCtxt<'a, 'tcx>, def_id: LocalDefId) -> Self {
+        Self::new_with_yield_handle(infcx, def_id, None)
+    }
+
     #[instrument(level = "debug", skip(self))]
     pub(super) fn register_predicate(&self, obligation: traits::PredicateObligation<'tcx>) {
         if obligation.has_escaping_bound_vars() {

compiler/rustc_typeck/src/check/mod.rs

@@ -104,6 +104,7 @@ pub use inherited::{Inherited, InheritedBuilder};
 use crate::astconv::AstConv;
 use crate::check::gather_locals::GatherLocalsVisitor;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_data_structures::steal::Steal;
 use rustc_errors::{
     pluralize, struct_span_err, Applicability, DiagnosticBuilder, EmissionGuarantee, MultiSpan,
 };
@@ -132,6 +133,7 @@ use rustc_trait_selection::traits::error_reporting::recursive_type_with_infinite
 use rustc_trait_selection::traits::error_reporting::suggestions::ReturnsVisitor;
 
 use std::cell::{Ref, RefCell, RefMut};
+use std::ops::{Generator, GeneratorState};
 
 use crate::require_c_abi_if_c_variadic;
 use crate::util::common::indenter;
@@ -244,6 +246,7 @@ pub fn provide(providers: &mut Providers) {
     method::provide(providers);
     *providers = Providers {
         typeck_item_bodies,
+        typeck_generator,
         typeck_const_arg,
         typeck,
         diagnostic_only_typeck,
@@ -319,20 +322,57 @@ fn used_trait_imports(tcx: TyCtxt<'_>, def_id: LocalDefId) -> &FxHashSet<LocalDe
     &*tcx.typeck(def_id).used_trait_imports
 }
 
+fn typeck_generator<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    (def_id, step): (LocalDefId, u32),
+) -> (
+    &'tcx Steal<ty::TypeckResultGenerator<'tcx>>,
+    GeneratorState<(LocalDefId, DefId), &'tcx ty::TypeckResults<'tcx>>,
+) {
+    use std::mem::ManuallyDrop;
+
+    let mut gen = if step == 0 {
+        ty::TypeckResultGenerator(ManuallyDrop::new(Box::pin(
+            stackful::generator::StackfulGenerator::new(move |y, ()| {
+                let fallback = move || tcx.type_of(def_id.to_def_id());
+                typeck_with_fallback(tcx, def_id, fallback, Some(y))
+            }),
+        )))
+    } else {
+        tcx.typeck_generator((def_id, step - 1)).0.steal()
+    };
+
+    let state = std::pin::Pin::new(&mut *gen.0).resume(());
+    let steal = tcx.arena.alloc(Steal::new(gen));
+    if let GeneratorState::Complete(_) = state {
+        drop(ManuallyDrop::into_inner(steal.steal().0));
+    }
+    (steal, state)
+}
+
 fn typeck_const_arg<'tcx>(
     tcx: TyCtxt<'tcx>,
     (did, param_did): (LocalDefId, DefId),
 ) -> &ty::TypeckResults<'tcx> {
     let fallback = move || tcx.type_of(param_did);
-    typeck_with_fallback(tcx, did, fallback)
+    typeck_with_fallback(tcx, did, fallback, None)
 }
 
 fn typeck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &ty::TypeckResults<'tcx> {
     if let Some(param_did) = tcx.opt_const_param_of(def_id) {
         tcx.typeck_const_arg((def_id, param_did))
     } else {
-        let fallback = move || tcx.type_of(def_id.to_def_id());
-        typeck_with_fallback(tcx, def_id, fallback)
+        let mut step = 0;
+        loop {
+            match tcx.typeck_generator((def_id, step)).1 {
+                GeneratorState::Yielded(y) => {
+                    // eprintln!(">>> {:?}", y);
+                    let _ = y;
+                }
+                GeneratorState::Complete(v) => break v,
+            }
+            step += 1;
+        }
     }
 }
 
@@ -343,14 +383,15 @@ fn diagnostic_only_typeck<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &ty::T
         let span = tcx.hir().span(tcx.hir().local_def_id_to_hir_id(def_id));
         tcx.ty_error_with_message(span, "diagnostic only typeck table used")
     };
-    typeck_with_fallback(tcx, def_id, fallback)
+    typeck_with_fallback(tcx, def_id, fallback, None)
 }
 
-#[instrument(skip(tcx, fallback))]
+#[instrument(skip(tcx, fallback, yield_handle))]
 fn typeck_with_fallback<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: LocalDefId,
     fallback: impl Fn() -> Ty<'tcx> + 'tcx,
+    yield_handle: Option<&stackful::generator::YieldHandle<(LocalDefId, DefId)>>,
 ) -> &'tcx ty::TypeckResults<'tcx> {
     // Closures' typeck results come from their outermost function,
     // as they are part of the same "inference environment".
@@ -368,7 +409,7 @@ fn typeck_with_fallback<'tcx>(
     });
     let body = tcx.hir().body(body_id);
 
-    let typeck_results = Inherited::build(tcx, def_id).enter(|inh| {
+    let typeck_results = Inherited::build(tcx, def_id).enter_with_yield_handle(yield_handle, |inh| {
         let param_env = tcx.param_env(def_id);
         let (fcx, wf_tys) = if let Some(hir::FnSig { header, decl, .. }) = fn_sig {
             let fn_sig = if crate::collect::get_infer_ret_ty(&decl.output).is_some() {

compiler/rustc_typeck/src/collect/type_of.rs

@@ -100,7 +100,21 @@ pub(super) fn opt_const_param_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<
             ..
         }) => {
             let body_owner = tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
-            let tables = tcx.typeck(body_owner);
+
+            let mut step = 0;
+            let tables = loop {
+                use std::ops::GeneratorState;
+                match tcx.typeck_generator((body_owner, step)).1 {
+                    GeneratorState::Yielded((const_did, const_param)) => {
+                        if const_did == def_id {
+                            return Some(const_param);
+                        }
+                    }
+                    GeneratorState::Complete(v) => break v,
+                }
+                step += 1;
+            };
+
             // This may fail in case the method/path does not actually exist.
             // As there is no relevant param for `def_id`, we simply return
             // `None` here.
@@ -134,7 +148,21 @@ pub(super) fn opt_const_param_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<
                     ..
                 }) => {
                     let body_owner = tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
-                    let _tables = tcx.typeck(body_owner);
+
+                    let mut step = 0;
+                    loop {
+                        use std::ops::GeneratorState;
+                        match tcx.typeck_generator((body_owner, step)).1 {
+                            GeneratorState::Yielded((const_did, _)) => {
+                                if const_did == def_id {
+                                    break;
+                                }
+                            }
+                            GeneratorState::Complete(_) => break,
+                        }
+                        step += 1;
+                    }
+
                     &*path
                 }
                 Node::Pat(pat) => {

compiler/rustc_typeck/src/lib.rs

@@ -73,6 +73,7 @@ This API is completely unstable and subject to change.
 #![feature(once_cell)]
 #![feature(slice_partition_dedup)]
 #![feature(try_blocks)]
+#![feature(generator_trait)]
 #![recursion_limit = "256"]
 
 #[macro_use]