naga 29.0.1

use alloc::{boxed::Box, vec::Vec};

use crate::{
    front::wgsl::{
        error::Error,
        lower::{ExpressionContext, Lowerer, Result},
        parse::{ast, conv},
    },
    ir, Handle, Span,
};

/// Iterator over a template list.
///
/// All functions will attempt to consume an element in the list.
///
/// Function variants prefixed with "maybe" will not return an error if there
/// are no more elements left in the list.
pub struct TemplateListIter<'iter, 'source> {
    ident_span: Span,
    template_list: core::slice::Iter<'iter, Handle<ast::Expression<'source>>>,
}

impl<'iter, 'source> TemplateListIter<'iter, 'source> {
    pub fn new(ident_span: Span, template_list: &'iter [Handle<ast::Expression<'source>>]) -> Self {
        Self {
            ident_span,
            template_list: template_list.iter(),
        }
    }

    pub fn finish(self, ctx: &ExpressionContext<'source, '_, '_>) -> Result<'source, ()> {
        let unused_args: Vec<Span> = self
            .template_list
            .map(|expr| ctx.ast_expressions.get_span(*expr))
            .collect();
        if unused_args.is_empty() {
            Ok(())
        } else {
            Err(Box::new(Error::UnusedArgsForTemplate(unused_args)))
        }
    }

    fn expect_next(
        &mut self,
        description: &'static str,
    ) -> Result<'source, Handle<ast::Expression<'source>>> {
        if let Some(expr) = self.template_list.next() {
            Ok(*expr)
        } else {
            Err(Box::new(Error::MissingTemplateArg {
                span: self.ident_span,
                description,
            }))
        }
    }

    pub fn ty(
        &mut self,
        lowerer: &mut Lowerer<'source, '_>,
        ctx: &mut ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, Handle<ir::Type>> {
        let expr = self.expect_next("`T`, a type")?;
        lowerer.type_expression(expr, ctx)
    }

    /// Lower the next template list element as a type, and return its span.
    ///
    /// This returns the span of the template list element. This is generally
    /// different from the span of the returned `Handle<ir::Type>`, as the
    /// latter may refer to the type's definition, not its use in the template list.
    pub fn ty_with_span(
        &mut self,
        lowerer: &mut Lowerer<'source, '_>,
        ctx: &mut ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, (Handle<ir::Type>, Span)> {
        let expr = self.expect_next("`T`, a type")?;
        let span = ctx.ast_expressions.get_span(expr);
        let ty = lowerer.type_expression(expr, ctx)?;
        Ok((ty, span))
    }

    pub fn scalar_ty(
        &mut self,
        lowerer: &mut Lowerer<'source, '_>,
        ctx: &mut ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, (ir::Scalar, Span)> {
        let expr = self.expect_next("`T`, a scalar type")?;
        let ty = lowerer.type_expression(expr, ctx)?;
        let span = ctx.ast_expressions.get_span(expr);
        match ctx.module.types[ty].inner {
            ir::TypeInner::Scalar(scalar) => Ok((scalar, span)),
            _ => Err(Box::new(Error::UnknownScalarType(span))),
        }
    }

    pub fn maybe_array_size(
        &mut self,
        lowerer: &mut Lowerer<'source, '_>,
        ctx: &mut ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, ir::ArraySize> {
        if let Some(expr) = self.template_list.next() {
            lowerer.array_size(*expr, ctx)
        } else {
            Ok(ir::ArraySize::Dynamic)
        }
    }

    pub fn address_space(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, ir::AddressSpace> {
        let expr = self.expect_next("`AS`, an address space")?;
        let (enumerant, span) = ctx.enumerant(expr)?;
        conv::map_address_space(enumerant, span, &ctx.enable_extensions)
    }
    pub fn maybe_address_space(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, Option<ir::AddressSpace>> {
        let Some(expr) = self.template_list.next() else {
            return Ok(None);
        };

        let (enumerant, span) = ctx.enumerant(*expr)?;
        Ok(Some(conv::map_address_space(
            enumerant,
            span,
            &ctx.enable_extensions,
        )?))
    }

    pub fn access_mode(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, ir::StorageAccess> {
        let expr = self.expect_next("`Access`, an access mode")?;
        let (enumerant, span) = ctx.enumerant(expr)?;
        conv::map_access_mode(enumerant, span)
    }

    /// Update `space` with an access mode from `self`, if appropriate.
    ///
    /// If `space` is [`Storage`], and there is another template parameter in
    /// `self`, parse it as a storage mode and mutate `space` accordingly. If
    /// there are no more template parameters, treat that like `read`.
    ///
    /// If `space` is some other address space, don't consume any template
    /// parameters.
    ///
    /// [`Storage`]: ir::AddressSpace::Storage
    pub fn maybe_access_mode(
        &mut self,
        space: &mut ir::AddressSpace,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, ()> {
        if let &mut ir::AddressSpace::Storage { ref mut access } = space {
            if let Some(expr) = self.template_list.next() {
                let (enumerant, span) = ctx.enumerant(*expr)?;
                let access_mode = conv::map_access_mode(enumerant, span)?;
                *access = access_mode;
            } else {
                // defaulting to `read`
                *access = ir::StorageAccess::LOAD
            }
        }
        Ok(())
    }

    pub fn storage_format(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, ir::StorageFormat> {
        let expr = self.expect_next("`Format`, a texel format")?;
        let (enumerant, span) = ctx.enumerant(expr)?;
        conv::map_storage_format(enumerant, span)
    }

    pub fn maybe_vertex_return(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, bool> {
        let Some(expr) = self.template_list.next() else {
            return Ok(false);
        };

        let (enumerant, span) = ctx.enumerant(*expr)?;
        conv::map_ray_flag(&ctx.enable_extensions, enumerant, span)?;
        Ok(true)
    }

    pub fn cooperative_role(
        &mut self,
        ctx: &ExpressionContext<'source, '_, '_>,
    ) -> Result<'source, crate::CooperativeRole> {
        let role_expr = self.expect_next("`Role`, a cooperative matrix role")?;
        let (enumerant, span) = ctx.enumerant(role_expr)?;
        let role = conv::map_cooperative_role(enumerant, span)?;
        Ok(role)
    }
}