zirium 0.1.0

//! Shared JSON registry configuration for library callers and the CLI.

use std::{
    collections::{BTreeMap, BTreeSet},
    fmt, fs, io,
    path::{Path, PathBuf},
};

use super::{DeclarativeRegistryError, DialectRegistry, OperationShape};

/// A complete registry: selected built-ins plus caller-named operation shapes.
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RegistryConfig {
    #[serde(default)]
    pub presets: Vec<String>,
    pub builtins: Vec<String>,
    pub operation_shapes: Vec<OperationShapeConfig>,
    #[serde(default)]
    pub operation_formats: Vec<OperationFormatConfig>,
    #[serde(default)]
    pub operation_alternatives: Vec<OperationAlternativesConfig>,
}

/// One named operation using an existing custom grammar.
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OperationShapeConfig {
    pub name: String,
    pub shape: OperationShape,
    #[serde(default)]
    pub callee_attribute: Option<String>,
}

/// One named operation using a validated format description.
#[derive(Clone, Debug, serde::Deserialize)]
#[serde(deny_unknown_fields)]
pub struct OperationFormatConfig {
    pub name: String,
    pub format: String,
}

#[derive(Clone, Debug, serde::Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct OperationAlternativesConfig {
    pub name: String,
    pub alternatives: Vec<OperationGrammarConfig>,
}

#[derive(Clone, Debug, serde::Deserialize, PartialEq, Eq)]
#[serde(deny_unknown_fields)]
pub struct OperationGrammarConfig {
    #[serde(default)]
    pub shape: Option<OperationShape>,
    #[serde(default)]
    pub format: Option<String>,
}

type ExpandedRegistry = (
    BTreeSet<String>,
    BTreeMap<String, OperationShape>,
    BTreeMap<String, String>,
    BTreeMap<String, Vec<OperationGrammarConfig>>,
    BTreeMap<String, String>,
);

const PRESET_NAMES: &[&str] = &[
    "stablehlo",
    "tosa",
    "scf",
    "linalg",
    "acc",
    "affine",
    "amdgpu",
    "amx",
    "arith",
    "arm_neon",
    "arm_sme",
    "arm_sve",
    "async",
    "bufferization",
    "cf",
    "complex",
    "dlti",
    "emitc",
    "func",
    "gpu",
    "index",
    "irdl",
    "llvm",
    "math",
    "memref",
    "ml_program",
    "mpi",
    "shard",
    "nvgpu",
    "nvvm",
    "omp",
    "pdl",
    "pdl_interp",
    "ptr",
    "quant",
    "rocdl",
    "shape",
    "sparse_tensor",
    "smt",
    "spirv",
    "tensor",
    "transform",
    "ub",
    "vector",
    "wasmssa",
    "x86vector",
    "xegpu",
    "xevm",
];

fn preset_json(name: &str) -> Option<&'static str> {
    match name {
        "stablehlo" => Some(include_str!("../../registries/stablehlo.json")),
        "tosa" => Some(include_str!("../../registries/tosa.json")),
        "scf" => Some(include_str!("../../registries/scf.json")),
        "linalg" => Some(include_str!("../../registries/linalg.json")),
        "acc" => Some(include_str!("../../registries/acc.json")),
        "affine" => Some(include_str!("../../registries/affine.json")),
        "amdgpu" => Some(include_str!("../../registries/amdgpu.json")),
        "amx" => Some(include_str!("../../registries/amx.json")),
        "arith" => Some(include_str!("../../registries/arith.json")),
        "arm_neon" => Some(include_str!("../../registries/arm_neon.json")),
        "arm_sme" => Some(include_str!("../../registries/arm_sme.json")),
        "arm_sve" => Some(include_str!("../../registries/arm_sve.json")),
        "async" => Some(include_str!("../../registries/async.json")),
        "bufferization" => Some(include_str!("../../registries/bufferization.json")),
        "cf" => Some(include_str!("../../registries/cf.json")),
        "complex" => Some(include_str!("../../registries/complex.json")),
        "dlti" => Some(include_str!("../../registries/dlti.json")),
        "emitc" => Some(include_str!("../../registries/emitc.json")),
        "func" => Some(include_str!("../../registries/func.json")),
        "gpu" => Some(include_str!("../../registries/gpu.json")),
        "index" => Some(include_str!("../../registries/index.json")),
        "irdl" => Some(include_str!("../../registries/irdl.json")),
        "llvm" => Some(include_str!("../../registries/llvm.json")),
        "math" => Some(include_str!("../../registries/math.json")),
        "memref" => Some(include_str!("../../registries/memref.json")),
        "ml_program" => Some(include_str!("../../registries/ml_program.json")),
        "mpi" => Some(include_str!("../../registries/mpi.json")),
        "shard" => Some(include_str!("../../registries/shard.json")),
        "nvgpu" => Some(include_str!("../../registries/nvgpu.json")),
        "nvvm" => Some(include_str!("../../registries/nvvm.json")),
        "omp" => Some(include_str!("../../registries/omp.json")),
        "pdl" => Some(include_str!("../../registries/pdl.json")),
        "pdl_interp" => Some(include_str!("../../registries/pdl_interp.json")),
        "ptr" => Some(include_str!("../../registries/ptr.json")),
        "quant" => Some(include_str!("../../registries/quant.json")),
        "rocdl" => Some(include_str!("../../registries/rocdl.json")),
        "shape" => Some(include_str!("../../registries/shape.json")),
        "sparse_tensor" => Some(include_str!("../../registries/sparse_tensor.json")),
        "smt" => Some(include_str!("../../registries/smt.json")),
        "spirv" => Some(include_str!("../../registries/spirv.json")),
        "tensor" => Some(include_str!("../../registries/tensor.json")),
        "transform" => Some(include_str!("../../registries/transform.json")),
        "ub" => Some(include_str!("../../registries/ub.json")),
        "vector" => Some(include_str!("../../registries/vector.json")),
        "wasmssa" => Some(include_str!("../../registries/wasmssa.json")),
        "x86vector" => Some(include_str!("../../registries/x86vector.json")),
        "xegpu" => Some(include_str!("../../registries/xegpu.json")),
        "xevm" => Some(include_str!("../../registries/xevm.json")),
        _ => None,
    }
}

