wgsl-rs 0.1.0-beta.4

//! Runtime wgpu linkage analysis for WGSL modules.
//!
//! Given an [`ir::Module`] (produced by a `#[wgsl]` source's
//! `WGSL_SOURCE.ir_constructor`, or by the macro-emitted
//! `instantiate::<...>()` function for templates), this module extracts
//! the binding and entry-point information needed to build wgpu
//! pipelines.
//!
//! # Why at runtime?
//!
//! The pre-[#120](https://github.com/schell/wgsl-rs/issues/120)
//! implementation generated all of this at proc-macro expansion time by walking
//! the parse tree in `wgsl-rs-macros/src/linkage.rs`. That had two drawbacks:
//!
//! 1. Template modules (`#[wgsl]` modules with type parameters) couldn't get
//!    any linkage — the generated WGSL was a template with `__TP{name}__`
//!    placeholders, so a `wgpu::ShaderModule` couldn't be built from it.
//! 2. The proc-macro walked a parse tree that duplicated information already
//!    present in `wgsl-rs-ir`.
//!
//! Walking the runtime IR unifies the two cases. After the macro-emitted
//! `instantiate::<...>()` function produces a concrete `ir::Module`, the
//! same analyzer used for non-template sources works on the result.
//!
//! # Example
//!
//! ```ignore
//! use wgsl_rs::linkage::wgpu;
//!
//! let linkage = wgpu::analyze_wgsl_module(&hello_triangle::WGSL_SOURCE).unwrap();
//! let shader_module = linkage.shader_module(&device);
//! let bg_layout = linkage.bind_group(0).unwrap().layout(&device);
//! let bg = linkage.bind_group(0).unwrap().create(
//!     &device, &bg_layout, &[frame_uniform.as_entire_binding()],
//! );
//! let vtx = linkage.vertex_entries.iter().find(|e| e.name == "vtx_main").unwrap()
//!     .vertex_state(&shader_module);
//! ```

use std::{
    borrow::Cow,
    collections::{BTreeMap, HashMap},
};

use snafu::prelude::*;
use wgsl_rs_ir as ir;

use crate::Source;

/// Errors thrown by the WGPU linkage module.
#[derive(Snafu, Debug)]
pub enum Error {
    #[snafu(display(
        "Unable to resolve template '{template_name}' for type args {:?}; available templates: \
         {:?}",
        mangled_type_args,
        available
    ))]
    TemplateResolution {
        template_name: String,
        mangled_type_args: Vec<String>,
        available: Vec<String>,
    },

    #[snafu(display(
        "{module}::bind_group_{group}: missing binding resource for declared name {name:?}"
    ))]
    MissingBinding {
        module: &'static str,
        group: u32,
        name: String,
    },

    #[snafu(display(
        "{module}::bind_group_{group}: resource provided for unknown binding name {name:?} \
         (declared names: {declared:?})"
    ))]
    UnknownBindingName {
        module: &'static str,
        group: u32,
        name: String,
        declared: Vec<String>,
    },

    #[snafu(display(
        "{module}::bind_group_{group}: binding name {name:?} appears more than once in resources"
    ))]
    DuplicateBindingName {
        module: &'static str,
        group: u32,
        name: String,
    },

    #[snafu(display(
        "{module}: no bind group @group({group}) declared; available groups: {available:?}"
    ))]
    NoSuchBindGroup {
        module: &'static str,
        group: u32,
        available: Vec<u32>,
    },
}

impl From<crate::SourceError<'_>> for Error {
    fn from(err: crate::SourceError<'_>) -> Self {
        match err {
            crate::SourceError::TemplateWgsl {
                uninstantiated_source,
            } => Error::TemplateResolution {
                template_name: format!("<source {}>", uninstantiated_source.name),
                mangled_type_args: uninstantiated_source
                    .module_type_params
                    .iter()
                    .map(|s| (*s).to_string())
                    .collect(),
                available: Vec::new(),
            },
            crate::SourceError::TemplateNotFound {
                template_name,
                mangled_type_args,
                available_templates,
            } => Error::TemplateResolution {
                template_name,
                mangled_type_args,
                available: available_templates,
            },
            crate::SourceError::AmbiguousTemplate {
                template_name,
                mangled_type_args,
                matching_source_names,
                available_templates: _,
            } => Error::TemplateResolution {
                template_name,
                mangled_type_args,
                available: matching_source_names,
            },
        }
    }
}

/// All the wgpu linkage information extracted from a WGSL module.
///
/// Produced by [`analyze_ir_module`] or [`analyze_wgsl_module`]. Owns the
/// string labels that the various `wgpu::*Descriptor` types borrow, so
/// hold this struct for the lifetime of any borrowed descriptor.
#[derive(Clone, Debug)]
pub struct WgpuLinkage {
    /// The module's name (used as a default label for descriptors).
    pub module_label: &'static str,

    /// The concrete, fully-assembled IR the linkage was built from. This
    /// is the IR rendered by [`Self::wgsl_source`] to produce the shader
    /// source's WGSL text. For sources built from [`crate::Source`] it
    /// is the result of flattening imports and template instantiations;
    /// for modules built from [`ir::Module`] (e.g. an instantiated
    /// template) it is that module directly. Always concrete: no
    /// `Type::TypeParam` nodes.
    ir: ir::Module,

    /// Bind groups keyed by `@group(N)` index. Each group is sorted by
    /// binding number.
    pub bind_groups: BTreeMap<u32, BindGroupInfo>,

    /// Lazily-built `wgpu::BindGroupLayout` per `@group(N)` index. Empty
    /// until a cache-populating method is called. Populating one does
    /// not populate the others; calling
    /// [`WgpuLinkage::pipeline_layout`] populates all groups at once.
    bind_group_layouts: HashMap<u32, wgpu::BindGroupLayout>,

    /// Lazily-built `wgpu::PipelineLayout` referencing all bind group
    /// layouts. `None` until [`WgpuLinkage::pipeline_layout`] is called.
    pipeline_layout: Option<wgpu::PipelineLayout>,

    /// Vertex shader entry points declared in the module.
    pub vertex_entries: Vec<EntryPointInfo>,

    /// Fragment shader entry points declared in the module.
    pub fragment_entries: Vec<EntryPointInfo>,

    /// Compute shader entry points declared in the module.
    pub compute_entries: Vec<ComputeEntryInfo>,

    /// Buffer descriptors for `uniform!` and `storage!` declarations, in
    /// declaration order. Indexed by binding name via [`Self::buffer`].
    pub buffers: Vec<BufferDescriptorInfo>,
}

impl core::fmt::Display for WgpuLinkage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let WgpuLinkage {
            module_label,
            ir,
            bind_groups,
            bind_group_layouts,
            pipeline_layout,
            vertex_entries,
            fragment_entries,
            compute_entries,
            buffers,
        } = self;

        writeln!(
            f,
            "linkage for '{module_label}' ({} IR items)",
            ir.items.len()
        )?;

        writeln!(f, "  entry points:")?;
        if vertex_entries.is_empty() && fragment_entries.is_empty() && compute_entries.is_empty() {
            writeln!(f, "    (none)")?;
        }
        for e in vertex_entries {
            writeln!(f, "    vertex: {}", e.name)?;
        }
        for e in fragment_entries {
            writeln!(f, "    fragment: {}", e.name)?;
        }
        for e in compute_entries {
            let (x, y, z) = e.workgroup_size;
            writeln!(f, "    compute: {} @workgroup_size({x}, {y}, {z})", e.name)?;
        }

        writeln!(f, "  bind groups:")?;
        if bind_groups.is_empty() {
            writeln!(f, "    (none)")?;
        }
        for (group, info) in bind_groups {
            writeln!(f, "    @group({group}) {}", info.label)?;
            for meta in &info.bindings {
                writeln!(
                    f,
                    "      @binding({}) {} ({})",
                    meta.binding,
                    meta.name,
                    binding_kind_label(&meta.kind)
                )?;
            }
        }

        writeln!(f, "  buffers:")?;
        if buffers.is_empty() {
            writeln!(f, "    (none)")?;
        }
        for b in buffers {
            // `size` is 0 for runtime-sized storage buffers; the buffer's
            // length is then whatever the caller chooses to allocate.
            let size = if b.size == 0 {
                "runtime-sized".to_string()
            } else {
                format!("{} bytes", b.size)
            };
            writeln!(
                f,
                "      @group({}) @binding({}): {} ({}, {size}, {ty})",
                b.group,
                b.binding,
                b.binding_name,
                buffer_kind_label(&b.kind),
                ty = b.ty
            )?;
        }

        // The layouts are built lazily; show how much of the cache is
        // populated. Calling `pipeline_layout` fills every group at once.
        write!(
            f,
            "  cache: {} of {} bind group layouts built, pipeline layout {}",
            bind_group_layouts.len(),
            bind_groups.len(),
            if pipeline_layout.is_some() {
                "built"
            } else {
                "not built"
            }
        )?;

        Ok(())
    }
}

/// A human-readable label for a [`BindingKind`].
fn binding_kind_label(kind: &BindingKind) -> &'static str {
    match kind {
        BindingKind::Uniform => "uniform",
        BindingKind::Storage { read_only: true } => "storage, read-only",
        BindingKind::Storage { read_only: false } => "storage, read-write",
        BindingKind::Sampler { comparison: true } => "comparison sampler",
        BindingKind::Sampler { comparison: false } => "sampler",
        BindingKind::Texture => "texture",
        BindingKind::DepthTexture => "depth texture",
        BindingKind::StorageTexture { read_write: true } => "storage texture, read-write",
        BindingKind::StorageTexture { read_write: false } => "storage texture",
    }
}

