khive_types/

entity_type.rs

//! EntityTypeRegistry — validates and normalises `(EntityKind, entity_type)` pairs.
//!
//! Canonical, dependency-light home for the closed entity-subtype taxonomy
//! (#571): both KG create validation and schedule replay validation resolve
//! against this single registry instead of maintaining separate copies.

use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use core::fmt;

use crate::entity::EntityKind;

/// Normalise a raw `entity_type` string to canonical snake_case.
///
/// Pipeline: trim → lowercase → runs of separators (space, hyphen, underscore)
/// collapsed to a single `_` → leading/trailing `_` stripped.
///
/// This implements the ADR-001:106 write-time normalisation step that precedes
/// alias resolution. `pub` (not `pub(crate)`) so callers outside this crate —
/// e.g. `khive-pack-kg`'s alias-substitution reporting — can compare a raw
/// value against this SAME cosmetic normalisation instead of maintaining a
/// separate, drift-prone copy of the rule.
pub fn to_snake_case(s: &str) -> String {
    let mut out = String::with_capacity(s.len());
    let mut prev_sep = true; // treat start as separator so leading _ are stripped
    for ch in s.chars() {
        if ch == ' ' || ch == '-' || ch == '_' {
            if !prev_sep && !out.is_empty() {
                out.push('_');
                prev_sep = true;
            }
        } else {
            out.push(ch.to_ascii_lowercase());
            prev_sep = false;
        }
    }
    if out.ends_with('_') {
        out.pop();
    }
    out
}

/// One entry in the registry: a canonical subtype name for a specific kind,
/// together with any accepted aliases.
#[derive(Clone, Debug)]
pub struct EntityTypeDef {
    /// The entity kind this subtype belongs to.
    pub kind: EntityKind,
    /// Canonical name that is written to the DB.
    pub type_name: &'static str,
    /// Alternative spellings that are accepted at the wire level but
    /// normalised to `type_name` before storage.
    pub aliases: &'static [&'static str],
}

/// Static table of built-in subtypes (non-exhaustive; packs may extend).
static BUILTIN_DEFS: &[EntityTypeDef] = &[
    // ── Document ────────────────────────────────────────────────────────────
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "paper",
        aliases: &["preprint", "article"],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "report",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "blog_post",
        aliases: &["blog"],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "book",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "specification",
        aliases: &["spec"],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "documentation",
        aliases: &["docs"],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "thesis",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "page",
        aliases: &["web_page"],
    },
    // ADR-191: any fetched body other than an HTML page — robots.txt, a
    // sitemap, a feed, a JSON body, a PDF, an alternate rendering served at
    // its own URL.
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "resource",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Document,
        type_name: "research_report",
        aliases: &[],
    },
    // ── Concept ─────────────────────────────────────────────────────────────
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "algorithm",
        aliases: &["algo"],
    },
    // ── Formal math ─────────────────────────────────────────────────────────
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "theorem",
        aliases: &["lemma", "proposition", "corollary"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "definition",
        aliases: &["def"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "structure",
        aliases: &["inductive", "struct", "class"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "instance",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "axiom",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "goal",
        aliases: &["proof_goal"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "technique",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "architecture",
        aliases: &["arch"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        // "model" is the alias; canonical name is "model_family"
        // to distinguish from a Dataset or Artifact trained model instance.
        type_name: "model_family",
        aliases: &["model"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "theory",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "research_gap",
        aliases: &["gap"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "design_pattern",
        aliases: &["pattern"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "mathematical_operation",
        aliases: &["math_op"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "metric",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "objective",
        aliases: &["loss"],
    },
    // ── Code (ADR-085) ──────────────────────────────────────────────────────
    // "struct" and "class" are not registered as aliases here: formal-math
    // "structure" already owns them (see above), and ADR-085 requires
    // ingesters to write the canonical code tokens instead.
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "module",
        aliases: &["mod", "namespace"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "function",
        aliases: &["fn", "func", "method"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "datatype",
        aliases: &["enum", "record", "type_alias"],
    },
    EntityTypeDef {
        kind: EntityKind::Concept,
        type_name: "interface",
        aliases: &["trait", "protocol"],
    },
    // ── Dataset ──────────────────────────────────────────────────────────────
    // benchmark belongs to Dataset, not Concept (it evaluates models, it is not itself a concept).
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "benchmark",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "corpus",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "training_set",
        aliases: &["train_set"],
    },
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "evaluation_set",
        aliases: &["eval_set"],
    },
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "test_set",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Dataset,
        type_name: "synthetic_dataset",
        aliases: &["synthetic"],
    },
    // ── Project ─────────────────────────────────────────────────────────────
    // Project subtypes: library, framework, tool, application, repository.
    // "service" and "svc" are omitted — EntityKind::Service handles running
    // instances; a service codebase repo is Project + application or tool.
    EntityTypeDef {
        kind: EntityKind::Project,
        type_name: "library",
        aliases: &["lib"],
    },
    EntityTypeDef {
        kind: EntityKind::Project,
        type_name: "framework",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Project,
        type_name: "tool",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Project,
        type_name: "application",
        aliases: &["app"],
    },
    EntityTypeDef {
        kind: EntityKind::Project,
        type_name: "repository",
        aliases: &["repo"],
    },
    // ── Org ──────────────────────────────────────────────────────────────────
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "academic_institution",
        aliases: &["university", "uni"],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "company",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "research_lab",
        aliases: &["lab"],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "nonprofit",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "government_agency",
        aliases: &["gov_agency"],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "consortium",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Org,
        type_name: "standards_body",
        aliases: &[],
    },
    // ── Artifact ─────────────────────────────────────────────────────────────
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "checkpoint",
        aliases: &["ckpt"],
    },
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "snapshot",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "export",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "embedding_index",
        aliases: &["embed_index"],
    },
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "state_bundle",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Artifact,
        type_name: "profile",
        aliases: &[],
    },
    // ── Service ──────────────────────────────────────────────────────────────
    // Service subtypes: inference_engine, retrieval_engine,
    // embedding_engine, api, database, search_engine, mcp_server, site.
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "inference_engine",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "retrieval_engine",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "embedding_engine",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "api",
        aliases: &["endpoint"],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "database",
        aliases: &["db"],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "search_engine",
        aliases: &[],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "mcp_server",
        aliases: &["mcp"],
    },
    EntityTypeDef {
        kind: EntityKind::Service,
        type_name: "site",
        aliases: &["origin"],
    },
    // Person  — no standard subtypes (roles are metadata, not subtypes).
];

