Skip to main content

spirv_webgpu_transform/
splitdref.rs

1use super::*;
2
3#[derive(Debug, Clone, Copy, PartialEq, Eq)]
4enum OperationVariant {
5    Regular,
6    Dref,
7}
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10enum LoadType {
11    Variable,
12    FunctionArgument,
13}
14
15#[derive(Debug, Clone, Copy, PartialEq, Eq)]
16enum MixState {
17    Mixed,
18    PotentiallyMixed,
19}
20
21trait IsIndexOrId {}
22impl IsIndexOrId for u32 {}
23impl IsIndexOrId for usize {}
24
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
26enum PatchObjectType<T: IsIndexOrId> {
27    Sampler(T),
28    Image(T),
29}
30
31impl<T> PatchObjectType<T>
32where
33    T: IsIndexOrId,
34{
35    fn next<N: IsIndexOrId>(self, next_id: N) -> PatchObjectType<N> {
36        match self {
37            PatchObjectType::Sampler(_) => PatchObjectType::Sampler(next_id),
38            PatchObjectType::Image(_) => PatchObjectType::Image(next_id),
39        }
40    }
41
42    fn inner(self) -> T {
43        match self {
44            PatchObjectType::Sampler(v) => v,
45            PatchObjectType::Image(v) => v,
46        }
47    }
48}
49
50/// Perform the operation on a `Vec<u32>`.
51/// Use [u8_slice_to_u32_vec] to convert a `&[u8]` into a `Vec<u32>`.
52pub fn drefsplitter(
53    in_spv: &[u32],
54    corrections: &mut CorrectionMap,
55) -> Result<Vec<u32>, ()> {
56    let spv = in_spv.to_owned();
57
58    let mut instruction_bound = spv[SPV_HEADER_INSTRUCTION_BOUND_OFFSET];
59    let magic_number = spv[SPV_HEADER_MAGIC_NUM_OFFSET];
60
61    let spv_header = spv[0..SPV_HEADER_LENGTH].to_owned();
62
63    assert_eq!(magic_number, SPV_HEADER_MAGIC);
64
65    let mut instruction_inserts: Vec<InstructionInsert> = vec![];
66    let mut word_inserts: Vec<WordInsert> = vec![];
67
68    let spv = spv.into_iter().skip(SPV_HEADER_LENGTH).collect::<Vec<_>>();
69    let mut new_spv = spv.clone();
70
71    // 1. Find locations instructions we need
72    let mut op_dref_operation_idxs = vec![];
73    let mut op_sampled_operation_idxs = vec![];
74    let mut op_sampled_image_idxs = vec![];
75    let mut op_load_idxs = vec![];
76    let mut op_variable_idxs = vec![];
77    let mut op_decorate_idxs = vec![];
78    let mut op_type_image_idxs = vec![];
79    let mut op_type_pointer_idxs = vec![];
80    let mut op_type_function_idxs = vec![];
81    let mut op_function_idxs = vec![];
82    let mut op_function_call_idxs = vec![];
83    let mut op_function_parameter_idxs = vec![];
84
85    let mut first_op_type_sampler_id = None;
86    let mut first_op_type_pointer_sampler_id = None;
87
88    let mut spv_idx = 0;
89    while spv_idx < spv.len() {
90        let op = spv[spv_idx];
91        let word_count = hiword(op);
92        let instruction = loword(op);
93
94        match instruction {
95            SPV_INSTRUCTION_OP_SAMPLED_IMAGE => op_sampled_image_idxs.push(spv_idx),
96            SPV_INSTRUCTION_OP_LOAD => op_load_idxs.push(spv_idx),
97            SPV_INSTRUCTION_OP_IMAGE_SAMPLE_DREF_IMPLICIT_LOD
98            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_DREF_EXPLICIT_LOD
99            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_PROJ_DREF_IMPLICIT_LOD
100            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_PROJ_DREF_EXPLICIT_LOD
101            | SPV_INSTRUCTION_OP_IMAGE_DREF_GATHER
102            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_SAMPLE_DREF_IMPLICIT_LOD
103            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_SAMPLE_DREF_EXPLICIT_LOD
104            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_DREF_GATHER => op_dref_operation_idxs.push(spv_idx),
105            SPV_INSTRUCTION_OP_IMAGE_SAMPLE_IMPLICIT_LOD
106            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_EXPLICIT_LOD
107            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_PROJ_IMPLICIT_LOD
108            | SPV_INSTRUCTION_OP_IMAGE_SAMPLE_PROJ_EXPLICIT_LOD
109            | SPV_INSTRUCTION_OP_IMAGE_GATHER
110            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_SAMPLE_IMPLICIT_LOD
111            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_SAMPLE_EXPLICIT_LOD
112            | SPV_INSTRUCTION_OP_IMAGE_SPARSE_GATHER => op_sampled_operation_idxs.push(spv_idx),
113            SPV_INSTRUCTION_OP_VARIABLE => op_variable_idxs.push(spv_idx),
114            SPV_INSTRUCTION_OP_DECORATE => op_decorate_idxs.push(spv_idx),
115            SPV_INSTRUCTION_OP_TYPE_IMAGE => op_type_image_idxs.push(spv_idx),
116            SPV_INSTRUCTION_OP_TYPE_SAMPLER => {
117                first_op_type_sampler_id.get_or_insert(spv[spv_idx + 1]);
118            }
119            SPV_INSTRUCTION_OP_TYPE_POINTER => {
120                if first_op_type_sampler_id == Some(spv[spv_idx + 3])
121                    && spv[spv_idx + 2] == SPV_STORAGE_CLASS_UNIFORM_CONSTANT
122                {
123                    first_op_type_pointer_sampler_id = Some(spv[spv_idx + 1]);
124                }
125                op_type_pointer_idxs.push(spv_idx)
126            }
127            SPV_INSTRUCTION_OP_TYPE_FUNCTION => op_type_function_idxs.push(spv_idx),
128            SPV_INSTRUCTION_OP_FUNCTION => op_function_idxs.push(spv_idx),
129            SPV_INSTRUCTION_OP_FUNCTION_CALL => op_function_call_idxs.push(spv_idx),
130            SPV_INSTRUCTION_OP_FUNCTION_PARAMETER => op_function_parameter_idxs.push(spv_idx),
131            _ => {}
132        }
133
134        spv_idx += word_count as usize;
135    }
136
137    let first_op_deocrate_idx = op_decorate_idxs.first().copied();
138
139    // If there is no OpTypeSampler, either this is invalid, or we do not need to do any patching at all.
140    let (Some(first_op_type_sampler_id), Some(first_op_type_pointer_sampler_id)) =
141        (first_op_type_sampler_id, first_op_type_pointer_sampler_id)
142    else {
143        return Ok(in_spv.to_vec());
144    };
145
146    // 2. Collect all the loaded sampled images of both operation types
147    // Conveniently, the offset for this value is always +3 for all of these operations
148    let loaded_sampled_image_ids = op_sampled_operation_idxs
149        .iter()
150        .map(|idx| (spv[idx + 3], OperationVariant::Regular))
151        .chain(
152            op_dref_operation_idxs
153                .iter()
154                .map(|idx| (spv[idx + 3], OperationVariant::Dref)),
155        )
156        .collect::<Vec<_>>();
157
158    // 3. Backtrace to find the OpSampledImage that resulted in our loaded sampled images
159    let loaded_variable_ids = op_sampled_image_idxs
160        .iter()
161        .filter_map(|idx| {
162            let sampled_result_id = spv[idx + 2];
163            let loaded_image_id = spv[idx + 3];
164            let loaded_sampler_id = spv[idx + 4];
165            loaded_sampled_image_ids
166                .iter()
167                .find_map(|(id, ty)| (*id == sampled_result_id).then_some(ty))
168                .map(|ty| {
169                    [
170                        (PatchObjectType::Image(loaded_image_id), ty),
171                        (PatchObjectType::Sampler(loaded_sampler_id), ty),
172                    ]
173                })
174        })
175        .flatten()
176        .collect::<Vec<_>>();
177
178    // 4. Backtrack to find the OpLoad that resulted in our loaded images
179    let object_ids = op_load_idxs
180        .iter()
181        .filter_map(|idx| {
182            let loaded_result_id = spv[idx + 2];
183            let original_image_or_sampler = spv[idx + 3];
184            loaded_variable_ids
185                .iter()
186                .find_map(|(id, ty)| (id.inner() == loaded_result_id).then_some((id, ty)))
187                .map(|(id, ty)| (id.next(original_image_or_sampler), idx, ty))
188        })
189        .collect::<Vec<_>>();
190
191    // 5. Find the images that mismatch operations
192    let mut object_flags = HashMap::new();
193    let mut patch_object_id_to_loads = HashMap::new();
194
195    for (id, load_idx, ty) in object_ids.iter().copied() {
196        let entry = object_flags.entry(id).or_insert((false, false));
197
198        match ty {
199            OperationVariant::Regular => entry.0 = true,
200            OperationVariant::Dref => {
201                entry.1 = true;
202            }
203        }
204        patch_object_id_to_loads
205            .entry(id)
206            .or_insert(vec![])
207            .push((load_idx, ty));
208    }
209
210    let mixed_object_ids = object_flags
211        .iter()
212        .filter_map(|(&id, &(uses_regular, uses_dref))| (uses_regular && uses_dref).then_some(id))
213        .collect::<Vec<_>>();
214
215    // 6. Mix object flags from across contexts (OpVariable + OpFunctionArgument)
216    // See `test_cross_dref.spv` and `test_hidden3_dref.spv`.
217    let mut aggregate_flags: HashMap<PatchObjectType<usize>, (bool, bool)> = HashMap::new();
218    for (&id, &flags) in object_flags.iter() {
219        if let Some(&v_idx) = op_variable_idxs
220            .iter()
221            .find(|&&idx| spv[idx + 2] == id.inner())
222        {
223            let entry = aggregate_flags
224                .entry(id.next(v_idx))
225                .or_insert((false, false));
226            entry.0 |= flags.0;
227            entry.1 |= flags.1;
228        }
229    }
230    for (&id, &flags) in object_flags.iter() {
231        if let Some(&fp_idx) = op_function_parameter_idxs
232            .iter()
233            .find(|&&idx| spv[idx + 2] == id.inner())
234        {
235            let entry = get_function_from_parameter(&spv, fp_idx);
236            let mut traced = vec![];
237            let variables =
238                trace_function_argument_to_variables(TraceFunctionArgumentToVariablesIn {
239                    spv: &spv,
240                    op_variable_idxs: &op_variable_idxs,
241                    op_function_parameter_idxs: &op_function_parameter_idxs,
242                    op_function_call_idxs: &op_function_call_idxs,
243                    entry,
244                    traced_function_call_idxs: &mut traced,
245                });
246            for v_idx in variables {
247                if let Some(entry) = aggregate_flags.get_mut(&id.next(v_idx)) {
248                    entry.0 |= flags.0;
249                    entry.1 |= flags.1;
250                }
251            }
252        }
253    }
254    let aggregate_mixed_variables = aggregate_flags
255        .into_iter()
256        .filter_map(|(g, (uses_regular, uses_dref))| (uses_regular && uses_dref).then_some(g))
257        .collect::<HashSet<_>>();
258
259    // 6. Find the OpVariable of the mismatched images
260    let patch_variable_idxs = op_variable_idxs
261        .iter()
262        .filter_map(|idx: &usize| {
263            let result_id = spv[*idx + 2];
264            mixed_object_ids
265                .iter()
266                .find(|id| id.inner() == result_id)
267                .map(|id| id.next(*idx))
268        })
269        .collect::<Vec<_>>();
270
271    // 7. Find OpFunctionParameter of ~~the mismatched~~ ALL images operations
272    // Later, we can keep the ones that trace to mismatched global variables
273    let patch_function_parameter_idxs = op_function_parameter_idxs
274        .iter()
275        .filter_map(|idx: &usize| {
276            let result_id = spv[*idx + 2];
277            mixed_object_ids
278                .iter()
279                .find_map(|id| {
280                    (id.inner() == result_id).then_some((id.next(*idx), MixState::Mixed))
281                })
282                .or(object_ids.iter().find_map(|(id, _, _)| {
283                    (id.inner() == result_id).then_some((id.next(*idx), MixState::PotentiallyMixed))
284                }))
285        })
286        .collect::<Vec<_>>();
287
288    // 8. Find the OpVariable that eventually reaches OpFunctionCall of our OpFunctions
289    // Because functions may be deeply nested, we'll have to account for other OpFunctionCalls
290    let function_patch_variables_with_calls = patch_function_parameter_idxs
291        .iter()
292        .map(|&(idx, mix_state)| {
293            let mut traced_function_calls = vec![];
294            let entry = get_function_from_parameter(&spv, idx.inner());
295            let variables =
296                trace_function_argument_to_variables(TraceFunctionArgumentToVariablesIn {
297                    spv: &spv,
298                    op_variable_idxs: &op_variable_idxs,
299                    op_function_parameter_idxs: &op_function_parameter_idxs,
300                    op_function_call_idxs: &op_function_call_idxs,
301                    entry,
302                    traced_function_call_idxs: &mut traced_function_calls,
303                });
304            (
305                variables
306                    .into_iter()
307                    .map(|v| idx.next(v))
308                    .collect::<Vec<_>>(),
309                traced_function_calls,
310                mix_state,
311            )
312        })
313        .collect::<Vec<_>>();
314
315    let function_patch_variables_with_calls = function_patch_variables_with_calls
316        .iter()
317        .cloned()
318        .filter_map(|(variables, calls, mix_state)| match mix_state {
319            MixState::Mixed => Some((variables, calls)),
320            MixState::PotentiallyMixed => (function_patch_variables_with_calls.iter().any(
321                |(mixed_variables, _, mix_state)| {
322                    *mix_state == MixState::Mixed
323                        && mixed_variables
324                            .iter()
325                            .any(|va| variables.iter().any(|vb| va == vb))
326                },
327            ) || patch_variable_idxs
328                .iter()
329                .any(|idx| variables.iter().any(|va| va == idx))
330                || variables
331                    .iter()
332                    .any(|va| aggregate_mixed_variables.contains(va)))
333            .then_some((variables, calls)),
334        })
335        .collect::<Vec<_>>();
336
337    let mut patch_variable_idxs = patch_variable_idxs
338        .iter()
339        .copied()
340        .map(|idx| (idx, LoadType::Variable))
341        .collect::<Vec<_>>();
342
343    let mut function_patch_variable_set = HashSet::new();
344    for (variables, _) in function_patch_variables_with_calls.iter() {
345        for variable in variables {
346            if !function_patch_variable_set.contains(variable) {
347                patch_variable_idxs.push((*variable, LoadType::FunctionArgument));
348                function_patch_variable_set.insert(*variable);
349            }
350        }
351    }
352
353    // 9. Find OpTypePointer that resulted in OpVariable
354    let patch_variable_idxs = patch_variable_idxs.into_iter().map(|(variable_idx, lty)| {
355        let type_pointer_id = spv[variable_idx.inner() + 1];
356        let maybe_tp_idx = op_type_pointer_idxs.iter().find(|&tp_idx| {
357            let tp_id = spv[tp_idx + 1];
358            type_pointer_id == tp_id
359        });
360        (variable_idx, lty, maybe_tp_idx.copied())
361    });
362
363    // 10. Find OpTypeImage that resulted in OpTypePointer
364    //    We also want to create an complement OpTypeImage (depth=!depth) (without duplicates) and
365    //    a respective OpTypePointer ~~and OpTypeSampledImage pair~~ (also no duplicates).
366    let mut existing_type_pointers_from_type_image = HashMap::new();
367    let mut existing_type_images_from_complement_instruction = HashMap::new();
368
369    let patch_variable_idxs = patch_variable_idxs
370        .map(|(variable_idx, lty, tp_idx)| {
371            match variable_idx {
372                v @ PatchObjectType::Sampler(variable_idx) => {
373                    (
374                        v.next(variable_idx),
375                        lty,
376                        first_op_type_sampler_id,
377                        first_op_type_pointer_sampler_id,
378                        first_op_type_sampler_id,
379                        // From the perspective of a SPIRV sampler variable, this doesn't matter
380                        OperationVariant::Dref,
381                    )
382                }
383                v @ PatchObjectType::Image(variable_idx) => {
384                    let variable_result_id = spv[variable_idx];
385                    let image_type_id = if let Some(tp_idx) = tp_idx {
386                        // type_image_id
387                        spv[tp_idx + 3]
388                    } else if let Some(load_idxs) =
389                        patch_object_id_to_loads.get(&PatchObjectType::Image(variable_result_id))
390                        && let Some(&(load_idx, _)) = load_idxs.first()
391                    {
392                        // We don't have a type pointer, let's find the OpTypeImage via our original OpLoad!
393                        // load_type_result_id
394                        spv[load_idx + 1]
395                    } else {
396                        unreachable!(
397                            "Our OpVariable image id should always point back to a OpLoad id"
398                        );
399                    };
400
401                    // Grab the existing type image
402                    let (ti_idx, ti_id) = op_type_image_idxs
403                        .iter()
404                        .find_map(|&ti_idx| {
405                            let result_id = spv[ti_idx + 1];
406                            (result_id == image_type_id).then_some((ti_idx, result_id))
407                        })
408                        .unwrap();
409
410                    // Try to find an type image with the complement properties or (re-)create one
411                    let ti_word_count = hiword(spv[ti_idx]) as usize;
412                    let mut ti_complement = spv[ti_idx + 2..ti_idx + ti_word_count].to_vec();
413                    let complement_ty = match ti_complement[2] {
414                        0 | 2 => {
415                            ti_complement[2] = 1;
416                            OperationVariant::Dref
417                        }
418                        1 => {
419                            ti_complement[2] = 0;
420                            OperationVariant::Regular
421                        }
422                        _ => panic!("depth field on valid spv can only be 0, 1, or 2"),
423                    };
424
425                    let mut new_instructions = vec![];
426
427                    let complement_ti_id = existing_type_images_from_complement_instruction
428                        .get(&ti_complement)
429                        .copied()
430                        .or(op_type_image_idxs.iter().find_map(|&idx| {
431                            let word_count = hiword(spv[idx]) as usize;
432                            let result_id = spv[idx + 1];
433                            // To have a consistent instruction ordering, we white-out the existing OpTypeImage
434                            if ti_complement == spv[idx + 2..idx + word_count] {
435                                for it in new_spv.iter_mut().skip(idx).take(word_count) {
436                                    *it = encode_word(1, SPV_INSTRUCTION_OP_NOP);
437                                }
438                                Some(result_id)
439                            } else {
440                                None
441                            }
442                        }));
443                    let complement_ti_id = {
444                        let new_type_image_id = complement_ti_id.unwrap_or_else(|| {
445                            instruction_bound += 1;
446                            instruction_bound - 1
447                        });
448                        if !existing_type_images_from_complement_instruction
449                            .contains_key(&ti_complement)
450                        {
451                            let mut new_instruction = vec![
452                                encode_word(
453                                    (ti_complement.len() + 2) as u16,
454                                    SPV_INSTRUCTION_OP_TYPE_IMAGE,
455                                ),
456                                new_type_image_id,
457                            ];
458                            existing_type_images_from_complement_instruction
459                                .insert(ti_complement.clone(), new_type_image_id);
460                            new_instruction.append(&mut ti_complement);
461                            drop(ti_complement);
462                            new_instructions.append(&mut new_instruction);
463                        }
464                        new_type_image_id
465                    };
466
467                    // Try to find a type id for complement type image or create one
468                    let complement_tp_id = existing_type_pointers_from_type_image
469                        .get(&complement_ti_id)
470                        .copied()
471                        .or(op_type_pointer_idxs.iter().find_map(|&idx| {
472                            let result_id = spv[idx + 1];
473                            let type_id = spv[idx + 3];
474                            if type_id == complement_ti_id {
475                                existing_type_pointers_from_type_image
476                                    .insert(complement_ti_id, result_id);
477                                Some(result_id)
478                            } else {
479                                None
480                            }
481                        }))
482                        .unwrap_or_else(|| {
483                            let new_type_pointer_id = instruction_bound;
484                            instruction_bound += 1;
485                            let mut new_instruction = vec![
486                                encode_word(4, SPV_INSTRUCTION_OP_TYPE_POINTER),
487                                new_type_pointer_id,
488                                SPV_STORAGE_CLASS_UNIFORM_CONSTANT,
489                                complement_ti_id,
490                            ];
491                            new_instructions.append(&mut new_instruction);
492                            existing_type_pointers_from_type_image
493                                .insert(complement_ti_id, new_type_pointer_id);
494                            new_type_pointer_id
495                        });
496
497                    instruction_inserts.push(InstructionInsert {
498                        previous_spv_idx: ti_idx,
499                        instruction: new_instructions,
500                    });
501
502                    (
503                        v.next(variable_idx),
504                        lty,
505                        ti_id,
506                        complement_tp_id,
507                        complement_ti_id,
508                        complement_ty,
509                    )
510                }
511            }
512        })
513        .collect::<Vec<_>>();
514
515    // 11. New OpVariable with a new_id, patch old OpLoads, and new depth=1 OpTypeImage.
516    // Map new function arguments to the correct instructions.
517    // NOTE: GENERALLY, with glslc, each OpImage* will get its own OpLoad, so we don't need to
518    // check that its result isn't used for both regular and dref operations!
519    let mut affected_variables = Vec::new();
520
521    // There may be a shared OpTypeFunction but not shared OpFunctionParameter
522    let mut patched_function_types = HashMap::new();
523    let mut patched_function_parameters = HashSet::new();
524
525    // We may patch ourselves a new OpTypeFunction multiple times.
526    // Maps function type id and function index to our new type.
527    let mut defered_new_function_types: HashMap<(u32, usize), InstructionInsert> = HashMap::new();
528
529    for (
530        variable_idx_typed,
531        lty,
532        original_ti_id,
533        complement_tp_id,
534        complement_ti_id,
535        complement_ty,
536    ) in patch_variable_idxs
537    {
538        let variable_idx = variable_idx_typed.inner();
539        // OpVariable
540        let word_count = hiword(spv[variable_idx]);
541        let new_variable_id = instruction_bound;
542        instruction_bound += 1;
543        let mut new_variable = Vec::new();
544        new_variable.extend_from_slice(&spv[variable_idx..variable_idx + word_count as usize]);
545        new_variable[1] = complement_tp_id;
546        new_variable[2] = new_variable_id;
547        instruction_inserts.push(InstructionInsert {
548            previous_spv_idx: variable_idx,
549            instruction: new_variable,
550        });
551
552        affected_variables.push(AffectedDecoration {
553            original_res_id: spv[variable_idx + 2],
554            new_res_ids: vec![new_variable_id],
555            correction_type: match complement_ty {
556                OperationVariant::Regular => CorrectionType::SplitDrefRegular,
557                OperationVariant::Dref => CorrectionType::SplitDrefComparison,
558            },
559        });
560
561        // OpLoad
562        match lty {
563            LoadType::Variable => {
564                let old_variable_id = spv[variable_idx + 2];
565                if let Some(op_load_idxs) =
566                    patch_object_id_to_loads.get(&variable_idx_typed.next(old_variable_id))
567                {
568                    for &(op_load_idx, ty) in op_load_idxs {
569                        if **ty == complement_ty {
570                            new_spv[op_load_idx + 1] = complement_ti_id;
571                            new_spv[op_load_idx + 3] = new_variable_id;
572                        } else {
573                            new_spv[op_load_idx + 1] = original_ti_id;
574                            new_spv[op_load_idx + 3] = old_variable_id;
575                        };
576                    }
577                }
578            }
579            LoadType::FunctionArgument => {
580                let mut function_id_and_index_to_new_parameter_id = HashMap::new();
581
582                // Patch function types, definition parameter, and final loads
583                for (variables, calls) in function_patch_variables_with_calls.iter() {
584                    if variables.contains(&variable_idx_typed.next(variable_idx)) {
585                        for &call in calls.iter().rev() {
586                            let function_id = spv[call.call_parameter.function_idx + 2];
587                            let type_function_id = spv[call.call_parameter.function_idx + 4];
588                            if !patched_function_parameters.contains(&(
589                                call.call_parameter.parameter_instruction_idx,
590                                spv[call.call_parameter.function_idx + 2],
591                            )) {
592                                let Some(type_function_idx) =
593                                    op_type_function_idxs.iter().find(|&&idx| {
594                                        let result_id = spv[idx + 1];
595                                        type_function_id == result_id
596                                    })
597                                else {
598                                    panic!(
599                                        "OpTypeFunction does not exist for function {}, type {}",
600                                        function_id, type_function_id
601                                    );
602                                };
603
604                                // To allow multiple patching we can either patch by taking
605                                // an instruction directly from the code, or by patching a
606                                // function we have already began
607                                // We save patching the function's type for later in case of
608                                // duplicate OpTypeFunction
609                                {
610                                    let (new_type_function_id, type_instruction_type_info) =
611                                        if let Some(new_function_type) = defered_new_function_types
612                                            .get_mut(&(
613                                                type_function_id,
614                                                call.call_parameter.function_idx,
615                                            ))
616                                        {
617                                            new_function_type.instruction.insert(
618                                                3 + call.call_parameter.parameter_instruction_idx
619                                                    + 1
620                                                    + 1,
621                                                complement_tp_id,
622                                            );
623                                            new_function_type.instruction[0] = encode_word(
624                                                new_function_type.instruction.len() as u16,
625                                                SPV_INSTRUCTION_OP_TYPE_FUNCTION,
626                                            );
627                                            (
628                                                new_function_type.instruction[1],
629                                                new_function_type.instruction[2..].to_vec(),
630                                            )
631                                        } else {
632                                            let new_function_type_id = instruction_bound;
633                                            instruction_bound += 1;
634
635                                            let word_count = hiword(spv[*type_function_idx]);
636                                            let mut type_function = Vec::new();
637                                            type_function.extend_from_slice(
638                                                &spv[*type_function_idx
639                                                    ..*type_function_idx + word_count as usize],
640                                            );
641                                            type_function[0] = encode_word(
642                                                word_count + 1,
643                                                SPV_INSTRUCTION_OP_TYPE_FUNCTION,
644                                            );
645                                            type_function[1] = new_function_type_id;
646                                            type_function.insert(
647                                                3 + call.call_parameter.parameter_instruction_idx
648                                                    + 1,
649                                                complement_tp_id,
650                                            );
651
652                                            let type_instruction_type_info =
653                                                type_function[2..].to_vec();
654
655                                            defered_new_function_types.insert(
656                                                (
657                                                    type_function_id,
658                                                    call.call_parameter.function_idx,
659                                                ),
660                                                InstructionInsert {
661                                                    previous_spv_idx: *type_function_idx,
662                                                    instruction: type_function,
663                                                },
664                                            );
665                                            (new_function_type_id, type_instruction_type_info)
666                                        };
667                                    let entry = patched_function_types
668                                        .entry(type_instruction_type_info)
669                                        .or_insert((new_type_function_id, vec![]));
670                                    entry
671                                        .1
672                                        .push((type_function_id, call.call_parameter.function_idx));
673                                }
674
675                                // Patch function parameter
676                                let new_parameter_id = instruction_bound;
677                                instruction_bound += 1;
678                                instruction_inserts.push(InstructionInsert {
679                                    previous_spv_idx: call.call_parameter.parameter_idx,
680                                    instruction: vec![
681                                        encode_word(3, SPV_INSTRUCTION_OP_FUNCTION_PARAMETER),
682                                        complement_tp_id,
683                                        new_parameter_id,
684                                    ],
685                                });
686
687                                // Use our new parameters to patch dependent OpLoads
688                                for load_idx in op_load_idxs.iter() {
689                                    let result_id = spv[load_idx + 2];
690                                    let ptr_id = spv[load_idx + 3];
691                                    let parameter_result_id =
692                                        spv[call.call_parameter.parameter_idx + 2];
693
694                                    // TODO: OPT Someone else can come by and rearrange these silly data
695                                    // structures later.
696                                    if ptr_id == parameter_result_id {
697                                        let ty = loaded_variable_ids
698                                            .iter()
699                                            .find_map(|&(id, ty)| {
700                                                (id.inner() == result_id).then_some(ty)
701                                            })
702                                            .unwrap();
703                                        if *ty == complement_ty {
704                                            new_spv[load_idx + 1] = complement_ti_id;
705                                            new_spv[load_idx + 3] = new_parameter_id;
706                                        }
707                                    }
708                                }
709
710                                let function_id = spv[call.function_call_idx + 3];
711                                function_id_and_index_to_new_parameter_id.insert(
712                                    (function_id, call.call_parameter.parameter_instruction_idx),
713                                    new_parameter_id,
714                                );
715                                patched_function_parameters.insert((
716                                    call.call_parameter.parameter_instruction_idx,
717                                    function_id,
718                                ));
719                            }
720                        }
721                    }
722                }
723
724                // Patch function calls that call other functions
725                for (variables, calls) in function_patch_variables_with_calls.iter() {
726                    if variables.contains(&variable_idx_typed.next(variable_idx)) {
727                        for &call in calls.iter().rev() {
728                            let function_idx = get_function_index_of_instruction_index(
729                                &spv,
730                                call.function_call_idx,
731                            );
732                            let function_id = spv[function_idx + 2];
733                            if let Some(parameter_word) = function_id_and_index_to_new_parameter_id
734                                .get(&(function_id, call.call_parameter.parameter_instruction_idx))
735                            {
736                                word_inserts.push(WordInsert {
737                                    idx: call.function_call_idx
738                                        + 4
739                                        + call.call_parameter.parameter_instruction_idx,
740                                    word: *parameter_word,
741                                    head_idx: call.function_call_idx,
742                                });
743                            } else {
744                                word_inserts.push(WordInsert {
745                                    idx: call.function_call_idx
746                                        + 4
747                                        + call.call_parameter.parameter_instruction_idx,
748                                    word: new_variable_id,
749                                    head_idx: call.function_call_idx,
750                                });
751                            }
752                        }
753                    }
754                }
755            }
756        }
757
758        // OpSampledImage
759        // NOTE: We did not patch in a new OpSampledImage and OpTypeSampledImage.
760        // Thankfully, it seems that `spirv-val`, `naga`, nor `tint` seem to care.
761    }
762
763    // 12. Remove duplicate function types and patch them into OpFunction
764    for (_, (new_type_function_id, functions)) in patched_function_types.into_iter() {
765        for (idx, &(type_function_id, function_idx)) in functions.iter().enumerate() {
766            if idx != 0 {
767                defered_new_function_types.remove(&(type_function_id, function_idx));
768            }
769            new_spv[function_idx + 4] = new_type_function_id;
770        }
771    }
772
773    // We now insert our new function types
774    for (_, new_instruction) in defered_new_function_types {
775        instruction_inserts.push(new_instruction)
776    }
777
778    // 13. Insert new OpDecorate
779    let DecorateOut {
780        descriptor_sets_to_correct,
781    } = util::decorate(DecorateIn {
782        spv: &spv,
783        instruction_inserts: &mut instruction_inserts,
784        first_op_deocrate_idx,
785        op_decorate_idxs: &op_decorate_idxs,
786        affected_decorations: &affected_variables,
787        corrections,
788    });
789
790    // 14. Insert New Instructions
791    insert_new_instructions(&spv, &mut new_spv, &word_inserts, &instruction_inserts);
792
793    // 15. Correct OpDecorate Bindings
794    util::correct_decorate(CorrectDecorateIn {
795        new_spv: &mut new_spv,
796        descriptor_sets_to_correct,
797    });
798
799    // 16. Remove Instructions that have been Whited Out.
800    prune_noops(&mut new_spv);
801
802    // 17. Write New Header and New Code
803    Ok(fuse_final(spv_header, new_spv, instruction_bound))
804}