Skip to main content

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::Block;
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(pub Vec<Block>);
28#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct Blocks(Vec<Block>);
30
31impl BlocksBuilder {
32    pub fn new() -> Self {
33        Self(vec![])
34    }
35    pub fn alloc(&mut self, block: Block) -> 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(Block::default());
44        id
45    }
46    /// Sets an already-allocated block.
47    pub fn set_block(&mut self, id: BlockId, block: Block) {
48        self.0[id.0] = block;
49    }
50
51    /// Returns a mutable reference to an already-allocated block.
52    pub fn get_mut_block(&mut self, id: BlockId) -> &mut Block {
53        &mut self.0[id.0]
54    }
55
56    pub fn len(&self) -> usize {
57        self.0.len()
58    }
59
60    pub fn is_empty(&self) -> bool {
61        self.0.is_empty()
62    }
63
64    pub fn build(self) -> Option<Blocks> {
65        require(!self.is_empty())?;
66        Some(Blocks(self.0))
67    }
68}
69impl Blocks {
70    pub fn new_errored(_diag_added: DiagnosticAdded) -> Self {
71        Self(vec![])
72    }
73
74    pub fn get(&self) -> &Vec<Block> {
75        &self.0
76    }
77
78    pub fn len(&self) -> usize {
79        self.0.len()
80    }
81
82    pub fn is_empty(&self) -> bool {
83        self.0.is_empty()
84    }
85
86    pub fn iter(&self) -> BlocksIter<'_> {
87        self.into_iter()
88    }
89
90    // Note: It is safe to create DiagnosticAdded here, since BlocksBuilder::build() guarantees to
91    // build a non empty Blocks. The only way to create an empty Blocks is using
92    // `new_errored(DiagnosticAdded)`.
93    pub fn root_block(&self) -> Maybe<&Block> {
94        if self.is_empty() { Err(DiagnosticAdded) } else { Ok(&self.0[0]) }
95    }
96
97    pub fn has_root(&self) -> Maybe<()> {
98        if self.is_empty() { Err(DiagnosticAdded) } else { Ok(()) }
99    }
100
101    pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, Block> {
102        self.0.iter_mut()
103    }
104
105    pub fn push(&mut self, block: Block) -> BlockId {
106        let id = BlockId(self.0.len());
107        self.0.push(block);
108        id
109    }
110
111    pub fn reset_block(&mut self, block_id: BlockId, block: Block) {
112        self.0[block_id.0] = block;
113    }
114}
115impl Index<BlockId> for Blocks {
116    type Output = Block;
117
118    fn index(&self, index: BlockId) -> &Self::Output {
119        &self.0[index.0]
120    }
121}
122impl IndexMut<BlockId> for Blocks {
123    fn index_mut(&mut self, index: BlockId) -> &mut Self::Output {
124        &mut self.0[index.0]
125    }
126}
127impl<'a> IntoIterator for &'a Blocks {
128    type Item = (BlockId, &'a Block);
129    type IntoIter = BlocksIter<'a>;
130
131    fn into_iter(self) -> Self::IntoIter {
132        BlocksIter { blocks: self, index: 0 }
133    }
134}
135pub struct BlocksIter<'a> {
136    pub blocks: &'a Blocks,
137    pub index: usize,
138}
139impl<'a> Iterator for BlocksIter<'a> {
140    type Item = (BlockId, &'a Block);
141
142    fn next(&mut self) -> Option<Self::Item> {
143        self.blocks.0.get(self.index).map(|b| {
144            let res = (BlockId(self.index), b);
145            self.index += 1;
146            res
147        })
148    }
149}