fugue_ir/
register.rs

1use ahash::AHashMap as Map;
2use iset::IntervalMap;
3use std::sync::Arc;
4use unsafe_unwrap::UnsafeUnwrap;
5use ustr::Ustr;
6
7use crate::space::AddressSpace;
8
9#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
10pub struct RegisterNames {
11    exact: Map<(u64, usize), Ustr>,
12    reversed: Map<Ustr, (u64, usize)>,
13    overlaps: IntervalMap<u64, Ustr>,
14    space: Arc<AddressSpace>,
15}
16
17impl RegisterNames {
18    pub fn new(space: Arc<AddressSpace>) -> Self {
19        Self {
20            exact: Map::default(),
21            reversed: Map::default(),
22            overlaps: IntervalMap::new(),
23            space,
24        }
25    }
26
27    pub fn insert(&mut self, offset: u64, size: usize, name: Ustr) {
28        self.exact.insert((offset, size), name.clone());
29        self.reversed.insert(name.clone(), (offset, size));
30        self.overlaps.insert(offset..offset + size as u64, name);
31    }
32
33    pub fn get(&self, offset: u64, size: usize) -> Option<&Ustr> {
34        if let Some(exact) = self.exact.get(&(offset, size)) {
35            return Some(exact);
36        }
37
38        let range = offset..offset + size as u64;
39        self.overlaps
40            .iter(range.clone())
41            .into_iter()
42            .find_map(|(r, v)| {
43                if r.start <= range.start && r.end >= range.end {
44                    Some(v)
45                } else {
46                    None
47                }
48            })
49    }
50
51    pub unsafe fn unchecked_get(&self, offset: u64, size: usize) -> &Ustr {
52        self.get(offset, size).unsafe_unwrap()
53    }
54
55    pub fn get_by_name<N>(&self, name: N) -> Option<(&Ustr, u64, usize)>
56    where
57        N: AsRef<str>,
58    {
59        self.reversed
60            .get_key_value(&name.as_ref().into())
61            .map(|(k, vv)| (k, vv.0, vv.1))
62    }
63
64    pub fn register_space(&self) -> &Arc<AddressSpace> {
65        &self.space
66    }
67
68    pub fn iter(&self) -> impl ExactSizeIterator<Item = (&(u64, usize), &Ustr)> {
69        self.exact.iter()
70    }
71}