/// A human-readable label for a [`BufferKind`].
fn buffer_kind_label(kind: &BufferKind) -> &'static str {
    match kind {
        BufferKind::Uniform => "uniform",
        BufferKind::Storage { read_only: true } => "storage, read-only",
        BufferKind::Storage { read_only: false } => "storage, read-write",
    }
}

impl WgpuLinkage {
    /// Returns the bind group for the given `@group(N)` index, if any.
    pub fn bind_group(&self, group: u32) -> Option<&BindGroupInfo> {
        self.bind_groups.get(&group)
    }

    /// Returns the buffer descriptor for a `uniform!` or `storage!`
    /// binding by its declared name, if any.
    pub fn buffer(&self, name: &str) -> Option<&BufferDescriptorInfo> {
        self.buffers.iter().find(|b| b.binding_name == name)
    }

    /// Returns the buffer descriptor for a `uniform!` or `storage!`
    /// binding by its declared name, if any - removing it from the linkage.
    ///
    /// Use this to ensure that you provide data for all buffers in your
    /// program.
    pub fn take_buffer(&mut self, name: &str) -> Option<BufferDescriptorInfo> {
        let mut info = None;
        self.buffers.retain(|bi| {
            if bi.binding_name == name {
                info = Some(bi.clone());
                false
            } else {
                true
            }
        });
        info
    }

    /// Returns the WGSL source for this linkage's source, rendered from
    /// the concrete IR the linkage was built from.
    ///
    /// The linkage's IR is concrete (no `Type::TypeParam`s), so the
    /// rendered text is free of `__TP{name}__` placeholders. For a
    /// linkage built from a template `Source` via
    /// [`analyze_ir_module`], this is the instantiated source.
    pub fn wgsl_source(&self) -> String {
        ir::render_module(&self.ir)
    }

    /// Builds a `wgpu::ShaderModuleDescriptor` for this module using
    /// the concrete IR the linkage was built from. The source string
    /// is owned by the descriptor, so no borrow of `self` is required
    /// to build a wgpu pipeline.
    pub fn shader_module_descriptor(&self) -> wgpu::ShaderModuleDescriptor<'static> {
        wgpu::ShaderModuleDescriptor {
            label: Some(self.module_label),
            source: wgpu::ShaderSource::Wgsl(std::borrow::Cow::Owned(self.wgsl_source())),
        }
    }

    /// Creates a `wgpu::ShaderModule` from the concrete IR the linkage
    /// was built from.
    pub fn shader_module(&self, device: &wgpu::Device) -> wgpu::ShaderModule {
        device.create_shader_module(self.shader_module_descriptor())
    }

    /// Returns the first vertex entry matching `name`, or `None`.
    pub fn vertex_entry(&self, name: &str) -> Option<&EntryPointInfo> {
        self.vertex_entries.iter().find(|e| e.name == name)
    }

    /// Returns the first fragment entry matching `name`, or `None`.
    pub fn fragment_entry(&self, name: &str) -> Option<&EntryPointInfo> {
        self.fragment_entries.iter().find(|e| e.name == name)
    }

    /// Returns the first compute entry matching `name`, or `None`.
    pub fn compute_entry(&self, name: &str) -> Option<&ComputeEntryInfo> {
        self.compute_entries.iter().find(|e| e.name == name)
    }

    /// Returns the cached `wgpu::BindGroupLayout` for the given
    /// `@group(N)` index, building and caching it on first access.
    /// Returns `None` if the linkage has no such bind group.
    ///
    /// The returned layout is an owned clone of the cached value; the
    /// cache is not invalidated. Multiple calls are cheap after the
    /// first.
    pub fn bind_group_layout(
        &mut self,
        group: u32,
        device: &wgpu::Device,
    ) -> Option<wgpu::BindGroupLayout> {
        if let Some(layout) = self.bind_group_layouts.get(&group) {
            return Some(layout.clone());
        }
        let info = self.bind_groups.get(&group)?;
        let layout = info.layout(device);
        self.bind_group_layouts.insert(group, layout.clone());
        Some(layout)
    }

    /// Creates a bind group for `@group(N)` by matching each binding's
    /// declared name to a resource. The order of `resources` does not
    /// have to match the binding declaration order; the produced
    /// `BindGroupEntry` list is sorted by binding number, matching the
    /// cached layout.
    ///
    /// Looks up (and builds+caches on first access) the bind group
    /// layout for `group`, then delegates to
    /// [`BindGroupInfo::create_named`].
    ///
    /// # Errors
    ///
    /// - [`Error::NoSuchBindGroup`] if no bind group with the given index is
    ///   declared in this module.
    /// - [`Error::DuplicateBindingName`] if a name appears twice in
    ///   `resources`.
    /// - [`Error::MissingBinding`] if a declared binding has no resource in
    ///   `resources`.
    /// - [`Error::UnknownBindingName`] if `resources` contains a name that
    ///   isn't declared in this bind group.
    pub fn create_bind_group_named<'a>(
        &mut self,
        group: u32,
        device: &wgpu::Device,
        resources: &[(&str, wgpu::BindingResource<'a>)],
    ) -> Result<wgpu::BindGroup, Error> {
        // Build (or reuse cached) the layout for this group first.
        // The cache lookup is cheap and avoids a redundant layout
        // rebuild if the caller has already populated it.
        if !self.bind_group_layouts.contains_key(&group) {
            let Some(info) = self.bind_groups.get(&group) else {
                let mut available: Vec<u32> = self.bind_groups.keys().copied().collect();
                available.sort_unstable();
                return NoSuchBindGroupSnafu {
                    module: self.module_label,
                    group,
                    available,
                }
                .fail();
            };
            self.bind_group_layouts.insert(group, info.layout(device));
        }
        let info = self
            .bind_groups
            .get(&group)
            .expect("just inserted or found");
        let layout = self
            .bind_group_layouts
            .get(&group)
            .expect("just inserted or found");
        info.create_named(self.module_label, device, layout, resources)
    }

    /// Returns the cached `wgpu::PipelineLayout` for this module,
    /// building and caching it (and all bind group layouts it
    /// references) on first access. Subsequent calls return a cheap
    /// Arc-backed clone.
    ///
    /// For modules with no bind groups (e.g. depth-only pipelines), the
    /// cached layout references an empty `bind_group_layouts` slice.
    ///
    /// **Infallible.** The bind groups are already analyzed; this only
    /// drives `wgpu::Device` constructors.
    pub fn pipeline_layout(
        &mut self,
        device: &wgpu::Device,
        label: Option<&str>,
    ) -> wgpu::PipelineLayout {
        if let Some(layout) = &self.pipeline_layout {
            return layout.clone();
        }
        // Build (and cache) every bind group layout we don't already have.
        for (&group, info) in &self.bind_groups {
            self.bind_group_layouts
                .entry(group)
                .or_insert_with(|| info.layout(device));
        }
        // `wgpu::PipelineLayoutDescriptor::bind_group_layouts` is indexed
        // by `@group(N)`: position N must hold group N's layout. Sort the
        // cached layouts by group index so the slice lines up. Gaps (e.g.
        // a module declaring only group 0 and group 2) are filled with
        // `None`, which wgpu accepts as "no bind group at this index".
        let max_group = self.bind_group_layouts.keys().copied().max().unwrap_or(0);
        let mut bg_layouts: Vec<Option<&wgpu::BindGroupLayout>> = (0..=max_group)
            .map(|g| self.bind_group_layouts.get(&g))
            .collect();
        // Trim trailing `None`s so we don't pass unnecessary slots.
        while bg_layouts.last() == Some(&None) {
            bg_layouts.pop();
        }
        let layout = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label,
            bind_group_layouts: &bg_layouts,
            immediate_size: 0,
        });
        self.pipeline_layout = Some(layout.clone());
        layout
    }
}

// ===== ir::Module extension trait =====

/// Extension trait that adds `wgsl-rs` operations to [`wgsl_rs_ir::Module`].
///
/// [`wgsl_rs_ir::Module`] lives in a separate crate and so cannot have
/// inherent methods that reference `wgsl-rs` types. The methods that
/// need `WgpuLinkage` or naga validation live here instead. The trait
/// is re-exported from `wgsl_rs::*` so `m.generate_linkage()` works
/// after `use wgsl_rs;` without an explicit import.
pub trait IrModuleExt {
    /// Builds a [`WgpuLinkage`] from this concrete IR module.
    ///
    /// Equivalent to [`analyze_ir_module`](self).
    fn generate_linkage(&self) -> WgpuLinkage;

    /// Validates the WGSL source rendered from this IR module using naga.
    ///
    /// Equivalent to the free function [`crate::validate_wgsl_source`]
    /// called on the output of [`wgsl_rs_ir::Module::wgsl_source`].
    ///
    /// Only available with the `validation` feature enabled.
    #[cfg(feature = "validation")]
    fn validate(&self) -> Result<(), String>;
}

impl IrModuleExt for wgsl_rs_ir::Module {
    fn generate_linkage(&self) -> WgpuLinkage {
        analyze_ir_module(self.clone())
    }

    #[cfg(feature = "validation")]
    fn validate(&self) -> Result<(), String> {
        crate::validate_wgsl_source(&self.wgsl_source())
    }
}

// ===== Bind groups =====

/// A single bind group, including its layout entries and metadata for
/// resource construction.
#[derive(Clone, Debug)]
pub struct BindGroupInfo {
    /// The `@group(N)` index.
    pub group: u32,
    /// The bind group layout entries, sorted by binding number.
    pub entries: Vec<wgpu::BindGroupLayoutEntry>,
    /// Per-binding metadata: the declared name, the binding number, and
    /// the kind of resource the binding expects.
    pub bindings: Vec<BindingMeta>,
    /// A default label for descriptors (`"<module>::bind_group_<N>"`).
    pub label: String,
}