impl RegistryConfig {
    /// Reads JSON, rejecting missing/unknown fields and unsupported shapes.
    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
        // Deserialize first so duplicate fields and normal schema errors retain
        // their source locations. Serde structs also accept arrays; this file
        // format deliberately requires JSON objects for its records.
        let config = serde_json::from_str(json)?;
        let value: serde_json::Value = serde_json::from_str(json)?;
        if !value.is_object()
            || value["operation_shapes"]
                .as_array()
                .is_some_and(|entries| entries.iter().any(|entry| !entry.is_object()))
            || value["operation_formats"]
                .as_array()
                .is_some_and(|entries| entries.iter().any(|entry| !entry.is_object()))
            || value["operation_alternatives"]
                .as_array()
                .is_some_and(|entries| {
                    entries.iter().any(|entry| {
                        !entry.is_object()
                            || entry["alternatives"].as_array().is_some_and(|grammars| {
                                grammars.iter().any(|grammar| !grammar.is_object())
                            })
                    })
                })
        {
            return Err(<serde_json::Error as serde::de::Error>::custom(
                "registry and operation records must be JSON objects",
            ));
        }
        Ok(config)
    }

    /// Combines complete configurations. Identical entries across configurations
    /// are shared; duplicate entries within a configuration still fail validation.
    pub fn build_many(configs: &[Self]) -> Result<DialectRegistry, DeclarativeRegistryError> {
        let mut builtins = BTreeSet::new();
        let mut shapes = BTreeMap::new();
        let mut formats = BTreeMap::new();
        let mut alternatives = BTreeMap::new();
        let mut call_target_attributes = BTreeMap::new();
        for config in configs {
            let (
                config_builtins,
                config_shapes,
                config_formats,
                config_alternatives,
                config_call_target_attributes,
            ) = config.expanded()?;
            builtins.extend(config_builtins);
            for (name, shape) in config_shapes {
                if let Some(previous) = shapes.insert(name.clone(), shape)
                    && previous != shape
                {
                    return Err(DeclarativeRegistryError::ConflictingShape(name));
                }
            }
            for (name, format) in config_formats {
                if let Some(previous) = formats.insert(name.clone(), format.clone())
                    && previous != format
                {
                    return Err(DeclarativeRegistryError::ConflictingFormat(name));
                }
            }
            for (name, grammars) in config_alternatives {
                if let Some(previous) = alternatives.insert(name.clone(), grammars.clone())
                    && previous != grammars
                {
                    return Err(DeclarativeRegistryError::ConflictingAlternatives(name));
                }
            }
            for (name, attribute) in config_call_target_attributes {
                if let Some(previous) =
                    call_target_attributes.insert(name.clone(), attribute.clone())
                    && previous != attribute
                {
                    return Err(DeclarativeRegistryError::ConflictingCallTargetAttribute(
                        name,
                    ));
                }
            }
        }
        Self::build_entries(
            builtins,
            shapes,
            formats,
            alternatives,
            call_target_attributes,
        )
    }

    /// Validates all registrations and constructs an owned registry.
    pub fn build(&self) -> Result<DialectRegistry, DeclarativeRegistryError> {
        let (builtins, shapes, formats, alternatives, call_target_attributes) = self.expanded()?;
        Self::build_entries(
            builtins,
            shapes,
            formats,
            alternatives,
            call_target_attributes,
        )
    }

    fn expanded(&self) -> Result<ExpandedRegistry, DeclarativeRegistryError> {
        let mut seen_presets = BTreeSet::new();
        let mut builtins = BTreeSet::new();
        let mut shapes = BTreeMap::new();
        let mut formats = BTreeMap::new();
        let mut alternatives = BTreeMap::new();
        let mut call_target_attributes = BTreeMap::new();
        for name in &self.presets {
            if !seen_presets.insert(name.as_str()) {
                return Err(DeclarativeRegistryError::DuplicatePreset(name.clone()));
            }
            let json = preset_json(name)
                .ok_or_else(|| DeclarativeRegistryError::UnknownPreset(name.clone()))?;
            let preset = Self::from_json(json).expect("bundled registry preset must be valid");
            let (
                preset_builtins,
                preset_shapes,
                preset_formats,
                preset_alternatives,
                preset_call_target_attributes,
            ) = preset.expanded()?;
            builtins.extend(preset_builtins);
            for (name, shape) in preset_shapes {
                if let Some(previous) = shapes.insert(name.clone(), shape)
                    && previous != shape
                {
                    return Err(DeclarativeRegistryError::ConflictingShape(name));
                }
            }
            for (name, format) in preset_formats {
                formats.insert(name, format);
            }
            alternatives.extend(preset_alternatives);
            for (name, attribute) in preset_call_target_attributes {
                if let Some(previous) =
                    call_target_attributes.insert(name.clone(), attribute.clone())
                    && previous != attribute
                {
                    return Err(DeclarativeRegistryError::ConflictingCallTargetAttribute(
                        name,
                    ));
                }
            }
        }

        let mut explicit_builtins = BTreeSet::new();
        for name in &self.builtins {
            if !explicit_builtins.insert(name.as_str()) {
                return Err(DeclarativeRegistryError::DuplicateOperation(name.clone()));
            }
            builtins.insert(name.clone());
        }
        let mut explicit_shapes = BTreeSet::new();
        for operation in &self.operation_shapes {
            if !explicit_shapes.insert(operation.name.as_str()) {
                return Err(DeclarativeRegistryError::DuplicateOperation(
                    operation.name.clone(),
                ));
            }
            if let Some(previous) = shapes.insert(operation.name.clone(), operation.shape)
                && previous != operation.shape
            {
                return Err(DeclarativeRegistryError::ConflictingShape(
                    operation.name.clone(),
                ));
            }
            if let Some(attribute) = operation.callee_attribute.as_deref()
                && (operation.shape != OperationShape::CallLike || !valid_attribute_name(attribute))
            {
                return Err(DeclarativeRegistryError::InvalidCallTargetAttribute(
                    operation.name.clone(),
                ));
            }
            if operation.shape == OperationShape::CallLike {
                let attribute = operation.callee_attribute.as_deref().unwrap_or("callee");
                if let Some(previous) =
                    call_target_attributes.insert(operation.name.clone(), attribute.to_owned())
                    && previous != attribute
                {
                    return Err(DeclarativeRegistryError::ConflictingCallTargetAttribute(
                        operation.name.clone(),
                    ));
                }
            }
        }
        let mut explicit_formats = BTreeSet::new();
        for operation in &self.operation_formats {
            if !explicit_formats.insert(operation.name.as_str())
                || explicit_shapes.contains(operation.name.as_str())
            {
                return Err(DeclarativeRegistryError::DuplicateOperation(
                    operation.name.clone(),
                ));
            }
            if let Some(previous) = formats.insert(operation.name.clone(), operation.format.clone())
                && previous != operation.format
            {
                return Err(DeclarativeRegistryError::ConflictingFormat(
                    operation.name.clone(),
                ));
            }
        }
        let mut explicit_alternatives = BTreeSet::new();
        for operation in &self.operation_alternatives {
            if !explicit_alternatives.insert(operation.name.as_str())
                || explicit_shapes.contains(operation.name.as_str())
                || explicit_formats.contains(operation.name.as_str())
            {
                return Err(DeclarativeRegistryError::DuplicateOperation(
                    operation.name.clone(),
                ));
            }
            if operation.alternatives.len() < 2
                || operation.alternatives.iter().any(|grammar| {
                    matches!(
                        (&grammar.shape, &grammar.format),
                        (None, None) | (Some(_), Some(_))
                    )
                })
            {
                return Err(DeclarativeRegistryError::InvalidOperationAlternatives(
                    operation.name.clone(),
                ));
            }
            if let Some(previous) =
                alternatives.insert(operation.name.clone(), operation.alternatives.clone())
                && previous != operation.alternatives
            {
                return Err(DeclarativeRegistryError::ConflictingAlternatives(
                    operation.name.clone(),
                ));
            }
        }
        Ok((
            builtins,
            shapes,
            formats,
            alternatives,
            call_target_attributes,
        ))
    }

    fn build_entries(
        builtins: BTreeSet<String>,
        shapes: BTreeMap<String, OperationShape>,
        formats: BTreeMap<String, String>,
        alternatives: BTreeMap<String, Vec<OperationGrammarConfig>>,
        call_target_attributes: BTreeMap<String, String>,
    ) -> Result<DialectRegistry, DeclarativeRegistryError> {
        let builtins = builtins.iter().map(String::as_str).collect::<Vec<_>>();
        let shapes = shapes
            .iter()
            .map(|(name, shape)| (name.as_str(), *shape))
            .collect::<Vec<_>>();
        let formats = formats
            .iter()
            .map(|(name, format)| (name.as_str(), format.as_str()))
            .collect::<Vec<_>>();
        DialectRegistry::declarative(&builtins)?
            .extend_operation_shapes(&shapes)?
            .extend_operation_formats(&formats)
            .and_then(|registry| {
                registry
                    .extend_operation_alternatives(&alternatives.into_iter().collect::<Vec<_>>())
            })
            .map(|registry| {
                registry.with_call_target_attributes(call_target_attributes.into_iter().collect())
            })
    }
}

