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