/*! A representation for highly regular overload sets common in Naga IR.

Many Naga builtin functions' overload sets have a highly regular
structure. For example, many arithmetic functions can be applied to
any floating-point type, or any vector thereof. This module defines a
handful of types for representing such simple overload sets that is
simple and efficient.

*/

use crate::common::{DiagnosticDebug, ForDebugWithTypes};
use crate::ir;
use crate::proc::overloads::constructor_set::{ConstructorSet, ConstructorSize};
use crate::proc::overloads::rule::{Conclusion, Rule};
use crate::proc::overloads::scalar_set::ScalarSet;
use crate::proc::overloads::OverloadSet;
use crate::proc::{GlobalCtx, TypeResolution};
use crate::UniqueArena;

use alloc::vec::Vec;
use core::fmt;

/// Overload sets represented as sets of scalars and constructors.
///
/// This type represents an [`OverloadSet`] using a bitset of scalar
/// types and a bitset of type constructors that might be applied to
/// those scalars. The overload set contains a rule for every possible
/// combination of scalars and constructors, essentially the cartesian
/// product of the two sets.
///
/// For example, if the arity is 2, set of scalars is { AbstractFloat,
/// `f32` }, and the set of constructors is { `vec2`, `vec3` }, then
/// that represents the set of overloads:
///
/// - (`vec2<AbstractFloat>`, `vec2<AbstractFloat>`) -> `vec2<AbstractFloat>`
/// - (`vec2<f32>`, `vec2<f32>`) -> `vec2<f32>`
/// - (`vec3<AbstractFloat>`, `vec3<AbstractFloat>`) -> `vec3<AbstractFloat>`
/// - (`vec3<f32>`, `vec3<f32>`) -> `vec3<f32>`
///
/// The `conclude` value says how to determine the return type from
/// the argument type.
///
/// Restrictions:
///
/// - All overloads must take the same number of arguments.
///
/// - For any given overload, all its arguments must have the same
///   type.
#[derive(Clone)]
pub(in crate::proc::overloads) struct Regular {
    /// The number of arguments in the rules.
    pub arity: usize,

    /// The set of type constructors to apply.
    pub constructors: ConstructorSet,

    /// The set of scalars to apply them to.
    pub scalars: ScalarSet,

    /// How to determine a member rule's return type given the type of
    /// its arguments.
    pub conclude: ConclusionRule,
}

impl Regular {
    pub(in crate::proc::overloads) const EMPTY: Regular = Regular {
        arity: 0,
        constructors: ConstructorSet::empty(),
        scalars: ScalarSet::empty(),
        conclude: ConclusionRule::ArgumentType,
    };

    /// Return an iterator over all the argument types allowed by `self`.
    ///
    /// Return an iterator that produces, for each overload in `self`, the
    /// constructor and scalar of its argument types and return type.
    ///
    /// A [`Regular`] value can only represent overload sets where, in
    /// each overload, all the arguments have the same type, and the
    /// return type is always going to be a determined by the argument
    /// types, so giving the constructor and scalar is sufficient to
    /// characterize the entire rule.
    fn members(&self) -> impl Iterator<Item = (ConstructorSize, ir::Scalar)> {
        let scalars = self.scalars;
        self.constructors.members().flat_map(move |constructor| {
            let size = constructor.size();
            // Technically, we don't need the "most general" `TypeInner` here,
            // but since `ScalarSet::members` only produces singletons anyway,
            // the effect is the same.
            scalars
                .members()
                .map(move |singleton| (size, singleton.most_general_scalar()))
        })
    }

    fn rules(&self) -> impl Iterator<Item = Rule> {
        let arity = self.arity;
        let conclude = self.conclude;
        self.members()
            .map(move |(size, scalar)| make_rule(arity, size, scalar, conclude))
    }
}

impl OverloadSet for Regular {
    fn is_empty(&self) -> bool {
        self.constructors.is_empty() || self.scalars.is_empty()
    }

    fn min_arguments(&self) -> usize {
        assert!(!self.is_empty());
        self.arity
    }

    fn max_arguments(&self) -> usize {
        assert!(!self.is_empty());
        self.arity
    }

    fn arg(&self, i: usize, ty: &ir::TypeInner, types: &UniqueArena<ir::Type>) -> Self {
        if i >= self.arity {
            return Self::EMPTY;
        }

        let constructor = ConstructorSet::singleton(ty);

        let scalars = match ty.scalar_for_conversions(types) {
            Some(ty_scalar) => ScalarSet::convertible_from(ty_scalar),
            None => ScalarSet::empty(),
        };

        Self {
            arity: self.arity,

            // Constrain all member rules' constructors to match `ty`'s.
            constructors: self.constructors & constructor,

            // Constrain all member rules' arguments to be something
            // that `ty` can be converted to.
            scalars: self.scalars & scalars,

            conclude: self.conclude,
        }
    }