/// Resolved output of [`EntityTypeRegistry::resolve`].
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResolvedEntityType {
    /// The canonical `EntityKind` for the resolved combination.
    pub kind: EntityKind,
    /// The canonical type name to store, or `None` when no `entity_type` was
    /// supplied and none was inferred.
    pub entity_type: Option<String>,
}

/// Error produced by [`EntityTypeRegistry::resolve`].
///
/// Dependency-light: `khive-types` cannot depend on `khive-runtime`, so
/// callers map this to their own error type at the pack boundary.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EntityTypeError {
    /// The raw type name exists, but belongs to a different `EntityKind`.
    WrongKind {
        raw_type: String,
        actual_kind: &'static str,
        expected_kind: &'static str,
        valid: String,
    },
    /// The raw type name is not registered for any kind.
    UnknownType {
        raw_type: String,
        kind: &'static str,
        valid: String,
    },
}

impl fmt::Display for EntityTypeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::WrongKind {
                raw_type,
                actual_kind,
                expected_kind,
                valid,
            } => write!(
                f,
                "entity_type {raw_type:?} belongs to {actual_kind:?}, not {expected_kind:?}; \
                 valid types for {expected_kind:?}: {valid}",
            ),
            Self::UnknownType {
                raw_type,
                kind,
                valid,
            } => write!(
                f,
                "unknown entity_type {raw_type:?} for {kind:?}; valid: {valid}"
            ),
        }
    }
}

#[cfg(feature = "std")]
impl std::error::Error for EntityTypeError {}

/// Registry for `(EntityKind, entity_type)` pair validation and alias normalisation.
#[derive(Clone)]
pub struct EntityTypeRegistry {
    /// `alias_or_name (lowercase) → def index`.  Covers both canonical names
    /// and all registered aliases.
    lookup: BTreeMap<String, usize>,
    defs: Vec<EntityTypeDef>,
}