/// Metadata for a single binding inside a bind group.
#[derive(Clone, Debug)]
pub struct BindingMeta {
    /// The declared Rust name (e.g. `"FRAME"`, `"DIFFUSE_TEX"`).
    pub name: String,
    /// The `@binding(N)` index.
    pub binding: u32,
    /// What kind of resource the binding expects.
    pub kind: BindingKind,
}

/// The kind of resource a binding expects.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BindingKind {
    /// A uniform buffer.
    Uniform,
    /// A storage buffer. `read_only` mirrors the WGSL access mode.
    Storage { read_only: bool },
    /// A (filtering) sampler.
    Sampler { comparison: bool },
    /// A sampled texture.
    Texture,
    /// A depth texture.
    DepthTexture,
    /// A storage texture. The `read_write` flag distinguishes
    /// `read_write` access from `read`-only or `write`-only.
    StorageTexture { read_write: bool },
}

impl BindGroupInfo {
    /// Builds a `wgpu::BindGroupLayoutDescriptor` borrowing this struct.
    /// The returned descriptor borrows from `self`, so `self` must
    /// outlive the descriptor and any wgpu object built from it.
    pub fn layout_descriptor<'a>(&'a self) -> wgpu::BindGroupLayoutDescriptor<'a> {
        wgpu::BindGroupLayoutDescriptor {
            label: Some(&self.label),
            entries: &self.entries,
        }
    }

    /// Creates the bind group layout on the given device.
    pub fn layout(&self, device: &wgpu::Device) -> wgpu::BindGroupLayout {
        device.create_bind_group_layout(&self.layout_descriptor())
    }

    /// Creates a bind group with one entry per binding, in `self.entries`
    /// order — which is sorted by binding number (see
    /// [`BindGroupInfo::entries`]). The caller must pass exactly one
    /// `BindingResource` per entry, in that same (binding-number) order.
    pub fn create<'a>(
        &self,
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        resources: &[wgpu::BindingResource<'a>],
    ) -> wgpu::BindGroup {
        assert_eq!(
            resources.len(),
            self.entries.len(),
            "{}: expected {} binding resource(s), got {}",
            self.label,
            self.entries.len(),
            resources.len(),
        );
        let entries: Vec<wgpu::BindGroupEntry> = self
            .entries
            .iter()
            .zip(resources.iter())
            .map(|(layout_entry, resource)| wgpu::BindGroupEntry {
                binding: layout_entry.binding,
                resource: resource.clone(),
            })
            .collect();
        device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some(&self.label),
            layout,
            entries: &entries,
        })
    }

    /// Creates a bind group by matching each binding's declared name
    /// to a resource. The order of `resources` does not have to match
    /// the binding declaration order; the produced `BindGroupEntry`
    /// list is sorted by binding number, matching the layout.
    ///
    /// `module_label` is the linkage's `module_label`, used to build
    /// helpful error messages. It is not borrowed into the returned
    /// bind group.
    ///
    /// # Errors
    ///
    /// - [`Error::DuplicateBindingName`] if a name appears twice in
    ///   `resources`.
    /// - [`Error::MissingBinding`] if a declared name has no resource in
    ///   `resources`.
    /// - [`Error::UnknownBindingName`] if `resources` contains a name that
    ///   isn't declared in this bind group.
    pub fn create_named<'a>(
        &self,
        module_label: &'static str,
        device: &wgpu::Device,
        layout: &wgpu::BindGroupLayout,
        resources: &[(&str, wgpu::BindingResource<'a>)],
    ) -> Result<wgpu::BindGroup, Error> {
        use std::collections::HashMap;

        // Build a name → resource map. Detect duplicates as we go.
        let mut by_name: HashMap<&str, wgpu::BindingResource<'a>> =
            HashMap::with_capacity(resources.len());
        for &(name, ref resource) in resources {
            if by_name.insert(name, resource.clone()).is_some() {
                return DuplicateBindingNameSnafu {
                    module: module_label,
                    group: self.group,
                    name: name.to_string(),
                }
                .fail();
            }
        }

        // Every declared binding must be present.
        for binding in &self.bindings {
            if !by_name.contains_key(binding.name.as_str()) {
                return MissingBindingSnafu {
                    module: module_label,
                    group: self.group,
                    name: binding.name.clone(),
                }
                .fail();
            }
        }

        // Every supplied name must be declared. We check this by
        // building the entries from the declared list (not from the
        // supplied slice) and tracking which supplied names we hit.
        let mut hit: HashMap<&str, ()> = HashMap::with_capacity(by_name.len());
        let mut entries: Vec<wgpu::BindGroupEntry> = Vec::with_capacity(self.bindings.len());
        for (layout_entry, binding) in self.entries.iter().zip(self.bindings.iter()) {
            let resource = by_name
                .get(binding.name.as_str())
                .expect("checked above")
                .clone();
            hit.insert(binding.name.as_str(), ());
            entries.push(wgpu::BindGroupEntry {
                binding: layout_entry.binding,
                resource,
            });
        }
        for name in by_name.keys() {
            if !hit.contains_key(name) {
                let declared: Vec<String> = self.bindings.iter().map(|b| b.name.clone()).collect();
                return UnknownBindingNameSnafu {
                    module: module_label,
                    group: self.group,
                    name: (*name).to_string(),
                    declared,
                }
                .fail();
            }
        }

        Ok(device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: Some(&self.label),
            layout,
            entries: &entries,
        }))
    }
}

// ===== Entry points =====

/// A vertex or fragment shader entry point.
#[derive(Clone, Debug)]
pub struct EntryPointInfo {
    /// The entry point name. For non-monomorphized functions this is a
    /// `Cow::Borrowed` of a `stringify!`-emitted `'static` literal;
    /// for monomorphized instances (e.g. `id` → `id_f32`) it is a
    /// `Cow::Owned` holding the runtime-computed name.
    pub name: Cow<'static, str>,
}

impl EntryPointInfo {
    /// Builds a `wgpu::VertexState` borrowing this entry's name.
    pub fn vertex_state<'a>(&'a self, module: &'a wgpu::ShaderModule) -> wgpu::VertexState<'a> {
        wgpu::VertexState {
            module,
            entry_point: Some(&self.name),
            buffers: &[],
            compilation_options: Default::default(),
        }
    }

    /// Builds a `wgpu::FragmentState` borrowing this entry's name.
    pub fn fragment_state<'a>(
        &'a self,
        module: &'a wgpu::ShaderModule,
        targets: &'a [Option<wgpu::ColorTargetState>],
    ) -> wgpu::FragmentState<'a> {
        wgpu::FragmentState {
            module,
            entry_point: Some(&self.name),
            targets,
            compilation_options: Default::default(),
        }
    }
}

/// A compute shader entry point.
#[derive(Clone, Debug)]
pub struct ComputeEntryInfo {
    /// The entry point name (see [`EntryPointInfo::name`] for the
    /// `Cow<'static, str>` rationale).
    pub name: Cow<'static, str>,
    /// The `@workgroup_size(X, Y, Z)` dimensions.
    pub workgroup_size: (u32, u32, u32),
    /// A default label for descriptors (`"<module>::<name>"`).
    pub label: String,
}

// `ComputeEntryInfo` currently exposes only public fields. Helper
// methods for building compute pipelines live in caller code (the
// `wgpu::ComputePipelineDescriptor` shape is simple enough to assemble
// at the call site, and the previous `compute_pipeline_descriptor` /
// `compute_pipeline` helpers had zero in-tree call sites).

// ===== Buffer descriptors =====

/// A pre-built `wgpu::BufferDescriptor` for a `uniform!` or `storage!`
/// declaration. Sizing follows WGSL §14.4.1 ("Alignment and Size").
///
/// `size` is `0` for runtime-sized storage buffers, which signals the
/// caller to choose a buffer size appropriate for the workload (the GPU
/// side stores a runtime array whose length is determined by the buffer
/// size).
#[derive(Clone, Debug)]
pub struct BufferDescriptorInfo {
    /// The declared Rust name of the linkage variable.
    pub binding_name: String,
    /// The `@group(N)` index.
    pub group: u32,
    /// The `@binding(N)` index.
    pub binding: u32,
    /// What kind of buffer this is.
    pub kind: BufferKind,
    /// The byte size, computed per WGSL §14.4.1. `0` for runtime arrays.
    pub size: u64,
    /// The pre-built descriptor's usage flags, ready to plug into a
    /// `wgpu::BufferDescriptor` along with the borrow of
    /// [`Self::binding_name`].
    pub usage: wgpu::BufferUsages,
    /// The buffer's IR type.
    ///
    /// Here because it's useful for users to see the type from the runtime/CPU
    /// side.
    pub ty: crate::ir::Type,
}

impl BufferDescriptorInfo {
    /// Builds a `wgpu::BufferDescriptor` borrowing from
    /// [`Self::binding_name`]. The caller must keep this
    /// `BufferDescriptorInfo` alive for the duration of any wgpu call
    /// using the returned descriptor.
    pub fn descriptor(&self) -> wgpu::BufferDescriptor<'_> {
        wgpu::BufferDescriptor {
            label: Some(&self.binding_name),
            size: self.size,
            usage: self.usage,
            mapped_at_creation: false,
        }
    }

    /// Creates an empty buffer on the given device. The caller is
    /// responsible for populating the buffer (typically via
    /// `queue.write_buffer`).
    pub fn create_buffer(&self, device: &wgpu::Device) -> wgpu::Buffer {
        device.create_buffer(&self.descriptor())
    }
}

/// What kind of GPU buffer a [`BufferDescriptorInfo`] describes.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BufferKind {
    /// A uniform buffer (`var<uniform>`).
    Uniform,
    /// A storage buffer (`var<storage, read[_write]>`).
    Storage { read_only: bool },
}

