naga/compact/
expressions.rs

1use super::{HandleMap, HandleSet, ModuleMap};
2use crate::arena::{Arena, Handle};
3
4pub struct ExpressionTracer<'tracer> {
5    pub constants: &'tracer Arena<crate::Constant>,
6    pub overrides: &'tracer Arena<crate::Override>,
7
8    /// The arena in which we are currently tracing expressions.
9    pub expressions: &'tracer Arena<crate::Expression>,
10
11    /// The used map for `types`.
12    pub types_used: &'tracer mut HandleSet<crate::Type>,
13
14    /// The used map for global variables.
15    pub global_variables_used: &'tracer mut HandleSet<crate::GlobalVariable>,
16
17    /// The used map for `constants`.
18    pub constants_used: &'tracer mut HandleSet<crate::Constant>,
19
20    /// The used map for `overrides`.
21    pub overrides_used: &'tracer mut HandleSet<crate::Override>,
22
23    /// The used set for `arena`.
24    ///
25    /// This points to whatever arena holds the expressions we are
26    /// currently tracing: either a function's expression arena, or
27    /// the module's constant expression arena.
28    pub expressions_used: &'tracer mut HandleSet<crate::Expression>,
29
30    /// The used set for the module's `global_expressions` arena.
31    ///
32    /// If `None`, we are already tracing the constant expressions,
33    /// and `expressions_used` already refers to their handle set.
34    pub global_expressions_used: Option<&'tracer mut HandleSet<crate::Expression>>,
35}
36
37impl ExpressionTracer<'_> {
38    /// Propagate usage through `self.expressions`, starting with `self.expressions_used`.
39    ///
40    /// Treat `self.expressions_used` as the initial set of "known
41    /// live" expressions, and follow through to identify all
42    /// transitively used expressions.
43    ///
44    /// Mark types, constants, and constant expressions used directly
45    /// by `self.expressions` as used. Items used indirectly are not
46    /// marked.
47    ///
48    /// [fe]: crate::Function::expressions
49    /// [ce]: crate::Module::global_expressions
50    pub fn trace_expressions(&mut self) {
51        log::trace!(
52            "entering trace_expression of {}",
53            if self.global_expressions_used.is_some() {
54                "function expressions"
55            } else {
56                "const expressions"
57            }
58        );
59
60        // We don't need recursion or a work list. Because an
61        // expression may only refer to other expressions that precede
62        // it in the arena, it suffices to make a single pass over the
63        // arena from back to front, marking the referents of used
64        // expressions as used themselves.
65        for (handle, expr) in self.expressions.iter().rev() {
66            // If this expression isn't used, it doesn't matter what it uses.
67            if !self.expressions_used.contains(handle) {
68                continue;
69            }
70
71            log::trace!("tracing new expression {expr:?}");
72            self.trace_expression(expr);
73        }
74    }
75
76    pub fn trace_expression(&mut self, expr: &crate::Expression) {
77        use crate::Expression as Ex;
78        match *expr {
79            // Expressions that do not contain handles that need to be traced.
80            Ex::Literal(_)
81            | Ex::FunctionArgument(_)
82            | Ex::LocalVariable(_)
83            | Ex::SubgroupBallotResult
84            | Ex::RayQueryProceedResult => {}
85
86            // Expressions can refer to constants and overrides, which can refer
87            // in turn to expressions, which complicates our nice one-pass
88            // algorithm. But since constants and overrides don't refer to each
89            // other directly, only via expressions, we can get around this by
90            // looking *through* each constant/override and marking its
91            // initializer expression as used immediately. Since `expr` refers
92            // to the constant/override, which then refers to the initializer,
93            // the initializer must precede `expr` in the arena, so we know we
94            // have yet to visit the initializer, so it's not too late to mark
95            // it.
96            Ex::Constant(handle) => {
97                self.constants_used.insert(handle);
98                let constant = &self.constants[handle];
99                self.types_used.insert(constant.ty);
100                match self.global_expressions_used {
101                    Some(ref mut used) => used.insert(constant.init),
102                    None => self.expressions_used.insert(constant.init),
103                };
104            }
105            Ex::Override(handle) => {
106                self.overrides_used.insert(handle);
107                let r#override = &self.overrides[handle];
108                self.types_used.insert(r#override.ty);
109                if let Some(init) = r#override.init {
110                    match self.global_expressions_used {
111                        Some(ref mut used) => used.insert(init),
112                        None => self.expressions_used.insert(init),
113                    };
114                }
115            }
116            Ex::ZeroValue(ty) => {
117                self.types_used.insert(ty);
118            }
119            Ex::Compose { ty, ref components } => {
120                self.types_used.insert(ty);
121                self.expressions_used
122                    .insert_iter(components.iter().cloned());
123            }
124            Ex::Access { base, index } => self.expressions_used.insert_iter([base, index]),
125            Ex::AccessIndex { base, index: _ } => {
126                self.expressions_used.insert(base);
127            }
128            Ex::Splat { size: _, value } => {
129                self.expressions_used.insert(value);
130            }
131            Ex::Swizzle {
132                size: _,
133                vector,
134                pattern: _,
135            } => {
136                self.expressions_used.insert(vector);
137            }
138            Ex::GlobalVariable(handle) => {
139                self.global_variables_used.insert(handle);
140            }
141            Ex::Load { pointer } => {
142                self.expressions_used.insert(pointer);
143            }
144            Ex::ImageSample {
145                image,
146                sampler,
147                gather: _,
148                coordinate,
149                array_index,
150                offset,
151                ref level,
152                depth_ref,
153                clamp_to_edge: _,
154            } => {
155                self.expressions_used
156                    .insert_iter([image, sampler, coordinate]);
157                self.expressions_used.insert_iter(array_index);
158                self.expressions_used.insert_iter(offset);
159                use crate::SampleLevel as Sl;
160                match *level {
161                    Sl::Auto | Sl::Zero => {}
162                    Sl::Exact(expr) | Sl::Bias(expr) => {
163                        self.expressions_used.insert(expr);
164                    }
165                    Sl::Gradient { x, y } => self.expressions_used.insert_iter([x, y]),
166                }
167                self.expressions_used.insert_iter(depth_ref);
168            }
169            Ex::ImageLoad {
170                image,
171                coordinate,
172                array_index,
173                sample,
174                level,
175            } => {
176                self.expressions_used.insert(image);
177                self.expressions_used.insert(coordinate);
178                self.expressions_used.insert_iter(array_index);
179                self.expressions_used.insert_iter(sample);
180                self.expressions_used.insert_iter(level);
181            }
182            Ex::ImageQuery { image, ref query } => {
183                self.expressions_used.insert(image);
184                use crate::ImageQuery as Iq;
185                match *query {
186                    Iq::Size { level } => self.expressions_used.insert_iter(level),
187                    Iq::NumLevels | Iq::NumLayers | Iq::NumSamples => {}
188                }
189            }
190            Ex::RayQueryVertexPositions {
191                query,
192                committed: _,
193            } => {
194                self.expressions_used.insert(query);
195            }
196            Ex::Unary { op: _, expr } => {
197                self.expressions_used.insert(expr);
198            }
199            Ex::Binary { op: _, left, right } => {
200                self.expressions_used.insert_iter([left, right]);
201            }
202            Ex::Select {
203                condition,
204                accept,
205                reject,
206            } => self
207                .expressions_used
208                .insert_iter([condition, accept, reject]),
209            Ex::Derivative {
210                axis: _,
211                ctrl: _,
212                expr,
213            } => {
214                self.expressions_used.insert(expr);
215            }
216            Ex::Relational { fun: _, argument } => {
217                self.expressions_used.insert(argument);
218            }
219            Ex::Math {
220                fun: _,
221                arg,
222                arg1,
223                arg2,
224                arg3,
225            } => {
226                self.expressions_used.insert(arg);
227                self.expressions_used.insert_iter(arg1);
228                self.expressions_used.insert_iter(arg2);
229                self.expressions_used.insert_iter(arg3);
230            }
231            Ex::As {
232                expr,
233                kind: _,
234                convert: _,
235            } => {
236                self.expressions_used.insert(expr);
237            }
238            Ex::ArrayLength(expr) => {
239                self.expressions_used.insert(expr);
240            }
241            // `CallResult` expressions do contain a function handle, but any used
242            // `CallResult` expression should have an associated `ir::Statement::Call`
243            // that we will trace.
244            Ex::CallResult(_) => {}
245            Ex::AtomicResult { ty, comparison: _ }
246            | Ex::WorkGroupUniformLoadResult { ty }
247            | Ex::SubgroupOperationResult { ty } => {
248                self.types_used.insert(ty);
249            }
250            Ex::RayQueryGetIntersection {
251                query,
252                committed: _,
253            } => {
254                self.expressions_used.insert(query);
255            }
256        }
257    }
258}
259
260impl ModuleMap {
261    /// Fix up all handles in `expr`.
262    ///
263    /// Use the expression handle remappings in `operand_map`, and all
264    /// other mappings from `self`.
265    pub fn adjust_expression(
266        &self,
267        expr: &mut crate::Expression,
268        operand_map: &HandleMap<crate::Expression>,
269    ) {
270        let adjust = |expr: &mut Handle<crate::Expression>| {
271            operand_map.adjust(expr);
272        };
273
274        use crate::Expression as Ex;
275        match *expr {
276            // Expressions that do not contain handles that need to be adjusted.
277            Ex::Literal(_)
278            | Ex::FunctionArgument(_)
279            | Ex::LocalVariable(_)
280            | Ex::SubgroupBallotResult
281            | Ex::RayQueryProceedResult => {}
282
283            // Expressions that contain handles that need to be adjusted.
284            Ex::Constant(ref mut constant) => self.constants.adjust(constant),
285            Ex::Override(ref mut r#override) => self.overrides.adjust(r#override),
286            Ex::ZeroValue(ref mut ty) => self.types.adjust(ty),
287            Ex::Compose {
288                ref mut ty,
289                ref mut components,
290            } => {
291                self.types.adjust(ty);
292                for component in components {
293                    adjust(component);
294                }
295            }
296            Ex::Access {
297                ref mut base,
298                ref mut index,
299            } => {
300                adjust(base);
301                adjust(index);
302            }
303            Ex::AccessIndex {
304                ref mut base,
305                index: _,
306            } => adjust(base),
307            Ex::Splat {
308                size: _,
309                ref mut value,
310            } => adjust(value),
311            Ex::Swizzle {
312                size: _,
313                ref mut vector,
314                pattern: _,
315            } => adjust(vector),
316            Ex::GlobalVariable(ref mut handle) => self.globals.adjust(handle),
317            Ex::Load { ref mut pointer } => adjust(pointer),
318            Ex::ImageSample {
319                ref mut image,
320                ref mut sampler,
321                gather: _,
322                ref mut coordinate,
323                ref mut array_index,
324                ref mut offset,
325                ref mut level,
326                ref mut depth_ref,
327                clamp_to_edge: _,
328            } => {
329                adjust(image);
330                adjust(sampler);
331                adjust(coordinate);
332                operand_map.adjust_option(array_index);
333                operand_map.adjust_option(offset);
334                self.adjust_sample_level(level, operand_map);
335                operand_map.adjust_option(depth_ref);
336            }
337            Ex::ImageLoad {
338                ref mut image,
339                ref mut coordinate,
340                ref mut array_index,
341                ref mut sample,
342                ref mut level,
343            } => {
344                adjust(image);
345                adjust(coordinate);
346                operand_map.adjust_option(array_index);
347                operand_map.adjust_option(sample);
348                operand_map.adjust_option(level);
349            }
350            Ex::ImageQuery {
351                ref mut image,
352                ref mut query,
353            } => {
354                adjust(image);
355                self.adjust_image_query(query, operand_map);
356            }
357            Ex::Unary {
358                op: _,
359                ref mut expr,
360            } => adjust(expr),
361            Ex::Binary {
362                op: _,
363                ref mut left,
364                ref mut right,
365            } => {
366                adjust(left);
367                adjust(right);
368            }
369            Ex::Select {
370                ref mut condition,
371                ref mut accept,
372                ref mut reject,
373            } => {
374                adjust(condition);
375                adjust(accept);
376                adjust(reject);
377            }
378            Ex::Derivative {
379                axis: _,
380                ctrl: _,
381                ref mut expr,
382            } => adjust(expr),
383            Ex::Relational {
384                fun: _,
385                ref mut argument,
386            } => adjust(argument),
387            Ex::Math {
388                fun: _,
389                ref mut arg,
390                ref mut arg1,
391                ref mut arg2,
392                ref mut arg3,
393            } => {
394                adjust(arg);
395                operand_map.adjust_option(arg1);
396                operand_map.adjust_option(arg2);
397                operand_map.adjust_option(arg3);
398            }
399            Ex::As {
400                ref mut expr,
401                kind: _,
402                convert: _,
403            } => adjust(expr),
404            Ex::CallResult(ref mut function) => {
405                self.functions.adjust(function);
406            }
407            Ex::AtomicResult {
408                ref mut ty,
409                comparison: _,
410            } => self.types.adjust(ty),
411            Ex::WorkGroupUniformLoadResult { ref mut ty } => self.types.adjust(ty),
412            Ex::SubgroupOperationResult { ref mut ty } => self.types.adjust(ty),
413            Ex::ArrayLength(ref mut expr) => adjust(expr),
414            Ex::RayQueryGetIntersection {
415                ref mut query,
416                committed: _,
417            } => adjust(query),
418            Ex::RayQueryVertexPositions {
419                ref mut query,
420                committed: _,
421            } => adjust(query),
422        }
423    }
424
425    fn adjust_sample_level(
426        &self,
427        level: &mut crate::SampleLevel,
428        operand_map: &HandleMap<crate::Expression>,
429    ) {
430        let adjust = |expr: &mut Handle<crate::Expression>| operand_map.adjust(expr);
431
432        use crate::SampleLevel as Sl;
433        match *level {
434            Sl::Auto | Sl::Zero => {}
435            Sl::Exact(ref mut expr) => adjust(expr),
436            Sl::Bias(ref mut expr) => adjust(expr),
437            Sl::Gradient {
438                ref mut x,
439                ref mut y,
440            } => {
441                adjust(x);
442                adjust(y);
443            }
444        }
445    }
446
447    fn adjust_image_query(
448        &self,
449        query: &mut crate::ImageQuery,
450        operand_map: &HandleMap<crate::Expression>,
451    ) {
452        use crate::ImageQuery as Iq;
453
454        match *query {
455            Iq::Size { ref mut level } => operand_map.adjust_option(level),
456            Iq::NumLevels | Iq::NumLayers | Iq::NumSamples => {}
457        }
458    }
459}