Skip to main content

spirv_webgpu_transform/
splitbindingarray.rs

1use super::*;
2
3fn inc(ib: &mut u32) -> u32 {
4    *ib += 1;
5    *ib - 1
6}
7
8// For the purposes of this patch, I consider an OpTypeImage and OpTypeSampler to be opaque.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10struct OpaqueArrayType;
11
12mod rechain_instructions;
13mod select_template;
14
15use rechain_instructions::*;
16use select_template::*;
17
18/// Perform the operation on a `Vec<u32>`.
19/// Use [u8_slice_to_u32_vec] to convert a `&[u8]` into a `Vec<u32>`
20///
21/// Assumed SPIR-V properties for this patch:
22///
23/// TODO:
24/// - No nested
25/// - No additional capabilities (SparseResidency or ImageQuery)
26///
27/// wgpu Properties:
28///
29/// - The only opaque types that can be in an array are `OpTypeImage` and `OpTypeSampler`
30///
31/// SPIR-V Properties (These should always be true):
32/// - No opaque types in structures
33/// - All UBOs and SSBO hold a structure and therefore are accessed with `OpAccessChain*` first.
34///
35pub fn splitbindingarray(in_spv: &[u32], corrections: &mut CorrectionMap) -> Result<Vec<u32>, ()> {
36    let spv = in_spv.to_owned();
37
38    let mut instruction_bound = spv[SPV_HEADER_INSTRUCTION_BOUND_OFFSET];
39    let magic_number = spv[SPV_HEADER_MAGIC_NUM_OFFSET];
40
41    let spv_header = spv[0..SPV_HEADER_LENGTH].to_owned();
42
43    assert_eq!(magic_number, SPV_HEADER_MAGIC);
44
45    let mut instruction_inserts = vec![];
46    let word_inserts = vec![];
47
48    let spv = spv.into_iter().skip(SPV_HEADER_LENGTH).collect::<Vec<_>>();
49    let mut new_spv = spv.clone();
50
51    let mut op_type_int_idxs = vec![];
52    let mut op_type_array_idxs = vec![];
53    let mut op_type_pointer_idxs = vec![];
54    let mut op_type_image_idxs = vec![];
55    let mut op_type_sampler_idxs = vec![];
56    let mut op_constant_idxs = vec![];
57    let mut op_variable_idxs = vec![];
58    let mut op_access_chain_idxs = vec![];
59    let mut op_in_bounds_access_chain_idxs = vec![];
60    let mut op_load_idxs = vec![];
61    let mut op_store_idxs = vec![];
62    let mut op_copy_memory_idxs = vec![];
63    let mut op_type_function_idxs = vec![];
64    let mut op_function_parameter_idxs = vec![];
65    let mut op_function_call_idxs = vec![];
66    let mut op_function_end_idxs = vec![];
67    let mut op_decorate_idxs = vec![];
68    let mut op_name_idxs = vec![];
69    let mut op_sampled_image_idxs = vec![];
70
71    // 1. Find locations instructions we need
72    let mut spv_idx = 0;
73    while spv_idx < spv.len() {
74        let op = spv[spv_idx];
75        let word_count = hiword(op);
76        let instruction = loword(op);
77
78        match instruction {
79            SPV_INSTRUCTION_OP_TYPE_INT => op_type_int_idxs.push(spv_idx),
80            SPV_INSTRUCTION_OP_TYPE_ARRAY => op_type_array_idxs.push(spv_idx),
81            SPV_INSTRUCTION_OP_TYPE_POINTER => op_type_pointer_idxs.push(spv_idx),
82            SPV_INSTRUCTION_OP_TYPE_IMAGE => op_type_image_idxs.push(spv_idx),
83            SPV_INSTRUCTION_OP_TYPE_SAMPLER => op_type_sampler_idxs.push(spv_idx),
84            SPV_INSTRUCTION_OP_CONSTANT => op_constant_idxs.push(spv_idx),
85            SPV_INSTRUCTION_OP_VARIABLE => op_variable_idxs.push(spv_idx),
86            SPV_INSTRUCTION_OP_ACCESS_CHAIN => op_access_chain_idxs.push(spv_idx),
87            SPV_INSTRUCTION_OP_IN_BOUNDS_ACCESS_CHAIN => {
88                op_in_bounds_access_chain_idxs.push(spv_idx)
89            }
90            SPV_INSTRUCTION_OP_LOAD => op_load_idxs.push(spv_idx),
91            SPV_INSTRUCTION_OP_STORE => op_store_idxs.push(spv_idx),
92            SPV_INSTRUCTION_OP_COPY_MEMORY => op_copy_memory_idxs.push(spv_idx),
93            SPV_INSTRUCTION_OP_TYPE_FUNCTION => op_type_function_idxs.push(spv_idx),
94            SPV_INSTRUCTION_OP_FUNCTION_PARAMETER => op_function_parameter_idxs.push(spv_idx),
95            SPV_INSTRUCTION_OP_FUNCTION_CALL => op_function_call_idxs.push(spv_idx),
96            SPV_INSTRUCTION_OP_FUNCTION_END => op_function_end_idxs.push(spv_idx),
97            SPV_INSTRUCTION_OP_DECORATE => op_decorate_idxs.push(spv_idx),
98            SPV_INSTRUCTION_OP_NAME => op_name_idxs.push(spv_idx),
99            SPV_INSTRUCTION_OP_SAMPLED_IMAGE => op_sampled_image_idxs.push(spv_idx),
100
101            _ => {}
102        }
103
104        spv_idx += word_count as usize;
105    }
106
107    // 2. OpTypeArray -> OpTypePointer
108    //      -> OpVariable
109    //      -> OpFunctionParameter
110    let array_tp_ta_idxs = op_type_pointer_idxs
111        .iter()
112        .filter_map(|&tp_idx| {
113            let tp_storage_class = spv[tp_idx + 2];
114            let tp_underlying_id = spv[tp_idx + 3];
115
116            if tp_storage_class != SPV_STORAGE_CLASS_UNIFORM_CONSTANT
117                && tp_storage_class != SPV_STORAGE_CLASS_UNIFORM
118            {
119                return None;
120            }
121
122            op_type_array_idxs
123                .iter()
124                .find(|&ta_idx| {
125                    let ta_res_id = spv[ta_idx + 1];
126
127                    ta_res_id == tp_underlying_id
128                })
129                .map(|&ta_idx| {
130                    let array_type = op_type_image_idxs
131                        .iter()
132                        .chain(op_type_sampler_idxs.iter())
133                        .any(|&t_idx| spv[t_idx + 1] == spv[ta_idx + 2])
134                        .then_some(OpaqueArrayType);
135
136                    (tp_idx, ta_idx, array_type)
137                })
138        })
139        .collect::<Vec<_>>();
140
141    // Contains ((OpVariable or OpFunctionParameter), OpTypePointer, Option<OpaqueArrayType>)
142    // OpVariable is a subtype of OpFunctionParameter over the first three words.
143    let array_vfp_ta_idxs = op_variable_idxs
144        .iter()
145        .chain(op_function_parameter_idxs.iter())
146        .filter_map(|&vfp_idx| {
147            let variable_type_id = spv[vfp_idx + 1];
148            array_tp_ta_idxs
149                .iter()
150                .find(|&(tp_idx, _, _)| {
151                    let tp_res_id = spv[tp_idx + 1];
152                    tp_res_id == variable_type_id
153                })
154                .map(|&(_, ta_idx, array_type)| (vfp_idx, ta_idx, array_type))
155        })
156        .collect::<Vec<_>>();
157
158    // TODO: Implement for nested arrays.
159    for ta_idx in op_type_array_idxs.iter() {
160        let ta_underlying_id = spv[ta_idx + 2];
161        for (_, ta_jdx, _) in array_tp_ta_idxs.iter() {
162            let j_result_id = spv[ta_jdx + 1];
163            if j_result_id == ta_underlying_id && ta_idx != ta_jdx {
164                unimplemented!("How dare you use nested arrays! (Unimplemented)");
165            }
166        }
167    }
168
169    // 3. Build mapping of lengths
170    let length_map = array_vfp_ta_idxs
171        .iter()
172        .map(|(_, ta_idx, _)| {
173            let length_id = spv[ta_idx + 3];
174            let Some(length) = op_constant_idxs.iter().find_map(|&constant_idx| {
175                (spv[constant_idx + 2] == length_id).then_some(spv[constant_idx + 3])
176            }) else {
177                panic!("Missing OpConstant")
178            };
179            (ta_idx, length)
180        })
181        .collect::<HashMap<_, _>>();
182
183    // 4. Unroll array variables
184    let types_header_position = last_of_indices!(op_type_int_idxs, op_type_pointer_idxs);
185    let mut types_header_insert = InstructionInsert {
186        previous_spv_idx: types_header_position.unwrap(),
187        instruction: vec![],
188    };
189    let mut new_vfp_map = HashMap::new();
190    let mut function_type_changes = HashMap::new();
191    let mut affected_decorations = vec![];
192
193    for &(vfp_idx, ta_idx, array_type) in array_vfp_ta_idxs.iter() {
194        new_spv[vfp_idx..vfp_idx + hiword(spv[vfp_idx]) as usize]
195            .fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
196
197        let mut new_type_instructions = vec![];
198
199        let instruction = loword(spv[vfp_idx]);
200        let underlying_type_id = spv[ta_idx + 2];
201        let type_pointer_id = ensure_type_pointer(
202            &spv,
203            &op_type_pointer_idxs,
204            &mut instruction_bound,
205            &mut new_type_instructions,
206            match array_type {
207                Some(OpaqueArrayType) => SPV_STORAGE_CLASS_UNIFORM_CONSTANT,
208                _ => SPV_STORAGE_CLASS_UNIFORM,
209            },
210            underlying_type_id,
211        );
212
213        let length = length_map[&ta_idx];
214
215        let base_id = instruction_bound;
216        instruction_bound += length;
217
218        match instruction {
219            SPV_INSTRUCTION_OP_VARIABLE => {
220                for i in 0..length {
221                    new_type_instructions.append(&mut vec![
222                        encode_word(4, SPV_INSTRUCTION_OP_VARIABLE),
223                        type_pointer_id,
224                        base_id + i,
225                        match array_type {
226                            Some(OpaqueArrayType) => SPV_STORAGE_CLASS_UNIFORM_CONSTANT,
227                            _ => SPV_STORAGE_CLASS_UNIFORM,
228                        },
229                    ]);
230                }
231                // Ordering issues with this, let's keep it after all other type pointers.
232                //
233                // instruction_inserts.push(InstructionInsert {
234                //     previous_spv_idx: v_idx,
235                //     instruction: new_instruction,
236                // });
237                types_header_insert
238                    .instruction
239                    .append(&mut new_type_instructions);
240                let old_result_id = spv[vfp_idx + 2];
241
242                // We manually correct the base variable to reuse the original decorations.
243                // That way, we can output `N-1` correction bindings.
244                for &d_idx in op_decorate_idxs.iter() {
245                    if spv[d_idx + 1] == old_result_id {
246                        new_spv[d_idx + 1] = base_id;
247                    }
248                }
249                for &n_idx in op_name_idxs.iter() {
250                    if spv[n_idx + 1] == old_result_id {
251                        new_spv[n_idx + 1] = base_id;
252                    }
253                }
254
255                // We only want `N-1` correction bindings.
256                let new_ids = (base_id + 1..base_id + length).collect::<Vec<_>>();
257                affected_decorations.push(AffectedDecoration {
258                    original_res_id: old_result_id,
259                    new_res_ids: new_ids,
260                    correction_type: CorrectionType::SplitBindingArray,
261                });
262            }
263            SPV_INSTRUCTION_OP_FUNCTION_PARAMETER => {
264                let mut new_param_instructions = vec![];
265                for i in 0..length {
266                    new_param_instructions.append(&mut vec![
267                        encode_word(3, SPV_INSTRUCTION_OP_FUNCTION_PARAMETER),
268                        type_pointer_id,
269                        base_id + i,
270                    ]);
271                }
272                instruction_inserts.push(InstructionInsert {
273                    previous_spv_idx: vfp_idx,
274                    instruction: new_param_instructions,
275                });
276
277                let entry = get_function_from_parameter(&spv, vfp_idx);
278                let function_type_id = spv[entry.function_idx + 4];
279
280                // `entry.parameter_instruction_idx` is the 0-based ordinal of the parameter
281                // within the function; step 5 compares it against the loop variable `i`.
282                function_type_changes
283                    .entry(function_type_id)
284                    .or_insert(vec![])
285                    .push((entry.parameter_instruction_idx, type_pointer_id, length));
286            }
287            _ => unreachable!("Expected OpVariable or OpFunctionParameter"),
288        };
289
290        new_vfp_map.insert(vfp_idx, (base_id, ta_idx));
291    }
292
293    // 5. Change affected OpTypeFunction
294    for &tf_idx in op_type_function_idxs.iter() {
295        let tf_result_id = spv[tf_idx + 1];
296
297        let Some(changes) = function_type_changes.get(&tf_result_id) else {
298            continue;
299        };
300
301        let tf_wc = hiword(spv[tf_idx]) as usize;
302        let num_params = tf_wc - 3;
303
304        let mut new_params: Vec<u32> = vec![];
305        let mut change_i = 0;
306        for i in 0..num_params {
307            if change_i < changes.len() && changes[change_i].0 == i {
308                let (_, type_ptr, length) = changes[change_i];
309                for _ in 0..length {
310                    new_params.push(type_ptr);
311                }
312                change_i += 1;
313            } else {
314                new_params.push(spv[tf_idx + 3 + i]);
315            }
316        }
317
318        new_spv[tf_idx..tf_idx + tf_wc].fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
319
320        let new_wc = (3 + new_params.len()) as u16;
321        let mut new_tf = vec![
322            encode_word(new_wc, SPV_INSTRUCTION_OP_TYPE_FUNCTION),
323            tf_result_id,
324            spv[tf_idx + 2], // return type (unchanged)
325        ];
326        new_tf.extend_from_slice(&new_params);
327        types_header_insert.instruction.extend_from_slice(&new_tf);
328    }
329
330    let access_idxs = op_access_chain_idxs
331        .iter()
332        .chain(op_in_bounds_access_chain_idxs.iter())
333        .filter_map(|&ac_idx| {
334            let base_id = spv[ac_idx + 3];
335            array_vfp_ta_idxs
336                .iter()
337                .find(|&(vfp_idx, _, _)| {
338                    let result_id = spv[*vfp_idx + 2];
339                    result_id == base_id
340                })
341                .map(|(vfp_idx, ta_idx, array_type)| (ac_idx, vfp_idx, ta_idx, array_type))
342        })
343        .collect::<Vec<_>>();
344
345    // 6. Trace array samplers into a map
346    // Arrayed samplers turn our neat trace tree into a DAG.
347    // To keep things simple, we handle samplers separately.
348    // See `opaque_trace.rs` for details.
349    let mut arrayed_sampler_map = HashMap::new();
350    for &(ac_idx, &vfp_idx, ta_idx, &array_type) in access_idxs.iter() {
351        let access_result_id = spv[ac_idx + 2];
352        if let Some(OpaqueArrayType) = array_type {
353            for &load_idx in op_load_idxs.iter() {
354                let result_id = spv[load_idx + 2];
355                let pointer_id = spv[load_idx + 3];
356                if pointer_id == access_result_id {
357                    for &sampled_image_idx in op_sampled_image_idxs.iter() {
358                        let sampler_id = spv[sampled_image_idx + 4];
359                        if sampler_id == result_id {
360                            arrayed_sampler_map
361                                .insert(sampled_image_idx, (ac_idx, vfp_idx, ta_idx));
362                        }
363                    }
364                }
365            }
366        }
367    }
368
369    // 7. Replace OpAccessChain with selection function
370    for &(ac_idx, vfp_idx, ta_idx, array_type) in access_idxs.iter() {
371        let ac_word_count = hiword(spv[ac_idx]) as usize;
372        new_spv[ac_idx..ac_idx + ac_word_count].fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
373
374        let old_result_id = spv[ac_idx + 2];
375        let index_0_id = spv[ac_idx + 4];
376
377        let length = length_map[&ta_idx];
378
379        let (base_id, _) = new_vfp_map[vfp_idx];
380
381        if let Some(OpaqueArrayType) = *array_type {
382            // When both a texture array and a sampler array feed the same OpSampledImage,
383            // the texture AC's processing already generates the correct nested switch
384            // (outer = texture index, inner = sampler index via `maybe_sampler_array_data`).
385            //
386            // Detect this by checking whether this AC is already stored as the sampler
387            // dimension in `arrayed_sampler_map`.  If so, just NOP its dependent loads
388            // (they reference the now-undefined AC result) and skip switch generation.
389            let is_inner_sampler_ac = arrayed_sampler_map
390                .values()
391                .any(|&(map_ac_idx, _, _)| map_ac_idx == ac_idx);
392            if is_inner_sampler_ac {
393                for &load_idx in op_load_idxs.iter() {
394                    if spv[load_idx + 3] == old_result_id {
395                        let wc = hiword(spv[load_idx]) as usize;
396                        new_spv[load_idx..load_idx + wc]
397                            .fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
398                    }
399                }
400                continue;
401            }
402
403            let load_idxs = op_load_idxs
404                .iter()
405                .filter(|&idx| {
406                    let pointer = spv[idx + 3];
407                    pointer == old_result_id
408                })
409                .copied()
410                .collect::<Vec<_>>();
411            let dependent_traces = trace_loaded_opaques(&spv, &load_idxs);
412            for trace in dependent_traces {
413                let maybe_sampler_array_data = match trace.next {
414                    OpaqueImageOp::Sampled(sampled_image_op) => {
415                        arrayed_sampler_map.get(&sampled_image_op.idx)
416                    }
417                    _ => None,
418                };
419
420                let switch_instructions =
421                    reconstruct_opaque_trace_and_overwrite(&spv, &mut new_spv, &trace);
422                let underlying_type_and_target_id =
423                    get_last_instruction_result_type_and_id(&switch_instructions);
424                let rotate_image_sampler = matches!(
425                    trace.next,
426                    OpaqueImageOp::Sampled(SampledImageOp {
427                        parent: SampledImageParent::Sampler,
428                        ..
429                    })
430                );
431                let rechain_instructions = |ib: &mut u32, target_id: u32| {
432                    let (instructions, output) = rechain_instructions_with_target_id(
433                        ib,
434                        &switch_instructions,
435                        target_id,
436                        false,
437                        rotate_image_sampler,
438                    );
439                    (instructions, output.map(|(_, id)| id))
440                };
441                // Track inner merge labels per outer case so we can fix the outer phi after select_template_spv runs.
442                // `select_template_spv` puts the outer case labels in the phi, but with a nested inner switch the actual predecessor
443                // of the outer merge is the inner merge block, not the outer case block.
444                let mut inner_merge_labels: Vec<u32> = vec![];
445
446                let builder = |ib: &mut u32, target_id: u32| {
447                    if let Some((sampler_array_ac_idx, sampler_array_v_idx, sampler_array_ta_idx)) =
448                        maybe_sampler_array_data
449                    {
450                        let (sampler_base_id, _) = new_vfp_map[sampler_array_v_idx];
451                        let sampler_index_0_id = spv[sampler_array_ac_idx + 4];
452                        let sampler_length = length_map[sampler_array_ta_idx] as usize;
453
454                        // Rechain only the image load for this outer case.
455                        // switch_instructions = [image_load, OpSampledImage, ...]
456                        // target_id is the split image variable for this outer case.
457                        let image_load_wc = hiword(switch_instructions[0]) as usize;
458                        let (image_load_instrs, image_out) = rechain_instructions_with_target_id(
459                            ib,
460                            &switch_instructions[..image_load_wc],
461                            target_id,
462                            false,
463                            false,
464                        );
465                        let (_, new_image_id) =
466                            image_out.expect("image load must produce a result");
467
468                        // Locate OpSampledImage and any instructions that follow it.
469                        let si_wc = hiword(switch_instructions[image_load_wc]) as usize;
470                        let after_si = &switch_instructions[image_load_wc + si_wc..];
471                        let sampler_type_id = spv[op_type_sampler_idxs[0] + 1];
472
473                        // Inner builder: per sampler variable j, emit sampler load + OpSampledImage + trailing instructions.
474                        // The image load is placed before the inner switch.
475                        let inner_builder = |ib: &mut u32, inner_target_id: u32| {
476                            let mut instrs = vec![];
477
478                            let new_sampler_result = inc(ib);
479                            instrs.extend_from_slice(&[
480                                encode_word(4, SPV_INSTRUCTION_OP_LOAD),
481                                sampler_type_id,
482                                new_sampler_result,
483                                inner_target_id,
484                            ]);
485
486                            let new_si_result = inc(ib);
487                            let mut si_patched =
488                                switch_instructions[image_load_wc..image_load_wc + si_wc].to_vec();
489                            si_patched[2] = new_si_result;
490                            si_patched[3] = new_image_id;
491                            si_patched[4] = new_sampler_result;
492                            instrs.extend_from_slice(&si_patched);
493
494                            if !after_si.is_empty() {
495                                let (chained, output) = rechain_instructions_with_target_id(
496                                    ib,
497                                    after_si,
498                                    new_si_result,
499                                    false,
500                                    false,
501                                );
502                                instrs.extend_from_slice(&chained);
503                                return (instrs, output.map(|(_, id)| id));
504                            }
505                            (instrs, Some(new_si_result))
506                        };
507
508                        let mut inner_switch = select_template_spv(
509                            ib,
510                            sampler_base_id,
511                            sampler_index_0_id,
512                            sampler_length,
513                            inner_builder,
514                            underlying_type_and_target_id,
515                        );
516
517                        // Find the inner merge label (last OpLabel before the inner phi).
518                        let phi_idx = get_last_instruction_index(&inner_switch);
519                        {
520                            let mut idx = 0;
521                            let mut label = 0u32;
522                            while idx < phi_idx {
523                                if loword(inner_switch[idx]) == SPV_INSTRUCTION_OP_LABEL {
524                                    label = inner_switch[idx + 1];
525                                }
526                                idx += hiword(inner_switch[idx]) as usize;
527                            }
528                            inner_merge_labels.push(label);
529                        }
530
531                        // Patch the inner phi's result id to a fresh id for the outer phi.
532                        let output_id = (loword(inner_switch[phi_idx]) == SPV_INSTRUCTION_OP_PHI)
533                            .then(|| {
534                                let new_id = inc(ib);
535                                inner_switch[phi_idx + 2] = new_id;
536                                new_id
537                            });
538
539                        // Emit: image load once for this outer case, then the inner switch.
540                        let mut result = image_load_instrs;
541                        result.extend_from_slice(&inner_switch);
542                        (result, output_id)
543                    } else {
544                        rechain_instructions(ib, target_id)
545                    }
546                };
547
548                let mut switch = select_template_spv(
549                    &mut instruction_bound,
550                    base_id,
551                    index_0_id,
552                    length as usize,
553                    builder,
554                    underlying_type_and_target_id,
555                );
556
557                // Patch the outer phi's predecessor labels.
558                // select_template_spv filled them with the outer case labels,
559                // but each outer case now ends at its inner merge block, not at the outer case label.
560                // phi layout: [opword, type, result_id, val0, pred0, val1, pred1, ...]
561                if !inner_merge_labels.is_empty() {
562                    let phi_idx = get_last_instruction_index(&switch);
563                    if loword(switch[phi_idx]) == SPV_INSTRUCTION_OP_PHI {
564                        for (i, &label) in inner_merge_labels.iter().enumerate() {
565                            switch[phi_idx + 4 + 2 * i] = label;
566                        }
567                    }
568                }
569
570                instruction_inserts.push(InstructionInsert {
571                    previous_spv_idx: trace.last_result_id(),
572                    instruction: switch,
573                });
574            }
575        } else {
576            // For concreate types, find all dependent operations afterwards and replace each instruction with index switch
577            for &spv_idx in op_load_idxs
578                .iter()
579                .chain(op_store_idxs.iter())
580                .chain(op_access_chain_idxs.iter())
581                .chain(op_in_bounds_access_chain_idxs.iter())
582                .chain(op_copy_memory_idxs.iter())
583            {
584                let word_count = hiword(spv[spv_idx]) as usize;
585                let instruction = loword(spv[spv_idx]);
586
587                let mut flip_store_into = false;
588                let is_dependent = match instruction {
589                    SPV_INSTRUCTION_OP_STORE | SPV_INSTRUCTION_OP_COPY_MEMORY => {
590                        // We need to handle cases where buffers are stored from and to.
591                        let source_id = spv[spv_idx + 1];
592                        let dest_id = spv[spv_idx + 2];
593
594                        // OpStore: %result = %a
595                        if dest_id == old_result_id {
596                            flip_store_into = true;
597                        }
598
599                        source_id == old_result_id || dest_id == old_result_id
600                    }
601                    SPV_INSTRUCTION_OP_LOAD
602                    | SPV_INSTRUCTION_OP_ACCESS_CHAIN
603                    | SPV_INSTRUCTION_OP_IN_BOUNDS_ACCESS_CHAIN => {
604                        let source_id = spv[spv_idx + 3];
605                        source_id == old_result_id
606                    }
607                    _ => unreachable!("Unexpected instruction {} while matching", instruction),
608                };
609
610                if is_dependent && ac_idx != spv_idx {
611                    if instruction == SPV_INSTRUCTION_OP_ACCESS_CHAIN
612                        || instruction == SPV_INSTRUCTION_OP_IN_BOUNDS_ACCESS_CHAIN
613                    {
614                        unimplemented!(
615                            "Nested OpAccessChain / OpInBoundsAccessChain on binding array (Unimplemented)"
616                        );
617                    }
618
619                    // We don't want to fully overwrite the access chain since UBOs and SSBOs
620                    // accesses will always be followed by these.
621                    let mut new_instructions = [
622                        &spv[ac_idx..ac_idx + 4],
623                        &spv[ac_idx + 5..ac_idx + ac_word_count],
624                        &spv[spv_idx..spv_idx + word_count],
625                    ]
626                    .concat();
627                    new_instructions[0] =
628                        encode_word(ac_word_count as u16 - 1, SPV_INSTRUCTION_OP_ACCESS_CHAIN);
629
630                    new_spv[spv_idx..spv_idx + word_count]
631                        .fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
632
633                    let builder = &|ib: &mut u32, target_id: u32| {
634                        let (instructions, output) = rechain_instructions_with_target_id(
635                            ib,
636                            &new_instructions,
637                            target_id,
638                            flip_store_into,
639                            false,
640                        );
641                        (instructions, output.map(|(_, id)| id))
642                    };
643
644                    let underlying_type_and_target_id =
645                        get_last_instruction_result_type_and_id(&new_instructions);
646                    let switch = select_template_spv(
647                        &mut instruction_bound,
648                        base_id,
649                        index_0_id,
650                        length as usize,
651                        builder,
652                        underlying_type_and_target_id,
653                    );
654                    instruction_inserts.push(InstructionInsert {
655                        previous_spv_idx: spv_idx,
656                        instruction: switch,
657                    });
658                }
659            }
660        }
661    }
662
663    // 8. Replace all OpFunctionCall references of arrayed resources
664    let new_vfp_id_map = new_vfp_map
665        .iter()
666        .map(|(&vfp_idx, &v)| {
667            let result_id = spv[vfp_idx + 2];
668            (result_id, v)
669        })
670        .collect::<HashMap<_, _>>();
671    for &function_call_idx in op_function_call_idxs.iter() {
672        const ARGUMENT_OFFSET: usize = 4;
673        let word_count = hiword(spv[function_call_idx]) as usize;
674        let mut arguments = vec![];
675        for &argument_id in spv
676            .iter()
677            .take(function_call_idx + word_count)
678            .skip(function_call_idx + ARGUMENT_OFFSET)
679        {
680            if let Some(&(base_id, ta_idx)) = new_vfp_id_map.get(&argument_id) {
681                let length = length_map[&ta_idx];
682                for i in 0..length {
683                    arguments.push(base_id + i);
684                }
685            } else {
686                arguments.push(argument_id)
687            }
688        }
689
690        if arguments.len() != word_count - ARGUMENT_OFFSET {
691            new_spv[function_call_idx..function_call_idx + word_count]
692                .fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
693            let new_instruction = [
694                &[encode_word(
695                    (arguments.len() + ARGUMENT_OFFSET) as u16,
696                    SPV_INSTRUCTION_OP_FUNCTION_CALL,
697                )],
698                &spv[function_call_idx + 1..function_call_idx + ARGUMENT_OFFSET],
699                arguments.as_slice(),
700            ]
701            .concat();
702            instruction_inserts.push(InstructionInsert {
703                previous_spv_idx: function_call_idx,
704                instruction: new_instruction,
705            });
706        }
707    }
708
709    // 9. Copy Old OpDecorate to New Variables.
710    for (&vfp_idx, &(base_id, ta_idx)) in new_vfp_map.iter() {
711        let old_id = spv[vfp_idx + 2];
712        let length = length_map[&ta_idx];
713
714        for i in 1..length {
715            let new_id = base_id + i;
716            copy_decorate(CopyDecorateIn {
717                spv: &spv,
718                op_decorate_idxs: &op_decorate_idxs,
719                instruction_inserts: &mut instruction_inserts,
720                old_id,
721                new_id,
722            });
723        }
724    }
725
726    // 10. Find OpDecorate / OpName to OpVariable
727    let unused_decorate_idxs = op_decorate_idxs
728        .iter()
729        .filter(|&&idx| {
730            let target = spv[idx + 1];
731            if new_spv[idx + 1] != target {
732                return false;
733            }
734            new_vfp_map.iter().any(|(vfp_idx, _)| {
735                let result_id = spv[vfp_idx + 2];
736                target == result_id
737            })
738        })
739        .copied()
740        .collect::<Vec<_>>();
741    let unused_name_idxs = op_name_idxs
742        .iter()
743        .filter(|&&idx| {
744            let target = spv[idx + 1];
745            if new_spv[idx + 1] != target {
746                return false;
747            }
748            new_vfp_map.iter().any(|(vfp_idx, _)| {
749                let result_id = spv[vfp_idx + 2];
750                target == result_id
751            })
752        })
753        .copied()
754        .collect::<Vec<_>>();
755
756    // 11. Remove Instructions that have been Whited Out.
757    for &spv_idx in unused_decorate_idxs.iter().chain(unused_name_idxs.iter()) {
758        let op = spv[spv_idx];
759        let word_count = hiword(op) as usize;
760
761        new_spv[spv_idx..spv_idx + word_count].fill(encode_word(1, SPV_INSTRUCTION_OP_NOP));
762    }
763
764    // 12. OpDecorate
765    let DecorateOut {
766        descriptor_sets_to_correct,
767    } = util::decorate(DecorateIn {
768        spv: &spv,
769        instruction_inserts: &mut instruction_inserts,
770        first_op_deocrate_idx: op_decorate_idxs.first().copied(),
771        op_decorate_idxs: &op_decorate_idxs,
772        affected_decorations: &affected_decorations,
773        corrections,
774    });
775
776    // 13. Insert New Instructions
777    instruction_inserts.insert(0, types_header_insert);
778    insert_new_instructions(&spv, &mut new_spv, &word_inserts, &instruction_inserts);
779
780    // 14. Correct OpDecorate Bindings
781    util::correct_decorate(CorrectDecorateIn {
782        new_spv: &mut new_spv,
783        descriptor_sets_to_correct,
784    });
785    prune_noops(&mut new_spv);
786
787    // 15. Write New Header and New Code
788    Ok(fuse_final(spv_header, new_spv, instruction_bound))
789}