/// The wgpu shader stage of an entry-point fn, or `None` for helper fns.
fn stage_of(fn_attrs: &ir::FnAttrs) -> Option<wgpu::ShaderStages> {
    match fn_attrs {
        ir::FnAttrs::Vertex => Some(wgpu::ShaderStages::VERTEX),
        ir::FnAttrs::Fragment => Some(wgpu::ShaderStages::FRAGMENT),
        ir::FnAttrs::Compute { .. } => Some(wgpu::ShaderStages::COMPUTE),
        ir::FnAttrs::None => None,
    }
}

/// Value identifiers and call edges collected from a fn body, kept
/// separate so a local (or parameter) that shadows a helper-fn name is
/// not mistaken for a call (PR #190 review finding).
#[derive(Default)]
struct FnRefs {
    /// `Expr::Ident` references: bindings, locals, params.
    idents: std::collections::HashSet<String>,
    /// `Expr::FnCall` callee paths; method calls are mangled
    /// (`Type_method`). Only these are traversed as call-graph edges.
    callees: std::collections::HashSet<String>,
}

/// Records a function's call-graph entry: its (possibly mangled) name
/// and the value identifiers and call edges its body references (see
/// [`collect_refs_in_block`]).
fn record_fn_idents(
    f: &ir::ItemFn,
    name: String,
    fn_names: &mut std::collections::HashSet<String>,
    fn_refs: &mut std::collections::HashMap<String, FnRefs>,
) {
    let mut refs = FnRefs::default();
    collect_refs_in_block(&mut refs, &f.block);
    fn_names.insert(name.clone());
    fn_refs.insert(name, refs);
}

/// Analyzes an IR module and returns its wgpu linkage.
///
/// This is the core entry point. It expects a concrete IR module (no
/// `Type::TypeParam` nodes); for template modules, call the
/// macro-emitted `instantiate::<...>()` function first and then pass the
/// result here.
///
/// The passed IR is consumed and stored on the returned linkage so that
/// [`WgpuLinkage::wgsl_source`] and [`WgpuLinkage::shader_module`] can
/// render the same source the linkage was analyzed from — even when
/// the linkage was built from an instantiated template.
///
/// **Infallible.** The input must be concrete (no `Type::TypeParam`
/// nodes); the proc-macro's `instantiate` enforces this at the type
/// level. The infallibility is honest: every failure mode would be a
/// bug in IR construction, not a user-facing error.
pub fn analyze_ir_module(mut ir_module: wgsl_rs_ir::Module) -> WgpuLinkage {
    let module_label = ir_module.name;
    ir::deshadow_module(&mut ir_module);
    ir::suffix_module(&mut ir_module);
    let mut linkage = WgpuLinkage {
        module_label,
        ir: ir_module,
        bind_groups: BTreeMap::new(),
        bind_group_layouts: HashMap::new(),
        pipeline_layout: None,
        vertex_entries: Vec::new(),
        fragment_entries: Vec::new(),
        compute_entries: Vec::new(),
        buffers: Vec::new(),
    };

    // First pass: collect every binding declaration, and record the
    // identifiers used by every function — entry points and helpers
    // alike. Per-binding `ShaderStages` visibility is then computed
    // transitively in a second pass by walking each entry point's call
    // graph, so bindings referenced only inside helper functions
    // called from an entry point still carry that entry point's stage
    // (wgsl-rs#177).
    let mut binding_names: std::collections::HashSet<String> = std::collections::HashSet::new();
    // Every function name in the module, including impl methods under
    // their mangled render names (`Type_method`).
    let mut fn_names: std::collections::HashSet<String> = std::collections::HashSet::new();
    // Per-fn reference sets: value identifiers and call edges, kept
    // separate (see [`FnRefs`]) so a local shadowing a helper-fn name
    // is not mistaken for a call.
    let mut fn_refs: std::collections::HashMap<String, FnRefs> = std::collections::HashMap::new();
    let mut entry_points: Vec<(String, wgpu::ShaderStages)> = Vec::new();

    for item in &linkage.ir.items {
        match item {
            ir::Item::Uniform(u) => {
                binding_names.insert(u.name.clone());
            }
            ir::Item::Storage(s) => {
                binding_names.insert(s.name.clone());
            }
            ir::Item::Sampler(s) => {
                binding_names.insert(s.name.clone());
            }
            ir::Item::Texture(t) => {
                binding_names.insert(t.name.clone());
            }
            ir::Item::Fn(f) => {
                record_fn_idents(f, f.name.to_string(), &mut fn_names, &mut fn_refs);
                if let Some(stage) = stage_of(&f.fn_attrs) {
                    entry_points.push((f.name.to_string(), stage));
                }
            }
            ir::Item::Impl(imp) => {
                for impl_item in &imp.items {
                    let ir::ImplItem::Fn(f) = impl_item else {
                        continue;
                    };
                    // Impl methods are invoked via their mangled render
                    // names (`Type_method`), so key the call graph by the
                    // same mangled name.
                    let mangled = ir::mangle::mangle(&[&imp.self_ty, &f.name]);
                    record_fn_idents(f, mangled.clone(), &mut fn_names, &mut fn_refs);
                    if let Some(stage) = stage_of(&f.fn_attrs) {
                        entry_points.push((mangled, stage));
                    }
                }
            }
            _ => {}
        }
    }

    // Second pass: resolve per-binding stage visibility transitively.
    // Starting from each entry point, walk the module call graph and
    // union every reachable function's referenced bindings into the
    // entry point's stage. The visited set guards against (spec-
    // illegal) recursive call cycles.
    let mut referenced: std::collections::HashMap<String, wgpu::ShaderStages> =
        std::collections::HashMap::new();
    for (entry_name, stage) in &entry_points {
        let mut visited: std::collections::HashSet<&String> = std::collections::HashSet::new();
        let mut stack: Vec<&String> = vec![entry_name];
        while let Some(fname) = stack.pop() {
            if !visited.insert(fname) {
                continue;
            }
            let Some(refs) = fn_refs.get(fname) else {
                continue;
            };
            // Bindings are attributed from value identifiers only; call
            // edges come from `callees` only, so a same-named local can
            // never masquerade as a call.
            for name in &refs.idents {
                if binding_names.contains(name) {
                    referenced
                        .entry(name.clone())
                        .and_modify(|s| *s |= *stage)
                        .or_insert(*stage);
                }
            }
            for callee in &refs.callees {
                if fn_names.contains(callee) {
                    stack.push(callee);
                }
            }
        }
    }

    for item in &linkage.ir.items {
        match item {
            ir::Item::Uniform(u) => {
                let size = type_byte_size_or_zero(&u.ty, &linkage.ir);
                let entry = wgpu::BindGroupLayoutEntry {
                    binding: u.binding,
                    visibility: visibility_for(&referenced, &u.name),
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Uniform,
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                };
                linkage
                    .bind_groups
                    .entry(u.group)
                    .or_insert_with(|| BindGroupInfo {
                        group: u.group,
                        entries: Vec::new(),
                        bindings: Vec::new(),
                        label: format!("{}::bind_group_{}", linkage.module_label, u.group),
                    })
                    .entries
                    .push(entry);
                linkage
                    .bind_groups
                    .get_mut(&u.group)
                    .unwrap()
                    .bindings
                    .push(BindingMeta {
                        name: u.name.clone(),
                        binding: u.binding,
                        kind: BindingKind::Uniform,
                    });
                linkage.buffers.push(BufferDescriptorInfo {
                    binding_name: u.name.clone(),
                    group: u.group,
                    binding: u.binding,
                    kind: BufferKind::Uniform,
                    size,
                    usage: wgpu::BufferUsages::UNIFORM
                        | wgpu::BufferUsages::COPY_DST
                        | wgpu::BufferUsages::COPY_SRC,
                    ty: u.ty.clone(),
                });
            }
            ir::Item::Storage(s) => {
                let read_only = matches!(s.access, ir::StorageAccess::Read);
                let size = type_byte_size_or_zero(&s.ty, &linkage.ir);
                let entry = wgpu::BindGroupLayoutEntry {
                    binding: s.binding,
                    visibility: visibility_for(&referenced, &s.name),
                    ty: wgpu::BindingType::Buffer {
                        ty: wgpu::BufferBindingType::Storage { read_only },
                        has_dynamic_offset: false,
                        min_binding_size: None,
                    },
                    count: None,
                };
                linkage
                    .bind_groups
                    .entry(s.group)
                    .or_insert_with(|| BindGroupInfo {
                        group: s.group,
                        entries: Vec::new(),
                        bindings: Vec::new(),
                        label: format!("{}::bind_group_{}", linkage.module_label, s.group),
                    })
                    .entries
                    .push(entry);
                linkage
                    .bind_groups
                    .get_mut(&s.group)
                    .unwrap()
                    .bindings
                    .push(BindingMeta {
                        name: s.name.clone(),
                        binding: s.binding,
                        kind: BindingKind::Storage { read_only },
                    });
                linkage.buffers.push(BufferDescriptorInfo {
                    binding_name: s.name.clone(),
                    group: s.group,
                    binding: s.binding,
                    kind: BufferKind::Storage { read_only },
                    size,
                    usage: wgpu::BufferUsages::STORAGE
                        | wgpu::BufferUsages::COPY_DST
                        | wgpu::BufferUsages::COPY_SRC,
                    ty: s.ty.clone(),
                });
            }
            ir::Item::Sampler(s) => {
                let comparison = matches!(s.ty, ir::Type::SamplerComparison);
                let sampler_ty = if comparison {
                    wgpu::SamplerBindingType::Comparison
                } else if !s.filterable {
                    wgpu::SamplerBindingType::NonFiltering
                } else {
                    wgpu::SamplerBindingType::Filtering
                };
                let entry = wgpu::BindGroupLayoutEntry {
                    binding: s.binding,
                    visibility: visibility_for(&referenced, &s.name),
                    ty: wgpu::BindingType::Sampler(sampler_ty),
                    count: None,
                };
                linkage
                    .bind_groups
                    .entry(s.group)
                    .or_insert_with(|| BindGroupInfo {
                        group: s.group,
                        entries: Vec::new(),
                        bindings: Vec::new(),
                        label: format!("{}::bind_group_{}", linkage.module_label, s.group),
                    })
                    .entries
                    .push(entry);
                linkage
                    .bind_groups
                    .get_mut(&s.group)
                    .unwrap()
                    .bindings
                    .push(BindingMeta {
                        name: s.name.clone(),
                        binding: s.binding,
                        kind: BindingKind::Sampler { comparison },
                    });
            }
            ir::Item::Texture(t) => {
                let (mut entry, kind) = texture_layout_entry(t);
                entry.visibility = visibility_for(&referenced, &t.name);
                linkage
                    .bind_groups
                    .entry(t.group)
                    .or_insert_with(|| BindGroupInfo {
                        group: t.group,
                        entries: Vec::new(),
                        bindings: Vec::new(),
                        label: format!("{}::bind_group_{}", linkage.module_label, t.group),
                    })
                    .entries
                    .push(entry);
                linkage
                    .bind_groups
                    .get_mut(&t.group)
                    .unwrap()
                    .bindings
                    .push(BindingMeta {
                        name: t.name.clone(),
                        binding: t.binding,
                        kind,
                    });
            }
            ir::Item::Fn(f) => match &f.fn_attrs {
                ir::FnAttrs::Vertex => {
                    linkage.vertex_entries.push(EntryPointInfo {
                        name: f.name.clone(),
                    });
                }
                ir::FnAttrs::Fragment => {
                    linkage.fragment_entries.push(EntryPointInfo {
                        name: f.name.clone(),
                    });
                }
                ir::FnAttrs::Compute { workgroup_size } => {
                    let x = workgroup_size.x;
                    let y = workgroup_size.y.unwrap_or(1);
                    let z = workgroup_size.z.unwrap_or(1);
                    linkage.compute_entries.push(ComputeEntryInfo {
                        name: f.name.clone(),
                        workgroup_size: (x, y, z),
                        label: format!("{}::{}", linkage.module_label, f.name),
                    });
                }
                ir::FnAttrs::None => {}
            },
            // Structs, impls, enums, consts, workgroup vars don't directly
            // produce wgpu objects.
            ir::Item::Struct(_)
            | ir::Item::Impl(_)
            | ir::Item::Enum(_)
            | ir::Item::Const(_)
            | ir::Item::Workgroup(_) => {}
        }
    }

    // Sort each bind group's entries by binding number for determinism
    // (matches the compile-time codegen behavior).
    for bg in linkage.bind_groups.values_mut() {
        let mut paired: Vec<_> = bg.entries.drain(..).zip(bg.bindings.drain(..)).collect();
        paired.sort_by_key(|(e, _m)| e.binding);
        bg.entries = paired.iter().map(|(e, _)| *e).collect();
        bg.bindings = paired.into_iter().map(|(_, m)| m).collect();
    }

    linkage
}

