naga/common/wgsl/

types.rs

//! Code for formatting Naga IR types as WGSL source code.

use super::{address_space_str, ToWgsl, TryToWgsl};
use crate::common;
use crate::proc::TypeResolution;
use crate::{Handle, Scalar, TypeInner};

use alloc::string::String;
use core::fmt::Write;

/// A context for printing Naga IR types as WGSL.
///
/// This trait's default methods [`write_type`] and
/// [`write_type_inner`] do the work of formatting types as WGSL.
/// Implementors must provide the remaining methods, to customize
/// behavior for the context at hand.
///
/// For example, the WGSL backend would provide an implementation of
/// [`type_name`] that handles hygienic renaming, whereas the WGSL
/// front end would simply show the name that was given in the source.
///
/// [`write_type`]: TypeContext::write_type
/// [`write_type_inner`]: TypeContext::write_type_inner
/// [`type_name`]: TypeContext::type_name
pub trait TypeContext {
    /// Return the [`Type`] referred to by `handle`.
    ///
    /// [`Type`]: crate::Type
    fn lookup_type(&self, handle: Handle<crate::Type>) -> &crate::Type;

    /// Return the name to be used for the type referred to by
    /// `handle`.
    fn type_name(&self, handle: Handle<crate::Type>) -> &str;

    /// Write the WGSL form of `override` to `out`.
    fn write_override<W: Write>(
        &self,
        r#override: Handle<crate::Override>,
        out: &mut W,
    ) -> core::fmt::Result;

    /// Write a [`TypeInner`] that has no representation as WGSL source,
    /// even including Naga extensions.
    ///
    /// A backend might implement this with a call to the [`unreachable!`]
    /// macro, since backends are allowed to assume that the module has passed
    /// validation.
    ///
    /// The default implementation is appropriate for generating type names to
    /// appear in error messages. It punts to `TypeInner`'s [`core::fmt::Debug`]
    /// implementation, since it's probably best to show the user something they
    /// can act on.
    fn write_non_wgsl_inner<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result {
        write!(out, "{{non-WGSL Naga type {inner:?}}}")
    }

    /// Write a [`Scalar`] that has no representation as WGSL source,
    /// even including Naga extensions.
    ///
    /// A backend might implement this with a call to the [`unreachable!`]
    /// macro, since backends are allowed to assume that the module has passed
    /// validation.
    ///
    /// The default implementation is appropriate for generating type names to
    /// appear in error messages. It punts to `Scalar`'s [`core::fmt::Debug`]
    /// implementation, since it's probably best to show the user something they
    /// can act on.
    fn write_non_wgsl_scalar<W: Write>(&self, scalar: Scalar, out: &mut W) -> core::fmt::Result {
        match scalar.kind {
            crate::ScalarKind::Sint
            | crate::ScalarKind::Uint
            | crate::ScalarKind::Float
            | crate::ScalarKind::Bool => write!(out, "{{non-WGSL Naga scalar {scalar:?}}}"),

            // The abstract types are kind of an odd quasi-WGSL category:
            // they are definitely part of the spec, but they are not expressible
            // in WGSL itself. So we want to call them out by name in error messages,
            // but the WGSL backend should never generate these.
            crate::ScalarKind::AbstractInt => out.write_str("{AbstractInt}"),
            crate::ScalarKind::AbstractFloat => out.write_str("{AbstractFloat}"),
        }
    }

    /// Write the type `ty` as it would appear in a value's declaration.
    ///
    /// Write the type referred to by `ty` in `module` as it would appear in
    /// a `var`, `let`, etc. declaration, or in a function's argument list.
    fn write_type<W: Write>(&self, handle: Handle<crate::Type>, out: &mut W) -> core::fmt::Result {
        let ty = self.lookup_type(handle);
        match ty.inner {
            TypeInner::Struct { .. } => out.write_str(self.type_name(handle))?,
            ref other => self.write_type_inner(other, out)?,
        }

        Ok(())
    }