    fn concrete_only(self, _types: &UniqueArena<ir::Type>) -> Self {
        Self {
            scalars: self.scalars & ScalarSet::CONCRETE,
            ..self
        }
    }

    fn most_preferred(&self) -> Rule {
        assert!(!self.is_empty());

        // If there is more than one constructor allowed, then we must
        // not have had any arguments supplied at all. In any case, we
        // don't have any unambiguously preferred candidate.
        assert!(self.constructors.is_singleton());

        let size = self.constructors.size();
        let scalar = self.scalars.most_general_scalar();
        make_rule(self.arity, size, scalar, self.conclude)
    }

    fn overload_list(&self, _gctx: &GlobalCtx<'_>) -> Vec<Rule> {
        self.rules().collect()
    }

    fn allowed_args(&self, i: usize, _gctx: &GlobalCtx<'_>) -> Vec<TypeResolution> {
        if i >= self.arity {
            return Vec::new();
        }
        self.members()
            .map(|(size, scalar)| TypeResolution::Value(size.to_inner(scalar)))
            .collect()
    }

    fn for_debug(&self, types: &UniqueArena<ir::Type>) -> impl fmt::Debug {
        DiagnosticDebug((self, types))
    }
}

/// Construct a [`Regular`] member [`Rule`] for the given arity and type.
///
/// [`Regular`] can only represent rules where all the argument types and the
/// return type are the same, so just knowing `arity` and `inner` is sufficient.
///
/// [`Rule`]: crate::proc::overloads::Rule
fn make_rule(
    arity: usize,
    size: ConstructorSize,
    scalar: ir::Scalar,
    conclusion_rule: ConclusionRule,
) -> Rule {
    let inner = size.to_inner(scalar);
    let arg = TypeResolution::Value(inner.clone());
    Rule {
        arguments: core::iter::repeat_n(arg.clone(), arity).collect(),
        conclusion: conclusion_rule.conclude(size, scalar),
    }
}

/// Conclusion-computing rules.
#[derive(Clone, Copy, Debug)]
#[repr(u8)]
pub(in crate::proc::overloads) enum ConclusionRule {
    ArgumentType,
    Scalar,
    Frexp,
    Modf,
    U32,
    I32,
    Vec2F,
    Vec4F,
    Vec4I,
    Vec4U,
}

impl ConclusionRule {
    fn conclude(self, size: ConstructorSize, scalar: ir::Scalar) -> Conclusion {
        match self {
            Self::ArgumentType => Conclusion::Value(size.to_inner(scalar)),
            Self::Scalar => Conclusion::Value(ir::TypeInner::Scalar(scalar)),
            Self::Frexp => Conclusion::for_frexp_modf(ir::MathFunction::Frexp, size, scalar),
            Self::Modf => Conclusion::for_frexp_modf(ir::MathFunction::Modf, size, scalar),
            Self::U32 => Conclusion::Value(ir::TypeInner::Scalar(ir::Scalar::U32)),
            Self::I32 => Conclusion::Value(ir::TypeInner::Scalar(ir::Scalar::I32)),
            Self::Vec2F => Conclusion::Value(ir::TypeInner::Vector {
                size: ir::VectorSize::Bi,
                scalar: ir::Scalar::F32,
            }),
            Self::Vec4F => Conclusion::Value(ir::TypeInner::Vector {
                size: ir::VectorSize::Quad,
                scalar: ir::Scalar::F32,
            }),
            Self::Vec4I => Conclusion::Value(ir::TypeInner::Vector {
                size: ir::VectorSize::Quad,
                scalar: ir::Scalar::I32,
            }),
            Self::Vec4U => Conclusion::Value(ir::TypeInner::Vector {
                size: ir::VectorSize::Quad,
                scalar: ir::Scalar::U32,
            }),
        }
    }
}

impl fmt::Debug for DiagnosticDebug<(&Regular, &UniqueArena<ir::Type>)> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (regular, types) = self.0;
        let rules: Vec<Rule> = regular.rules().collect();
        f.debug_struct("List")
            .field("rules", &rules.for_debug(types))
            .field("conclude", &regular.conclude)
            .finish()
    }
}

impl ForDebugWithTypes for &Regular {}

impl fmt::Debug for DiagnosticDebug<(&[Rule], &UniqueArena<ir::Type>)> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let (rules, types) = self.0;
        f.debug_list()
            .entries(rules.iter().map(|rule| rule.for_debug(types)))
            .finish()
    }
}