/// Returns the `ShaderStages` visibility for a binding given the
/// per-binding reference map. If no entry-point function references
/// the binding (e.g. the binding is only declared but unused, or the
/// reference walker missed it), we default to `COMPUTE` — the most
/// common stage for the wgsl-rs roundtrip-test use case. The previous
/// `ShaderStages::all()` default would silently demand the
/// `VERTEX_WRITABLE_STORAGE` feature for any read_write storage, which
/// is almost never what the user wants.
fn visibility_for(
    referenced: &std::collections::HashMap<String, wgpu::ShaderStages>,
    name: &str,
) -> wgpu::ShaderStages {
    referenced
        .get(name)
        .copied()
        .unwrap_or(wgpu::ShaderStages::COMPUTE)
}

/// Walks an `ir::Block`, collecting value identifiers and call edges
/// into `refs` (see [`FnRefs`]).
fn collect_refs_in_block(refs: &mut FnRefs, block: &ir::Block) {
    for stmt in &block.stmts {
        collect_refs_in_stmt(refs, stmt);
    }
}

fn collect_refs_in_stmt(refs: &mut FnRefs, stmt: &ir::Stmt) {
    match stmt {
        ir::Stmt::Local(l) => {
            if let Some(init) = &l.init {
                collect_refs_in_expr(refs, init);
            }
        }
        ir::Stmt::Const(c) => {
            collect_refs_in_expr(refs, &c.expr);
        }
        ir::Stmt::Assignment { lhs, rhs } => {
            collect_refs_in_expr(refs, lhs);
            collect_refs_in_expr(refs, rhs);
        }
        ir::Stmt::CompoundAssignment { lhs, rhs, .. } => {
            collect_refs_in_expr(refs, lhs);
            collect_refs_in_expr(refs, rhs);
        }
        ir::Stmt::While { condition, body } => {
            collect_refs_in_expr(refs, condition);
            collect_refs_in_block(refs, body);
        }
        ir::Stmt::Loop { body } => {
            collect_refs_in_block(refs, body);
        }
        ir::Stmt::Expr { expr, .. } => {
            collect_refs_in_expr(refs, expr);
        }
        ir::Stmt::If(s) => {
            collect_refs_in_expr(refs, &s.condition);
            collect_refs_in_block(refs, &s.then_block);
            if let Some(else_branch) = &s.else_branch {
                match else_branch {
                    ir::ElseBranch::Block(b) => collect_refs_in_block(refs, b),
                    ir::ElseBranch::If(i) => {
                        collect_refs_in_expr(refs, &i.condition);
                        collect_refs_in_block(refs, &i.then_block);
                        if let Some(else_branch) = &i.else_branch {
                            match else_branch {
                                ir::ElseBranch::Block(b) => collect_refs_in_block(refs, b),
                                ir::ElseBranch::If(_) => {
                                    // Shouldn't recurse infinitely; the IR is
                                    // finite.
                                }
                            }
                        }
                    }
                }
            }
        }
        ir::Stmt::Return(Some(e)) => collect_refs_in_expr(refs, e),
        ir::Stmt::For(f) => {
            collect_refs_in_expr(refs, &f.from);
            collect_refs_in_expr(refs, &f.to);
            collect_refs_in_block(refs, &f.body);
        }
        ir::Stmt::Switch(s) => {
            collect_refs_in_expr(refs, &s.selector);
            for arm in &s.arms {
                for sel in &arm.selectors {
                    match sel {
                        ir::CaseSelector::Literal(_) => {}
                        ir::CaseSelector::Expr(e) => collect_refs_in_expr(refs, e),
                        ir::CaseSelector::Default => {}
                    }
                }
                collect_refs_in_block(refs, &arm.body);
            }
        }
        ir::Stmt::Block(b) => collect_refs_in_block(refs, b),
        ir::Stmt::SlabCopy {
            src,
            src_offset,
            dest,
            dest_offset,
            size,
        } => {
            collect_refs_in_expr(refs, src);
            collect_refs_in_expr(refs, src_offset);
            collect_refs_in_expr(refs, dest);
            collect_refs_in_expr(refs, dest_offset);
            collect_refs_in_expr(refs, size);
        }
        ir::Stmt::Break
        | ir::Stmt::Continue
        | ir::Stmt::Discard
        | ir::Stmt::Return(None)
        | ir::Stmt::Macro { .. } => {}
    }
}

fn collect_refs_in_expr(refs: &mut FnRefs, expr: &ir::Expr) {
    match expr {
        ir::Expr::Ident(name) => {
            refs.idents.insert(name.clone());
        }
        ir::Expr::Lit(_) => {}
        ir::Expr::Array { elems } => {
            for e in elems {
                collect_refs_in_expr(refs, e);
            }
        }
        ir::Expr::Paren(e) => collect_refs_in_expr(refs, e),
        ir::Expr::Binary { lhs, rhs, .. } => {
            collect_refs_in_expr(refs, lhs);
            collect_refs_in_expr(refs, rhs);
        }
        ir::Expr::Unary { expr, .. } => collect_refs_in_expr(refs, expr),
        ir::Expr::ArrayIndexing { lhs, index } => {
            collect_refs_in_expr(refs, lhs);
            collect_refs_in_expr(refs, index);
        }
        ir::Expr::Swizzle { lhs, params, .. } => {
            collect_refs_in_expr(refs, lhs);
            if let Some(args) = params {
                for a in args {
                    collect_refs_in_expr(refs, a);
                }
            }
        }
        ir::Expr::Cast { lhs, .. } => collect_refs_in_expr(refs, lhs),
        ir::Expr::FnCall { path, params, .. } => {
            // Record the callee so the linkage analysis can resolve
            // transitive binding use through helper functions
            // (wgsl-rs#177). Method calls are recorded under their
            // mangled render names, matching how impl methods are
            // keyed in the call graph.
            match path {
                ir::FnPath::Ident(name) => {
                    refs.callees.insert(name.clone());
                }
                ir::FnPath::TypeMethod { ty, method } => {
                    refs.callees.insert(ir::mangle::mangle(&[ty, method]));
                }
            }
            for p in params {
                collect_refs_in_expr(refs, p);
            }
        }
        ir::Expr::Struct { fields, .. } => {
            for fv in fields {
                collect_refs_in_expr(refs, &fv.expr);
            }
        }
        ir::Expr::FieldAccess { base, .. } => collect_refs_in_expr(refs, base),
        ir::Expr::TypePath { .. } => {}
        ir::Expr::Reference(e) => collect_refs_in_expr(refs, e),
        ir::Expr::ZeroValueArray { len, .. } => collect_refs_in_expr(refs, len),
    }
}

