mikros/rust
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Auto merge of #108820 - cjgillot:ensure-on-disk, r=oli-obk

Browse Source 
Ensure value is on the on-disk cache before returning from `ensure()`.

The current logic for `ensure()` a query just checks that the node is green in the dependency graph.
However, a lot of places use `ensure()` to prevent the query from being called later. This is the case before stealing a query result.

If the query is actually green but the value is not available in the on-disk cache, `ensure` would return, but a subsequent call to the full query would run the code, and attempt to read from a stolen value.

This PR conforms the query system to the usage by checking whether the queried value is loadable from disk before returning.

Sadly, I can't manage to craft a proper test...

Should fix all instances of "attempted to read from stolen value".
This commit is contained in:
bors 2023-03-12 14:00:28 +00:00
commit f41927f309
10 changed files with 116 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
compiler/rustc_ast_lowering/src/lib.rs
View File

@ -435,8 +435,9 @@ fn compute_hir_hash(
pub fn lower_to_hir(tcx: TyCtxt<'_>, (): ()) -> hir::Crate<'_> {
let sess = tcx.sess;
tcx.ensure().output_filenames(());
let _ = tcx.early_lint_checks(()); // Borrows `resolver_for_lowering`.
// Queries that borrow `resolver_for_lowering`.
tcx.ensure_with_value().output_filenames(());
tcx.ensure_with_value().early_lint_checks(());
let (mut resolver, krate) = tcx.resolver_for_lowering(()).steal();
let ast_index = index_crate(&resolver.node_id_to_def_id, &krate);

6
compiler/rustc_metadata/src/rmeta/encoder.rs
View File

@ -2050,13 +2050,13 @@ fn prefetch_mir(tcx: TyCtxt<'_>) {
let (encode_const, encode_opt) = should_encode_mir(tcx, def_id);
if encode_const {
tcx.ensure().mir_for_ctfe(def_id);
tcx.ensure_with_value().mir_for_ctfe(def_id);
}
if encode_opt {
tcx.ensure().optimized_mir(def_id);
tcx.ensure_with_value().optimized_mir(def_id);
}
if encode_opt || encode_const {
tcx.ensure().promoted_mir(def_id);
tcx.ensure_with_value().promoted_mir(def_id);
}
})
}

40
compiler/rustc_middle/src/ty/query.rs
View File

@ -99,6 +99,11 @@ pub struct TyCtxtEnsure<'tcx> {
pub tcx: TyCtxt<'tcx>,
}
#[derive(Copy, Clone)]
pub struct TyCtxtEnsureWithValue<'tcx> {
pub tcx: TyCtxt<'tcx>,
}
impl<'tcx> TyCtxt<'tcx> {
/// Returns a transparent wrapper for `TyCtxt`, which ensures queries
/// are executed instead of just returning their results.
@ -107,6 +112,15 @@ pub fn ensure(self) -> TyCtxtEnsure<'tcx> {
TyCtxtEnsure { tcx: self }
}
/// Returns a transparent wrapper for `TyCtxt`, which ensures queries
/// are executed instead of just returning their results.
///
/// This version verifies that the computed result exists in the cache before returning.
#[inline(always)]
pub fn ensure_with_value(self) -> TyCtxtEnsureWithValue<'tcx> {
TyCtxtEnsureWithValue { tcx: self }
}
/// Returns a transparent wrapper for `TyCtxt` which uses
/// `span` as the location of queries performed through it.
#[inline(always)]
@ -314,7 +328,31 @@ pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
Some(_) => return,
None => self.tcx.queries.$name(self.tcx, DUMMY_SP, key, QueryMode::Ensure),
None => self.tcx.queries.$name(
self.tcx,
DUMMY_SP,
key,
QueryMode::Ensure { check_cache: false },
),
};
})*
}
impl<'tcx> TyCtxtEnsureWithValue<'tcx> {
$($(#[$attr])*
#[inline(always)]
pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
let key = key.into_query_param();
opt_remap_env_constness!([$($modifiers)*][key]);
match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
Some(_) => return,
None => self.tcx.queries.$name(
self.tcx,
DUMMY_SP,
key,
QueryMode::Ensure { check_cache: true },
),
};
})*
}

11
compiler/rustc_mir_build/src/build/mod.rs
View File

