cairo_lang_lowering/objects/
blocks.rs

1use std::ops::{Index, IndexMut};
2
3use cairo_lang_diagnostics::{DiagnosticAdded, Maybe};
4use cairo_lang_utils::require;
5
6use crate::FlatBlock;
7
8#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
9pub struct BlockId(pub usize);
10impl BlockId {
11    pub fn root() -> Self {
12        Self(0)
13    }
14
15    pub fn is_root(&self) -> bool {
16        self.0 == 0
17    }
18
19    pub fn next_block_id(&self) -> BlockId {
20        BlockId(self.0 + 1)
21    }
22}
23
24/// A convenient wrapper around a vector of blocks.
25/// This is used instead of id_arena, since the latter is harder to clone and modify.
26#[derive(Clone, Debug, Default, PartialEq, Eq)]
27pub struct BlocksBuilder<T>(pub Vec<T>);
28#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct Blocks<T>(Vec<T>);
30
31impl<T: Default> BlocksBuilder<T> {
32    pub fn new() -> Self {
33        Self(vec![])
34    }
35    pub fn alloc(&mut self, block: T) -> BlockId {
36        let id = BlockId(self.0.len());
37        self.0.push(block);
38        id
39    }
40    /// Allocate a new block ID. The block itself should be populated later.
41    pub fn alloc_empty(&mut self) -> BlockId {
42        let id = BlockId(self.0.len());
43        self.0.push(T::default());
44        id
45    }
46    /// Sets an already-allocated block.
47    pub fn set_block(&mut self, id: BlockId, block: T) {
48        self.0[id.0] = block;
49    }
50
51    pub fn len(&self) -> usize {
52        self.0.len()
53    }
54
55    pub fn is_empty(&self) -> bool {
56        self.0.is_empty()
57    }
58
59    pub fn build(self) -> Option<Blocks<T>> {
60        require(!self.is_empty())?;
61        Some(Blocks(self.0))
62    }
63}
64impl<T: Default> Blocks<T> {
65    pub fn new_errored(_diag_added: DiagnosticAdded) -> Self {
66        Self(vec![])
67    }
68
69    pub fn get(&self) -> &Vec<T> {
70        &self.0
71    }
72
73    pub fn len(&self) -> usize {
74        self.0.len()
75    }
76
77    pub fn is_empty(&self) -> bool {
78        self.0.is_empty()
79    }
80
81    pub fn iter(&self) -> BlocksIter<'_, T> {
82        self.into_iter()
83    }
84
85    // Note: It is safe to create DiagnosticAdded here, since BlocksBuilder::build() guarantees to
86    // build a non empty Blocks. The only way to create an empty Blocks is using
87    // `new_errored(DiagnosticAdded)`.
88    pub fn root_block(&self) -> Maybe<&T> {
89        if self.is_empty() { Err(DiagnosticAdded) } else { Ok(&self.0[0]) }
90    }
91
92    pub fn has_root(&self) -> Maybe<()> {
93        if self.is_empty() { Err(DiagnosticAdded) } else { Ok(()) }
94    }
95
96    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
97        self.0.iter_mut()
98    }
99
100    pub fn push(&mut self, block: T) -> BlockId {
101        let id = BlockId(self.0.len());
102        self.0.push(block);
103        id
104    }
105
106    pub fn reset_block(&mut self, block_id: BlockId, block: T) {
107        self.0[block_id.0] = block;
108    }
109}
110impl<T> Index<BlockId> for Blocks<T> {
111    type Output = T;
112
113    fn index(&self, index: BlockId) -> &Self::Output {
114        &self.0[index.0]
115    }
116}
117impl<T> IndexMut<BlockId> for Blocks<T> {
118    fn index_mut(&mut self, index: BlockId) -> &mut Self::Output {
119        &mut self.0[index.0]
120    }
121}
122impl<'a, T> IntoIterator for &'a Blocks<T> {
123    type Item = (BlockId, &'a T);
124    type IntoIter = BlocksIter<'a, T>;
125
126    fn into_iter(self) -> Self::IntoIter {
127        BlocksIter { blocks: self, index: 0 }
128    }
129}
130pub struct BlocksIter<'a, T> {
131    pub blocks: &'a Blocks<T>,
132    pub index: usize,
133}
134impl<'a, T> Iterator for BlocksIter<'a, T> {
135    type Item = (BlockId, &'a T);
136
137    fn next(&mut self) -> Option<Self::Item> {
138        self.blocks.0.get(self.index).map(|b| {
139            let res = (BlockId(self.index), b);
140            self.index += 1;
141            res
142        })
143    }
144}
145
146pub type FlatBlocksBuilder = BlocksBuilder<FlatBlock>;
147pub type FlatBlocks = Blocks<FlatBlock>;