    /// Write the [`TypeInner`] `inner` as it would appear in a value's declaration.
    ///
    /// Write `inner` as it would appear in a `var`, `let`, etc.
    /// declaration, or in a function's argument list.
    ///
    /// Note that this cannot handle writing [`Struct`] types: those
    /// must be referred to by name, but the name isn't available in
    /// [`TypeInner`].
    ///
    /// [`Struct`]: TypeInner::Struct
    fn write_type_inner<W: Write>(&self, inner: &TypeInner, out: &mut W) -> core::fmt::Result {
        match try_write_type_inner(self, inner, out) {
            Ok(()) => Ok(()),
            Err(WriteTypeError::Format(err)) => Err(err),
            Err(WriteTypeError::NonWgsl) => self.write_non_wgsl_inner(inner, out),
        }
    }

    /// Write the [`Scalar`] `scalar` as a WGSL type.
    fn write_scalar<W: Write>(&self, scalar: Scalar, out: &mut W) -> core::fmt::Result {
        match scalar.try_to_wgsl() {
            Some(string) => out.write_str(string),
            None => self.write_non_wgsl_scalar(scalar, out),
        }
    }

    /// Write the [`TypeResolution`] `resolution` as a WGSL type.
    fn write_type_resolution<W: Write>(
        &self,
        resolution: &TypeResolution,
        out: &mut W,
    ) -> core::fmt::Result {
        match *resolution {
            TypeResolution::Handle(handle) => self.write_type(handle, out),
            TypeResolution::Value(ref inner) => self.write_type_inner(inner, out),
        }
    }

    fn type_to_string(&self, handle: Handle<crate::Type>) -> String {
        let mut buf = String::new();
        self.write_type(handle, &mut buf).unwrap();
        buf
    }

    fn type_inner_to_string(&self, inner: &TypeInner) -> String {
        let mut buf = String::new();
        self.write_type_inner(inner, &mut buf).unwrap();
        buf
    }

    fn type_resolution_to_string(&self, resolution: &TypeResolution) -> String {
        let mut buf = String::new();
        self.write_type_resolution(resolution, &mut buf).unwrap();
        buf
    }
}

