use alloc::collections::BTreeMap;
use alloc::format;
use alloc::string::{String, ToString};
use alloc::vec::Vec;
use core::fmt;
use crate::entity::EntityKind;
pub fn to_snake_case(s: &str) -> String {
let mut out = String::with_capacity(s.len());
let mut prev_sep = true; 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
}
#[derive(Clone, Debug)]
pub struct EntityTypeDef {
pub kind: EntityKind,
pub type_name: &'static str,
pub aliases: &'static [&'static str],
}
static BUILTIN_DEFS: &[EntityTypeDef] = &[
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"],
},
EntityTypeDef {
kind: EntityKind::Document,
type_name: "resource",
aliases: &[],
},
EntityTypeDef {
kind: EntityKind::Document,
type_name: "research_report",
aliases: &[],
},
EntityTypeDef {
kind: EntityKind::Concept,
type_name: "algorithm",
aliases: &["algo"],
},
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,
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"],
},
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"],
},
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"],
},
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"],
},
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: &[],
},
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: &[],
},
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"],
},
];
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ResolvedEntityType {
pub kind: EntityKind,
pub entity_type: Option<String>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EntityTypeError {
WrongKind {
raw_type: String,
actual_kind: &'static str,
expected_kind: &'static str,
valid: String,
},
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 {}
#[derive(Clone)]
pub struct EntityTypeRegistry {
lookup: BTreeMap<String, usize>,
defs: Vec<EntityTypeDef>,
}
impl EntityTypeRegistry {
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);
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 }
}
pub fn builtin() -> Self {
Self::new(BUILTIN_DEFS.iter().cloned())
}
pub fn with_extra(extra: impl IntoIterator<Item = EntityTypeDef>) -> Self {
let defs: Vec<EntityTypeDef> = BUILTIN_DEFS.iter().cloned().chain(extra).collect();
Self::new(defs)
}
pub fn definitions(&self) -> &[EntityTypeDef] {
&self.defs
}
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(())
}
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);
}
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());
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()),
});
}
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()),
});
}
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),
});
}
Err(EntityTypeError::UnknownType {
raw_type: raw_type.to_string(),
kind: kind.name(),
valid: self.valid_types_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(" | ")
}
}
}
#[cfg(feature = "std")]
static GLOBAL_REGISTRY: std::sync::OnceLock<EntityTypeRegistry> = std::sync::OnceLock::new();
#[cfg(feature = "std")]
impl EntityTypeRegistry {
pub fn global() -> &'static EntityTypeRegistry {
GLOBAL_REGISTRY.get_or_init(EntityTypeRegistry::builtin)
}
}
#[cfg(test)]
#[path = "entity_type_tests.rs"]
mod tests;