spade_common/
loc_map.rs

1use std::collections::{BTreeMap, HashMap};
2
3use spade_codespan::{ByteIndex, Span};
4
5use crate::location_info::Loc;
6
7/// A map that allows quick (O(log(n))) lookup of things by source code location.
8#[derive(Debug)]
9pub struct LocMap<T> {
10    inner: HashMap<usize, BTreeMap<ByteIndex, BTreeMap<ByteIndex, Vec<T>>>>,
11}
12
13impl<T> LocMap<T> {
14    pub fn new() -> Self {
15        Self {
16            inner: HashMap::new(),
17        }
18    }
19
20    pub fn insert(&mut self, thing: Loc<T>) {
21        self.inner
22            .entry(thing.file_id)
23            .or_insert(Default::default())
24            .entry(thing.span.start())
25            .or_insert(Default::default())
26            .entry(thing.span.end())
27            .or_insert(Default::default())
28            .push(thing.inner)
29    }
30
31    pub fn around(&self, loc: &Loc<()>) -> Vec<Loc<&T>> {
32        self.inner
33            .get(&loc.file_id)
34            .map(|starts| {
35                let after_start = starts.range(..=loc.span.start());
36                after_start
37                    .map(|(start, ends)| {
38                        ends.range(loc.span.end()..)
39                            .map(|(end, things)| {
40                                let with_locs = things.iter().map(|thing| {
41                                    Loc::new(thing, Span::new(*start, *end), loc.file_id)
42                                });
43                                with_locs
44                            })
45                            .flatten()
46                    })
47                    .flatten()
48                    .rev()
49                    .collect::<Vec<_>>()
50            })
51            .unwrap_or_default()
52    }
53}
54
55#[test]
56fn around_tests() {
57    let mut test = LocMap::new();
58
59    test.insert(Loc::new("a", Span::new(100, 105), 1));
60    test.insert(Loc::new("b", Span::new(262, 265), 1));
61    test.insert(Loc::new("c", Span::new(300, 305), 1));
62    test.insert(Loc::new("d", Span::new(400, 405), 1));
63    test.insert(Loc::new("e", Span::new(500, 500), 1));
64
65    test.insert(Loc::new("0", Span::new(0, 100), 1));
66    test.insert(Loc::new("1", Span::new(30, 70), 1));
67    test.insert(Loc::new("2", Span::new(40, 60), 1));
68    test.insert(Loc::new("3", Span::new(40, 70), 1));
69    test.insert(Loc::new("4", Span::new(45, 50), 1));
70
71    println!("{test:?}");
72
73    assert_eq!(
74        test.around(&Loc::new((), Span::new(263, 263), 1))
75            .iter()
76            .map(|l| *l.inner)
77            .collect::<Vec<_>>(),
78        vec!["b"]
79    );
80
81    assert_eq!(
82        test.around(&Loc::new((), Span::new(500, 500), 1))
83            .iter()
84            .map(|l| *l.inner)
85            .collect::<Vec<_>>(),
86        vec!["e"]
87    );
88
89    assert_eq!(
90        test.around(&Loc::new((), Span::new(50, 50), 1))
91            .iter()
92            .map(|l| *l.inner)
93            .collect::<Vec<_>>(),
94        vec!["4", "3", "2", "1", "0"],
95    )
96}
spade_common/loc_map.rs

spade_common/
loc_map.rs