impl EntityTypeRegistry {
    /// Build a fresh registry from the supplied definitions.
    ///
    /// Every lookup key (kind-qualified and bare, canonical and alias) is
    /// routed through `to_snake_case` before insertion — the same
    /// normalisation `resolve` applies to incoming wire values (ADR-001:106).
    /// Two differently-cased/separated spellings that normalise to the same
    /// key therefore land on the same registry slot instead of silently
    /// shadowing one another under distinct raw keys.
    pub fn new(defs: impl IntoIterator<Item = EntityTypeDef>) -> Self {
        let defs: Vec<EntityTypeDef> = defs.into_iter().collect();
        let mut lookup: BTreeMap<String, usize> = BTreeMap::new();
        for (idx, def) in defs.iter().enumerate() {
            let canonical_type = to_snake_case(def.type_name);
            let canonical_key = format!("{}:{}", def.kind.name(), canonical_type);
            lookup.insert(canonical_key, idx);
            // Bare name without kind prefix — only insert when unambiguous.
            lookup.entry(canonical_type).or_insert(idx);
            for alias in def.aliases {
                let normalised_alias = to_snake_case(alias);
                let kind_alias_key = format!("{}:{}", def.kind.name(), normalised_alias);
                lookup.insert(kind_alias_key, idx);
                lookup.entry(normalised_alias).or_insert(idx);
            }
        }
        Self { lookup, defs }
    }

    /// Return the built-in registry with all static subtype definitions.
    pub fn builtin() -> Self {
        Self::new(BUILTIN_DEFS.iter().cloned())
    }

    /// Derive a registry that includes all built-in subtypes plus the
    /// caller-supplied extras (used by packs that extend the vocabulary).
    pub fn with_extra(extra: impl IntoIterator<Item = EntityTypeDef>) -> Self {
        let defs: Vec<EntityTypeDef> = BUILTIN_DEFS.iter().cloned().chain(extra).collect();
        Self::new(defs)
    }

    /// Definitions in this composed registry, including canonical names and aliases.
    pub fn definitions(&self) -> &[EntityTypeDef] {
        &self.defs
    }

