rewrite_combined_image_samplers.rs

use std::borrow::Cow;
use std::collections::hash_map::Entry;
use rspirv::dr::{Builder, InsertPoint, Instruction, Module, Operand};
use rustc_hash::{FxHashMap, FxHashSet};
use spirv::{Decoration, Op, StorageClass, Word};

pub struct Pass<'a> {
    pub builder: &'a mut Builder,
}

#[derive(Clone, Debug)]
pub struct CombinedImageSampler {
    texture_variable: spirv::Word,
    sampler_variable: spirv::Word,
    sampler_pointer_type: spirv::Word,
    original_uniform_pointer_type_id: spirv::Word,
    original_uniform_type: Instruction,

    target_texture_type_id: spirv::Word,
    target_texture_pointer_type_id: spirv::Word,

    // Always OpTypeImage
    base_type: Instruction,
}

#[derive(Debug)]
pub struct OpAccessChain<'a> {
    sampled_image: &'a CombinedImageSampler,
    index: Operand,
    original_result_type: Word,
    target_pointer_type: Word,
    pub result_id: Word,
}

#[derive(Debug)]
pub struct OpFunctionCallParam<'a> {
    sampled_image: &'a CombinedImageSampler,
    type_handle: spirv::Word,
}

impl<'a> Pass<'a> {
    pub fn new(builder: &'a mut Builder) -> Self {
        let mut val = Self { builder };

        val
    }

    fn do_pass(&mut self) {
        let mut combined_image_samplers = self.collect_global_sampled_images();
        self.retype_combined_image_sampler_uniforms(&combined_image_samplers);
        self.put_variables_to_end();

        // First rewrite global loads
        self.rewrite_loads(&combined_image_samplers);

        let mut op_access_chains = self.collect_op_access_chain(&combined_image_samplers);
        self.rewrite_global_op_access_chain_loads(&op_access_chains);

        // Collect functions that reference combined image samplers indirectly (todo) (also rewrite the found OpFunctionCall)


        while !combined_image_samplers.is_empty() {
            let op_functions = self.rewrite_function_calls(&op_access_chains, &combined_image_samplers);

            combined_image_samplers = self.rewrite_functions_definitions(&op_functions);
            self.rewrite_loads(&combined_image_samplers);

            op_access_chains = self.collect_op_access_chain(&combined_image_samplers);

            eprintln!("{:?}", &op_access_chains);
            self.rewrite_global_op_access_chain_loads(&op_access_chains);
        }

        // notes: OpFunctionCall %void %assign_s21_ %tex (i.e for scalar no OpLoad precedes, OpLoad and OpAccessChain is in function), this is true even when passing array of samplers
        //        However, if passing single sampler from array, OpAccessChain happens in outer scope.
        // Rewrite functions that do so by adding sampler param
        // rewrite loads + op_access_chains for functions
    }

    fn put_variables_to_end(&mut self) {
        // this is easier than doing proper topo sort.
        let mut vars = Vec::new();

        self.builder.module_mut()
            .types_global_values.retain(|instr| {
            if instr.class.opcode == spirv::Op::Variable {
                vars.push(instr.clone());
                return false;
            };
            true
        });

        self.builder.module_mut()
            .types_global_values.append(&mut vars);
    }
    fn ensure_op_type_sampler(&mut self) {
        self.builder.type_sampler();
    }

    fn find_instruction(&self, word: Word) -> Option<&Instruction> {
        self.builder
            .module_ref()
            .all_inst_iter()
            .find(|i| i.result_id == Some(word))
    }

    fn find_global_instruction(&self, word: Word) -> Option<&Instruction> {
        self.builder
            .module_ref()
            .global_inst_iter()
            .find(|i| i.result_id == Some(word))
    }

    fn find_global_instruction_mut(&mut self, word: Word) -> Option<&mut Instruction> {
        self.builder
            .module_mut()
            .global_inst_iter_mut()
            .find(|i| i.result_id == Some(word))
    }

