Skip to content

Commit

Permalink
coverage: Split out SpanFromMir from CoverageSpan
Browse files Browse the repository at this point in the history
This draws a clear distinction between the fields/methods that are needed by
initial span extraction and preprocessing, and those that are needed by the
main "refinement" loop.
  • Loading branch information
Zalathar committed Dec 23, 2023
1 parent 86beddc commit 062c4d0
Show file tree
Hide file tree
Showing 2 changed files with 80 additions and 66 deletions.
59 changes: 4 additions & 55 deletions compiler/rustc_mir_transform/src/coverage/spans.rs
Original file line number Diff line number Diff line change
@@ -1,11 +1,9 @@
use std::cell::OnceCell;

use rustc_data_structures::graph::WithNumNodes;
use rustc_index::IndexVec;
use rustc_middle::mir;
use rustc_span::{BytePos, ExpnKind, MacroKind, Span, Symbol, DUMMY_SP};
use rustc_span::{BytePos, Span, DUMMY_SP};

use super::graph::{BasicCoverageBlock, CoverageGraph, START_BCB};
use super::graph::{BasicCoverageBlock, CoverageGraph};
use crate::coverage::ExtractedHirInfo;

mod from_mir;
Expand Down Expand Up @@ -66,8 +64,6 @@ impl CoverageSpans {
#[derive(Debug, Clone)]
struct CoverageSpan {
span: Span,
expn_span: Span,
current_macro_or_none: OnceCell<Option<Symbol>>,
bcb: BasicCoverageBlock,
/// List of all the original spans from MIR that have been merged into this
/// span. Mainly used to precisely skip over gaps when truncating a span.
Expand All @@ -76,24 +72,8 @@ struct CoverageSpan {
}

impl CoverageSpan {
pub fn for_fn_sig(fn_sig_span: Span) -> Self {
Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
}

pub(super) fn new(
span: Span,
expn_span: Span,
bcb: BasicCoverageBlock,
is_closure: bool,
) -> Self {
Self {
span,
expn_span,
current_macro_or_none: Default::default(),
bcb,
merged_spans: vec![span],
is_closure,
}
fn new(span: Span, bcb: BasicCoverageBlock, is_closure: bool) -> Self {
Self { span, bcb, merged_spans: vec![span], is_closure }
}

pub fn merge_from(&mut self, other: &Self) {
Expand All @@ -118,37 +98,6 @@ impl CoverageSpan {
pub fn is_in_same_bcb(&self, other: &Self) -> bool {
self.bcb == other.bcb
}

/// If the span is part of a macro, returns the macro name symbol.
pub fn current_macro(&self) -> Option<Symbol> {
self.current_macro_or_none
.get_or_init(|| {
if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
self.expn_span.ctxt().outer_expn_data().kind
{
return Some(current_macro);
}
None
})
.map(|symbol| symbol)
}

/// If the span is part of a macro, and the macro is visible (expands directly to the given
/// body_span), returns the macro name symbol.
pub fn visible_macro(&self, body_span: Span) -> Option<Symbol> {
let current_macro = self.current_macro()?;
let parent_callsite = self.expn_span.parent_callsite()?;

// In addition to matching the context of the body span, the parent callsite
// must also be the source callsite, i.e. the parent must have no parent.
let is_visible_macro =
parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span);
is_visible_macro.then_some(current_macro)
}

pub fn is_macro_expansion(&self) -> bool {
self.current_macro().is_some()
}
}

/// Converts the initial set of `CoverageSpan`s (one per MIR `Statement` or `Terminator`) into a
Expand Down
87 changes: 76 additions & 11 deletions compiler/rustc_mir_transform/src/coverage/spans/from_mir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,11 @@ use rustc_middle::mir::{
self, AggregateKind, FakeReadCause, Rvalue, Statement, StatementKind, Terminator,
TerminatorKind,
};
use rustc_span::Span;
use rustc_span::{ExpnKind, MacroKind, Span, Symbol};

use crate::coverage::graph::{BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph};
use crate::coverage::graph::{
BasicCoverageBlock, BasicCoverageBlockData, CoverageGraph, START_BCB,
};
use crate::coverage::spans::CoverageSpan;
use crate::coverage::ExtractedHirInfo;

Expand All @@ -17,7 +19,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
) -> Vec<CoverageSpan> {
let &ExtractedHirInfo { is_async_fn, fn_sig_span, body_span, .. } = hir_info;

let mut initial_spans = vec![CoverageSpan::for_fn_sig(fn_sig_span)];
let mut initial_spans = vec![SpanFromMir::for_fn_sig(fn_sig_span)];

if is_async_fn {
// An async function desugars into a function that returns a future,
Expand Down Expand Up @@ -57,7 +59,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
.then_with(|| Ord::cmp(&a.is_closure, &b.is_closure).reverse())
});

initial_spans
initial_spans.into_iter().map(SpanFromMir::into_coverage_span).collect::<Vec<_>>()
}

/// Macros that expand into branches (e.g. `assert!`, `trace!`) tend to generate
Expand All @@ -67,7 +69,7 @@ pub(super) fn mir_to_initial_sorted_coverage_spans(
///
/// (The input spans should be sorted in BCB dominator order, so that the
/// retained "first" span is likely to dominate the others.)
fn remove_redundant_macro_spans(initial_spans: &mut Vec<CoverageSpan>) {
fn remove_redundant_macro_spans(initial_spans: &mut Vec<SpanFromMir>) {
let mut seen_spans = FxHashSet::default();
initial_spans.retain(|covspan| {
// Ignore (retain) closure spans and non-macro-expansion spans.
Expand All @@ -84,7 +86,7 @@ fn remove_redundant_macro_spans(initial_spans: &mut Vec<CoverageSpan>) {
/// function body, split it into two parts. The first part covers just the
/// macro name plus `!`, and the second part covers the rest of the macro
/// invocation. This seems to give better results for code that uses macros.
fn split_visible_macro_spans(initial_spans: &mut Vec<CoverageSpan>, hir_info: &ExtractedHirInfo) {
fn split_visible_macro_spans(initial_spans: &mut Vec<SpanFromMir>, hir_info: &ExtractedHirInfo) {
let mut extra_spans = vec![];

initial_spans.retain(|covspan| {
Expand All @@ -105,8 +107,8 @@ fn split_visible_macro_spans(initial_spans: &mut Vec<CoverageSpan>, hir_info: &E
}

assert!(!covspan.is_closure);
extra_spans.push(CoverageSpan::new(before, covspan.expn_span, covspan.bcb, false));
extra_spans.push(CoverageSpan::new(after, covspan.expn_span, covspan.bcb, false));
extra_spans.push(SpanFromMir::new(before, covspan.expn_span, covspan.bcb, false));
extra_spans.push(SpanFromMir::new(after, covspan.expn_span, covspan.bcb, false));
false // Discard the original covspan that we just split.
});

Expand All @@ -125,22 +127,22 @@ fn bcb_to_initial_coverage_spans<'a, 'tcx>(
body_span: Span,
bcb: BasicCoverageBlock,
bcb_data: &'a BasicCoverageBlockData,
) -> impl Iterator<Item = CoverageSpan> + Captures<'a> + Captures<'tcx> {
) -> impl Iterator<Item = SpanFromMir> + Captures<'a> + Captures<'tcx> {
bcb_data.basic_blocks.iter().flat_map(move |&bb| {
let data = &mir_body[bb];

let statement_spans = data.statements.iter().filter_map(move |statement| {
let expn_span = filtered_statement_span(statement)?;
let span = unexpand_into_body_span(expn_span, body_span)?;

Some(CoverageSpan::new(span, expn_span, bcb, is_closure(statement)))
Some(SpanFromMir::new(span, expn_span, bcb, is_closure(statement)))
});

let terminator_span = Some(data.terminator()).into_iter().filter_map(move |terminator| {
let expn_span = filtered_terminator_span(terminator)?;
let span = unexpand_into_body_span(expn_span, body_span)?;

Some(CoverageSpan::new(span, expn_span, bcb, false))
Some(SpanFromMir::new(span, expn_span, bcb, false))
});

statement_spans.chain(terminator_span)
Expand Down Expand Up @@ -270,3 +272,66 @@ fn unexpand_into_body_span(span: Span, body_span: Span) -> Option<Span> {
let original_span = original_sp(span, body_span).with_ctxt(body_span.ctxt());
body_span.contains(original_span).then_some(original_span)
}

#[derive(Debug)]
struct SpanFromMir {
/// A copy of `expn_span` that has been "un-expanded" back to the current
/// function's `body_span`. After various intermediate processing steps,
/// this span is emitted as part of the final coverage mappings.
///
/// With the exception of `fn_sig_span`, this should always be contained
/// within `body_span`.
span: Span,
/// The actual span that was extracted from MIR, used to look up information
/// about macro expansions.
expn_span: Span,
current_macro_or_none: std::cell::OnceCell<Option<Symbol>>,
bcb: BasicCoverageBlock,
is_closure: bool,
}

impl SpanFromMir {
fn for_fn_sig(fn_sig_span: Span) -> Self {
Self::new(fn_sig_span, fn_sig_span, START_BCB, false)
}

fn new(span: Span, expn_span: Span, bcb: BasicCoverageBlock, is_closure: bool) -> Self {
Self { span, expn_span, current_macro_or_none: Default::default(), bcb, is_closure }
}

/// If the span is part of a macro, returns the macro name symbol.
fn current_macro(&self) -> Option<Symbol> {
self.current_macro_or_none
.get_or_init(|| {
if let ExpnKind::Macro(MacroKind::Bang, current_macro) =
self.expn_span.ctxt().outer_expn_data().kind
{
return Some(current_macro);
}
None
})
.map(|symbol| symbol)
}

/// If the span is part of a macro, and the macro is visible (expands directly to the given
/// body_span), returns the macro name symbol.
fn visible_macro(&self, body_span: Span) -> Option<Symbol> {
let current_macro = self.current_macro()?;
let parent_callsite = self.expn_span.parent_callsite()?;

// In addition to matching the context of the body span, the parent callsite
// must also be the source callsite, i.e. the parent must have no parent.
let is_visible_macro =
parent_callsite.parent_callsite().is_none() && parent_callsite.eq_ctxt(body_span);
is_visible_macro.then_some(current_macro)
}

fn is_macro_expansion(&self) -> bool {
self.current_macro().is_some()
}

fn into_coverage_span(self) -> CoverageSpan {
let Self { span, expn_span: _, current_macro_or_none: _, bcb, is_closure } = self;
CoverageSpan::new(span, bcb, is_closure)
}
}

0 comments on commit 062c4d0

Please sign in to comment.