fn try_write_type_inner<C, W>(ctx: &C, inner: &TypeInner, out: &mut W) -> Result<(), WriteTypeError>
where
    C: TypeContext + ?Sized,
    W: Write,
{
    match *inner {
        TypeInner::Vector { size, scalar } => {
            write!(out, "vec{}<", common::vector_size_str(size))?;
            ctx.write_scalar(scalar, out)?;
            out.write_str(">")?;
        }
        TypeInner::Sampler { comparison: false } => {
            write!(out, "sampler")?;
        }
        TypeInner::Sampler { comparison: true } => {
            write!(out, "sampler_comparison")?;
        }
        TypeInner::Image {
            dim,
            arrayed,
            class,
        } => {
            // More about texture types: https://gpuweb.github.io/gpuweb/wgsl/#sampled-texture-type
            use crate::ImageClass as Ic;

            let dim_str = dim.to_wgsl();
            let arrayed_str = if arrayed { "_array" } else { "" };
            match class {
                Ic::Sampled { kind, multi } => {
                    let multisampled_str = if multi { "multisampled_" } else { "" };
                    write!(out, "texture_{multisampled_str}{dim_str}{arrayed_str}<")?;
                    ctx.write_scalar(Scalar { kind, width: 4 }, out)?;
                    out.write_str(">")?;
                }
                Ic::Depth { multi } => {
                    let multisampled_str = if multi { "multisampled_" } else { "" };
                    write!(
                        out,
                        "texture_depth_{multisampled_str}{dim_str}{arrayed_str}"
                    )?;
                }
                Ic::Storage { format, access } => {
                    let format_str = format.to_wgsl();
                    let access_str = if access.contains(crate::StorageAccess::ATOMIC) {
                        ",atomic"
                    } else if access
                        .contains(crate::StorageAccess::LOAD | crate::StorageAccess::STORE)
                    {
                        ",read_write"
                    } else if access.contains(crate::StorageAccess::LOAD) {
                        ",read"
                    } else {
                        ",write"
                    };
                    write!(
                        out,
                        "texture_storage_{dim_str}{arrayed_str}<{format_str}{access_str}>"
                    )?;
                }
            }
        }
        TypeInner::Scalar(scalar) => {
            ctx.write_scalar(scalar, out)?;
        }
        TypeInner::Atomic(scalar) => {
            out.write_str("atomic<")?;
            ctx.write_scalar(scalar, out)?;
            out.write_str(">")?;
        }
        TypeInner::Array {
            base,
            size,
            stride: _,
        } => {
            // More info https://gpuweb.github.io/gpuweb/wgsl/#array-types
            // array<A, 3> -- Constant array
            // array<A> -- Dynamic array
            write!(out, "array<")?;
            match size {
                crate::ArraySize::Constant(len) => {
                    ctx.write_type(base, out)?;
                    write!(out, ", {len}")?;
                }
                crate::ArraySize::Pending(r#override) => {
                    ctx.write_override(r#override, out)?;
                }
                crate::ArraySize::Dynamic => {
                    ctx.write_type(base, out)?;
                }
            }
            write!(out, ">")?;
        }
        TypeInner::BindingArray { base, size } => {
            // More info https://github.com/gpuweb/gpuweb/issues/2105
            write!(out, "binding_array<")?;
            match size {
                crate::ArraySize::Constant(len) => {
                    ctx.write_type(base, out)?;
                    write!(out, ", {len}")?;
                }
                crate::ArraySize::Pending(r#override) => {
                    ctx.write_override(r#override, out)?;
                }
                crate::ArraySize::Dynamic => {
                    ctx.write_type(base, out)?;
                }
            }
            write!(out, ">")?;
        }
        TypeInner::Matrix {
            columns,
            rows,
            scalar,
        } => {
            write!(
                out,
                "mat{}x{}<",
                common::vector_size_str(columns),
                common::vector_size_str(rows),
            )?;
            ctx.write_scalar(scalar, out)?;
            out.write_str(">")?;
        }
        TypeInner::Pointer { base, space } => {
            let (address, maybe_access) = address_space_str(space);
            // Everything but `AddressSpace::Handle` gives us a `address` name, but
            // Naga IR never produces pointers to handles, so it doesn't matter much
            // how we write such a type. Just write it as the base type alone.
            if let Some(space) = address {
                write!(out, "ptr<{space}, ")?;
            }
            ctx.write_type(base, out)?;
            if address.is_some() {
                if let Some(access) = maybe_access {
                    write!(out, ", {access}")?;
                }
                write!(out, ">")?;
            }
        }
        TypeInner::ValuePointer {
            size: None,
            scalar,
            space,
        } => {
            let (address, maybe_access) = address_space_str(space);
            if let Some(space) = address {
                write!(out, "ptr<{}, ", space)?;
                ctx.write_scalar(scalar, out)?;
                if let Some(access) = maybe_access {
                    write!(out, ", {access}")?;
                }
                write!(out, ">")?;
            } else {
                return Err(WriteTypeError::NonWgsl);
            }
        }
        TypeInner::ValuePointer {
            size: Some(size),
            scalar,
            space,
        } => {
            let (address, maybe_access) = address_space_str(space);
            if let Some(space) = address {
                write!(out, "ptr<{}, vec{}<", space, common::vector_size_str(size),)?;
                ctx.write_scalar(scalar, out)?;
                out.write_str(">")?;
                if let Some(access) = maybe_access {
                    write!(out, ", {access}")?;
                }
                write!(out, ">")?;
            } else {
                return Err(WriteTypeError::NonWgsl);
            }
            write!(out, ">")?;
        }
        TypeInner::AccelerationStructure { vertex_return } => {
            let caps = if vertex_return { "<vertex_return>" } else { "" };
            write!(out, "acceleration_structure{}", caps)?
        }
        TypeInner::Struct { .. } => {
            unreachable!("structs can only be referenced by name in WGSL");
        }
        TypeInner::RayQuery { vertex_return } => {
            let caps = if vertex_return { "<vertex_return>" } else { "" };
            write!(out, "ray_query{}", caps)?
        }
    }

    Ok(())
}

/// Error type returned by `try_write_type_inner`.
///
/// This type is private to the module.
enum WriteTypeError {
    Format(core::fmt::Error),
    NonWgsl,
}

impl From<core::fmt::Error> for WriteTypeError {
    fn from(err: core::fmt::Error) -> Self {
        Self::Format(err)
    }
}