Skip to main content

midenc_hir/ir/effects/
interface.rs

1use smallvec::SmallVec;
2
3use super::*;
4use crate::{OperationRef, SymbolRef, ValueRef};
5
6/// Marker trait for ops with recursive effects of a given [Effect] type.
7///
8/// Ops with recursive effects are considered to have any of the effects of operations nested within
9/// its regions, in addition to any effects it declares on itself. Only when the operation and none
10/// of its nested operations carry effects of the given type, can it be assumed that the operation
11/// is free of that effect.
12pub trait HasRecursiveEffects<T: Effect> {}
13
14impl<T: HasRecursiveMemoryEffects> HasRecursiveEffects<MemoryEffect> for T {}
15impl<T: HasRecursiveAdviceEffects> HasRecursiveEffects<AdviceEffect> for T {}
16
17pub trait EffectOpInterface<T: Effect> {
18    /// Return the set all of the operation's effects
19    fn effects(&self) -> EffectIterator<T>;
20    /// Returns true if this operation has no effects
21    fn has_no_effect(&self) -> bool {
22        self.effects().is_empty()
23    }
24    /// Return the set of effect instances that operate on the provided value
25    fn effects_on_value(&self, value: ValueRef) -> ValueEffectIterator<T> {
26        EffectIterator::for_value(self.effects(), value)
27    }
28    /// Return the set of effect instances that operate on the provided symbol
29    fn effects_on_symbol(&self, symbol: SymbolRef) -> SymbolEffectIterator<T> {
30        EffectIterator::for_symbol(self.effects(), symbol)
31    }
32    /// Return the set of effect instances that operate on the provided resource
33    fn effects_on_resource<'a, 'b: 'a>(
34        &self,
35        resource: &'b dyn Resource,
36    ) -> ResourceEffectIterator<'b, T> {
37        EffectIterator::for_resource(self.effects(), resource)
38    }
39}
40
41impl<T: Effect> dyn EffectOpInterface<T> {
42    /// Return the set all of the operation's effects that correspond to effect type `T`
43    pub fn effects_of_type<E>(&self) -> impl Iterator<Item = EffectInstance<T>> + '_
44    where
45        E: Effect,
46    {
47        self.effects().filter(|instance| instance.effect().as_any().is::<E>())
48    }
49
50    /// Returns true if the operation exhibits the given effect.
51    pub fn has_effect<E>(&self) -> bool
52    where
53        E: Effect,
54    {
55        self.effects().any(|instance| instance.effect().as_any().is::<E>())
56    }
57
58    /// Returns true if the operation only exhibits the given effect.
59    pub fn only_has_effect<E>(&self) -> bool
60    where
61        E: Effect,
62    {
63        let mut effects = self.effects();
64        !effects.is_empty() && effects.all(|instance| instance.effect().as_any().is::<E>())
65    }
66}
67
68pub struct EffectIterator<T> {
69    effects: smallvec::IntoIter<[EffectInstance<T>; 4]>,
70}
71impl<T> EffectIterator<T> {
72    pub fn from_smallvec(effects: SmallVec<[EffectInstance<T>; 4]>) -> Self {
73        Self {
74            effects: effects.into_iter(),
75        }
76    }
77
78    pub fn new(effects: impl IntoIterator<Item = EffectInstance<T>>) -> Self {
79        let effects = effects.into_iter().collect::<SmallVec<[_; 4]>>();
80        Self {
81            effects: effects.into_iter(),
82        }
83    }
84
85    pub const fn for_value(effects: Self, value: ValueRef) -> ValueEffectIterator<T> {
86        ValueEffectIterator {
87            iter: effects,
88            value,
89        }
90    }
91
92    pub const fn for_symbol(effects: Self, symbol: SymbolRef) -> SymbolEffectIterator<T> {
93        SymbolEffectIterator {
94            iter: effects,
95            symbol,
96        }
97    }
98
99    pub const fn for_resource(
100        effects: Self,
101        resource: &dyn Resource,
102    ) -> ResourceEffectIterator<'_, T> {
103        ResourceEffectIterator {
104            iter: effects,
105            resource,
106        }
107    }
108
109    #[inline]
110    pub fn as_slice(&self) -> &[EffectInstance<T>] {
111        self.effects.as_slice()
112    }
113}
114impl<T> core::iter::FusedIterator for EffectIterator<T> {}
115impl<T> ExactSizeIterator for EffectIterator<T> {
116    fn is_empty(&self) -> bool {
117        self.effects.is_empty()
118    }
119
120    fn len(&self) -> usize {
121        self.effects.len()
122    }
123}
124impl<T> Iterator for EffectIterator<T> {
125    type Item = EffectInstance<T>;
126
127    #[inline]
128    fn next(&mut self) -> Option<Self::Item> {
129        self.effects.next()
130    }
131}
132
133pub struct ValueEffectIterator<T> {
134    iter: EffectIterator<T>,
135    value: ValueRef,
136}
137impl<T> core::iter::FusedIterator for ValueEffectIterator<T> {}
138impl<T> Iterator for ValueEffectIterator<T> {
139    type Item = EffectInstance<T>;
140
141    fn next(&mut self) -> Option<Self::Item> {
142        while let Some(instance) = self.iter.next() {
143            if instance.value().is_some_and(|v| v == self.value) {
144                return Some(instance);
145            }
146        }
147
148        None
149    }
150}
151
152pub struct SymbolEffectIterator<T> {
153    iter: EffectIterator<T>,
154    symbol: SymbolRef,
155}
156impl<T> core::iter::FusedIterator for SymbolEffectIterator<T> {}
157impl<T> Iterator for SymbolEffectIterator<T> {
158    type Item = EffectInstance<T>;
159
160    fn next(&mut self) -> Option<Self::Item> {
161        while let Some(instance) = self.iter.next() {
162            if instance.symbol().is_some_and(|s| s == self.symbol) {
163                return Some(instance);
164            }
165        }
166
167        None
168    }
169}
170
171pub struct ResourceEffectIterator<'a, T> {
172    iter: EffectIterator<T>,
173    resource: &'a dyn Resource,
174}
175impl<T> core::iter::FusedIterator for ResourceEffectIterator<'_, T> {}
176impl<T> Iterator for ResourceEffectIterator<'_, T> {
177    type Item = EffectInstance<T>;
178
179    fn next(&mut self) -> Option<Self::Item> {
180        #[allow(clippy::while_let_on_iterator)]
181        while let Some(instance) = self.iter.next() {
182            if instance.resource().dyn_eq(self.resource) {
183                return Some(instance);
184            }
185        }
186
187        None
188    }
189}
190
191/// An iterator over the recursive effects of an [crate::Operation].
192///
193/// The value produced by the iterator is `(OperationRef, Option<EffectInterface<T>>)`, where the
194/// operation reference is the effecting op, and the second element is the identified effect:
195///
196/// * `Some` represents an effect on the operation or one of its nested operations
197/// * `None` indicates that we have identified that the given operation has unknown effects, and
198///   thus the entire operation could have unknown effects.
199///
200/// Note that in the case of discovering that an operation has unknown effects, the iterator can
201/// continue to visit all effects recursively - it is up to the caller to stop iteration if the
202/// presence of unknown effects makes further search wasteful.
203pub struct RecursiveEffectIterator<T> {
204    buffer: crate::adt::SmallDeque<(OperationRef, Option<EffectInstance<T>>), 2>,
205    effecting_ops: SmallVec<[OperationRef; 4]>,
206}
207
208impl<T: Effect> RecursiveEffectIterator<T> {
209    /// Iterate over the recursive effects of `op`
210    pub fn new(op: OperationRef) -> Self {
211        Self {
212            buffer: Default::default(),
213            effecting_ops: SmallVec::from_iter([op]),
214        }
215    }
216}
217
218impl<T: Effect> core::iter::FusedIterator for RecursiveEffectIterator<T> {}
219
220impl<T: Effect> Iterator for RecursiveEffectIterator<T> {
221    type Item = (OperationRef, Option<EffectInstance<T>>);
222
223    fn next(&mut self) -> Option<Self::Item> {
224        loop {
225            if let Some(next) = self.buffer.pop_front() {
226                return Some(next);
227            }
228
229            if let Some(op) = self.effecting_ops.pop() {
230                let operation = op.borrow();
231
232                let has_recursive_effects = operation.implements::<dyn HasRecursiveEffects<T>>();
233                if has_recursive_effects {
234                    for region in operation.regions() {
235                        for block in region.body() {
236                            let mut next = block.body().front().as_pointer();
237                            while let Some(nested) = next.take() {
238                                next = nested.next();
239                                self.effecting_ops.push(nested);
240                            }
241                        }
242                    }
243                }
244
245                if let Some(effect_interface) = operation.as_trait::<dyn EffectOpInterface<T>>() {
246                    self.buffer.extend(effect_interface.effects().map(|eff| (op, Some(eff))));
247                } else if !has_recursive_effects {
248                    // The operation does not have recursive memory effects or implement
249                    // EffectOpInterface, so its effects are unknown.
250                    self.buffer.push_back((op, None));
251                }
252            } else {
253                break;
254            }
255        }
256
257        None
258    }
259}