fn valid_attribute_name(name: &str) -> bool {
    !name.is_empty()
        && name.split('.').all(|component| {
            let mut chars = component.chars();
            chars
                .next()
                .is_some_and(|ch| ch.is_ascii_alphabetic() || ch == '_')
                && chars.all(|ch| ch.is_ascii_alphanumeric() || ch == '_')
        })
}

/// Reading, deserializing, or constructing a registry failed.
#[derive(Debug)]
pub enum RegistryConfigError {
    Io {
        path: PathBuf,
        error: io::Error,
    },
    Json {
        path: PathBuf,
        error: serde_json::Error,
    },
    Registry(DeclarativeRegistryError),
}

impl fmt::Display for RegistryConfigError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Io { path, error } => write!(f, "{}: {error}", path.display()),
            Self::Json { path, error } => write!(f, "{}: {error}", path.display()),
            Self::Registry(error) => error.fmt(f),
        }
    }
}

impl std::error::Error for RegistryConfigError {
    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
        Some(match self {
            Self::Io { error, .. } => error,
            Self::Json { error, .. } => error,
            Self::Registry(error) => error,
        })
    }
}

impl DialectRegistry {
    /// Names of the registry presets bundled with this Zirium release.
    pub const fn preset_names() -> &'static [&'static str] {
        PRESET_NAMES
    }

    /// Builds one registry preset bundled with this Zirium release.
    pub fn from_name(name: &str) -> Result<Self, DeclarativeRegistryError> {
        RegistryConfig {
            presets: vec![name.to_owned()],
            builtins: Vec::new(),
            operation_shapes: Vec::new(),
            operation_formats: Vec::new(),
            operation_alternatives: Vec::new(),
        }
        .build()
    }

    /// Loads a complete registry from a UTF-8 JSON file.
    ///
    /// No default operations are implicitly added. I/O, JSON, and registration
    /// failures are distinguished in the returned error.
    pub fn from_config_file(path: impl AsRef<Path>) -> Result<Self, RegistryConfigError> {
        Self::from_config_files([path])
    }

    /// Combines JSON registry files, rejecting conflicting definitions.
    pub fn from_config_files<P: AsRef<Path>>(
        paths: impl IntoIterator<Item = P>,
    ) -> Result<Self, RegistryConfigError> {
        let mut configs = Vec::new();
        for path in paths {
            let path = path.as_ref();
            let json = fs::read_to_string(path).map_err(|error| RegistryConfigError::Io {
                path: path.to_owned(),
                error,
            })?;
            configs.push(RegistryConfig::from_json(&json).map_err(|error| {
                RegistryConfigError::Json {
                    path: path.to_owned(),
                    error,
                }
            })?);
        }
        RegistryConfig::build_many(&configs).map_err(RegistryConfigError::Registry)
    }
}