impl ForDebugWithTypes for &[Rule] {}

/// Construct a [`Regular`] [`OverloadSet`].
///
/// Examples:
///
/// - `regular!(2, SCALAR|VECN of FLOAT)`: An overload set whose rules take two
///   arguments of the same type: a floating-point scalar (possibly abstract) or
///   a vector of such. The return type is the same as the argument type.
///
/// - `regular!(1, VECN of FLOAT -> Scalar)`: An overload set whose rules take
///   one argument that is a vector of floats, and whose return type is the leaf
///   scalar type of the argument type.
///
/// The constructor values (before the `<` angle brackets `>`) are
/// constants from [`ConstructorSet`].
///
/// The scalar values (inside the `<` angle brackets `>`) are
/// constants from [`ScalarSet`].
///
/// When a return type identifier is given, it is treated as a variant
/// of the the [`ConclusionRule`] enum.
macro_rules! regular {
    // regular!(ARITY, CONSTRUCTOR of SCALAR)
    ( $arity:literal , $( $constr:ident )|* of $( $scalar:ident )|*) => {
        {
            use $crate::proc::overloads;
            use overloads::constructor_set::constructor_set;
            use overloads::regular::{Regular, ConclusionRule};
            use overloads::scalar_set::scalar_set;
            Regular {
                arity: $arity,
                constructors: constructor_set!( $( $constr )|* ),
                scalars: scalar_set!( $( $scalar )|* ),
                conclude: ConclusionRule::ArgumentType,
            }
        }
    };

    // regular!(ARITY, CONSTRUCTOR of SCALAR -> CONCLUSION_RULE)
    ( $arity:literal , $( $constr:ident )|* of $( $scalar:ident )|* -> $conclude:ident) => {
        {
            use $crate::proc::overloads;
            use overloads::constructor_set::constructor_set;
            use overloads::regular::{Regular, ConclusionRule};
            use overloads::scalar_set::scalar_set;
            Regular {
                arity: $arity,
                constructors:constructor_set!( $( $constr )|* ),
                scalars: scalar_set!( $( $scalar )|* ),
                conclude: ConclusionRule::$conclude,
            }
        }
    };
}

pub(in crate::proc::overloads) use regular;

#[cfg(test)]
mod test {
    use super::*;
    use crate::ir;

    const fn scalar(scalar: ir::Scalar) -> ir::TypeInner {
        ir::TypeInner::Scalar(scalar)
    }

    const fn vec2(scalar: ir::Scalar) -> ir::TypeInner {
        ir::TypeInner::Vector {
            scalar,
            size: ir::VectorSize::Bi,
        }
    }

    const fn vec3(scalar: ir::Scalar) -> ir::TypeInner {
        ir::TypeInner::Vector {
            scalar,
            size: ir::VectorSize::Tri,
        }
    }

    /// Assert that `set` has a most preferred candidate whose type
    /// conclusion is `expected`.
    #[track_caller]
    fn check_return_type(set: &Regular, expected: &ir::TypeInner, arena: &UniqueArena<ir::Type>) {
        assert!(!set.is_empty());

        let special_types = ir::SpecialTypes::default();

        let preferred = set.most_preferred();
        let conclusion = preferred.conclusion;
        let resolution = conclusion
            .into_resolution(&special_types)
            .expect("special types should have been pre-registered");
        let inner = resolution.inner_with(arena);

        assert!(
            inner.non_struct_equivalent(expected, arena),
            "Expected {:?}, got {:?}",
            expected.for_debug(arena),
            inner.for_debug(arena),
        );
    }

    #[test]
    fn unary_vec_or_scalar_numeric_scalar() {
        let arena = UniqueArena::default();

        let builtin = regular!(1, SCALAR of NUMERIC);

        let ok = builtin.arg(0, &scalar(ir::Scalar::U32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::U32), &arena);

        let err = builtin.arg(0, &scalar(ir::Scalar::BOOL), &arena);
        assert!(err.is_empty());
    }

    #[test]
    fn unary_vec_or_scalar_numeric_vector() {
        let arena = UniqueArena::default();

        let builtin = regular!(1, VECN|SCALAR of NUMERIC);

        let ok = builtin.arg(0, &vec3(ir::Scalar::F64), &arena);
        check_return_type(&ok, &vec3(ir::Scalar::F64), &arena);

        let err = builtin.arg(0, &vec3(ir::Scalar::BOOL), &arena);
        assert!(err.is_empty());
    }

