yz_basic_block/arena/
mod.rs

1use crate::bb::{BasicBlock, BasicBlockInner};
2use crate::jump::ForeachTarget;
3use crate::{BbId, Label};
4use alloc::collections::{btree_map::Entry as MapEntry, BTreeMap as Map};
5use alloc::{string::String, vec::Vec};
6use core::mem::{replace, take};
7
8#[cfg(feature = "serde")]
9use serde::{Deserialize, Serialize};
10
11mod check;
12mod optimize;
13
14type ABB<S, C> = BasicBlock<S, C, BbId>;
15type LabelMap = Map<String, BbId>;
16
17#[derive(Debug)]
18#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
19pub struct Arena<S, C> {
20    // invariant: every pointer to another BB should be valid inside the arena.
21    bbs: Map<BbId, ABB<S, C>>,
22    labels: LabelMap,
23
24    // cache earliest insert point, used to speed up 'push' calls.
25    #[cfg_attr(feature = "serde", serde(skip))]
26    cache_ins_start: usize,
27}
28
29#[derive(Clone, Debug)]
30#[cfg_attr(feature = "std", derive(thiserror::Error))]
31#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
32pub enum SetBbLabelError {
33    #[cfg_attr(feature = "std", error("got invalid basic block id {0}"))]
34    InvalidId(BbId),
35
36    #[cfg_attr(
37        feature = "std",
38        error("label already exists with target = {orig_target}")
39    )]
40    LabelAlreadyExists { orig_target: BbId },
41}
42
43#[derive(Clone, Debug)]
44#[cfg_attr(feature = "std", derive(thiserror::Error))]
45#[cfg_attr(
46    feature = "std",
47    error("got offending basic block ids (from -> to) {0:?}")
48)]
49#[cfg_attr(feature = "serde", derive(Deserialize, Serialize))]
50pub struct OffendingIds(pub Vec<(BbId, BbId)>);
51
52impl<S, C> Default for Arena<S, C> {
53    #[inline]
54    fn default() -> Self {
55        Self {
56            bbs: Map::new(),
57            labels: Map::new(),
58            cache_ins_start: 0,
59        }
60    }
61}
62
63fn labels_of_bb(labels: &LabelMap, bbid: BbId) -> impl Iterator<Item = &str> {
64    labels.iter().filter_map(move |(label, &curid)| {
65        if curid == bbid {
66            Some(label.as_str())
67        } else {
68            None
69        }
70    })
71}
72
73fn set_label(
74    labels: &mut LabelMap,
75    label: String,
76    target: BbId,
77    overwrite: bool,
78) -> Result<Option<BbId>, SetBbLabelError> {
79    match labels.entry(label) {
80        MapEntry::Occupied(mut e) if overwrite => Ok(Some(replace(e.get_mut(), target))),
81        MapEntry::Occupied(e) => Err(SetBbLabelError::LabelAlreadyExists {
82            orig_target: *e.get(),
83        }),
84        MapEntry::Vacant(e) => {
85            e.insert(target);
86            Ok(None)
87        }
88    }
89}
90
91impl<S, C> Arena<S, C> {
92    #[inline(always)]
93    pub fn new() -> Self {
94        Default::default()
95    }
96
97    #[inline(always)]
98    pub fn len(&self) -> usize {
99        self.bbs.len()
100    }
101
102    #[inline(always)]
103    pub fn is_empty(&self) -> bool {
104        self.bbs.is_empty()
105    }
106
107    #[inline(always)]
108    pub fn bbs(&self) -> &Map<BbId, ABB<S, C>> {
109        &self.bbs
110    }
111
112    #[inline(always)]
113    pub fn bbs_mut(&mut self) -> &mut Map<BbId, ABB<S, C>> {
114        &mut self.bbs
115    }
116
117    #[inline(always)]
118    pub fn labels(&self) -> &LabelMap {
119        &self.labels
120    }
121
122    #[inline(always)]
123    pub fn labels_of_bb(&self, bbid: BbId) -> impl Iterator<Item = &str> {
124        labels_of_bb(&self.labels, bbid)
125    }
126
127    pub fn label2bb(&self, label: &str) -> Option<(BbId, &ABB<S, C>)> {
128        if let Some(bbid) = self.labels.get(label) {
129            if let Some(bb) = self.bbs.get(bbid) {
130                return Some((*bbid, bb));
131            }
132        }
133        None
134    }
135
136    /// If this call replaced the current label->BB-ID association,
137    /// then the old associated BBID is returned.
138    pub fn set_label(
139        &mut self,
140        label: Label,
141        target: BbId,
142        overwrite: bool,
143    ) -> Result<Option<BbId>, SetBbLabelError> {
144        if self.bbs.get(&target).is_none() {
145            return Err(SetBbLabelError::InvalidId(target));
146        }
147        set_label(&mut self.labels, label.into_owned(), target, overwrite)
148    }
149
150    pub fn shrink_to_fit(&mut self) {
151        for i in self.bbs.values_mut() {
152            if let BasicBlockInner::Concrete { statements, .. } = &mut i.inner {
153                statements.shrink_to_fit();
154            }
155        }
156    }
157}
158
159impl<S, C> ForeachTarget for Arena<S, C>
160where
161    ABB<S, C>: ForeachTarget<JumpTarget = BbId>,
162{
163    type JumpTarget = BbId;
164
165    fn foreach_target<F>(&self, mut f: F)
166    where
167        F: FnMut(&Self::JumpTarget),
168    {
169        for i in self.bbs.values() {
170            i.foreach_target(&mut f);
171        }
172        for i in self.labels.values() {
173            f(i);
174        }
175    }
176
177    fn foreach_target_mut<F>(&mut self, mut f: F)
178    where
179        F: FnMut(&mut Self::JumpTarget),
180    {
181        for i in self.bbs.values_mut() {
182            i.foreach_target_mut(&mut f);
183        }
184        for i in self.labels.values_mut() {
185            f(i);
186        }
187    }
188}
yz_basic_block/arena/mod.rs

yz_basic_block/arena/
mod.rs