@ -48,17 +48,14 @@ pub(crate) fn mir_built(
/// Construct the MIR for a given `DefId`.
fn mir_build(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> Body<'_> {
// Ensure unsafeck and abstract const building is ran before we steal the THIR.
// We can't use `ensure()` for `thir_abstract_const` as it doesn't compute the query
// if inputs are green. This can cause ICEs when calling `thir_abstract_const` after
// THIR has been stolen if we haven't computed this query yet.
match def {
ty::WithOptConstParam { did, const_param_did: Some(const_param_did) } => {
tcx.ensure().thir_check_unsafety_for_const_arg((did, const_param_did));
drop(tcx.thir_abstract_const_of_const_arg((did, const_param_did)));
tcx.ensure_with_value().thir_check_unsafety_for_const_arg((did, const_param_did));
tcx.ensure_with_value().thir_abstract_const_of_const_arg((did, const_param_did));
}
ty::WithOptConstParam { did, const_param_did: None } => {
tcx.ensure().thir_check_unsafety(did);
drop(tcx.thir_abstract_const(did));
tcx.ensure_with_value().thir_check_unsafety(did);
tcx.ensure_with_value().thir_abstract_const(did);
}
}

12
compiler/rustc_mir_transform/src/lib.rs
View File

@ -278,14 +278,14 @@ fn mir_const(tcx: TyCtxt<'_>, def: ty::WithOptConstParam<LocalDefId>) -> &Steal<
// Unsafety check uses the raw mir, so make sure it is run.
if !tcx.sess.opts.unstable_opts.thir_unsafeck {
if let Some(param_did) = def.const_param_did {
tcx.ensure().unsafety_check_result_for_const_arg((def.did, param_did));
tcx.ensure_with_value().unsafety_check_result_for_const_arg((def.did, param_did));
} else {
tcx.ensure().unsafety_check_result(def.did);
tcx.ensure_with_value().unsafety_check_result(def.did);
}
}
// has_ffi_unwind_calls query uses the raw mir, so make sure it is run.
tcx.ensure().has_ffi_unwind_calls(def.did);
tcx.ensure_with_value().has_ffi_unwind_calls(def.did);
let mut body = tcx.mir_built(def).steal();
@ -433,7 +433,7 @@ fn mir_drops_elaborated_and_const_checked(
if tcx.sess.opts.unstable_opts.drop_tracking_mir
&& let DefKind::Generator = tcx.def_kind(def.did)
{
tcx.ensure().mir_generator_witnesses(def.did);
tcx.ensure_with_value().mir_generator_witnesses(def.did);
}
let mir_borrowck = tcx.mir_borrowck_opt_const_arg(def);
@ -444,7 +444,7 @@ fn mir_drops_elaborated_and_const_checked(
// Do not compute the mir call graph without said call graph actually being used.
if inline::Inline.is_enabled(&tcx.sess) {
let _ = tcx.mir_inliner_callees(ty::InstanceDef::Item(def));
tcx.ensure_with_value().mir_inliner_callees(ty::InstanceDef::Item(def));
}
}
@ -613,7 +613,7 @@ fn inner_optimized_mir(tcx: TyCtxt<'_>, did: LocalDefId) -> Body<'_> {
// Run the `mir_for_ctfe` query, which depends on `mir_drops_elaborated_and_const_checked`
// which we are going to steal below. Thus we need to run `mir_for_ctfe` first, so it
// computes and caches its result.
Some(hir::ConstContext::ConstFn) => tcx.ensure().mir_for_ctfe(did),
Some(hir::ConstContext::ConstFn) => tcx.ensure_with_value().mir_for_ctfe(did),
None => {}
Some(other) => panic!("do not use `optimized_mir` for constants: {:?}", other),
}

1
compiler/rustc_query_impl/src/lib.rs
View File

@ -31,6 +31,7 @@
#[macro_use]
mod plumbing;
pub use plumbing::QueryCtxt;
use rustc_query_system::dep_graph::SerializedDepNodeIndex;
use rustc_query_system::query::*;
#[cfg(parallel_compiler)]
pub use rustc_query_system::query::{deadlock, QueryContext};

14
compiler/rustc_query_impl/src/on_disk_cache.rs
View File

@ -388,6 +388,12 @@ pub fn store_side_effects(&self, dep_node_index: DepNodeIndex, side_effects: Que
debug_assert!(prev.is_none());
}
/// Return whether the cached query result can be decoded.
pub fn loadable_from_disk(&self, dep_node_index: SerializedDepNodeIndex) -> bool {
self.query_result_index.contains_key(&dep_node_index)
// with_decoder is infallible, so we can stop here
}
/// Returns the cached query result if there is something in the cache for
/// the given `SerializedDepNodeIndex`; otherwise returns `None`.
pub fn try_load_query_result<'tcx, T>(
@ -398,7 +404,9 @@ pub fn try_load_query_result<'tcx, T>(
where
T: for<'a> Decodable<CacheDecoder<'a, 'tcx>>,
{
self.load_indexed(tcx, dep_node_index, &self.query_result_index)
let opt_value = self.load_indexed(tcx, dep_node_index, &self.query_result_index);
debug_assert_eq!(opt_value.is_some(), self.loadable_from_disk(dep_node_index));
opt_value
}
/// Stores side effect emitted during computation of an anonymous query.
@ -428,8 +436,8 @@ fn load_indexed<'tcx, T>(
T: for<'a> Decodable<CacheDecoder<'a, 'tcx>>,
{
let pos = index.get(&dep_node_index).cloned()?;
self.with_decoder(tcx, pos, |decoder| Some(decode_tagged(decoder, dep_node_index)))
let value = self.with_decoder(tcx, pos, |decoder| decode_tagged(decoder, dep_node_index));
Some(value)
}
fn with_decoder<'a, 'tcx, T, F: for<'s> FnOnce(&mut CacheDecoder<'s, 'tcx>) -> T>(

23
compiler/rustc_query_impl/src/plumbing.rs
View File

@ -364,6 +364,14 @@ fn try_load_from_on_disk_cache<'tcx, Q>(query: Q, tcx: TyCtxt<'tcx>, dep_node: D
}
}
pub(crate) fn loadable_from_disk<'tcx>(tcx: QueryCtxt<'tcx>, id: SerializedDepNodeIndex) -> bool {
if let Some(cache) = tcx.on_disk_cache().as_ref() {
cache.loadable_from_disk(id)
} else {
false
}
}
pub(crate) fn try_load_from_disk<'tcx, V>(
tcx: QueryCtxt<'tcx>,
id: SerializedDepNodeIndex,
@ -535,6 +543,21 @@ fn try_load_from_disk(
})
}
#[inline]
fn loadable_from_disk(
self,
_qcx: QueryCtxt<'tcx>,
_key: &Self::Key,
_index: SerializedDepNodeIndex,
) -> bool {
should_ever_cache_on_disk!([$($modifiers)*] {
self.cache_on_disk(_qcx.tcx, _key) &&
$crate::plumbing::loadable_from_disk(_qcx, _index)
} {
false
})
}
#[inline(always)]
fn anon(self) -> bool {
is_anon!([$($modifiers)*])

2
compiler/rustc_query_system/src/query/config.rs
View File

@ -43,6 +43,8 @@ fn query_cache<'a>(self, tcx: Qcx) -> &'a Self::Cache
fn try_load_from_disk(self, qcx: Qcx, idx: &Self::Key) -> TryLoadFromDisk<Qcx, Self::Value>;
fn loadable_from_disk(self, qcx: Qcx, key: &Self::Key, idx: SerializedDepNodeIndex) -> bool;
fn anon(self) -> bool;
fn eval_always(self) -> bool;
fn depth_limit(self) -> bool;

32
compiler/rustc_query_system/src/query/plumbing.rs
View File

@ -564,10 +564,17 @@ fn try_load_from_disk_and_cache_in_memory<Q, Qcx>(
// can be forced from `DepNode`.
debug_assert!(
!qcx.dep_context().fingerprint_style(dep_node.kind).reconstructible(),
"missing on-disk cache entry for {dep_node:?}"
"missing on-disk cache entry for reconstructible {dep_node:?}"
);
}
// Sanity check for the logic in `ensure`: if the node is green and the result loadable,
// we should actually be able to load it.
debug_assert!(
!query.loadable_from_disk(qcx, &key, prev_dep_node_index),
"missing on-disk cache entry for loadable {dep_node:?}"
);
// We could not load a result from the on-disk cache, so
// recompute.
let prof_timer = qcx.dep_context().profiler().query_provider();
@ -718,6 +725,7 @@ fn ensure_must_run<Q, Qcx>(
query: Q,
qcx: Qcx,
key: &Q::Key,
check_cache: bool,
) -> (bool, Option<DepNode<Qcx::DepKind>>)
where
Q: QueryConfig<Qcx>,
@ -733,7 +741,7 @@ fn ensure_must_run<Q, Qcx>(
let dep_node = query.construct_dep_node(*qcx.dep_context(), key);
let dep_graph = qcx.dep_context().dep_graph();
match dep_graph.try_mark_green(qcx, &dep_node) {
let serialized_dep_node_index = match dep_graph.try_mark_green(qcx, &dep_node) {
None => {
// A None return from `try_mark_green` means that this is either
// a new dep node or that the dep node has already been marked red.
@ -741,20 +749,28 @@ fn ensure_must_run<Q, Qcx>(
// DepNodeIndex. We must invoke the query itself. The performance cost
// this introduces should be negligible as we'll immediately hit the
// in-memory cache, or another query down the line will.
(true, Some(dep_node))
return (true, Some(dep_node));
}
Some((_, dep_node_index)) => {
Some((serialized_dep_node_index, dep_node_index)) => {
dep_graph.read_index(dep_node_index);
qcx.dep_context().profiler().query_cache_hit(dep_node_index.into());
(false, None)
serialized_dep_node_index
}
};
// We do not need the value at all, so do not check the cache.
if !check_cache {
return (false, None);
}
let loadable = query.loadable_from_disk(qcx, key, serialized_dep_node_index);
(!loadable, Some(dep_node))
}
#[derive(Debug)]
pub enum QueryMode {
Get,
Ensure,
Ensure { check_cache: bool },
}
#[inline(always)]
@ -769,8 +785,8 @@ pub fn get_query<Q, Qcx>(
Q: QueryConfig<Qcx>,
Qcx: QueryContext,
{
let dep_node = if let QueryMode::Ensure = mode {
let (must_run, dep_node) = ensure_must_run(query, qcx, &key);
let dep_node = if let QueryMode::Ensure { check_cache } = mode {
let (must_run, dep_node) = ensure_must_run(query, qcx, &key, check_cache);
if !must_run {
return None;
}