    #[test]
    fn unary_vec_or_scalar_numeric_matrix() {
        let arena = UniqueArena::default();

        let builtin = regular!(1, VECN|SCALAR of NUMERIC);

        let err = builtin.arg(
            0,
            &ir::TypeInner::Matrix {
                columns: ir::VectorSize::Tri,
                rows: ir::VectorSize::Tri,
                scalar: ir::Scalar::F32,
            },
            &arena,
        );
        assert!(err.is_empty());
    }

    #[test]
    #[rustfmt::skip]
    fn binary_vec_or_scalar_numeric_scalar() {
        let arena = UniqueArena::default();

        let builtin = regular!(2, VECN|SCALAR of NUMERIC);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::F32), &arena)
            .arg(1, &scalar(ir::Scalar::F32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::F32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::ABSTRACT_FLOAT), &arena)
            .arg(1, &scalar(ir::Scalar::F32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::F32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::F32), &arena)
            .arg(1, &scalar(ir::Scalar::ABSTRACT_INT), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::F32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::U32), &arena)
            .arg(1, &scalar(ir::Scalar::U32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::U32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::U32), &arena)
            .arg(1, &scalar(ir::Scalar::ABSTRACT_INT), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::U32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::ABSTRACT_INT), &arena)
            .arg(1, &scalar(ir::Scalar::U32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::U32), &arena);

        // Not numeric.
        let err = builtin
            .arg(0, &scalar(ir::Scalar::BOOL), &arena)
            .arg(1, &scalar(ir::Scalar::BOOL), &arena);
        assert!(err.is_empty());

        // Different floating-point types.
        let err = builtin
            .arg(0, &scalar(ir::Scalar::F32), &arena)
            .arg(1, &scalar(ir::Scalar::F64), &arena);
        assert!(err.is_empty());

        // Different constructor.
        let err = builtin
            .arg(0, &scalar(ir::Scalar::F32), &arena)
            .arg(1, &vec2(ir::Scalar::F32), &arena);
        assert!(err.is_empty());

        // Different vector size
        let err = builtin
            .arg(0, &vec2(ir::Scalar::F32), &arena)
            .arg(1, &vec3(ir::Scalar::F32), &arena);
        assert!(err.is_empty());
    }

    #[test]
    #[rustfmt::skip]
    fn binary_vec_or_scalar_numeric_vector() {
        let arena = UniqueArena::default();

        let builtin = regular!(2, VECN|SCALAR of NUMERIC);

        let ok = builtin
            .arg(0, &vec3(ir::Scalar::F32), &arena)
            .arg(1, &vec3(ir::Scalar::F32), &arena);
        check_return_type(&ok, &vec3(ir::Scalar::F32), &arena);

        // Different vector sizes.
        let err = builtin
            .arg(0, &vec2(ir::Scalar::F32), &arena)
            .arg(1, &vec3(ir::Scalar::F32), &arena);
        assert!(err.is_empty());

        // Different vector scalars.
        let err = builtin
            .arg(0, &vec3(ir::Scalar::F32), &arena)
            .arg(1, &vec3(ir::Scalar::F64), &arena);
        assert!(err.is_empty());

        // Mix of vectors and scalars.
        let err = builtin
            .arg(0, &scalar(ir::Scalar::F32), &arena)
            .arg(1, &vec3(ir::Scalar::F32), &arena);
        assert!(err.is_empty());
    }

    #[test]
    #[rustfmt::skip]
    fn binary_vec_or_scalar_numeric_vector_abstract() {
        let arena = UniqueArena::default();

        let builtin = regular!(2, VECN|SCALAR of NUMERIC);

        let ok = builtin
            .arg(0, &vec2(ir::Scalar::ABSTRACT_INT), &arena)
            .arg(1, &vec2(ir::Scalar::U32), &arena);
        check_return_type(&ok, &vec2(ir::Scalar::U32), &arena);

        let ok = builtin
            .arg(0, &vec3(ir::Scalar::ABSTRACT_INT), &arena)
            .arg(1, &vec3(ir::Scalar::F32), &arena);
        check_return_type(&ok, &vec3(ir::Scalar::F32), &arena);

        let ok = builtin
            .arg(0, &scalar(ir::Scalar::ABSTRACT_FLOAT), &arena)
            .arg(1, &scalar(ir::Scalar::F32), &arena);
        check_return_type(&ok, &scalar(ir::Scalar::F32), &arena);

        let err = builtin
            .arg(0, &scalar(ir::Scalar::ABSTRACT_FLOAT), &arena)
            .arg(1, &scalar(ir::Scalar::U32), &arena);
        assert!(err.is_empty());

        let err = builtin
            .arg(0, &scalar(ir::Scalar::I32), &arena)
            .arg(1, &scalar(ir::Scalar::U32), &arena);
        assert!(err.is_empty());
    }
}