naga_oil/
derive.rs

1use indexmap::IndexMap;
2use naga::{
3    Arena, AtomicFunction, Block, Constant, EntryPoint, Expression, Function, FunctionArgument,
4    FunctionResult, GatherMode, GlobalVariable, Handle, ImageQuery, LocalVariable, Module,
5    Override, SampleLevel, Span, Statement, StructMember, SwitchCase, Type, TypeInner, UniqueArena,
6};
7use std::{cell::RefCell, rc::Rc};
8
9#[derive(Debug, Default)]
10pub struct DerivedModule<'a> {
11    shader: Option<&'a Module>,
12    span_offset: usize,
13
14    /// Maps the original type handle to the the mangled type handle.
15    type_map: IndexMap<Handle<Type>, Handle<Type>>,
16    /// Maps the original const handle to the the mangled const handle.
17    const_map: IndexMap<Handle<Constant>, Handle<Constant>>,
18    /// Maps the original pipeline override handle to the the mangled pipeline override handle.
19    pipeline_override_map: IndexMap<Handle<Override>, Handle<Override>>,
20    /// Contains both const expressions and pipeline override constant expressions.
21    /// The expressions are stored together because that's what Naga expects.
22    global_expressions: Rc<RefCell<Arena<Expression>>>,
23    /// Maps the original expression handle to the new expression handle for const expressions and pipeline override expressions.
24    /// The expressions are stored together because that's what Naga expects.
25    global_expression_map: Rc<RefCell<IndexMap<Handle<Expression>, Handle<Expression>>>>,
26    global_map: IndexMap<Handle<GlobalVariable>, Handle<GlobalVariable>>,
27    function_map: IndexMap<String, Handle<Function>>,
28    types: UniqueArena<Type>,
29    constants: Arena<Constant>,
30    globals: Arena<GlobalVariable>,
31    functions: Arena<Function>,
32    pipeline_overrides: Arena<Override>,
33    special_types: naga::SpecialTypes,
34}
35
36impl<'a> DerivedModule<'a> {
37    // set source context for import operations
38    pub fn set_shader_source(&mut self, shader: &'a Module, span_offset: usize) {
39        self.clear_shader_source();
40        self.shader = Some(shader);
41        self.span_offset = span_offset;
42
43        // eagerly import special types
44        if let Some(h_special_type) = shader.special_types.ray_desc.as_ref() {
45            if let Some(derived_special_type) = self.special_types.ray_desc.as_ref() {
46                self.type_map.insert(*h_special_type, *derived_special_type);
47            } else {
48                self.special_types.ray_desc = Some(self.import_type(h_special_type));
49            }
50        }
51        if let Some(h_special_type) = shader.special_types.ray_intersection.as_ref() {
52            if let Some(derived_special_type) = self.special_types.ray_intersection.as_ref() {
53                self.type_map.insert(*h_special_type, *derived_special_type);
54            } else {
55                self.special_types.ray_intersection = Some(self.import_type(h_special_type));
56            }
57        }
58        if let Some(h_special_type) = shader.special_types.ray_vertex_return.as_ref() {
59            if let Some(derived_special_type) = self.special_types.ray_vertex_return.as_ref() {
60                self.type_map.insert(*h_special_type, *derived_special_type);
61            } else {
62                self.special_types.ray_vertex_return = Some(self.import_type(h_special_type));
63            }
64        }
65        for (predeclared, h_predeclared_type) in shader.special_types.predeclared_types.iter() {
66            if let Some(derived_special_type) =
67                self.special_types.predeclared_types.get(predeclared)
68            {
69                self.type_map
70                    .insert(*h_predeclared_type, *derived_special_type);
71            } else {
72                let new_h = self.import_type(h_predeclared_type);
73                self.special_types
74                    .predeclared_types
75                    .insert(predeclared.clone(), new_h);
76            }
77        }
78    }
79
80    // detach source context
81    pub fn clear_shader_source(&mut self) {
82        self.shader = None;
83        self.type_map.clear();
84        self.const_map.clear();
85        self.global_map.clear();
86        self.global_expression_map.borrow_mut().clear();
87        self.pipeline_override_map.clear();
88    }
89
90    pub fn map_span(&self, span: Span) -> Span {
91        let span = span.to_range();
92        match span {
93            Some(rng) => Span::new(
94                (rng.start + self.span_offset) as u32,
95                (rng.end + self.span_offset) as u32,
96            ),
97            None => Span::UNDEFINED,
98        }
99    }
100
101    // remap a type from source context into our derived context
102    pub fn import_type(&mut self, h_type: &Handle<Type>) -> Handle<Type> {
103        self.rename_type(h_type, None)
104    }
105
106    // remap a type from source context into our derived context, and rename it
107    pub fn rename_type(&mut self, h_type: &Handle<Type>, name: Option<String>) -> Handle<Type> {
108        self.type_map.get(h_type).copied().unwrap_or_else(|| {
109            let ty = self
110                .shader
111                .as_ref()
112                .unwrap()
113                .types
114                .get_handle(*h_type)
115                .unwrap();
116
117            let name = match name {
118                Some(name) => Some(name),
119                None => ty.name.clone(),
120            };
121
122            let new_type = Type {
123                name,
124                inner: match &ty.inner {
125                    TypeInner::Scalar { .. }
126                    | TypeInner::Vector { .. }
127                    | TypeInner::Matrix { .. }
128                    | TypeInner::ValuePointer { .. }
129                    | TypeInner::Image { .. }
130                    | TypeInner::Sampler { .. }
131                    | TypeInner::Atomic { .. }
132                    | TypeInner::AccelerationStructure { .. }
133                    | TypeInner::RayQuery { .. } => ty.inner.clone(),
134                    TypeInner::Pointer { base, space } => TypeInner::Pointer {
135                        base: self.import_type(base),
136                        space: *space,
137                    },
138                    TypeInner::Struct { members, span } => {
139                        let members = members
140                            .iter()
141                            .map(|m| StructMember {
142                                name: m.name.clone(),
143                                ty: self.import_type(&m.ty),
144                                binding: m.binding.clone(),
145                                offset: m.offset,
146                            })
147                            .collect();
148                        TypeInner::Struct {
149                            members,
150                            span: *span,
151                        }
152                    }
153                    TypeInner::Array { base, size, stride } => TypeInner::Array {
154                        base: self.import_type(base),
155                        size: *size,
156                        stride: *stride,
157                    },
158                    TypeInner::BindingArray { base, size } => TypeInner::BindingArray {
159                        base: self.import_type(base),
160                        size: *size,
161                    },
162                    TypeInner::CooperativeMatrix {
163                        columns,
164                        rows,
165                        scalar,
166                        role,
167                    } => TypeInner::CooperativeMatrix {
168                        columns: *columns,
169                        rows: *rows,
170                        scalar: *scalar,
171                        role: *role,
172                    },
173                },
174            };
175            let span = self.shader.as_ref().unwrap().types.get_span(*h_type);
176            let new_h = self.types.insert(new_type, self.map_span(span));
177            self.type_map.insert(*h_type, new_h);
178            new_h
179        })
180    }
181
182    // remap a const from source context into our derived context
183    pub fn import_const(&mut self, h_const: &Handle<Constant>) -> Handle<Constant> {
184        self.const_map.get(h_const).copied().unwrap_or_else(|| {
185            let c = self
186                .shader
187                .as_ref()
188                .unwrap()
189                .constants
190                .try_get(*h_const)
191                .unwrap();
192
193            let new_const = Constant {
194                name: c.name.clone(),
195                ty: self.import_type(&c.ty),
196                init: self.import_global_expression(c.init),
197            };
198
199            let span = self.shader.as_ref().unwrap().constants.get_span(*h_const);
200            let new_h = self
201                .constants
202                .fetch_or_append(new_const, self.map_span(span));
203            self.const_map.insert(*h_const, new_h);
204            new_h
205        })
206    }
207
208    // remap a global from source context into our derived context
209    pub fn import_global(&mut self, h_global: &Handle<GlobalVariable>) -> Handle<GlobalVariable> {
210        self.global_map.get(h_global).copied().unwrap_or_else(|| {
211            let gv = self
212                .shader
213                .as_ref()
214                .unwrap()
215                .global_variables
216                .try_get(*h_global)
217                .unwrap();
218
219            let new_global = GlobalVariable {
220                name: gv.name.clone(),
221                space: gv.space,
222                binding: gv.binding,
223                ty: self.import_type(&gv.ty),
224                init: gv.init.map(|c| self.import_global_expression(c)),
225                memory_decorations: gv.memory_decorations,
226            };
227
228            let span = self
229                .shader
230                .as_ref()
231                .unwrap()
232                .global_variables
233                .get_span(*h_global);
234            let new_h = self
235                .globals
236                .fetch_or_append(new_global, self.map_span(span));
237            self.global_map.insert(*h_global, new_h);
238            new_h
239        })
240    }
241
242    // remap either a const or pipeline override expression from source context into our derived context
243    pub fn import_global_expression(&mut self, h_expr: Handle<Expression>) -> Handle<Expression> {
244        self.import_expression(
245            h_expr,
246            &self.shader.as_ref().unwrap().global_expressions,
247            self.global_expression_map.clone(),
248            self.global_expressions.clone(),
249            false,
250            true,
251        )
252    }
253
254    // remap a pipeline override from source context into our derived context
255    pub fn import_pipeline_override(&mut self, h_override: &Handle<Override>) -> Handle<Override> {
256        self.pipeline_override_map
257            .get(h_override)
258            .copied()
259            .unwrap_or_else(|| {
260                let pipeline_override = self
261                    .shader
262                    .as_ref()
263                    .unwrap()
264                    .overrides
265                    .try_get(*h_override)
266                    .unwrap();
267
268                let new_override = Override {
269                    name: pipeline_override.name.clone(),
270                    id: pipeline_override.id,
271                    ty: self.import_type(&pipeline_override.ty),
272                    init: pipeline_override
273                        .init
274                        .map(|init| self.import_global_expression(init)),
275                };
276
277                let span = self
278                    .shader
279                    .as_ref()
280                    .unwrap()
281                    .overrides
282                    .get_span(*h_override);
283                let new_h = self
284                    .pipeline_overrides
285                    .fetch_or_append(new_override, self.map_span(span));
286                self.pipeline_override_map.insert(*h_override, new_h);
287                new_h
288            })
289    }
290
291    // remap a block
292    fn import_block(
293        &mut self,
294        block: &Block,
295        old_expressions: &Arena<Expression>,
296        already_imported: Rc<RefCell<IndexMap<Handle<Expression>, Handle<Expression>>>>,
297        new_expressions: Rc<RefCell<Arena<Expression>>>,
298    ) -> Block {
299        macro_rules! map_expr {
300            ($e:expr) => {
301                self.import_expression(
302                    *$e,
303                    old_expressions,
304                    already_imported.clone(),
305                    new_expressions.clone(),
306                    false,
307                    false,
308                )
309            };
310        }
311
312        macro_rules! map_expr_opt {
313            ($e:expr) => {
314                $e.as_ref().map(|expr| map_expr!(expr))
315            };
316        }
317
318        macro_rules! map_block {
319            ($b:expr) => {
320                self.import_block(
321                    $b,
322                    old_expressions,
323                    already_imported.clone(),
324                    new_expressions.clone(),
325                )
326            };
327        }
328
329        let statements = block
330            .iter()
331            .map(|stmt| {
332                match stmt {
333                    Statement::Call {
334                        function,
335                        arguments,
336                        result,
337                    } => Statement::Call {
338                        function: self.map_function_handle(function),
339                        arguments: arguments.iter().map(|expr| map_expr!(expr)).collect(),
340                        result: result.as_ref().map(|result| map_expr!(result)),
341                    },
342                    Statement::Block(b) => Statement::Block(map_block!(b)),
343                    Statement::If {
344                        condition,
345                        accept,
346                        reject,
347                    } => Statement::If {
348                        condition: map_expr!(condition),
349                        accept: map_block!(accept),
350                        reject: map_block!(reject),
351                    },
352                    Statement::Switch { selector, cases } => Statement::Switch {
353                        selector: map_expr!(selector),
354                        cases: cases
355                            .iter()
356                            .map(|case| SwitchCase {
357                                value: case.value,
358                                body: map_block!(&case.body),
359                                fall_through: case.fall_through,
360                            })
361                            .collect(),
362                    },
363                    Statement::Loop {
364                        body,
365                        continuing,
366                        break_if,
367                    } => Statement::Loop {
368                        body: map_block!(body),
369                        continuing: map_block!(continuing),
370                        break_if: map_expr_opt!(break_if),
371                    },
372                    Statement::Emit(exprs) => {
373                        // iterate once to add expressions that should NOT be part of the emit statement
374                        for expr in exprs.clone() {
375                            self.import_expression(
376                                expr,
377                                old_expressions,
378                                already_imported.clone(),
379                                new_expressions.clone(),
380                                true,
381                                false,
382                            );
383                        }
384                        let old_length = new_expressions.borrow().len();
385                        // iterate again to add expressions that should be part of the emit statement
386                        for expr in exprs.clone() {
387                            map_expr!(&expr);
388                        }
389
390                        Statement::Emit(new_expressions.borrow().range_from(old_length))
391                    }
392                    Statement::Store { pointer, value } => Statement::Store {
393                        pointer: map_expr!(pointer),
394                        value: map_expr!(value),
395                    },
396                    Statement::ImageStore {
397                        image,
398                        coordinate,
399                        array_index,
400                        value,
401                    } => Statement::ImageStore {
402                        image: map_expr!(image),
403                        coordinate: map_expr!(coordinate),
404                        array_index: map_expr_opt!(array_index),
405                        value: map_expr!(value),
406                    },
407                    Statement::Atomic {
408                        pointer,
409                        fun,
410                        value,
411                        result,
412                    } => {
413                        let fun = match fun {
414                            AtomicFunction::Exchange {
415                                compare: Some(compare_expr),
416                            } => AtomicFunction::Exchange {
417                                compare: Some(map_expr!(compare_expr)),
418                            },
419                            fun => *fun,
420                        };
421                        Statement::Atomic {
422                            pointer: map_expr!(pointer),
423                            fun,
424                            value: map_expr!(value),
425                            result: map_expr_opt!(result),
426                        }
427                    }
428                    Statement::WorkGroupUniformLoad { pointer, result } => {
429                        Statement::WorkGroupUniformLoad {
430                            pointer: map_expr!(pointer),
431                            result: map_expr!(result),
432                        }
433                    }
434                    Statement::Return { value } => Statement::Return {
435                        value: map_expr_opt!(value),
436                    },
437                    Statement::RayQuery { query, fun } => Statement::RayQuery {
438                        query: map_expr!(query),
439                        fun: match fun {
440                            naga::RayQueryFunction::Initialize {
441                                acceleration_structure,
442                                descriptor,
443                            } => naga::RayQueryFunction::Initialize {
444                                acceleration_structure: map_expr!(acceleration_structure),
445                                descriptor: map_expr!(descriptor),
446                            },
447                            naga::RayQueryFunction::Proceed { result } => {
448                                naga::RayQueryFunction::Proceed {
449                                    result: map_expr!(result),
450                                }
451                            }
452                            naga::RayQueryFunction::Terminate => naga::RayQueryFunction::Terminate,
453                            naga::RayQueryFunction::GenerateIntersection { hit_t } => {
454                                naga::RayQueryFunction::GenerateIntersection {
455                                    hit_t: map_expr!(hit_t),
456                                }
457                            }
458                            naga::RayQueryFunction::ConfirmIntersection => {
459                                naga::RayQueryFunction::ConfirmIntersection
460                            }
461                        },
462                    },
463                    Statement::SubgroupBallot { result, predicate } => Statement::SubgroupBallot {
464                        result: map_expr!(result),
465                        predicate: map_expr_opt!(predicate),
466                    },
467                    Statement::SubgroupGather {
468                        mut mode,
469                        argument,
470                        result,
471                    } => {
472                        match mode {
473                            GatherMode::BroadcastFirst => (),
474                            GatherMode::Broadcast(ref mut h_src)
475                            | GatherMode::QuadBroadcast(ref mut h_src)
476                            | GatherMode::Shuffle(ref mut h_src)
477                            | GatherMode::ShuffleDown(ref mut h_src)
478                            | GatherMode::ShuffleUp(ref mut h_src)
479                            | GatherMode::ShuffleXor(ref mut h_src) => *h_src = map_expr!(h_src),
480                            GatherMode::QuadSwap(_) => (),
481                        };
482                        Statement::SubgroupGather {
483                            mode,
484                            argument: map_expr!(argument),
485                            result: map_expr!(result),
486                        }
487                    }
488                    Statement::SubgroupCollectiveOperation {
489                        op,
490                        collective_op,
491                        argument,
492                        result,
493                    } => Statement::SubgroupCollectiveOperation {
494                        op: *op,
495                        collective_op: *collective_op,
496                        argument: map_expr!(argument),
497                        result: map_expr!(result),
498                    },
499                    Statement::ImageAtomic {
500                        image,
501                        coordinate,
502                        array_index,
503                        fun,
504                        value,
505                    } => {
506                        let fun = match fun {
507                            AtomicFunction::Exchange {
508                                compare: Some(compare_expr),
509                            } => AtomicFunction::Exchange {
510                                compare: Some(map_expr!(compare_expr)),
511                            },
512                            fun => *fun,
513                        };
514                        Statement::ImageAtomic {
515                            image: map_expr!(image),
516                            coordinate: map_expr!(coordinate),
517                            array_index: map_expr_opt!(array_index),
518                            fun,
519                            value: map_expr!(value),
520                        }
521                    }
522                    Statement::Break
523                    | Statement::Continue
524                    | Statement::Kill
525                    | Statement::MemoryBarrier(_)
526                    | Statement::ControlBarrier(_) => stmt.clone(),
527                    Statement::RayPipelineFunction(ray_pipeline_function) => {
528                        Statement::RayPipelineFunction(*ray_pipeline_function)
529                    }
530                    Statement::CooperativeStore { target, data } => Statement::CooperativeStore {
531                        target: map_expr!(target),
532                        data: *data,
533                    },
534                }
535            })
536            .collect();
537
538        let mut new_block = Block::from_vec(statements);
539
540        for ((_, new_span), (_, old_span)) in new_block.span_iter_mut().zip(block.span_iter()) {
541            *new_span.unwrap() = self.map_span(*old_span);
542        }
543
544        new_block
545    }
546
547    fn import_expression(
548        &mut self,
549        h_expr: Handle<Expression>,
550        old_expressions: &Arena<Expression>,
551        already_imported: Rc<RefCell<IndexMap<Handle<Expression>, Handle<Expression>>>>,
552        new_expressions: Rc<RefCell<Arena<Expression>>>,
553        non_emitting_only: bool, // only brings items that should NOT be emitted into scope
554        unique: bool,            // ensure expressions are unique with custom comparison
555    ) -> Handle<Expression> {
556        if let Some(h_new) = already_imported.borrow().get(&h_expr) {
557            return *h_new;
558        }
559
560        macro_rules! map_expr {
561            ($e:expr) => {
562                self.import_expression(
563                    *$e,
564                    old_expressions,
565                    already_imported.clone(),
566                    new_expressions.clone(),
567                    non_emitting_only,
568                    unique,
569                )
570            };
571        }
572
573        macro_rules! map_expr_opt {
574            ($e:expr) => {
575                $e.as_ref().map(|expr| {
576                    self.import_expression(
577                        *expr,
578                        old_expressions,
579                        already_imported.clone(),
580                        new_expressions.clone(),
581                        non_emitting_only,
582                        unique,
583                    )
584                })
585            };
586        }
587
588        let mut is_external = false;
589        let expr = old_expressions.try_get(h_expr).unwrap();
590        let expr = match expr {
591            Expression::Literal(_) => {
592                is_external = true;
593                expr.clone()
594            }
595            Expression::ZeroValue(zv) => {
596                is_external = true;
597                Expression::ZeroValue(self.import_type(zv))
598            }
599            Expression::CallResult(f) => Expression::CallResult(self.map_function_handle(f)),
600            Expression::Constant(c) => {
601                is_external = true;
602                Expression::Constant(self.import_const(c))
603            }
604            Expression::Compose { ty, components } => Expression::Compose {
605                ty: self.import_type(ty),
606                components: components.iter().map(|expr| map_expr!(expr)).collect(),
607            },
608            Expression::GlobalVariable(gv) => {
609                is_external = true;
610                Expression::GlobalVariable(self.import_global(gv))
611            }
612            Expression::ImageSample {
613                image,
614                sampler,
615                gather,
616                coordinate,
617                array_index,
618                offset,
619                level,
620                depth_ref,
621                clamp_to_edge,
622            } => Expression::ImageSample {
623                image: map_expr!(image),
624                sampler: map_expr!(sampler),
625                gather: *gather,
626                coordinate: map_expr!(coordinate),
627                array_index: map_expr_opt!(array_index),
628                offset: map_expr_opt!(offset),
629                level: match level {
630                    SampleLevel::Auto | SampleLevel::Zero => *level,
631                    SampleLevel::Exact(expr) => SampleLevel::Exact(map_expr!(expr)),
632                    SampleLevel::Bias(expr) => SampleLevel::Bias(map_expr!(expr)),
633                    SampleLevel::Gradient { x, y } => SampleLevel::Gradient {
634                        x: map_expr!(x),
635                        y: map_expr!(y),
636                    },
637                },
638                depth_ref: map_expr_opt!(depth_ref),
639                clamp_to_edge: *clamp_to_edge,
640            },
641            Expression::Access { base, index } => Expression::Access {
642                base: map_expr!(base),
643                index: map_expr!(index),
644            },
645            Expression::AccessIndex { base, index } => Expression::AccessIndex {
646                base: map_expr!(base),
647                index: *index,
648            },
649            Expression::Splat { size, value } => Expression::Splat {
650                size: *size,
651                value: map_expr!(value),
652            },
653            Expression::Swizzle {
654                size,
655                vector,
656                pattern,
657            } => Expression::Swizzle {
658                size: *size,
659                vector: map_expr!(vector),
660                pattern: *pattern,
661            },
662            Expression::Load { pointer } => Expression::Load {
663                pointer: map_expr!(pointer),
664            },
665            Expression::ImageLoad {
666                image,
667                coordinate,
668                array_index,
669                sample,
670                level,
671            } => Expression::ImageLoad {
672                image: map_expr!(image),
673                coordinate: map_expr!(coordinate),
674                array_index: map_expr_opt!(array_index),
675                sample: map_expr_opt!(sample),
676                level: map_expr_opt!(level),
677            },
678            Expression::ImageQuery { image, query } => Expression::ImageQuery {
679                image: map_expr!(image),
680                query: match query {
681                    ImageQuery::Size { level } => ImageQuery::Size {
682                        level: map_expr_opt!(level),
683                    },
684                    _ => *query,
685                },
686            },
687            Expression::Unary { op, expr } => Expression::Unary {
688                op: *op,
689                expr: map_expr!(expr),
690            },
691            Expression::Binary { op, left, right } => Expression::Binary {
692                op: *op,
693                left: map_expr!(left),
694                right: map_expr!(right),
695            },
696            Expression::Select {
697                condition,
698                accept,
699                reject,
700            } => Expression::Select {
701                condition: map_expr!(condition),
702                accept: map_expr!(accept),
703                reject: map_expr!(reject),
704            },
705            Expression::Derivative { axis, expr, ctrl } => Expression::Derivative {
706                axis: *axis,
707                expr: map_expr!(expr),
708                ctrl: *ctrl,
709            },
710            Expression::Relational { fun, argument } => Expression::Relational {
711                fun: *fun,
712                argument: map_expr!(argument),
713            },
714            Expression::Math {
715                fun,
716                arg,
717                arg1,
718                arg2,
719                arg3,
720            } => Expression::Math {
721                fun: *fun,
722                arg: map_expr!(arg),
723                arg1: map_expr_opt!(arg1),
724                arg2: map_expr_opt!(arg2),
725                arg3: map_expr_opt!(arg3),
726            },
727            Expression::As {
728                expr,
729                kind,
730                convert,
731            } => Expression::As {
732                expr: map_expr!(expr),
733                kind: *kind,
734                convert: *convert,
735            },
736            Expression::ArrayLength(expr) => Expression::ArrayLength(map_expr!(expr)),
737            Expression::LocalVariable(_) | Expression::FunctionArgument(_) => {
738                is_external = true;
739                expr.clone()
740            }
741            Expression::AtomicResult { ty, comparison } => Expression::AtomicResult {
742                ty: self.import_type(ty),
743                comparison: *comparison,
744            },
745            Expression::WorkGroupUniformLoadResult { ty } => {
746                Expression::WorkGroupUniformLoadResult {
747                    ty: self.import_type(ty),
748                }
749            }
750            Expression::RayQueryProceedResult => expr.clone(),
751            Expression::RayQueryGetIntersection { query, committed } => {
752                Expression::RayQueryGetIntersection {
753                    query: map_expr!(query),
754                    committed: *committed,
755                }
756            }
757            Expression::RayQueryVertexPositions { query, committed } => {
758                Expression::RayQueryVertexPositions {
759                    query: map_expr!(query),
760                    committed: *committed,
761                }
762            }
763            Expression::Override(h_override) => {
764                is_external = true;
765                Expression::Override(self.import_pipeline_override(h_override))
766            }
767            Expression::SubgroupBallotResult => expr.clone(),
768            Expression::SubgroupOperationResult { ty } => Expression::SubgroupOperationResult {
769                ty: self.import_type(ty),
770            },
771            Expression::CooperativeLoad {
772                columns,
773                rows,
774                role,
775                data,
776            } => Expression::CooperativeLoad {
777                columns: *columns,
778                rows: *rows,
779                role: *role,
780                data: *data,
781            },
782            Expression::CooperativeMultiplyAdd { a, b, c } => Expression::CooperativeMultiplyAdd {
783                a: map_expr!(a),
784                b: map_expr!(b),
785                c: map_expr!(c),
786            },
787        };
788
789        if !non_emitting_only || is_external {
790            let span = old_expressions.get_span(h_expr);
791            let h_new = if unique {
792                new_expressions.borrow_mut().fetch_if_or_append(
793                    expr,
794                    self.map_span(span),
795                    |lhs, rhs| lhs == rhs,
796                )
797            } else {
798                new_expressions
799                    .borrow_mut()
800                    .append(expr, self.map_span(span))
801            };
802
803            already_imported.borrow_mut().insert(h_expr, h_new);
804            h_new
805        } else {
806            h_expr
807        }
808    }
809
810    // remap function global references (global vars, consts, types) into our derived context
811    pub fn localize_function(&mut self, func: &Function) -> Function {
812        let arguments = func
813            .arguments
814            .iter()
815            .map(|arg| FunctionArgument {
816                name: arg.name.clone(),
817                ty: self.import_type(&arg.ty),
818                binding: arg.binding.clone(),
819            })
820            .collect();
821
822        let result = func.result.as_ref().map(|r| FunctionResult {
823            ty: self.import_type(&r.ty),
824            binding: r.binding.clone(),
825        });
826
827        let expressions = Rc::new(RefCell::new(Arena::new()));
828        let expr_map = Rc::new(RefCell::new(IndexMap::new()));
829
830        let mut local_variables = Arena::new();
831        for (h_l, l) in func.local_variables.iter() {
832            let new_local = LocalVariable {
833                name: l.name.clone(),
834                ty: self.import_type(&l.ty),
835                init: l.init.map(|c| {
836                    self.import_expression(
837                        c,
838                        &func.expressions,
839                        expr_map.clone(),
840                        expressions.clone(),
841                        false,
842                        true,
843                    )
844                }),
845            };
846            let span = func.local_variables.get_span(h_l);
847            let new_h = local_variables.append(new_local, self.map_span(span));
848            assert_eq!(h_l, new_h);
849        }
850
851        let body = self.import_block(
852            &func.body,
853            &func.expressions,
854            expr_map.clone(),
855            expressions.clone(),
856        );
857
858        let named_expressions = func
859            .named_expressions
860            .iter()
861            .flat_map(|(h_expr, name)| {
862                expr_map
863                    .borrow()
864                    .get(h_expr)
865                    .map(|new_h| (*new_h, name.clone()))
866            })
867            .collect::<IndexMap<_, _, std::hash::BuildHasherDefault<rustc_hash::FxHasher>>>();
868
869        Function {
870            name: func.name.clone(),
871            arguments,
872            result,
873            local_variables,
874            expressions: Rc::try_unwrap(expressions).unwrap().into_inner(),
875            named_expressions,
876            body,
877            diagnostic_filter_leaf: None,
878        }
879    }
880
881    // import a function defined in the source shader context.
882    // func name may be already defined, the returned handle will refer to the new function.
883    // the previously defined function will still be valid.
884    pub fn import_function(&mut self, func: &Function, span: Span) -> Handle<Function> {
885        let name = func.name.as_ref().unwrap().clone();
886        let mapped_func = self.localize_function(func);
887        let new_span = self.map_span(span);
888        let new_h = self.functions.append(mapped_func, new_span);
889        self.function_map.insert(name, new_h);
890        new_h
891    }
892
893    // get the derived handle corresponding to the given source function handle
894    // requires func to be named
895    pub fn map_function_handle(&mut self, h_func: &Handle<Function>) -> Handle<Function> {
896        let functions = &self.shader.as_ref().unwrap().functions;
897        let func = functions.try_get(*h_func).unwrap();
898        let name = func.name.as_ref().unwrap();
899        self.function_map.get(name).copied().unwrap_or_else(|| {
900            let span = functions.get_span(*h_func);
901            self.import_function(func, span)
902        })
903    }
904
905    /// swap an already imported function for a new one.
906    /// note span cannot be updated
907    pub fn import_function_if_new(&mut self, func: &Function, span: Span) -> Handle<Function> {
908        let name = func.name.as_ref().unwrap().clone();
909        if let Some(h) = self.function_map.get(&name) {
910            return *h;
911        }
912
913        self.import_function(func, span)
914    }
915
916    /// get any required special types for this module
917    pub fn has_required_special_types(&self) -> bool {
918        !self.special_types.predeclared_types.is_empty()
919            || self.special_types.ray_desc.is_some()
920            || self.special_types.ray_intersection.is_some()
921            || self.special_types.ray_vertex_return.is_some()
922    }
923
924    pub fn into_module_with_entrypoints(mut self) -> naga::Module {
925        let entry_points = self
926            .shader
927            .unwrap()
928            .entry_points
929            .iter()
930            .map(|ep| EntryPoint {
931                name: ep.name.clone(),
932                stage: ep.stage,
933                early_depth_test: ep.early_depth_test,
934                workgroup_size: ep.workgroup_size,
935                function: self.localize_function(&ep.function),
936                workgroup_size_overrides: ep.workgroup_size_overrides,
937                mesh_info: ep.mesh_info.clone(),
938                task_payload: ep.task_payload,
939                incoming_ray_payload: ep.incoming_ray_payload,
940            })
941            .collect();
942
943        naga::Module {
944            entry_points,
945            ..self.into()
946        }
947    }
948}
949
950impl From<DerivedModule<'_>> for naga::Module {
951    fn from(derived: DerivedModule) -> Self {
952        naga::Module {
953            types: derived.types,
954            constants: derived.constants,
955            global_variables: derived.globals,
956            global_expressions: Rc::try_unwrap(derived.global_expressions)
957                .unwrap()
958                .into_inner(),
959            functions: derived.functions,
960            special_types: derived.special_types,
961            entry_points: Default::default(),
962            overrides: derived.pipeline_overrides,
963            diagnostic_filters: Default::default(),
964            diagnostic_filter_leaf: None,
965            doc_comments: None,
966        }
967    }
968}
naga_oil/derive.rs

naga_oil/
derive.rs