    fn create_sampler_name(&mut self, word: Word) -> Option<String> {
        self.builder.module_ref().debug_names.iter().find_map(|i| {
            if i.class.opcode != spirv::Op::Name {
                return None;
            }

            let Some(&Operand::IdRef(target)) = &i.operands.get(0) else {
                return None;
            };

            if target != word {
                return None;
            }

            let Some(Operand::LiteralString(string)) = &i.operands.get(1) else {
                return None;
            };

            return Some(format!("_{string}_sampler"));
        })
    }

    fn get_base_type_for_sampled_image(
        &'a self,
        mut inst: &'a Instruction,
    ) -> Option<&'a Instruction> {
        if inst.class.opcode != spirv::Op::TypeSampledImage {
            return None;
        }

        let Some(&Operand::IdRef(referand)) = inst.operands.get(0) else {
            return None;
        };

        self.find_global_instruction(referand)
    }

    // Create a sampler OpVariable at the selected block
    fn create_sampler_uniform(
        &mut self,
        uniform_type: spirv::Word,
        combined_image_sampler: spirv::Word,
    ) -> (spirv::Word, spirv::Word) {
        let sampler_pointer_type =
            self.builder
                .type_pointer(None, StorageClass::UniformConstant, uniform_type);

        let sampler_uniform = self.builder.variable(
            sampler_pointer_type,
            None,
            StorageClass::UniformConstant,
            None,
        );

        let decorations: Vec<Instruction> = self
            .builder
            .module_ref()
            .annotations
            .iter()
            .filter_map(|f| {
                if f.class.opcode == spirv::Op::Decorate
                    && f.operands[0] == Operand::IdRef(combined_image_sampler)
                {
                    Some(f.clone())
                } else {
                    None
                }
            })
            .collect();

        if let Some(name) = self.create_sampler_name(combined_image_sampler) {
            self.builder.name(sampler_uniform, name);
        }

        // Clone decorations to the created sampler
        for decoration in decorations {
            let Operand::Decoration(decoration_type) = decoration.operands[1] else {
                continue;
            };

            self.builder.decorate(
                sampler_uniform,
                decoration_type,
                decoration.operands[2..].iter().map(|f| f.clone()),
            )
        }

        (sampler_pointer_type, sampler_uniform)
    }
    fn collect_global_sampled_images(&mut self) -> FxHashMap<spirv::Word, CombinedImageSampler> {
        let mut image_sampler_cadidates = Vec::new();
        let mut image_sampler_types = FxHashMap::default();

        for global in self.builder.module_ref().types_global_values.iter() {
            if global.class.opcode == spirv::Op::Variable
                && global.operands[0] == Operand::StorageClass(StorageClass::UniformConstant)
            {
                let pointer_type = global.result_type;
                let Some(pointer_type) = pointer_type else {
                    continue;
                };
                image_sampler_cadidates.push((pointer_type, global.result_id))
            }
        }

        for (pointer_type_id, global_variable) in image_sampler_cadidates {
            let Some(pointer_type) = self.find_global_instruction(pointer_type_id).cloned() else {
                continue;
            };

            if pointer_type.class.opcode == spirv::Op::TypePointer
                && pointer_type.operands[0] == Operand::StorageClass(StorageClass::UniformConstant)
            {
                let Some(&Operand::IdRef(sampled_image_type)) = pointer_type.operands.get(1) else {
                    continue;
                };

                let Some(uniform_type) = self.find_global_instruction(sampled_image_type).cloned()
                else {
                    continue;
                };

                if uniform_type.class.opcode == spirv::Op::TypeSampledImage {
                    let Some(base_type) =
                        self.get_base_type_for_sampled_image(&uniform_type).cloned()
                    else {
                        continue;
                    };

                    let Some(combined_image_sampler) = global_variable else {
                        continue;
                    };

                    // insert the sampler
                    if base_type.class.opcode != spirv::Op::TypeImage {
                        continue;
                    }

                    let Some(base_type_id) = base_type.result_id else {
                        continue;
                    };

                    let sampler_type = self.builder.type_sampler();

                    let (sampler_pointer_type, sampler_uniform) =
                        self.create_sampler_uniform(sampler_type, combined_image_sampler);

                    image_sampler_types.insert(
                        combined_image_sampler,
                        CombinedImageSampler {
                            texture_variable: combined_image_sampler,
                            sampler_variable: sampler_uniform,
                            original_uniform_type: uniform_type,
                            target_texture_type_id: base_type_id,
                            original_uniform_pointer_type_id: pointer_type_id,
                            base_type,
                            sampler_pointer_type,
                            target_texture_pointer_type_id: pointer_type_id,
                        },
                    );

                    continue;
                }

                if uniform_type.class.opcode == spirv::Op::TypeArray {
                    let Some(&Operand::IdRef(array_base_type)) = uniform_type.operands.get(0)
                    else {
                        continue;
                    };

                    let Some(&Operand::IdRef(array_length)) = uniform_type.operands.get(1) else {
                        continue;
                    };

                    let Some(sampled_image_type) =
                        self.find_global_instruction(array_base_type).cloned()
                    else {
                        continue;
                    };

                    let Some(base_type) = self
                        .get_base_type_for_sampled_image(&sampled_image_type)
                        .cloned()
                    else {
                        continue;
                    };

                    let Some(combined_image_sampler) = global_variable else {
                        continue;
                    };

                    // insert the sampler
                    if base_type.class.opcode != spirv::Op::TypeImage {
                        continue;
                    }

                    let sampler_type = self.builder.type_sampler();
                    let sampler_array_type = self.builder.type_array(sampler_type, array_length);

                    let (sampler_pointer_type, sampler_uniform) =
                        self.create_sampler_uniform(sampler_array_type, combined_image_sampler);


                    // insert target types
                    let Some(base_type_id) = base_type.result_id else {
                        continue;
                    };

                    let target_texture_type_id = self.builder
                        .type_array(base_type_id, array_length);

                    let target_texture_pointer_type_id = self.builder
                        .type_pointer(None, StorageClass::UniformConstant, target_texture_type_id);

                    image_sampler_types.insert(
                        combined_image_sampler,
                        CombinedImageSampler {
                            texture_variable: combined_image_sampler,
                            sampler_variable: sampler_uniform,
                            original_uniform_type: uniform_type,
                            target_texture_type_id,
                            original_uniform_pointer_type_id: pointer_type_id,
                            base_type,
                            sampler_pointer_type,
                            target_texture_pointer_type_id,
                        },
                    );
                    continue;
                }
            }
        }

        image_sampler_types
    }

    fn retype_combined_image_sampler_uniforms(
        &mut self,
        combined_image_samplers: &FxHashMap<spirv::Word, CombinedImageSampler>,
    ) {
        // Need to rebuild the global instructions because we need to insert new types...
        let mut instructions = Vec::new();
        for instr in self.builder.module_ref().types_global_values.clone() {
            let Some(result_id) = instr.result_id else {
                instructions.push(instr);
                continue;
            };

            let Some(sampled_image) = combined_image_samplers.get(&result_id) else {
                // We need to fix..
                instructions.push(instr);
                continue;
            };

            let Some(base_type_id) = sampled_image.base_type.result_id else {
                instructions.push(instr);
                continue;
            };

            // If it's a OpTypeSampledImage, we want to change the variable type to &TypeImage.
            if sampled_image.original_uniform_type.class.opcode == spirv::Op::TypeSampledImage {
                // keep labels in sync
                let mut op_variable = instr;
                op_variable.result_type = Some(sampled_image.target_texture_pointer_type_id);

                instructions.push(op_variable);
                continue;
            }

            // Re-type array globals.
            // We don't need to worry about the pointer type of the load, as
            // we can instantiate that later.
            if sampled_image.original_uniform_type.class.opcode == spirv::Op::TypeArray {
                // let Some(&Operand::IdRef(array_length)) =
                //     sampled_image.original_uniform_type.operands.get(1)
                // else {
                //     instructions.push(instr);
                //     continue;
                // };

                let mut op_variable = instr;
                op_variable.result_type = Some(sampled_image.target_texture_pointer_type_id);

                instructions.push(op_variable);
            }
        }

        // replace
        self.builder.module_mut().types_global_values = instructions;
    }

    fn rewrite_loads(
        &mut self,
        combined_image_samplers: &FxHashMap<spirv::Word, CombinedImageSampler>,
    ) {
        let op_type_sampler = self.builder.type_sampler();

        // need to clone
        let mut functions = self.builder.module_ref().functions.clone();

        for function in functions.iter_mut() {
            for block in function.blocks.iter_mut() {
                let mut instructions = Vec::new();
                for instr in block.instructions.drain(..) {
                    if instr.class.opcode != Op::Load {
                        instructions.push(instr);
                        continue;
                    }

                    // This doesn't affect array loads because array loads load the result of the OpAccessChain which can be done in a separate pass.
                    let Some(Operand::IdRef(op_variable_id)) = &instr.operands.get(0) else {
                        instructions.push(instr);
                        continue;
                    };

                    let Some(combined_image_sampler) = combined_image_samplers.get(op_variable_id)
                    else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let op_load_texture_id = self.builder.id();
                    let op_load_sampler_id = self.builder.id();

                    let mut load_instr = instr.clone();

                    load_instr.result_type = combined_image_sampler.base_type.result_id;
                    load_instr.result_id = Some(op_load_texture_id);
                    instructions.push(load_instr);

                    // load the sampler
                    instructions.push(Instruction {
                        class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::Load),
                        result_type: Some(op_type_sampler),
                        result_id: Some(op_load_sampler_id),
                        operands: vec![Operand::IdRef(combined_image_sampler.sampler_variable)],
                    });

                    // reuse the old id for the OpSampleImage
                    instructions.push(Instruction {
                        class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::SampledImage),
                        result_type: combined_image_sampler.original_uniform_type.result_id,
                        result_id: instr.result_id,

                        operands: vec![
                            Operand::IdRef(op_load_texture_id),
                            Operand::IdRef(op_load_sampler_id),
                        ],
                    });
                }
                block.instructions = instructions;
            }
        }

        self.builder.module_mut().functions = functions;
    }

    fn collect_op_access_chain<'b>(
        &mut self,
        combined_image_samplers: &'b FxHashMap<spirv::Word, CombinedImageSampler>,
    ) -> FxHashMap<spirv::Word, OpAccessChain<'b>> {
        // need to clone
        let mut functions = self.builder.module_ref().functions.clone();

        let mut seen_op_access_chain = FxHashMap::default();

        for function in functions.iter_mut() {
            for block in function.blocks.iter_mut() {
                let mut instructions = Vec::new();

                // This needs to be done in two passes, first to find the OpAccessChains and
                // update their type to the pointer type of the scalar value,

                // Then to update the OpLoads.
                for instr in block.instructions.clone() {
                    if instr.class.opcode != Op::AccessChain {
                        instructions.push(instr);
                        continue;
                    }

                    // Need to refer to it later
                    let Some(result_id) = instr.result_id else {
                        instructions.push(instr);
                        continue;
                    };

                    // This doesn't affect array loads because array loads load the result of the OpAccessChain which can be done in a separate pass.
                    let Some(Operand::IdRef(op_variable)) = &instr.operands.get(0) else {
                        instructions.push(instr);
                        continue;
                    };

                    // Ensure that this is an OpAccessChain of our combined image sampler.
                    let Some(sampled_image) = combined_image_samplers.get(op_variable) else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let Some(original_result_type) = instr
                        .result_type else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let Some(index) = instr
                        .operands.get(1).cloned() else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let Some(base_type_id) = sampled_image.base_type.result_id else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };



                    let op_pointer_type_id = self.builder.type_pointer(
                        None,
                        StorageClass::UniformConstant,
                        base_type_id,
                    );

                    let mut op_access_chain = instr;

                    let Some(result_id) = op_access_chain.result_id else {
                        instructions.push(op_access_chain);
                        continue;
                    };

                    op_access_chain.result_type = Some(op_pointer_type_id);

                    instructions.push(op_access_chain);

                    seen_op_access_chain.insert(
                        result_id,
                        OpAccessChain {
                            sampled_image,
                            index,
                            original_result_type,
                            target_pointer_type: op_pointer_type_id,
                            result_id
                        },
                    );
                }

                block.instructions = instructions;
            }
        }

        self.builder.module_mut().functions = functions;
        seen_op_access_chain
    }

    fn rewrite_global_op_access_chain_loads(
        &mut self,
        op_access_chains: &FxHashMap<spirv::Word, OpAccessChain>,
    ) {
        let op_type_sampler = self.builder.type_sampler();
        let op_type_sampler_pointer =
            self.builder
                .type_pointer(None, StorageClass::UniformConstant, op_type_sampler);

        let mut functions = self.builder.module_ref().functions.clone();

        for function in functions.iter_mut() {
            for block in function.blocks.iter_mut() {
                let mut instructions = Vec::new();

                // This needs to be done in two passes, first to find the OpAccessChains and
                // update their type to the pointer type of the scalar value,

                // Then to update the OpLoads.
                for instr in block.instructions.clone() {
                    if instr.class.opcode != Op::Load {
                        instructions.push(instr);
                        continue;
                    }

                    let Some(Operand::IdRef(op_variable)) = &instr.operands.get(0) else {
                        instructions.push(instr);
                        continue;
                    };

                    // Ensure that this is an OpLoad of the previous OpAccessChain
                    let Some(&OpAccessChain {
                        sampled_image,
                        ref index,
                        ..
                    }) = op_access_chains.get(op_variable)
                    else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let Some(base_type_id) = sampled_image.base_type.result_id else {
                        // Ignore what we didn't process before.
                        instructions.push(instr);
                        continue;
                    };

                    let op_load_texture_id = self.builder.id();
                    let op_load_sampler_id = self.builder.id();
                    let op_access_chain_sampler_id = self.builder.id();

                    let op_type_sampled_image = self.builder.type_sampled_image(base_type_id);

                    let mut load_instr = instr.clone();

                    eprintln!("ADDING: %{op_load_texture_id} = OpLoad %{:?}", sampled_image.base_type.result_id);
                    eprintln!("{:?}", sampled_image.base_type);

                    load_instr.result_type = sampled_image.base_type.result_id;
                    load_instr.result_id = Some(op_load_texture_id);
                    instructions.push(load_instr);

                    // load the sampler
                    instructions.push(Instruction {
                        class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::AccessChain),
                        result_type: Some(op_type_sampler_pointer),
                        result_id: Some(op_access_chain_sampler_id),
                        operands: vec![
                            Operand::IdRef(sampled_image.sampler_variable),
                            index.clone(),
                        ],
                    });

                    instructions.push(Instruction {
                        class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::Load),
                        result_type: Some(op_type_sampler),
                        result_id: Some(op_load_sampler_id),
                        operands: vec![Operand::IdRef(op_access_chain_sampler_id)],
                    });

                    // reuse the old id for the OpSampleImage
                    instructions.push(Instruction {
                        class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::SampledImage),
                        result_type: Some(op_type_sampled_image),
                        result_id: instr.result_id,

                        operands: vec![
                            Operand::IdRef(op_load_texture_id),
                            Operand::IdRef(op_load_sampler_id),
                        ],
                    });
                }

                block.instructions = instructions;
            }
        }

        self.builder.module_mut().functions = functions;
    }

    // Returns nested hashmap of { OpFunction: { OpType: CombinedImageSampler } },
    // indicating that the listed parameters should change
    fn rewrite_function_calls<'b>(
        &mut self,
        op_access_chains: &'b FxHashMap<spirv::Word, OpAccessChain>,
        combined_image_samplers: &'b FxHashMap<spirv::Word, CombinedImageSampler>,
    ) -> FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>> {

        let mut seen_functions: FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>> = FxHashMap::default();

        // First pass, rewrite function calls involving sampled image access
        for instr in self.builder.module_mut().all_inst_iter_mut() {
            if instr.class.opcode != spirv::Op::FunctionCall {
                continue;
            }

            let Some(&Operand::IdRef(function_id)) = instr.operands.get(0) else {
                continue;
            };

            if !instr.operands[1..].iter().any(|param| {
                let &Operand::IdRef(function_id) = param else {
                    return false;
                };

                combined_image_samplers.contains_key(&function_id)
            }) {
                continue;
            };

            let mut function_call_operands = Vec::with_capacity(instr.operands.len());

            for operand in instr.operands.drain(..) {
                let Operand::IdRef(op_ref_id) = operand else {
                    function_call_operands.push(operand);
                    continue;
                };

                let Some(sampled_image) = combined_image_samplers.get(&op_ref_id) else {
                    function_call_operands.push(operand);
                    continue;
                };

                // todo: assumes homogenous types, but need to check behaviour for array types and
                // doing op_load after.
                // Will need to do like rewrite_function_calls_with_global_op_access_chain
                function_call_operands.push(operand);
                function_call_operands.push(Operand::IdRef(sampled_image.sampler_variable));

                match seen_functions.entry(function_id) {
                    Entry::Occupied(mut vec) => {
                        vec.get_mut().insert(sampled_image.original_uniform_pointer_type_id, Cow::Borrowed(sampled_image));
                    }
                    Entry::Vacant(vec) => {
                        let mut map = FxHashMap::default();
                        // eprintln!("fn({}), {:#?}", function_id, sampled_image);
                        map.insert(sampled_image.original_uniform_pointer_type_id, Cow::Borrowed(sampled_image));
                        vec.insert(map);
                    }
                }
            }
            instr.operands = function_call_operands;
        }

        // Deal with OpAccessChains.

        let op_type_sampler = self.builder.type_sampler();
        let op_type_sampler_pointer =
            self.builder
                .type_pointer(None, StorageClass::UniformConstant, op_type_sampler);

        let mut functions = self.builder.module_ref().functions.clone();

        for function in functions.iter_mut() {
            for block in function.blocks.iter_mut() {
                let mut instructions = Vec::new();

                // This needs to be done in two passes, first to find the OpAccessChains and
                // update their type to the pointer type of the scalar value,

                // Then to update the OpLoads.
                for mut instr in block.instructions.clone() {
                    if instr.class.opcode != Op::FunctionCall {
                        instructions.push(instr);
                        continue;
                    }

                    let Some(&Operand::IdRef(function_id)) = instr.operands.get(0) else {
                        instructions.push(instr);
                        continue;
                    };

                    // Ensure that this is an OpFunctionCall involving some previous OpAccessChain
                    let relevant_operands = instr.operands[1..].iter().filter(|param| {
                        let &Operand::IdRef(op_ref_id) = param else {
                            return false;
                        };

                        op_access_chains.contains_key(&op_ref_id)
                    }).collect::<Vec<_>>();

                    if relevant_operands.is_empty() {
                        instructions.push(instr);
                        continue;
                    }

                    // Maps OpAccessChain of texture to OpAccessChain of sampler
                    let mut op_access_sampler_mapping = FxHashMap::default();
                    // Insert necessary OpAccessChains
                    for operand in relevant_operands {
                        let &Operand::IdRef(op_ref_id) = operand else {
                            continue;
                        };

                        let Some(&OpAccessChain {
                            sampled_image,
                            ref index,
                            ..
                        }) = op_access_chains.get(&op_ref_id) else {
                            continue;
                        };

                        let op_access_chain_sampler_id = self.builder.id();
                        // access chain the sampler
                        instructions.push(Instruction {
                            class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::AccessChain),
                            result_type: Some(op_type_sampler_pointer),
                            result_id: Some(op_access_chain_sampler_id),
                            operands: vec![
                                Operand::IdRef(sampled_image.sampler_variable),
                                index.clone(),
                            ],
                        });

                        op_access_sampler_mapping.insert(op_ref_id, op_access_chain_sampler_id);
                    }

                    let mut function_call_operands = Vec::with_capacity(instr.operands.len());
                    for operand in instr.operands.drain(..) {
                        let Operand::IdRef(op_ref_id) = operand else {
                            function_call_operands.push(operand);
                            continue;
                        };

                        let Some(&OpAccessChain {
                            sampled_image,
                            ref index,
                            original_result_type,
                            target_pointer_type,
                            result_id
                        }) = op_access_chains.get(&op_ref_id) else {
                            function_call_operands.push(operand);
                            continue;
                        };

                        let Some(base_type_id) = sampled_image.base_type.result_id else {
                            function_call_operands.push(operand);
                            continue;
                        };

                        let Some(&op_access_chain_sampler) = op_access_sampler_mapping
                            .get(&op_ref_id) else {
                            function_call_operands.push(operand);
                            continue;
                        };

                        function_call_operands.push(operand);
                        function_call_operands.push(Operand::IdRef(op_access_chain_sampler));

                        let original_type = self
                            .find_global_instruction(original_result_type)
                            .cloned()
                            .expect("huh");

                        let sampled_image = CombinedImageSampler {
                            texture_variable: result_id,
                            sampler_variable: op_access_chain_sampler,
                            sampler_pointer_type: op_type_sampler_pointer,
                            original_uniform_pointer_type_id: original_result_type,
                            original_uniform_type: original_type,
                            target_texture_type_id: base_type_id,
                            target_texture_pointer_type_id: target_pointer_type,
                            base_type: sampled_image.base_type.clone(),
                        };

                        match seen_functions.entry(function_id) {
                            Entry::Occupied(mut vec) => {
                                vec.get_mut().insert(original_result_type, Cow::Owned(sampled_image));
                            }
                            Entry::Vacant(vec) => {
                                let mut map = FxHashMap::default();
                                map.insert(original_result_type, Cow::Owned(sampled_image));
                                vec.insert(map);
                            }
                        }
                    }

                    instr.operands = function_call_operands;

                    // Restore the instructions
                    instructions.push(instr);
                }

                block.instructions = instructions;
            }
        }

        self.builder.module_mut().functions = functions;
        seen_functions
    }

    fn rewrite_functions_definitions<'b>(&mut self, mappings: &FxHashMap<spirv::Word, FxHashMap<spirv::Word, Cow<'b, CombinedImageSampler>>>)
                                         -> FxHashMap<spirv::Word, CombinedImageSampler>
    {
        let mut sampled_refs = FxHashMap::default();
        let mut functions = self.builder.module_ref().functions.clone();

        for function in functions.iter_mut() {
            let Some(def_id) = function.def_id() else {
                continue;
            };

            let Some(function_def) = function.def.clone() else {
                continue;
            };

            let Some(param_mappings) = mappings.get(&def_id) else {
                continue;
            };

            let &Some(function_return_type_id) = &function_def.result_type else {
                continue;
            };


            let Some(&Operand::IdRef(function_type_id)) = &function_def.operands.last()  else {
                continue;
            };


            let mut parameters = Vec::new();
            let mut param_types = Vec::new();
            for param in function.parameters.drain(..) {
                let Some(param_id) = param.result_id else {
                    parameters.push(param);
                    continue;
                };

                let Some(param_type) = param.result_type else {
                    parameters.push(param);
                    continue;
                };

                let Some(sampled_image) = param_mappings.get(&param_type) else {
                    parameters.push(param);
                    continue;
                };

                let op_function_parameter_sampler_id = self.builder.id();

                // todo: updaite param type to base type
                // parameters.push(param);


                parameters.push(Instruction {
                    class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::FunctionParameter),
                    result_type: Some(sampled_image.target_texture_pointer_type_id),
                    result_id: param.result_id,
                    operands: vec![],
                });

                parameters.push(Instruction {
                    class: rspirv::grammar::CoreInstructionTable::get(spirv::Op::FunctionParameter),
                    result_type: Some(sampled_image.sampler_pointer_type),
                    result_id: Some(op_function_parameter_sampler_id),
                    operands: vec![],
                });

                param_types.push(sampled_image.target_texture_pointer_type_id);
                param_types.push(sampled_image.sampler_pointer_type);

                sampled_refs.insert(param_id, CombinedImageSampler {
                    texture_variable: param_id,
                    sampler_variable: op_function_parameter_sampler_id,
                    sampler_pointer_type: sampled_image.sampler_pointer_type,
                    original_uniform_pointer_type_id: param_type,
                    original_uniform_type: sampled_image.original_uniform_type.clone(),
                    target_texture_type_id: sampled_image.target_texture_type_id,
                    target_texture_pointer_type_id: sampled_image.target_texture_pointer_type_id,
                    base_type: sampled_image.base_type.clone()
                });
            }

            let new_type = self.builder
                .type_function(function_return_type_id, param_types);

            if let Some(function) = &mut function.def {
                if let Some(function_type) = function.operands.last_mut() {
                    *function_type = Operand::IdRef(new_type);
                }
            }

            function.parameters = parameters;
        }

        self.builder.module_mut().functions = functions;

        sampled_refs
    }

}

MIT Licensed

Copyright 2023 Ronny Chan

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.