/// Analyzes a `wgsl_rs::Source` by first assembling its full IR
/// (including imports and template instantiations), then delegating to
/// [`analyze_ir_module`].
///
/// Errors if the source is a template (has unresolved
/// `Type::TypeParam`s). For templates, call the macro-emitted
/// `instantiate::<...>()` function first and pass the resulting
/// `ir::Module` to [`analyze_ir_module`].
///
/// **Fallible.** Returns
/// [`SourceError::TemplateWgsl`](crate::SourceError::TemplateWgsl)-style errors
/// for template sources (caller must `instantiate` first) and
/// [`Error::TemplateResolution`] if a cross-source template instantiation
/// is ambiguous. Callers can `?` this on a possibly-template source, or
/// `unwrap` on a known concrete one.
pub fn analyze_wgsl_module(wgsl_source: &Source) -> Result<WgpuLinkage, Error> {
    if wgsl_source.is_template() {
        return Err(crate::SourceError::TemplateWgsl {
            uninstantiated_source: wgsl_source,
        }
        .into());
    }
    let ir_module = assemble_ir(wgsl_source)?;
    Ok(analyze_ir_module(ir_module))
}

// ===== IR assembly (flattens imports) =====

/// Build a single owned `ir::Module` from a `wgsl_rs::Source` and all
/// its transitive imports and template instantiations. The result is
/// concrete (no `Type::TypeParam`s) and self-contained: every struct
/// referenced by a binding is defined within the assembled source.
///
/// This is the same traversal that `Source::wgsl_source()` performs,
/// but it stops at IR rather than rendering. It exists so the runtime
/// linkage analyzer can correctly size buffers whose types are defined
/// in imported sources.
pub(crate) fn assemble_ir(wgsl_source: &Source) -> Result<ir::Module, Error> {
    let mut visited: std::collections::HashSet<u64> = std::collections::HashSet::new();
    let mut seen: std::collections::HashSet<(u64, String, Vec<String>)> =
        std::collections::HashSet::new();
    let mut items: Vec<ir::Item> = Vec::new();
    collect_items(wgsl_source, &mut visited, &mut seen, &mut items, None)?;
    // Note: deshadowing is deferred to `analyze_ir_module`, which is the
    // sole consumer of this assembled IR. Deshadowing here too would be
    // redundant (the pass is idempotent but wastes work).
    Ok(ir::Module {
        name: wgsl_source.name,
        items,
        attrs: Vec::new(),
    })
}

fn collect_items(
    wgsl_source: &Source,
    visited: &mut std::collections::HashSet<u64>,
    seen: &mut std::collections::HashSet<(u64, String, Vec<String>)>,
    out: &mut Vec<ir::Item>,
    subst: Option<&std::collections::HashMap<String, ir::Type>>,
) -> Result<(), Error> {
    // Imports first (depth-first, deduplicated by source id).
    for m in wgsl_source.imports {
        if visited.insert(m.id) {
            collect_items(m, visited, seen, out, None)?;
        }
    }

    // This source's own items, optionally substituted.
    let mut ir_module = (wgsl_source.ir_constructor)();
    if let Some(s) = subst {
        ir::substitute_types(&mut ir_module, s);
    }
    out.extend(ir_module.items);

    // Cross-source template instantiations.
    for inst in wgsl_source.instantiations {
        let mangled: Vec<String> = inst
            .mangled_type_args
            .iter()
            .map(|s| (*s).to_string())
            .collect();
        let type_args = (inst.type_args_constructor)();
        collect_instantiation(
            inst.modules,
            inst.template_name,
            &mangled,
            &type_args,
            out,
            seen,
        )?;
    }
    Ok(())
}

fn collect_instantiation(
    sources: &[&Source],
    template_name: &str,
    mangled_type_args: &[String],
    type_args: &[ir::Type],
    out: &mut Vec<ir::Item>,
    seen: &mut std::collections::HashSet<(u64, String, Vec<String>)>,
) -> Result<(), Error> {
    let available: Vec<String> = sources
        .iter()
        .copied()
        .flat_map(|m| m.templates.iter().map(|t| t.name.to_string()))
        .collect();

    let mut matching: Vec<&Source> = sources
        .iter()
        .copied()
        .filter(|m| m.templates.iter().any(|t| t.name == template_name))
        .collect();

    snafu::ensure!(
        !matching.is_empty(),
        TemplateResolutionSnafu {
            template_name: template_name.to_string(),
            mangled_type_args: mangled_type_args.to_vec(),
            available,
        }
    );

    if matching.len() > 1 {
        let names: Vec<String> = matching.iter().map(|m| m.name.to_string()).collect();
        return TemplateResolutionSnafu {
            template_name: template_name.to_string(),
            mangled_type_args: mangled_type_args.to_vec(),
            available: names,
        }
        .fail();
    }

    let source = matching
        .pop()
        .expect("matching is non-empty after the check above");
    let Some(template) = source.templates.iter().find(|t| t.name == template_name) else {
        // Invariant: `matching` was built by filtering on
        // `templates.iter().any(|t| t.name == template_name)`, so the
        // sole surviving source must contain the template.
        unreachable!(
            "resolved source '{}' has no template '{}'; available: {:?}",
            source.name, template_name, available
        );
    };

    let key = (
        source.id,
        template_name.to_string(),
        mangled_type_args.to_vec(),
    );
    snafu::ensure!(
        seen.insert(key),
        TemplateResolutionSnafu {
            template_name: template_name.to_owned(),
            mangled_type_args: mangled_type_args.to_vec(),
            available
        }
    );

    // Recurse into dependencies first.
    for dep in template.dependencies {
        let dep_mangled: Vec<String> = dep
            .type_param_mapping
            .iter()
            .map(|&idx| mangled_type_args[idx].clone())
            .collect();
        let dep_args: Vec<ir::Type> = dep
            .type_param_mapping
            .iter()
            .map(|&idx| type_args[idx].clone())
            .collect();
        collect_instantiation(&[source], dep.callee, &dep_mangled, &dep_args, out, seen)?;
    }

    let mut subst: std::collections::HashMap<String, ir::Type> = std::collections::HashMap::new();
    for (param, arg) in template.type_params.iter().zip(type_args.iter()) {
        subst.insert((*param).to_string(), arg.clone());
    }

    let mut items = (template.ir_constructor)();
    ir::substitute_items(&mut items, &subst);

    let instance_name = if mangled_type_args.is_empty() {
        template.name.to_string()
    } else {
        let mut components: Vec<&str> = Vec::with_capacity(1 + mangled_type_args.len());
        components.push(template.name);
        for s in mangled_type_args {
            components.push(s.as_str());
        }
        ir::mangle(&components)
    };
    if instance_name != template.name {
        ir::rename_items(&mut items, template.name, &instance_name);
    }
    out.extend(items);

    Ok(())
}

// ===== Texture helpers =====

fn texture_layout_entry(item: &ir::ItemTexture) -> (wgpu::BindGroupLayoutEntry, BindingKind) {
    let view_dimension = match &item.ty {
        ir::Type::Texture { kind, .. } => texture_kind_view_dimension(*kind),
        ir::Type::TextureDepth { kind } => texture_depth_view_dimension(*kind),
        ir::Type::TextureStorage { kind, .. } => texture_storage_kind_view_dimension(*kind),
        _ => unreachable!("ItemTexture validated to be a texture type"),
    };
    let multisampled = match &item.ty {
        ir::Type::Texture { kind, .. } => matches!(kind, ir::TextureKind::TextureMultisampled2D),
        ir::Type::TextureDepth { kind } => {
            matches!(kind, ir::TextureDepthKind::DepthMultisampled2D)
        }
        ir::Type::TextureStorage { .. } => false,
        _ => unreachable!(),
    };
    match &item.ty {
        ir::Type::Texture { sampled_type, .. } => {
            let sample = match sampled_type {
                ir::ScalarType::F32 => wgpu::TextureSampleType::Float {
                    filterable: item.filterable,
                },
                ir::ScalarType::I32 => wgpu::TextureSampleType::Sint,
                ir::ScalarType::U32 => wgpu::TextureSampleType::Uint,
                ir::ScalarType::Bool => unreachable!("textures can't be bool-sampled"),
            };
            (
                wgpu::BindGroupLayoutEntry {
                    binding: item.binding,
                    visibility: wgpu::ShaderStages::all(),
                    ty: wgpu::BindingType::Texture {
                        sample_type: sample,
                        view_dimension,
                        multisampled,
                    },
                    count: None,
                },
                BindingKind::Texture,
            )
        }
        ir::Type::TextureDepth { .. } => (
            wgpu::BindGroupLayoutEntry {
                binding: item.binding,
                visibility: wgpu::ShaderStages::all(),
                ty: wgpu::BindingType::Texture {
                    sample_type: wgpu::TextureSampleType::Depth,
                    view_dimension,
                    multisampled,
                },
                count: None,
            },
            BindingKind::DepthTexture,
        ),
        ir::Type::TextureStorage { format, access, .. } => {
            let wgpu_format = texel_format_to_wgpu(*format);
            let wgpu_access = match access {
                ir::StorageTextureAccess::Read => wgpu::StorageTextureAccess::ReadOnly,
                ir::StorageTextureAccess::Write => wgpu::StorageTextureAccess::WriteOnly,
                ir::StorageTextureAccess::ReadWrite => wgpu::StorageTextureAccess::ReadWrite,
            };
            let read_write = matches!(access, ir::StorageTextureAccess::ReadWrite);
            (
                wgpu::BindGroupLayoutEntry {
                    binding: item.binding,
                    visibility: wgpu::ShaderStages::all(),
                    ty: wgpu::BindingType::StorageTexture {
                        access: wgpu_access,
                        format: wgpu_format,
                        view_dimension,
                    },
                    count: None,
                },
                BindingKind::StorageTexture { read_write },
            )
        }
        _ => unreachable!(),
    }
}

