spirt/
func_at.rs

1//! Traversal helpers for intra-function entities.
2//!
3//! [`FuncAt<P>`]/[`FuncAtMut<P>`] are like `(&FuncDefBody, P)`/`(&mut FuncDefBody, P`)
4//! (where `P` is some type describing a "position" in the function), except:
5//! * they only borrow the [`EntityDefs`] fields of [`FuncDefBody`]
6//!   * this can prevent borrow conflicts, especially when mutating other fields
7//!   * it also avoids accidentally accessing parts of the function definition
8//!     without going through `P` (as [`EntityDefs`] requires keys for any access)
9//! * they're dedicated types with inherent methods and trait `impl`s
10
11// NOTE(eddyb) wrong wrt lifetimes (https://github.com/rust-lang/rust-clippy/issues/5004).
12#![allow(clippy::should_implement_trait)]
13
14use crate::{
15    Context, ControlNode, ControlNodeDef, ControlRegion, ControlRegionDef, DataInst, DataInstDef,
16    EntityDefs, EntityList, EntityListIter, FuncDefBody, Type, Value,
17};
18
19/// Immutable traversal (i.e. visiting) helper for intra-function entities.
20///
21/// The point/position type `P` should be an entity or a shallow entity wrapper
22/// (e.g. [`EntityList<ControlNode>`]).
23#[derive(Copy, Clone)]
24pub struct FuncAt<'a, P: Copy> {
25    pub control_regions: &'a EntityDefs<ControlRegion>,
26    pub control_nodes: &'a EntityDefs<ControlNode>,
27    pub data_insts: &'a EntityDefs<DataInst>,
28
29    pub position: P,
30}
31
32impl<'a, P: Copy> FuncAt<'a, P> {
33    /// Reposition to `new_position`.
34    pub fn at<P2: Copy>(self, new_position: P2) -> FuncAt<'a, P2> {
35        FuncAt {
36            control_regions: self.control_regions,
37            control_nodes: self.control_nodes,
38            data_insts: self.data_insts,
39            position: new_position,
40        }
41    }
42}
43
44impl<'a> FuncAt<'a, ControlRegion> {
45    pub fn def(self) -> &'a ControlRegionDef {
46        &self.control_regions[self.position]
47    }
48
49    pub fn at_children(self) -> FuncAt<'a, EntityList<ControlNode>> {
50        self.at(self.def().children)
51    }
52}
53
54impl<'a> IntoIterator for FuncAt<'a, EntityList<ControlNode>> {
55    type IntoIter = FuncAt<'a, EntityListIter<ControlNode>>;
56    type Item = FuncAt<'a, ControlNode>;
57    fn into_iter(self) -> Self::IntoIter {
58        self.at(self.position.iter())
59    }
60}
61
62impl<'a> Iterator for FuncAt<'a, EntityListIter<ControlNode>> {
63    type Item = FuncAt<'a, ControlNode>;
64    fn next(&mut self) -> Option<Self::Item> {
65        let (next, rest) = self.position.split_first(self.control_nodes)?;
66        self.position = rest;
67        Some(self.at(next))
68    }
69}
70
71impl<'a> FuncAt<'a, ControlNode> {
72    pub fn def(self) -> &'a ControlNodeDef {
73        &self.control_nodes[self.position]
74    }
75}
76
77impl<'a> IntoIterator for FuncAt<'a, EntityList<DataInst>> {
78    type IntoIter = FuncAt<'a, EntityListIter<DataInst>>;
79    type Item = FuncAt<'a, DataInst>;
80    fn into_iter(self) -> Self::IntoIter {
81        self.at(self.position.iter())
82    }
83}
84
85impl<'a> Iterator for FuncAt<'a, EntityListIter<DataInst>> {
86    type Item = FuncAt<'a, DataInst>;
87    fn next(&mut self) -> Option<Self::Item> {
88        let (next, rest) = self.position.split_first(self.data_insts)?;
89        self.position = rest;
90        Some(self.at(next))
91    }
92}
93
94impl DoubleEndedIterator for FuncAt<'_, EntityListIter<DataInst>> {
95    fn next_back(&mut self) -> Option<Self::Item> {
96        let (prev, rest) = self.position.split_last(self.data_insts)?;
97        self.position = rest;
98        Some(self.at(prev))
99    }
100}
101
102impl<'a> FuncAt<'a, DataInst> {
103    pub fn def(self) -> &'a DataInstDef {
104        &self.data_insts[self.position]
105    }
106}
107
108impl FuncAt<'_, Value> {
109    /// Return the [`Type`] of this [`Value`] ([`Context`] used for [`Value::Const`]).
110    pub fn type_of(self, cx: &Context) -> Type {
111        match self.position {
112            Value::Const(ct) => cx[ct].ty,
113            Value::ControlRegionInput { region, input_idx } => {
114                self.at(region).def().inputs[input_idx as usize].ty
115            }
116            Value::ControlNodeOutput { control_node, output_idx } => {
117                self.at(control_node).def().outputs[output_idx as usize].ty
118            }
119            Value::DataInstOutput(inst) => cx[self.at(inst).def().form].output_type.unwrap(),
120        }
121    }
122}
123
124/// Mutable traversal (i.e. transforming) helper for intra-function entities.
125///
126/// The point/position type `P` should be an entity or a shallow entity wrapper
127/// (e.g. [`EntityList<ControlNode>`]).
128pub struct FuncAtMut<'a, P: Copy> {
129    pub control_regions: &'a mut EntityDefs<ControlRegion>,
130    pub control_nodes: &'a mut EntityDefs<ControlNode>,
131    pub data_insts: &'a mut EntityDefs<DataInst>,
132
133    pub position: P,
134}
135
136impl<'a, P: Copy> FuncAtMut<'a, P> {
137    /// Emulate a "reborrow", which is automatic only for `&mut` types.
138    pub fn reborrow(&mut self) -> FuncAtMut<'_, P> {
139        FuncAtMut {
140            control_regions: self.control_regions,
141            control_nodes: self.control_nodes,
142            data_insts: self.data_insts,
143            position: self.position,
144        }
145    }
146
147    /// Reposition to `new_position`.
148    pub fn at<P2: Copy>(self, new_position: P2) -> FuncAtMut<'a, P2> {
149        FuncAtMut {
150            control_regions: self.control_regions,
151            control_nodes: self.control_nodes,
152            data_insts: self.data_insts,
153            position: new_position,
154        }
155    }
156
157    /// Demote to a `FuncAt`, with the same `position`.
158    //
159    // FIXME(eddyb) maybe find a better name for this?
160    pub fn freeze(self) -> FuncAt<'a, P> {
161        let FuncAtMut { control_regions, control_nodes, data_insts, position } = self;
162        FuncAt { control_regions, control_nodes, data_insts, position }
163    }
164}
165
166impl<'a> FuncAtMut<'a, ControlRegion> {
167    pub fn def(self) -> &'a mut ControlRegionDef {
168        &mut self.control_regions[self.position]
169    }
170
171    pub fn at_children(mut self) -> FuncAtMut<'a, EntityList<ControlNode>> {
172        let children = self.reborrow().def().children;
173        self.at(children)
174    }
175}
176
177// HACK(eddyb) can't implement `IntoIterator` because `next` borrows `self`.
178impl<'a> FuncAtMut<'a, EntityList<ControlNode>> {
179    pub fn into_iter(self) -> FuncAtMut<'a, EntityListIter<ControlNode>> {
180        let iter = self.position.iter();
181        self.at(iter)
182    }
183}
184
185// HACK(eddyb) can't implement `Iterator` because `next` borrows `self`.
186impl FuncAtMut<'_, EntityListIter<ControlNode>> {
187    pub fn next(&mut self) -> Option<FuncAtMut<'_, ControlNode>> {
188        let (next, rest) = self.position.split_first(self.control_nodes)?;
189        self.position = rest;
190        Some(self.reborrow().at(next))
191    }
192}
193
194impl<'a> FuncAtMut<'a, ControlNode> {
195    pub fn def(self) -> &'a mut ControlNodeDef {
196        &mut self.control_nodes[self.position]
197    }
198}
199
200// HACK(eddyb) can't implement `IntoIterator` because `next` borrows `self`.
201impl<'a> FuncAtMut<'a, EntityList<DataInst>> {
202    pub fn into_iter(self) -> FuncAtMut<'a, EntityListIter<DataInst>> {
203        let iter = self.position.iter();
204        self.at(iter)
205    }
206}
207
208// HACK(eddyb) can't implement `Iterator` because `next` borrows `self`.
209impl FuncAtMut<'_, EntityListIter<DataInst>> {
210    pub fn next(&mut self) -> Option<FuncAtMut<'_, DataInst>> {
211        let (next, rest) = self.position.split_first(self.data_insts)?;
212        self.position = rest;
213        Some(self.reborrow().at(next))
214    }
215}
216
217impl<'a> FuncAtMut<'a, DataInst> {
218    pub fn def(self) -> &'a mut DataInstDef {
219        &mut self.data_insts[self.position]
220    }
221}
222
223impl FuncDefBody {
224    /// Start immutably traversing the function at `position`.
225    pub fn at<P: Copy>(&self, position: P) -> FuncAt<'_, P> {
226        FuncAt {
227            control_regions: &self.control_regions,
228            control_nodes: &self.control_nodes,
229            data_insts: &self.data_insts,
230            position,
231        }
232    }
233
234    /// Start mutably traversing the function at `position`.
235    pub fn at_mut<P: Copy>(&mut self, position: P) -> FuncAtMut<'_, P> {
236        FuncAtMut {
237            control_regions: &mut self.control_regions,
238            control_nodes: &mut self.control_nodes,
239            data_insts: &mut self.data_insts,
240            position,
241        }
242    }
243
244    /// Shorthand for `func_def_body.at(func_def_body.body)`.
245    pub fn at_body(&self) -> FuncAt<'_, ControlRegion> {
246        self.at(self.body)
247    }
248
249    /// Shorthand for `func_def_body.at_mut(func_def_body.body)`.
250    pub fn at_mut_body(&mut self) -> FuncAtMut<'_, ControlRegion> {
251        self.at_mut(self.body)
252    }
253}
spirt/func_at.rs

spirt/
func_at.rs
