onnx-runtime-shape-inference 0.1.0-dev.6

//! The extensible, opset-aware operator registry.
//!
//! Inference rules are keyed by `(domain, op_type)` and, within a key, by the
//! opset version at which they were introduced. Registration is *range-based*:
//! a rule registered at version `N` applies to every opset `>= N` until a later
//! registration supersedes it — mirroring how ONNX operator schemas evolve.
//! Unregistered operators are not an error; their outputs are simply left
//! unresolved (permissive behaviour).

use std::collections::HashMap;

use onnx_runtime_ir::{Node, normalize_domain};

use crate::context::{InferenceContext, MergePolicy, NodeIo, SymbolInterner};
use crate::error::ShapeInferError;

/// An operator inference rule: reads inputs from the [`InferenceContext`] and
/// sets its outputs' types (and, where applicable, shape-data).
pub type InferenceFn = fn(&mut InferenceContext) -> Result<(), ShapeInferError>;

/// A registry mapping `(domain, op_type, opset)` to an [`InferenceFn`].
#[derive(Default)]
pub struct InferenceRegistry {
    /// `(domain, op)` → ascending list of `(min_opset, rule)`.
    handlers: HashMap<(String, String), Vec<(u64, InferenceFn)>>,
}

impl InferenceRegistry {
    /// An empty registry (no rules).
    pub fn empty() -> Self {
        Self::default()
    }

    /// A registry populated with every built-in rule.
    pub fn default_registry() -> Self {
        let mut reg = Self::empty();
        crate::handlers::register_all(&mut reg);
        reg
    }

    /// Register `rule` for `(domain, op)` applying from opset `min_opset`
    /// upward. A later registration at a higher `min_opset` supersedes this one
    /// for those versions.
    pub fn register(&mut self, domain: &str, op: &str, min_opset: u64, rule: InferenceFn) {
        let key = (normalize_domain(domain).to_string(), op.to_string());
        let entry = self.handlers.entry(key).or_default();
        match entry.binary_search_by_key(&min_opset, |(v, _)| *v) {
            Ok(idx) => entry[idx] = (min_opset, rule), // replace same-version rule
            Err(idx) => entry.insert(idx, (min_opset, rule)),
        }
    }

    /// Look up the rule for `(domain, op)` effective at opset `version`: the
    /// registration with the greatest `min_opset <= version`.
    pub fn get(&self, domain: &str, op: &str, version: u64) -> Option<InferenceFn> {
        let key = (normalize_domain(domain).to_string(), op.to_string());
        let entry = self.handlers.get(&key)?;
        let mut chosen = None;
        for &(min_opset, rule) in entry {
            if min_opset <= version {
                chosen = Some(rule);
            } else {
                break;
            }
        }
        chosen
    }

    /// Number of distinct `(domain, operator)` keys in the registry.
    pub fn operator_count(&self) -> usize {
        self.handlers.len()
    }

    /// Number of opset-versioned inference rule entries.
    pub fn entry_count(&self) -> usize {
        self.handlers.values().map(Vec::len).sum()
    }

    /// Every registered rule as `(domain, operator, min_opset)`, sorted.
    ///
    /// This is the full identity of the catalog — its keys and opset floors,
    /// not its handler bodies — and the thing to pin. A swap of the
    /// [`InferenceFn`] behind an unchanged triple is deliberately out of scope
    /// here; that is what the behavioural rule tests cover. Neither count above
    /// can see a change that preserves them:
    ///
    /// - a **rename** drops one key and adds another, so `operator_count` and
    ///   `entry_count` both hold;
    /// - an **opset move** rewrites an existing entry's `min_opset` in place,
    ///   so `entry_count` holds too.
    ///
    /// Both are silent in production rather than loud: [`Self::get`] returns
    /// `None` for a key it does not know *and* for a version below every
    /// registration, and [`Self::infer_node`] treats `None` permissively —
    /// outputs are left unknown and the model still runs.
    pub fn operator_versions(&self) -> Vec<(&str, &str, u64)> {
        let mut rules: Vec<(&str, &str, u64)> = self
            .handlers
            .iter()
            .flat_map(|((domain, op), entries)| {
                entries
                    .iter()
                    .map(move |(min_opset, _)| (domain.as_str(), op.as_str(), *min_opset))
            })
            .collect();
        rules.sort_unstable();
        rules
    }

    /// Infer a single node's outputs.
    ///
    /// Returns one [`NodeIo`] per output slot. An unregistered op (or one whose
    /// rule declines to resolve an output) yields empty [`NodeIo`]s — the
    /// permissive "leave it unknown" outcome, never an error.
    pub fn infer_node(
        &self,
        node: &Node,
        opset_imports: &HashMap<String, u64>,
        inputs: Vec<NodeIo>,
        policy: MergePolicy,
        interner: &mut SymbolInterner,
    ) -> Result<Vec<NodeIo>, ShapeInferError> {
        let version = if let Some(version) = node.local_opset() {
            version
        } else {
            // Loaded IR is canonical (`normalize_domain` applied at load), so the
            // default domain is `""` for both node domains and opset-import keys.
            if node.is_default_domain() {
                opset_imports.get("").copied().unwrap_or(1)
            } else {
                opset_imports.get(&node.domain).copied().unwrap_or(1)
            }
        };
        let Some(rule) = self.get(&node.domain, &node.op_type, version) else {
            return Ok(vec![NodeIo::default(); node.outputs.len()]);
        };
        let mut ctx = InferenceContext::new(node, inputs, opset_imports, policy, interner);
        rule(&mut ctx)?;
        Ok(ctx.into_outputs())
    }
}