fn texture_storage_kind_view_dimension(kind: ir::TextureStorageKind) -> wgpu::TextureViewDimension {
    match kind {
        ir::TextureStorageKind::Storage1D => wgpu::TextureViewDimension::D1,
        ir::TextureStorageKind::Storage2D => wgpu::TextureViewDimension::D2,
        ir::TextureStorageKind::Storage2DArray => wgpu::TextureViewDimension::D2Array,
        ir::TextureStorageKind::Storage3D => wgpu::TextureViewDimension::D3,
    }
}

/// Map an IR texel format to the corresponding wgpu texture format.
fn texel_format_to_wgpu(f: ir::TexelFormat) -> wgpu::TextureFormat {
    use ir::TexelFormat::*;
    match f {
        Rgba8unorm => wgpu::TextureFormat::Rgba8Unorm,
        Rgba8snorm => wgpu::TextureFormat::Rgba8Snorm,
        Rgba8uint => wgpu::TextureFormat::Rgba8Uint,
        Rgba8sint => wgpu::TextureFormat::Rgba8Sint,
        Rgba16uint => wgpu::TextureFormat::Rgba16Uint,
        Rgba16sint => wgpu::TextureFormat::Rgba16Sint,
        Rgba16float => wgpu::TextureFormat::Rgba16Float,
        R32uint => wgpu::TextureFormat::R32Uint,
        R32sint => wgpu::TextureFormat::R32Sint,
        R32float => wgpu::TextureFormat::R32Float,
        Rg32uint => wgpu::TextureFormat::Rg32Uint,
        Rg32sint => wgpu::TextureFormat::Rg32Sint,
        Rg32float => wgpu::TextureFormat::Rg32Float,
        Rgba32uint => wgpu::TextureFormat::Rgba32Uint,
        Rgba32sint => wgpu::TextureFormat::Rgba32Sint,
        Rgba32float => wgpu::TextureFormat::Rgba32Float,
        Bgra8unorm => wgpu::TextureFormat::Bgra8Unorm,
        Rgba16unorm => wgpu::TextureFormat::Rgba16Unorm,
        Rgba16snorm => wgpu::TextureFormat::Rgba16Snorm,
        Rg8unorm => wgpu::TextureFormat::Rg8Unorm,
        Rg8snorm => wgpu::TextureFormat::Rg8Snorm,
        Rg8uint => wgpu::TextureFormat::Rg8Uint,
        Rg8sint => wgpu::TextureFormat::Rg8Sint,
        Rg16unorm => wgpu::TextureFormat::Rg16Unorm,
        Rg16snorm => wgpu::TextureFormat::Rg16Snorm,
        Rg16uint => wgpu::TextureFormat::Rg16Uint,
        Rg16sint => wgpu::TextureFormat::Rg16Sint,
        Rg16float => wgpu::TextureFormat::Rg16Float,
        R8unorm => wgpu::TextureFormat::R8Unorm,
        R8snorm => wgpu::TextureFormat::R8Snorm,
        R8uint => wgpu::TextureFormat::R8Uint,
        R8sint => wgpu::TextureFormat::R8Sint,
        R16unorm => wgpu::TextureFormat::R16Unorm,
        R16snorm => wgpu::TextureFormat::R16Snorm,
        R16uint => wgpu::TextureFormat::R16Uint,
        R16sint => wgpu::TextureFormat::R16Sint,
        R16float => wgpu::TextureFormat::R16Float,
        Rgb10a2unorm => wgpu::TextureFormat::Rgb10a2Unorm,
        Rgb10a2uint => wgpu::TextureFormat::Rgb10a2Uint,
        Rg11b10ufloat => wgpu::TextureFormat::Rg11b10Ufloat,
    }
}

fn texture_kind_view_dimension(kind: ir::TextureKind) -> wgpu::TextureViewDimension {
    match kind {
        ir::TextureKind::Texture1D => wgpu::TextureViewDimension::D1,
        ir::TextureKind::Texture2D | ir::TextureKind::TextureMultisampled2D => {
            wgpu::TextureViewDimension::D2
        }
        ir::TextureKind::Texture2DArray => wgpu::TextureViewDimension::D2Array,
        ir::TextureKind::Texture3D => wgpu::TextureViewDimension::D3,
        ir::TextureKind::TextureCube => wgpu::TextureViewDimension::Cube,
        ir::TextureKind::TextureCubeArray => wgpu::TextureViewDimension::CubeArray,
    }
}

fn texture_depth_view_dimension(kind: ir::TextureDepthKind) -> wgpu::TextureViewDimension {
    match kind {
        ir::TextureDepthKind::Depth2D | ir::TextureDepthKind::DepthMultisampled2D => {
            wgpu::TextureViewDimension::D2
        }
        ir::TextureDepthKind::Depth2DArray => wgpu::TextureViewDimension::D2Array,
        ir::TextureDepthKind::DepthCube => wgpu::TextureViewDimension::Cube,
        ir::TextureDepthKind::DepthCubeArray => wgpu::TextureViewDimension::CubeArray,
    }
}

// ===== WGSL §14.4.1 type sizing =====

/// Returns the byte size of a host-shareable WGSL type per §14.4.1.
/// `RuntimeArray<T>` returns `0` (caller chooses the buffer size based
/// on workload). User struct lookups require the defining module.
///
/// Returns `0` for non-shareable types (samplers, textures, pointers) and
/// for unresolvable cases (e.g. an array length that isn't a literal or
/// same-module const). In the latter case a panic with location info
/// would be more user-friendly, but we err on the side of "do something
/// reasonable" so that callers can still build pipelines.
pub fn type_byte_size(ty: &ir::Type, module: &ir::Module) -> u64 {
    let layout = type_layout(ty, module);
    layout.size
}

fn type_byte_size_or_zero(ty: &ir::Type, module: &ir::Module) -> u64 {
    type_byte_size(ty, module)
}

/// Layout for a single type (size + align) per WGSL §14.4.1.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TypeLayout {
    pub size: u64,
    pub align: u64,
}

fn align_up(val: u64, align: u64) -> u64 {
    if align == 0 {
        val
    } else {
        (val + (align - 1)) & !(align - 1)
    }
}

fn type_layout(ty: &ir::Type, module: &ir::Module) -> TypeLayout {
    match ty {
        ir::Type::Scalar(s) => scalar_layout(*s),
        ir::Type::Vector { elements, .. } => vector_layout(*elements),
        ir::Type::Matrix { columns, rows, .. } => matrix_layout(*columns, *rows),
        ir::Type::Array { elem, len } => {
            let inner = type_layout(elem, module);
            let n = eval_array_len(len, module).unwrap_or(0);
            TypeLayout {
                size: (n as u64) * align_up(inner.size, inner.align.max(1)),
                align: inner.align,
            }
        }
        ir::Type::RuntimeArray { elem } => {
            let inner = type_layout(elem, module);
            TypeLayout {
                size: 0,
                align: inner.align,
            }
        }
        ir::Type::Atomic { elem } => type_layout(elem, module),
        ir::Type::Struct { name, type_args } => struct_layout(name, type_args, module),
        ir::Type::Ptr { .. } => TypeLayout { size: 0, align: 1 },
        // Samplers, textures, and type parameters aren't host-shareable.
        // Return a zero layout so the analyzer doesn't choke.
        ir::Type::Sampler
        | ir::Type::SamplerComparison
        | ir::Type::Texture { .. }
        | ir::Type::TextureDepth { .. }
        | ir::Type::TextureStorage { .. }
        | ir::Type::TypeParam { .. }
        | ir::Type::Phantom { .. }
        | ir::Type::AssocType { .. } => TypeLayout { size: 0, align: 1 },
    }
}

fn scalar_layout(s: ir::ScalarType) -> TypeLayout {
    match s {
        ir::ScalarType::F32 | ir::ScalarType::I32 | ir::ScalarType::U32 | ir::ScalarType::Bool => {
            TypeLayout { size: 4, align: 4 }
        }
    }
}

fn vector_layout(elements: u8) -> TypeLayout {
    match elements {
        2 => TypeLayout { size: 8, align: 8 },
        3 => TypeLayout {
            size: 12,
            align: 16,
        },
        4 => TypeLayout {
            size: 16,
            align: 16,
        },
        _ => TypeLayout { size: 0, align: 1 },
    }
}

fn matrix_layout(columns: u8, rows: u8) -> TypeLayout {
    // WGSL §14.4.1: matCxR<T> has:
    //   Align = AlignOf(vecR<T>)
    //   Size  = SizeOf(array<vecR<T>, C>)
    // Each column is stored as a vecR<T> with stride roundUp(SizeOf(vecR<T>),
    // AlignOf(vecR<T>)).
    let row_vec = vector_layout(rows);
    let stride = align_up(row_vec.size, row_vec.align);
    TypeLayout {
        size: (columns as u64) * stride,
        align: row_vec.align,
    }
}

fn struct_layout(name: &str, type_args: &[ir::Type], module: &ir::Module) -> TypeLayout {
    let Some(s) = find_struct(name, type_args, module) else {
        return TypeLayout { size: 0, align: 1 };
    };
    let mut offset: u64 = 0;
    let mut struct_align: u64 = 1;
    for field in &s.fields {
        let fl = type_layout(&field.ty, module);
        struct_align = struct_align.max(fl.align);
        offset = align_up(offset, fl.align);
        offset += fl.size;
    }
    let size = align_up(offset, struct_align);
    TypeLayout {
        size,
        align: struct_align,
    }
}

