#![allow(unknown_lints, unnecessary_transmutes)]

use std::mem::transmute;

use crate::{
    SpirvCompiler, SpirvTarget,
    item::{Elem, Item},
    lookups::Array,
};
use cubecl_core::ir::{self, ConstantValue, Id};
use rspirv::{
    dr::Builder,
    spirv::{self, FPEncoding, StorageClass, Word},
};

#[derive(Debug, Clone, PartialEq)]
pub enum Variable {
    GlobalInputArray(Word, Item, u32),
    GlobalOutputArray(Word, Item, u32),
    GlobalScalar(Word, Elem),
    Constant(Word, ConstVal, Item),
    Local {
        id: Word,
        item: Item,
    },
    Versioned {
        id: (Id, u16),
        item: Item,
        variable: ir::Variable,
    },
    LocalBinding {
        id: Id,
        item: Item,
        variable: ir::Variable,
    },
    Raw(Word, Item),
    Slice {
        ptr: Box<Variable>,
        offset: Word,
        end: Word,
        const_len: Option<u32>,
        item: Item,
    },
    SharedArray(Word, Item, u32),
    Shared(Word, Item),
    ConstantArray(Word, Item, u32),
    LocalArray(Word, Item, u32),
    CoopMatrix(Id, Elem),
    Id(Word),
    Builtin(Word, Item),
}

