spade_common/
loc_map.rs

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

spade_common/
loc_map.rs
