1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use anyhow::Result;
use flowistry::infoflow::{self, Direction};
use itertools::Itertools;
use rustc_hir::BodyId;
use rustc_middle::ty::TyCtxt;
use rustc_span::Span;
use rustc_utils::{
  mir::borrowck_facts::get_body_with_borrowck_facts,
  source_map::{
    range::CharRange,
    spanner::{EnclosingHirSpans, Spanner},
  },
  SpanExt,
};
use serde::Serialize;

mod direct_influence;

#[derive(Debug, Serialize)]
pub struct PlaceInfo {
  pub range: CharRange,
  pub ranges: Vec<CharRange>,
  pub slice: Vec<CharRange>,
  pub direct_influence: Vec<CharRange>,
}

#[derive(Debug, Serialize)]
pub struct FocusOutput {
  pub place_info: Vec<PlaceInfo>,
  pub containers: Vec<CharRange>,
}

pub fn focus(tcx: TyCtxt, body_id: BodyId) -> Result<FocusOutput> {
  let def_id = tcx.hir().body_owner_def_id(body_id);
  let body_with_facts = get_body_with_borrowck_facts(tcx, def_id);
  let body = &body_with_facts.body;
  let results = &infoflow::compute_flow(tcx, body_id, body_with_facts);

  let source_map = tcx.sess.source_map();
  let spanner = Spanner::new(tcx, body_id, body);

  let grouped_spans = spanner
    .mir_span_tree
    .iter()
    .map(|mir_span| {
      (
        mir_span.span,
        mir_span
          .locations
          .iter()
          .map(|location| (mir_span.place, *location))
          .collect::<Vec<_>>(),
      )
    })
    .into_group_map()
    .into_iter()
    .map(|(k, vs)| (k, vs.concat()))
    .collect::<Vec<_>>();

  let targets = grouped_spans
    .iter()
    .map(|(_, target)| target.clone())
    .collect();

  let relevant =
    infoflow::compute_dependency_spans(results, targets, Direction::Both, &spanner);

  let direct =
    direct_influence::DirectInfluence::build(body, &results.analysis.place_info);

  let slices = grouped_spans
    .iter()
    .zip(relevant)
    .filter_map(|((mir_span, targets), relevant)| {
      log::debug!("Slice for {mir_span:?} is {relevant:#?}");

      let direct_influence = targets
        .iter()
        .flat_map(|(target, _)| direct.lookup(*target))
        .flat_map(|location| {
          spanner.location_to_spans(location, body, EnclosingHirSpans::None)
        })
        .filter(|span| relevant.iter().any(|slice_span| slice_span.contains(*span)))
        .collect::<Vec<_>>();

      let slice = relevant;

      let to_ranges = |v: Vec<Span>| {
        v.into_iter()
          .filter_map(|span| span.trim_leading_whitespace(source_map))
          .flatten()
          .filter_map(|span| CharRange::from_span(span, source_map).ok())
          .collect::<Vec<_>>()
      };

      log::debug!("{:#?}", to_ranges(slice.clone()));

      Some(PlaceInfo {
        range: CharRange::from_span(mir_span.span(), source_map).ok()?,
        ranges: to_ranges(vec![mir_span.span()]),
        slice: to_ranges(slice),
        direct_influence: to_ranges(direct_influence),
      })
    })
    .collect::<Vec<_>>();

  let body_range = CharRange::from_span(spanner.body_span, source_map)?;
  let ret_range = CharRange::from_span(spanner.ret_span, source_map)?;
  let mut containers = vec![body_range, ret_range];

  let hir_body = tcx.hir().body(body_id);
  let arg_span = hir_body
    .params
    .iter()
    .map(|param| param.span)
    .reduce(|s1, s2| s1.to(s2));
  if let Some(sp) = arg_span {
    containers.push(CharRange::from_span(sp, source_map)?);
  }

  Ok(FocusOutput {
    place_info: slices,
    containers,
  })
}