impl Variable {
    pub fn scope(&self) -> spirv::Scope {
        match self {
            Variable::GlobalInputArray(..)
            | Variable::GlobalOutputArray(..)
            | Variable::GlobalScalar(..) => spirv::Scope::Device,
            Variable::SharedArray(..) | Variable::Shared(..) => spirv::Scope::Workgroup,
            Variable::CoopMatrix(..) => spirv::Scope::Subgroup,
            Variable::Slice { ptr, .. } => ptr.scope(),
            Variable::Raw(..) => unimplemented!("Can't get scope of raw variable"),
            Variable::Id(_) => unimplemented!("Can't get scope of raw id"),
            _ => spirv::Scope::Invocation,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ConstVal {
    Bit32(u32),
    Bit64(u64),
}

impl ConstVal {
    pub fn as_u64(&self) -> u64 {
        match self {
            ConstVal::Bit32(val) => *val as u64,
            ConstVal::Bit64(val) => *val,
        }
    }

    pub fn as_u32(&self) -> u32 {
        match self {
            ConstVal::Bit32(val) => *val,
            ConstVal::Bit64(_) => panic!("Truncating 64 bit variable to 32 bit"),
        }
    }

    pub fn as_float(&self, width: u32, encoding: Option<FPEncoding>) -> f64 {
        match (width, encoding) {
            (64, _) => f64::from_bits(self.as_u64()),
            (32, _) => f32::from_bits(self.as_u32()) as f64,
            (16, None) => half::f16::from_bits(self.as_u32() as u16).to_f64(),
            (_, Some(FPEncoding::BFloat16KHR)) => {
                half::bf16::from_bits(self.as_u32() as u16).to_f64()
            }
            (_, Some(FPEncoding::Float8E4M3EXT)) => {
                cubecl_common::e4m3::from_bits(self.as_u32() as u8).to_f64()
            }
            (_, Some(FPEncoding::Float8E5M2EXT)) => {
                cubecl_common::e5m2::from_bits(self.as_u32() as u8).to_f64()
            }
            _ => unreachable!(),
        }
    }

    pub fn as_int(&self, width: u32) -> i64 {
        unsafe {
            match width {
                64 => transmute::<u64, i64>(self.as_u64()),
                32 => transmute::<u32, i32>(self.as_u32()) as i64,
                16 => transmute::<u16, i16>(self.as_u32() as u16) as i64,
                8 => transmute::<u8, i8>(self.as_u32() as u8) as i64,
                _ => unreachable!(),
            }
        }
    }

    pub fn from_float(value: f64, width: u32, encoding: Option<FPEncoding>) -> Self {
        match (width, encoding) {
            (64, _) => ConstVal::Bit64(value.to_bits()),
            (32, _) => ConstVal::Bit32((value as f32).to_bits()),
            (16, None) => ConstVal::Bit32(half::f16::from_f64(value).to_bits() as u32),
            (_, Some(FPEncoding::BFloat16KHR)) => {
                ConstVal::Bit32(half::bf16::from_f64(value).to_bits() as u32)
            }
            (_, Some(FPEncoding::Float8E4M3EXT)) => {
                ConstVal::Bit32(cubecl_common::e4m3::from_f64(value).to_bits() as u32)
            }
            (_, Some(FPEncoding::Float8E5M2EXT)) => {
                ConstVal::Bit32(cubecl_common::e5m2::from_f64(value).to_bits() as u32)
            }
            _ => unreachable!(),
        }
    }

    pub fn from_int(value: i64, width: u32) -> Self {
        match width {
            64 => ConstVal::Bit64(unsafe { transmute::<i64, u64>(value) }),
            32 => ConstVal::Bit32(unsafe { transmute::<i32, u32>(value as i32) }),
            16 => ConstVal::Bit32(unsafe { transmute::<i16, u16>(value as i16) } as u32),
            8 => ConstVal::Bit32(unsafe { transmute::<i8, u8>(value as i8) } as u32),
            _ => unreachable!(),
        }
    }

    pub fn from_uint(value: u64, width: u32) -> Self {
        match width {
            64 => ConstVal::Bit64(value),
            32 => ConstVal::Bit32(value as u32),
            16 => ConstVal::Bit32(value as u16 as u32),
            8 => ConstVal::Bit32(value as u8 as u32),
            _ => unreachable!(),
        }
    }

    pub fn from_bool(value: bool) -> Self {
        ConstVal::Bit32(value as u32)
    }
}

impl From<(ConstantValue, Item)> for ConstVal {
    fn from((value, ty): (ConstantValue, Item)) -> Self {
        let elem = ty.elem();
        let width = elem.size() * 8;
        match value {
            ConstantValue::Int(val) => ConstVal::from_int(val, width),
            ConstantValue::Float(val) => ConstVal::from_float(val, width, elem.float_encoding()),
            ConstantValue::UInt(val) => ConstVal::from_uint(val, width),
            ConstantValue::Bool(val) => ConstVal::from_bool(val),
        }
    }
}

impl From<u32> for ConstVal {
    fn from(value: u32) -> Self {
        ConstVal::Bit32(value)
    }
}

impl From<f32> for ConstVal {
    fn from(value: f32) -> Self {
        ConstVal::Bit32(value.to_bits())
    }
}

impl Variable {
    pub fn id<T: SpirvTarget>(&self, b: &mut SpirvCompiler<T>) -> Word {
        match self {
            Variable::GlobalInputArray(id, _, _) => *id,
            Variable::GlobalOutputArray(id, _, _) => *id,
            Variable::GlobalScalar(id, _) => *id,
            Variable::Constant(id, _, _) => *id,
            Variable::Local { id, .. } => *id,
            Variable::Versioned {
                id, variable: var, ..
            } => b.get_versioned(*id, var),
            Variable::LocalBinding {
                id, variable: var, ..
            } => b.get_binding(*id, var),
            Variable::Raw(id, _) => *id,
            Variable::Slice { ptr, .. } => ptr.id(b),
            Variable::SharedArray(id, _, _) => *id,
            Variable::Shared(id, _) => *id,
            Variable::ConstantArray(id, _, _) => *id,
            Variable::LocalArray(id, _, _) => *id,
            Variable::CoopMatrix(_, _) => unimplemented!("Can't get ID from matrix var"),
            Variable::Id(id) => *id,
            Variable::Builtin(id, ..) => *id,
        }
    }

    pub fn item(&self) -> Item {
        match self {
            Variable::GlobalInputArray(_, item, _) => item.clone(),
            Variable::GlobalOutputArray(_, item, _) => item.clone(),
            Variable::GlobalScalar(_, elem) => Item::Scalar(*elem),
            Variable::Constant(_, _, item) => item.clone(),
            Variable::Local { item, .. } => item.clone(),
            Variable::Versioned { item, .. } => item.clone(),
            Variable::LocalBinding { item, .. } => item.clone(),
            Variable::Slice { item, .. } => item.clone(),
            Variable::SharedArray(_, item, _) => item.clone(),
            Variable::Shared(_, item) => item.clone(),
            Variable::ConstantArray(_, item, _) => item.clone(),
            Variable::LocalArray(_, item, _) => item.clone(),
            Variable::CoopMatrix(_, elem) => Item::Scalar(*elem),
            Variable::Builtin(_, item) => item.clone(),
            Variable::Raw(_, item) => item.clone(),
            Variable::Id(_) => unimplemented!("Can't get item of raw ID"),
        }
    }

    pub fn indexed_item(&self) -> Item {
        match self {
            Variable::LocalBinding {
                item: Item::Vector(elem, _),
                ..
            } => Item::Scalar(*elem),
            Variable::Local {
                item: Item::Vector(elem, _),
                ..
            } => Item::Scalar(*elem),
            Variable::Versioned {
                item: Item::Vector(elem, _),
                ..
            } => Item::Scalar(*elem),
            other => other.item(),
        }
    }

    pub fn elem(&self) -> Elem {
        self.item().elem()
    }

    pub fn has_len(&self) -> bool {
        matches!(
            self,
            Variable::GlobalInputArray(_, _, _)
                | Variable::GlobalOutputArray(_, _, _)
                | Variable::Slice { .. }
                | Variable::SharedArray(_, _, _)
                | Variable::ConstantArray(_, _, _)
                | Variable::LocalArray(_, _, _)
        )
    }

    pub fn has_buffer_len(&self) -> bool {
        matches!(
            self,
            Variable::GlobalInputArray(_, _, _) | Variable::GlobalOutputArray(_, _, _)
        )
    }

    pub fn as_const(&self) -> Option<ConstVal> {
        match self {
            Self::Constant(_, val, _) => Some(*val),
            _ => None,
        }
    }

    pub fn as_binding(&self) -> Option<Id> {
        match self {
            Self::LocalBinding { id, .. } => Some(*id),
            _ => None,
        }
    }
}

#[derive(Debug)]
pub enum IndexedVariable {
    Pointer(Word, Item),
    Composite(Word, u32, Item),
    DynamicComposite(Word, u32, Item),
    Scalar(Variable),
}

impl<T: SpirvTarget> SpirvCompiler<T> {
    pub fn compile_variable(&mut self, variable: ir::Variable) -> Variable {
        let item = variable.ty;
        match variable.kind {
            ir::VariableKind::Constant(value) => {
                let item = self.compile_type(item);
                let const_val = (value, item.clone()).into();

                if let Some(existing) = self.state.constants.get(&(const_val, item.clone())) {
                    Variable::Constant(*existing, const_val, item)
                } else {
                    let id = item.constant(self, const_val);
                    self.state.constants.insert((const_val, item.clone()), id);
                    Variable::Constant(id, const_val, item)
                }
            }
            ir::VariableKind::GlobalInputArray(pos) => {
                let id = self.state.buffers[pos as usize];
                Variable::GlobalInputArray(id, self.compile_type(item), pos)
            }
            ir::VariableKind::GlobalOutputArray(pos) => {
                let id = self.state.buffers[pos as usize];
                Variable::GlobalOutputArray(id, self.compile_type(item), pos)
            }
            ir::VariableKind::GlobalScalar(id) => self.global_scalar(id, item.storage_type()),
            ir::VariableKind::LocalMut { id } => {
                let item = self.compile_type(item);
                let var = self.get_local(id, &item, variable);
                Variable::Local { id: var, item }
            }
            ir::VariableKind::Versioned { id, version } => {
                let item = self.compile_type(item);
                let id = (id, version);
                Variable::Versioned { id, item, variable }
            }
            ir::VariableKind::LocalConst { id } => {
                let item = self.compile_type(item);
                Variable::LocalBinding { id, item, variable }
            }
            ir::VariableKind::Builtin(builtin) => {
                let item = self.compile_type(item);
                self.compile_builtin(builtin, item)
            }
            ir::VariableKind::ConstantArray { id, length, .. } => {
                let item = self.compile_type(item);
                let id = self.state.const_arrays[id as usize].id;
                Variable::ConstantArray(id, item, length as u32)
            }
            ir::VariableKind::SharedArray { id, length, .. } => {
                let item = self.compile_type(item);
                let id = self.state.shared_arrays[&id].id;
                Variable::SharedArray(id, item, length as u32)
            }
            ir::VariableKind::Shared { id } => {
                let item = self.compile_type(item);
                let id = self.state.shared[&id].id;
                Variable::Shared(id, item)
            }
            ir::VariableKind::LocalArray {
                id,
                length,
                unroll_factor,
            } => {
                let item = self.compile_type(item);
                let id = if let Some(arr) = self.state.local_arrays.get(&id) {
                    arr.id
                } else {
                    let arr_ty = Item::Array(Box::new(item.clone()), length as u32);
                    let ptr_ty = Item::Pointer(StorageClass::Function, Box::new(arr_ty)).id(self);
                    let arr_id = self.declare_function_variable(ptr_ty);
                    self.debug_var_name(arr_id, variable);
                    let arr = Array {
                        id: arr_id,
                        item: item.clone(),
                        len: (length * unroll_factor) as u32,
                        var: variable,
                        alignment: None,
                    };
                    self.state.local_arrays.insert(id, arr);
                    arr_id
                };
                Variable::LocalArray(id, item, length as u32)
            }
            ir::VariableKind::Matrix { id, mat } => {
                let elem = self.compile_type(ir::Type::new(mat.storage)).elem();
                if self.state.matrices.contains_key(&id) {
                    Variable::CoopMatrix(id, elem)
                } else {
                    let matrix = self.init_coop_matrix(mat, variable);
                    self.state.matrices.insert(id, matrix);
                    Variable::CoopMatrix(id, elem)
                }
            }
            ir::VariableKind::Pipeline { .. } => panic!("Pipeline not supported."),
            ir::VariableKind::BarrierToken { .. } => {
                panic!("Barrier not supported.")
            }
            ir::VariableKind::TensorMapInput(_) => panic!("Tensor map not supported."),
            ir::VariableKind::TensorMapOutput(_) => panic!("Tensor map not supported."),
        }
    }

    pub fn read(&mut self, variable: &Variable) -> Word {
        match variable {
            Variable::Slice { ptr, .. } => self.read(ptr),
            Variable::Shared(id, item)
                if self.compilation_options.vulkan.supports_explicit_smem =>
            {
                let ty = item.id(self);
                let ptr_ty =
                    Item::Pointer(StorageClass::Workgroup, Box::new(item.clone())).id(self);
                let index = vec![self.const_u32(0)];
                let access = self.access_chain(ptr_ty, None, *id, index).unwrap();
                self.load(ty, None, access, None, []).unwrap()
            }
            Variable::Local { id, item } | Variable::Shared(id, item) => {
                let ty = item.id(self);
                self.load(ty, None, *id, None, []).unwrap()
            }
            ssa => ssa.id(self),
        }
    }

    pub fn read_as(&mut self, variable: &Variable, item: &Item) -> Word {
        if let Some(as_const) = variable.as_const() {
            self.static_cast(as_const, &variable.elem(), item).0
        } else {
            let id = self.read(variable);
            variable.item().cast_to(self, None, id, item)
        }
    }

    pub fn index(
        &mut self,
        variable: &Variable,
        index: &Variable,
        unchecked: bool,
    ) -> IndexedVariable {
        let access_chain = if unchecked {
            Builder::in_bounds_access_chain
        } else {
            Builder::access_chain
        };
        let index_id = self.read(index);
        match variable {
            Variable::GlobalInputArray(id, item, _) | Variable::GlobalOutputArray(id, item, _) => {
                let ptr_ty =
                    Item::Pointer(StorageClass::StorageBuffer, Box::new(item.clone())).id(self);
                let zero = self.const_u32(0);
                let id = access_chain(self, ptr_ty, None, *id, vec![zero, index_id]).unwrap();

                IndexedVariable::Pointer(id, item.clone())
            }
            Variable::Local {
                id,
                item: Item::Vector(elem, _),
            } => {
                let ptr_ty =
                    Item::Pointer(StorageClass::Function, Box::new(Item::Scalar(*elem))).id(self);
                let id = access_chain(self, ptr_ty, None, *id, vec![index_id]).unwrap();

                IndexedVariable::Pointer(id, Item::Scalar(*elem))
            }
            Variable::Shared(id, Item::Vector(elem, _)) => {
                let ptr_ty =
                    Item::Pointer(StorageClass::Workgroup, Box::new(Item::Scalar(*elem))).id(self);

                let mut index = vec![index_id];
                if self.compilation_options.vulkan.supports_explicit_smem {
                    index.insert(0, self.const_u32(0));
                }

                let id = access_chain(self, ptr_ty, None, *id, index).unwrap();

                IndexedVariable::Pointer(id, Item::Scalar(*elem))
            }
            Variable::LocalBinding {
                id,
                item: Item::Vector(elem, vec),
                variable,
            } if index.as_const().is_some() => IndexedVariable::Composite(
                self.get_binding(*id, variable),
                index.as_const().unwrap().as_u64() as u32,
                Item::Vector(*elem, *vec),
            ),
            Variable::LocalBinding {
                id,
                item: Item::Vector(elem, vec),
                variable,
            } => IndexedVariable::DynamicComposite(
                self.get_binding(*id, variable),
                index_id,
                Item::Vector(*elem, *vec),
            ),
            Variable::Versioned {
                id,
                item: Item::Vector(elem, vec),
                variable,
            } if index.as_const().is_some() => IndexedVariable::Composite(
                self.get_versioned(*id, variable),
                index.as_const().unwrap().as_u64() as u32,
                Item::Vector(*elem, *vec),
            ),
            Variable::Versioned {
                id,
                item: Item::Vector(elem, vec),
                variable,
            } => IndexedVariable::DynamicComposite(
                self.get_versioned(*id, variable),
                index_id,
                Item::Vector(*elem, *vec),
            ),
            Variable::Shared(id, item)
                if self.compilation_options.vulkan.supports_explicit_smem =>
            {
                let ptr_ty =
                    Item::Pointer(StorageClass::Workgroup, Box::new(item.clone())).id(self);
                let index = vec![self.const_u32(0)];
                let id = access_chain(self, ptr_ty, None, *id, index).unwrap();
                IndexedVariable::Pointer(id, item.clone())
            }
            Variable::Local { .. }
            | Variable::Shared(..)
            | Variable::LocalBinding { .. }
            | Variable::Versioned { .. } => IndexedVariable::Scalar(variable.clone()),
            Variable::Constant(_, val, item) => {
                let scalar_item = Item::Scalar(item.elem());
                let (id, val) = self.static_cast(*val, &item.elem(), &scalar_item);
                IndexedVariable::Scalar(Variable::Constant(id, val, scalar_item))
            }
            Variable::Slice { ptr, offset, .. } => {
                let item = Item::Scalar(Elem::Int(32, false));
                let int = item.id(self);
                let index = match index.as_const() {
                    Some(ConstVal::Bit32(0)) => *offset,
                    _ => self.i_add(int, None, *offset, index_id).unwrap(),
                };
                self.index(ptr, &Variable::Raw(index, item), unchecked)
            }
            Variable::SharedArray(id, item, _) => {
                let ptr_ty =
                    Item::Pointer(StorageClass::Workgroup, Box::new(item.clone())).id(self);
                let mut index = vec![index_id];
                if self.compilation_options.vulkan.supports_explicit_smem {
                    index.insert(0, self.const_u32(0));
                }
                let id = access_chain(self, ptr_ty, None, *id, index).unwrap();
                IndexedVariable::Pointer(id, item.clone())
            }
            Variable::ConstantArray(id, item, _) | Variable::LocalArray(id, item, _) => {
                let ptr_ty = Item::Pointer(StorageClass::Function, Box::new(item.clone())).id(self);
                let id = access_chain(self, ptr_ty, None, *id, vec![index_id]).unwrap();
                IndexedVariable::Pointer(id, item.clone())
            }
            var => unimplemented!("Can't index into {var:?}"),
        }
    }

    pub fn read_indexed(&mut self, out: &Variable, variable: &Variable, index: &Variable) -> Word {
        let always_in_bounds = is_always_in_bounds(variable, index);
        let indexed = self.index(variable, index, always_in_bounds);

        let read = |b: &mut Self| match indexed {
            IndexedVariable::Pointer(ptr, item) => {
                let ty = item.id(b);
                let out_id = b.write_id(out);
                b.load(ty, Some(out_id), ptr, None, vec![]).unwrap()
            }
            IndexedVariable::Composite(var, index, item) => {
                let elem = item.elem();
                let ty = elem.id(b);
                let out_id = b.write_id(out);
                b.composite_extract(ty, Some(out_id), var, vec![index])
                    .unwrap()
            }
            IndexedVariable::DynamicComposite(var, index, item) => {
                let elem = item.elem();
                let ty = elem.id(b);
                let out_id = b.write_id(out);
                b.vector_extract_dynamic(ty, Some(out_id), var, index)
                    .unwrap()
            }
            IndexedVariable::Scalar(var) => {
                let ty = out.item().id(b);
                let input = b.read(&var);
                let out_id = b.write_id(out);
                b.copy_object(ty, Some(out_id), input).unwrap();
                b.write(out, out_id);
                out_id
            }
        };

        read(self)
    }

    pub fn read_indexed_unchecked(
        &mut self,
        out: &Variable,
        variable: &Variable,
        index: &Variable,
    ) -> Word {
        let indexed = self.index(variable, index, true);

        match indexed {
            IndexedVariable::Pointer(ptr, item) => {
                let ty = item.id(self);
                let out_id = self.write_id(out);
                self.load(ty, Some(out_id), ptr, None, vec![]).unwrap()
            }
            IndexedVariable::Composite(var, index, item) => {
                let elem = item.elem();
                let ty = elem.id(self);
                let out_id = self.write_id(out);
                self.composite_extract(ty, Some(out_id), var, vec![index])
                    .unwrap()
            }
            IndexedVariable::DynamicComposite(var, index, item) => {
                let elem = item.elem();
                let ty = elem.id(self);
                let out_id = self.write_id(out);
                self.vector_extract_dynamic(ty, Some(out_id), var, index)
                    .unwrap()
            }
            IndexedVariable::Scalar(var) => {
                let ty = out.item().id(self);
                let input = self.read(&var);
                let out_id = self.write_id(out);
                self.copy_object(ty, Some(out_id), input).unwrap();
                self.write(out, out_id);
                out_id
            }
        }
    }

    pub fn index_ptr(&mut self, var: &Variable, index: &Variable) -> Word {
        let always_in_bounds = is_always_in_bounds(var, index);
        match self.index(var, index, always_in_bounds) {
            IndexedVariable::Pointer(ptr, _) => ptr,
            other => unreachable!("{other:?}"),
        }
    }

    pub fn write_id(&mut self, variable: &Variable) -> Word {
        match variable {
            Variable::LocalBinding { id, variable, .. } => self.get_binding(*id, variable),
            Variable::Versioned { id, variable, .. } => self.get_versioned(*id, variable),
            Variable::Local { .. } => self.id(),
            Variable::Shared(..) => self.id(),
            Variable::GlobalScalar(id, _) => *id,
            Variable::Raw(id, _) => *id,
            Variable::Constant(_, _, _) => panic!("Can't write to constant scalar"),
            Variable::GlobalInputArray(_, _, _)
            | Variable::GlobalOutputArray(_, _, _)
            | Variable::Slice { .. }
            | Variable::SharedArray(_, _, _)
            | Variable::ConstantArray(_, _, _)
            | Variable::LocalArray(_, _, _) => panic!("Can't write to unindexed array"),
            global => panic!("Can't write to builtin {global:?}"),
        }
    }

    pub fn write(&mut self, variable: &Variable, value: Word) {
        match variable {
            Variable::Shared(id, item)
                if self.compilation_options.vulkan.supports_explicit_smem =>
            {
                let ptr_ty =
                    Item::Pointer(StorageClass::Workgroup, Box::new(item.clone())).id(self);
                let index = vec![self.const_u32(0)];
                let access = self.access_chain(ptr_ty, None, *id, index).unwrap();
                self.store(access, value, None, []).unwrap()
            }
            Variable::Local { id, .. } | Variable::Shared(id, _) => {
                self.store(*id, value, None, []).unwrap()
            }

            Variable::Slice { ptr, .. } => self.write(ptr, value),
            _ => {}
        }
    }

    pub fn write_indexed(&mut self, out: &Variable, index: &Variable, value: Word) {
        let always_in_bounds = is_always_in_bounds(out, index);
        let variable = self.index(out, index, always_in_bounds);

        let write = |b: &mut Self| match variable {
            IndexedVariable::Pointer(ptr, _) => b.store(ptr, value, None, vec![]).unwrap(),
            IndexedVariable::Composite(var, index, item) => {
                let ty = item.id(b);
                let id = b
                    .composite_insert(ty, None, value, var, vec![index])
                    .unwrap();
                b.write(out, id);
            }
            IndexedVariable::DynamicComposite(var, index, item) => {
                let ty = item.id(b);
                let id = b
                    .vector_insert_dynamic(ty, None, value, var, index)
                    .unwrap();
                b.write(out, id);
            }
            IndexedVariable::Scalar(var) => b.write(&var, value),
        };

        write(self)
    }

    pub fn write_indexed_unchecked(&mut self, out: &Variable, index: &Variable, value: Word) {
        let variable = self.index(out, index, true);

        match variable {
            IndexedVariable::Pointer(ptr, _) => self.store(ptr, value, None, vec![]).unwrap(),
            IndexedVariable::Composite(var, index, item) => {
                let ty = item.id(self);
                let out_id = self
                    .composite_insert(ty, None, value, var, vec![index])
                    .unwrap();
                self.write(out, out_id);
            }
            IndexedVariable::DynamicComposite(var, index, item) => {
                let ty = item.id(self);
                let out_id = self
                    .vector_insert_dynamic(ty, None, value, var, index)
                    .unwrap();
                self.write(out, out_id);
            }
            IndexedVariable::Scalar(var) => self.write(&var, value),
        }
    }
}

fn is_always_in_bounds(var: &Variable, index: &Variable) -> bool {
    let len = match var {
        Variable::SharedArray(_, _, len)
        | Variable::ConstantArray(_, _, len)
        | Variable::LocalArray(_, _, len)
        | Variable::Slice {
            const_len: Some(len),
            ..
        } => *len,
        _ => return false,
    };

    let const_index = match index {
        Variable::Constant(_, value, _) => value.as_u64(),
        _ => return false,
    };

    const_index < len as u64
}