/// Look up a struct by name (and optional type args) in the module. The
/// struct is matched by base name; type args are currently ignored for
/// lookup (we look up the generic template, then resolve via the
/// `type_args` substitution if any). Generic structs share a base name
/// across instantiations, so callers that need a specific instantiation
/// must pre-substitute the IR before calling.
fn find_struct<'a>(
    name: &str,
    _type_args: &[ir::Type],
    module: &'a ir::Module,
) -> Option<&'a ir::ItemStruct> {
    module.items.iter().find_map(|i| match i {
        ir::Item::Struct(s) if s.name == name => Some(s),
        _ => None,
    })
}

/// Evaluate a fixed-size array length expression. Supports integer
/// literals and named integer constants defined at module scope. Other
/// expressions (binary ops, casts, etc.) return `None` and the caller
/// falls back to size 0 with a clearly-too-small buffer.
fn eval_array_len(expr: &ir::Expr, module: &ir::Module) -> Option<usize> {
    match expr {
        ir::Expr::Lit(ir::Lit::Int { digits, .. }) => digits.parse::<usize>().ok(),
        // WGSL array lengths must be non-negative const expressions, so a
        // unary negation is unresolvable — returning `None` falls back to
        // size 0 rather than producing a huge (wraparound) length.
        ir::Expr::Unary {
            op: ir::UnOp::Neg, ..
        } => None,
        ir::Expr::Ident(name) => {
            let val = module.items.iter().find_map(|i| match i {
                ir::Item::Const(c) if c.name == *name => Some(eval_const_int(&c.expr, module)),
                _ => None,
            });
            val.flatten()
        }
        _ => None,
    }
}

fn eval_const_int(expr: &ir::Expr, module: &ir::Module) -> Option<usize> {
    eval_array_len(expr, module)
}

#[cfg(test)]
mod binding_visibility {
    //! Tests for per-binding `ShaderStages` visibility, including the
    //! transitive helper-function case from wgsl-rs#177.

    use std::borrow::Cow;

    use super::*;

    fn f32_ty() -> ir::Type {
        ir::Type::Scalar(ir::ScalarType::F32)
    }

    fn zero() -> ir::Expr {
        ir::Expr::Lit(ir::Lit::Float {
            text: "0.0".to_string(),
        })
    }

    fn uniform(name: &str, group: u32, binding: u32) -> ir::Item {
        ir::Item::Uniform(ir::ItemUniform {
            group,
            binding,
            name: name.to_string(),
            ty: f32_ty(),
            attrs: Vec::new(),
        })
    }

    fn call(name: &str, params: Vec<ir::Expr>) -> ir::Expr {
        ir::Expr::FnCall {
            path: ir::FnPath::Ident(name.to_string()),
            type_args: Vec::new(),
            params,
        }
    }

    fn ident(name: &str) -> ir::Expr {
        ir::Expr::Ident(name.to_string())
    }

    fn fn_item(name: &str, fn_attrs: ir::FnAttrs, stmts: Vec<ir::Stmt>) -> ir::Item {
        ir::Item::Fn(ir::ItemFn {
            type_params: Vec::new(),
            const_params: Vec::new(),
            fn_attrs,
            name: Cow::Owned(name.to_string()),
            inputs: Vec::new(),
            return_type: ir::ReturnType::Default,
            block: ir::Block { stmts },
            attrs: Vec::new(),
        })
    }

    fn returns(expr: ir::Expr) -> ir::Stmt {
        ir::Stmt::Return(Some(expr))
    }

    fn analyze(items: Vec<ir::Item>) -> WgpuLinkage {
        analyze_ir_module(ir::Module {
            name: "test_module",
            items,
            attrs: Vec::new(),
        })
    }

    /// Looks up the analyzed visibility for the named binding.
    fn visibility_of(linkage: &WgpuLinkage, name: &str) -> wgpu::ShaderStages {
        for bind_group in linkage.bind_groups.values() {
            if let Some(meta) = bind_group.bindings.iter().find(|b| b.name == name) {
                let entry = bind_group
                    .entries
                    .iter()
                    .find(|e| e.binding == meta.binding)
                    .expect("binding layout entry exists");
                return entry.visibility;
            }
        }
        panic!("binding '{name}' not found in any bind group");
    }

    #[test]
    fn direct_use_gets_entry_stage() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![returns(ident("MY_UNIFORM"))],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::VERTEX
        );
    }

    /// The wgsl-rs#177 repro: the binding is used only inside a helper
    /// called from both entry points, so it must carry both stages.
    #[test]
    fn helper_only_use_carries_entry_point_stages() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item(
                "read_uniform",
                ir::FnAttrs::None,
                vec![returns(ident("MY_UNIFORM"))],
            ),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![returns(call("read_uniform", vec![]))],
            ),
            fn_item(
                "frag",
                ir::FnAttrs::Fragment,
                vec![returns(call("read_uniform", vec![]))],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT
        );
    }

    /// A helper called only from `frag` must not leak vertex visibility.
    #[test]
    fn frag_only_helper_stays_fragment() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item(
                "read_uniform",
                ir::FnAttrs::None,
                vec![returns(ident("MY_UNIFORM"))],
            ),
            fn_item(
                "frag",
                ir::FnAttrs::Fragment,
                vec![returns(call("read_uniform", vec![]))],
            ),
            fn_item("vert", ir::FnAttrs::Vertex, vec![returns(zero())]),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::FRAGMENT
        );
    }

    /// Mixed use: direct reference in one entry point, helper-only in
    /// the other.
    #[test]
    fn mixed_direct_and_helper_use() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item(
                "read_uniform",
                ir::FnAttrs::None,
                vec![returns(ident("MY_UNIFORM"))],
            ),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![returns(ident("MY_UNIFORM"))],
            ),
            fn_item(
                "frag",
                ir::FnAttrs::Fragment,
                vec![returns(call("read_uniform", vec![]))],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::VERTEX | wgpu::ShaderStages::FRAGMENT
        );
    }

    /// Reachability must be transitive: entry point -> helper -> helper
    /// -> binding.
    #[test]
    fn helper_chain_is_transitive() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item(
                "inner",
                ir::FnAttrs::None,
                vec![returns(ident("MY_UNIFORM"))],
            ),
            fn_item(
                "outer",
                ir::FnAttrs::None,
                vec![returns(call("inner", vec![]))],
            ),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![returns(call("outer", vec![]))],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::VERTEX
        );
    }

    /// A declared-but-unused binding keeps the documented COMPUTE
    /// default.
    #[test]
    fn unused_binding_defaults_to_compute() {
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            fn_item("vert", ir::FnAttrs::Vertex, vec![returns(zero())]),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::COMPUTE
        );
    }

    /// A local (or parameter) that shadows a helper-fn name is a value
    /// reference, not a call edge (PR #190 review finding): the helper
    /// is never called, so its bindings must not gain this entry
    /// point's stage.
    #[test]
    fn local_shadowing_fn_name_is_not_a_call_edge() {
        let linkage = analyze(vec![
            uniform("DATA", 0, 0),
            fn_item(
                "read_uniform",
                ir::FnAttrs::None,
                vec![returns(ident("DATA"))],
            ),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![
                    ir::Stmt::Local(ir::Local {
                        mutable: false,
                        name: "read_uniform".to_string(),
                        ty: None,
                        init: Some(zero()),
                    }),
                    returns(ident("read_uniform")),
                ],
            ),
        ]);
        // The helper is unreachable, so DATA falls back to the COMPUTE
        // default rather than spuriously gaining VERTEX.
        assert_eq!(visibility_of(&linkage, "DATA"), wgpu::ShaderStages::COMPUTE);
    }

    /// Same-shaped control: a real call from frag still creates the
    /// edge (FRAGMENT), while vert's same-named local grants nothing.
    #[test]
    fn call_edge_survives_local_shadowing_elsewhere() {
        let linkage = analyze(vec![
            uniform("DATA", 0, 0),
            fn_item(
                "read_uniform",
                ir::FnAttrs::None,
                vec![returns(ident("DATA"))],
            ),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![
                    ir::Stmt::Local(ir::Local {
                        mutable: false,
                        name: "read_uniform".to_string(),
                        ty: None,
                        init: Some(zero()),
                    }),
                    returns(ident("read_uniform")),
                ],
            ),
            fn_item(
                "frag",
                ir::FnAttrs::Fragment,
                vec![returns(call("read_uniform", vec![]))],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "DATA"),
            wgpu::ShaderStages::FRAGMENT
        );
    }

    /// Impl methods are called via mangled `Type_method` names; bindings
    /// used inside them must be reachable from entry points that call
    /// them.
    #[test]
    fn impl_method_helper_is_reachable() {
        let method = ir::ItemFn {
            type_params: Vec::new(),
            const_params: Vec::new(),
            fn_attrs: ir::FnAttrs::None,
            name: Cow::Owned("read_uniform".to_string()),
            inputs: Vec::new(),
            return_type: ir::ReturnType::Default,
            block: ir::Block {
                stmts: vec![returns(ident("MY_UNIFORM"))],
            },
            attrs: Vec::new(),
        };
        let linkage = analyze(vec![
            uniform("MY_UNIFORM", 0, 0),
            ir::Item::Impl(ir::ItemImpl {
                type_params: Vec::new(),
                const_params: Vec::new(),
                self_ty: "Foo".to_string(),
                items: vec![ir::ImplItem::Fn(method)],
                attrs: Vec::new(),
            }),
            fn_item(
                "vert",
                ir::FnAttrs::Vertex,
                vec![returns(ir::Expr::FnCall {
                    path: ir::FnPath::TypeMethod {
                        ty: "Foo".to_string(),
                        method: "read_uniform".to_string(),
                    },
                    type_args: Vec::new(),
                    params: Vec::new(),
                })],
            ),
        ]);
        assert_eq!(
            visibility_of(&linkage, "MY_UNIFORM"),
            wgpu::ShaderStages::VERTEX
        );
    }
}