    /// Boot-time collision check across the composed registry `with_extra`
    /// builds: built-in defs plus every caller-supplied `(owner, def)` extra.
    ///
    /// `new`/`register` populate a single kind-qualified lookup keyspace
    /// (`"{kind}:{canonical_name}"` and `"{kind}:{alias}"`) with a hard
    /// `insert` — the later write silently wins. Two independently loaded
    /// owners can therefore disagree on what a given key resolves to, and
    /// the winner depends on registration order. ADR-001's registry-ownership
    /// contract instead requires this to fail at boot: "same `(base_kind,
    /// canonical_name)` from two different packs = boot error" and "an alias
    /// collision = boot error." This walks the same keyspace `new`/`register`
    /// populate and returns the first collision found, naming both
    /// contributing owners, instead of silently picking a winner.
    ///
    /// `extras` is `(owner_label, def)` for pack-declared entries only — the
    /// built-in table is checked automatically, attributed to the label
    /// `"builtin"`.
    ///
    /// Every canonical and alias key is routed through the same
    /// `to_snake_case` normalisation `new`/`register` apply before
    /// insertion (ADR-001:104-107): two spellings that only differ in case
    /// or separator style (e.g. `"Architecture Decision Record"` vs.
    /// `"architecture-decision-record"`) collide here exactly as they would
    /// collide in the composed registry's lookup keyspace.
    pub fn check_extra_collisions<'a>(
        extras: impl IntoIterator<Item = (&'a str, &'a EntityTypeDef)>,
    ) -> Result<(), String> {
        let mut all: Vec<(&'a str, &'a EntityTypeDef)> =
            BUILTIN_DEFS.iter().map(|def| ("builtin", def)).collect();
        all.extend(extras);

        let mut seen: BTreeMap<String, (&'a str, String)> = BTreeMap::new();
        for (owner, def) in all {
            let canonical_type = to_snake_case(def.type_name);
            let canonical_key = format!("{}:{}", def.kind.name(), canonical_type);
            if let Some((first_owner, first_type)) = seen.get(&canonical_key) {
                if *first_owner != owner {
                    return Err(format!(
                        "duplicate entity_type {canonical_key:?}: claimed by both \
                         {first_owner:?} and {owner:?}"
                    ));
                }
                if *first_type != canonical_type {
                    return Err(format!(
                        "entity_type alias {canonical_key:?}: claimed by both {first_owner:?} \
                         (canonical {first_type:?}) and {owner:?} (canonical {canonical_type:?})"
                    ));
                }
            } else {
                seen.insert(canonical_key, (owner, canonical_type.clone()));
            }

            for alias in def.aliases {
                let normalised_alias = to_snake_case(alias);
                let alias_key = format!("{}:{}", def.kind.name(), normalised_alias);
                if let Some((first_owner, first_type)) = seen.get(&alias_key) {
                    if *first_owner != owner || *first_type != canonical_type {
                        return Err(format!(
                            "entity_type alias {alias_key:?}: claimed by both {first_owner:?} \
                             (canonical {first_type:?}) and {owner:?} (canonical {canonical_type:?})",
                        ));
                    }
                } else {
                    seen.insert(alias_key, (owner, canonical_type.clone()));
                }
            }
        }
        Ok(())
    }

    /// Register additional subtypes into an existing registry clone.
    ///
    /// Mirrors [`Self::new`]'s normalisation: every key is routed through
    /// `to_snake_case` before insertion.
    pub fn register(&mut self, def: EntityTypeDef) {
        let idx = self.defs.len();
        let canonical_type = to_snake_case(def.type_name);
        let canonical_key = format!("{}:{}", def.kind.name(), canonical_type);
        self.lookup.insert(canonical_key, idx);
        self.lookup.entry(canonical_type).or_insert(idx);
        for alias in def.aliases {
            let normalised_alias = to_snake_case(alias);
            let kind_alias_key = format!("{}:{}", def.kind.name(), normalised_alias);
            self.lookup.insert(kind_alias_key, idx);
            self.lookup.entry(normalised_alias).or_insert(idx);
        }
        self.defs.push(def);
    }

    /// Validate and normalise a `(kind, entity_type)` wire pair.
    pub fn resolve(
        &self,
        kind: EntityKind,
        entity_type: Option<&str>,
    ) -> Result<ResolvedEntityType, EntityTypeError> {
        let Some(raw_type) = entity_type else {
            return Ok(ResolvedEntityType {
                kind,
                entity_type: None,
            });
        };

        let normalised = to_snake_case(raw_type.trim());

        // Try kind-qualified lookup first (unambiguous).
        let kind_key = format!("{}:{}", kind.name(), normalised);
        if let Some(&idx) = self.lookup.get(&kind_key) {
            let def = &self.defs[idx];
            return Ok(ResolvedEntityType {
                kind,
                entity_type: Some(def.type_name.to_string()),
            });
        }

        // Try bare lookup — only valid when the bare name belongs to this kind.
        if let Some(&idx) = self.lookup.get(&normalised) {
            let def = &self.defs[idx];
            if def.kind == kind {
                return Ok(ResolvedEntityType {
                    kind,
                    entity_type: Some(def.type_name.to_string()),
                });
            }
            // The name exists but belongs to a different kind.
            return Err(EntityTypeError::WrongKind {
                raw_type: raw_type.to_string(),
                actual_kind: def.kind.name(),
                expected_kind: kind.name(),
                valid: self.valid_types_for(kind),
            });
        }

        // Not found at all.
        Err(EntityTypeError::UnknownType {
            raw_type: raw_type.to_string(),
            kind: kind.name(),
            valid: self.valid_types_for(kind),
        })
    }

    /// Comma-separated list of canonical type names valid for `kind`.
    pub fn valid_types_for(&self, kind: EntityKind) -> String {
        let mut names: Vec<&str> = self
            .defs
            .iter()
            .filter(|d| d.kind == kind)
            .map(|d| d.type_name)
            .collect();
        names.sort_unstable();
        if names.is_empty() {
            "(none registered)".to_string()
        } else {
            names.join(" | ")
        }
    }
}

// ── Module-level lazy global ─────────────────────────────────────────────────

#[cfg(feature = "std")]
static GLOBAL_REGISTRY: std::sync::OnceLock<EntityTypeRegistry> = std::sync::OnceLock::new();

#[cfg(feature = "std")]
impl EntityTypeRegistry {
    /// Return a reference to the module-level built-in registry.
    pub fn global() -> &'static EntityTypeRegistry {
        GLOBAL_REGISTRY.get_or_init(EntityTypeRegistry::builtin)
    }
}

// ── Tests ────────────────────────────────────────────────────────────────────

#[cfg(test)]
#[path = "entity_type_tests.rs"]
mod tests;