use crate::error::RuntimeError;
use crate::ids::{EntityId, ScopeKey};
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum EntityKind {
Person,
Code,
Project,
Concept,
Custom(String),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Entity {
pub id: EntityId,
pub canonical_name: String,
pub kind: EntityKind,
pub scope: ScopeKey,
pub aliases: Vec<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub metadata: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MatchQuality {
ExactCanonical,
ExactAlias,
ScopedFallback,
Unresolved,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResolveResult {
pub entity: Option<Entity>,
pub quality: MatchQuality,
pub mention: String,
pub alternatives: Vec<Entity>,
pub queried_scope: ScopeKey,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
struct ScopedName {
scope: ScopeKey,
lower_name: String,
}
#[derive(Debug)]
pub struct EntityRegistry {
entities: BTreeMap<EntityId, Entity>,
name_index: BTreeMap<ScopedName, EntityId>,
alias_index: BTreeMap<ScopedName, EntityId>,
max_aliases: usize,
max_entities: usize,
}
impl EntityRegistry {
pub fn new(max_aliases: usize, max_entities: usize) -> Self {
Self {
entities: BTreeMap::new(),
name_index: BTreeMap::new(),
alias_index: BTreeMap::new(),
max_aliases,
max_entities,
}
}
pub fn register(&mut self, entity: Entity) -> Result<(), RuntimeError> {
if self.entities.len() >= self.max_entities && !self.entities.contains_key(&entity.id) {
return Err(RuntimeError::RegistryFull {
capacity: self.max_entities,
});
}
if entity.aliases.len() > self.max_aliases {
return Err(RuntimeError::InvalidConfig {
field: "entity.aliases",
reason: format!(
"entity has {} aliases, max is {}",
entity.aliases.len(),
self.max_aliases
),
});
}
if let Some(old) = self.entities.remove(&entity.id) {
let old_key = ScopedName {
scope: old.scope.clone(),
lower_name: old.canonical_name.to_lowercase(),
};
self.name_index.remove(&old_key);
for alias in &old.aliases {
let alias_key = ScopedName {
scope: old.scope.clone(),
lower_name: alias.to_lowercase(),
};
self.alias_index.remove(&alias_key);
}
}
let name_key = ScopedName {
scope: entity.scope.clone(),
lower_name: entity.canonical_name.to_lowercase(),
};
if let Some(existing_id) = self.name_index.get(&name_key) {
if existing_id != &entity.id {
return Err(RuntimeError::EntityCollision {
name: entity.canonical_name.clone(),
existing_id: existing_id.clone(),
new_id: entity.id.clone(),
});
}
}
for alias in &entity.aliases {
let alias_key = ScopedName {
scope: entity.scope.clone(),
lower_name: alias.to_lowercase(),
};
if let Some(existing_id) = self.alias_index.get(&alias_key) {
if existing_id != &entity.id {
return Err(RuntimeError::EntityCollision {
name: alias.clone(),
existing_id: existing_id.clone(),
new_id: entity.id.clone(),
});
}
}
if let Some(existing_id) = self.name_index.get(&alias_key) {
if existing_id != &entity.id {
return Err(RuntimeError::EntityCollision {
name: alias.clone(),
existing_id: existing_id.clone(),
new_id: entity.id.clone(),
});
}
}
}
self.name_index.insert(name_key, entity.id.clone());
for alias in &entity.aliases {
let alias_key = ScopedName {
scope: entity.scope.clone(),
lower_name: alias.to_lowercase(),
};
self.alias_index.insert(alias_key, entity.id.clone());
}
self.entities.insert(entity.id.clone(), entity);
Ok(())
}
pub fn remove(&mut self, id: &EntityId) -> Option<Entity> {
if let Some(entity) = self.entities.remove(id) {
let name_key = ScopedName {
scope: entity.scope.clone(),
lower_name: entity.canonical_name.to_lowercase(),
};
self.name_index.remove(&name_key);
for alias in &entity.aliases {
let alias_key = ScopedName {
scope: entity.scope.clone(),
lower_name: alias.to_lowercase(),
};
self.alias_index.remove(&alias_key);
}
Some(entity)
} else {
None
}
}
pub fn get(&self, id: &EntityId) -> Option<&Entity> {
self.entities.get(id)
}
pub fn resolve(&self, mention: &str, scope: &ScopeKey) -> ResolveResult {
let lower = mention.to_lowercase();
let exact_name_key = ScopedName {
scope: scope.clone(),
lower_name: lower.clone(),
};
if let Some(id) = self.name_index.get(&exact_name_key) {
if let Some(entity) = self.entities.get(id) {
return ResolveResult {
entity: Some(entity.clone()),
quality: MatchQuality::ExactCanonical,
mention: mention.to_string(),
alternatives: Vec::new(),
queried_scope: scope.clone(),
};
}
}
let exact_alias_key = ScopedName {
scope: scope.clone(),
lower_name: lower.clone(),
};
if let Some(id) = self.alias_index.get(&exact_alias_key) {
if let Some(entity) = self.entities.get(id) {
return ResolveResult {
entity: Some(entity.clone()),
quality: MatchQuality::ExactAlias,
mention: mention.to_string(),
alternatives: Vec::new(),
queried_scope: scope.clone(),
};
}
}
let has_extra_dims =
scope.domain.is_some() || scope.workspace_id.is_some() || scope.repo_id.is_some();
if has_extra_dims {
let ns_only = ScopeKey::namespace_only(&scope.namespace);
let ns_name_key = ScopedName {
scope: ns_only.clone(),
lower_name: lower.clone(),
};
if let Some(id) = self.name_index.get(&ns_name_key) {
if let Some(entity) = self.entities.get(id) {
let alternatives = self.find_alternatives(mention, scope, Some(id));
return ResolveResult {
entity: Some(entity.clone()),
quality: MatchQuality::ScopedFallback,
mention: mention.to_string(),
alternatives,
queried_scope: scope.clone(),
};
}
}
let ns_alias_key = ScopedName {
scope: ns_only,
lower_name: lower.clone(),
};
if let Some(id) = self.alias_index.get(&ns_alias_key) {
if let Some(entity) = self.entities.get(id) {
let alternatives = self.find_alternatives(mention, scope, Some(id));
return ResolveResult {
entity: Some(entity.clone()),
quality: MatchQuality::ScopedFallback,
mention: mention.to_string(),
alternatives,
queried_scope: scope.clone(),
};
}
}
}
let alternatives = self.find_alternatives(mention, scope, None);
ResolveResult {
entity: None,
quality: MatchQuality::Unresolved,
mention: mention.to_string(),
alternatives,
queried_scope: scope.clone(),
}
}
fn find_alternatives(
&self,
mention: &str,
scope: &ScopeKey,
exclude_id: Option<&EntityId>,
) -> Vec<Entity> {
let lower = mention.to_lowercase();
let mut seen = std::collections::HashSet::new();
let mut alts = Vec::new();
for (key, id) in self.name_index.iter().chain(self.alias_index.iter()) {
if key.lower_name == lower
&& key.scope.namespace == scope.namespace
&& key.scope != *scope
&& (exclude_id != Some(id))
&& seen.insert(id.clone())
{
if let Some(entity) = self.entities.get(id) {
alts.push(entity.clone());
}
}
}
alts
}
pub fn len(&self) -> usize {
self.entities.len()
}
pub fn is_empty(&self) -> bool {
self.entities.is_empty()
}
pub fn iter(&self) -> impl Iterator<Item = &Entity> {
self.entities.values()
}
pub fn iter_scope<'a>(&'a self, scope: &'a ScopeKey) -> impl Iterator<Item = &'a Entity> + 'a {
self.entities.values().filter(move |e| &e.scope == scope)
}
pub fn clear_all(&mut self) {
self.entities.clear();
self.name_index.clear();
self.alias_index.clear();
}
pub fn clear_scope(&mut self, scope: &ScopeKey) {
let ids_to_remove: Vec<EntityId> = self
.entities
.values()
.filter(|e| &e.scope == scope)
.map(|e| e.id.clone())
.collect();
for id in ids_to_remove {
self.remove(&id);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_entity(name: &str, scope: ScopeKey, aliases: Vec<&str>) -> Entity {
Entity {
id: EntityId::new(format!("{}@{}", name.to_lowercase(), scope)),
canonical_name: name.to_string(),
kind: EntityKind::Person,
scope,
aliases: aliases.into_iter().map(String::from).collect(),
metadata: None,
}
}
#[test]
fn exact_name_match_in_scope() {
let scope = ScopeKey::namespace_only("test");
let mut reg = EntityRegistry::new(16, 100);
reg.register(make_entity("Alice", scope.clone(), vec![]))
.unwrap();
let result = reg.resolve("alice", &scope);
assert_eq!(result.quality, MatchQuality::ExactCanonical);
assert!(result.entity.is_some());
}
#[test]
fn alias_match_in_scope() {
let scope = ScopeKey::namespace_only("test");
let mut reg = EntityRegistry::new(16, 100);
reg.register(make_entity("Alice", scope.clone(), vec!["ally"]))
.unwrap();
let result = reg.resolve("ally", &scope);
assert_eq!(result.quality, MatchQuality::ExactAlias);
}
#[test]
fn unresolved_mention() {
let scope = ScopeKey::namespace_only("test");
let reg = EntityRegistry::new(16, 100);
let result = reg.resolve("unknown", &scope);
assert_eq!(result.quality, MatchQuality::Unresolved);
assert!(result.entity.is_none());
}
#[test]
fn same_mention_different_scopes_resolves_differently() {
let scope_a = ScopeKey {
namespace: "ns".into(),
domain: None,
workspace_id: None,
repo_id: Some("repo-a".into()),
};
let scope_b = ScopeKey {
namespace: "ns".into(),
domain: None,
workspace_id: None,
repo_id: Some("repo-b".into()),
};
let mut reg = EntityRegistry::new(16, 100);
reg.register(Entity {
id: EntityId::new("alice-a"),
canonical_name: "Alice".into(),
kind: EntityKind::Person,
scope: scope_a.clone(),
aliases: vec![],
metadata: None,
})
.unwrap();
reg.register(Entity {
id: EntityId::new("alice-b"),
canonical_name: "Alice".into(),
kind: EntityKind::Person,
scope: scope_b.clone(),
aliases: vec![],
metadata: None,
})
.unwrap();
let result_a = reg.resolve("Alice", &scope_a);
let result_b = reg.resolve("Alice", &scope_b);
assert_eq!(result_a.quality, MatchQuality::ExactCanonical);
assert_eq!(result_b.quality, MatchQuality::ExactCanonical);
assert_eq!(result_a.entity.unwrap().id.as_str(), "alice-a");
assert_eq!(result_b.entity.unwrap().id.as_str(), "alice-b");
}
#[test]
fn alias_only_works_in_correct_scope() {
let scope_a = ScopeKey {
namespace: "ns".into(),
domain: None,
workspace_id: None,
repo_id: Some("repo-a".into()),
};
let scope_b = ScopeKey {
namespace: "ns".into(),
domain: None,
workspace_id: None,
repo_id: Some("repo-b".into()),
};
let mut reg = EntityRegistry::new(16, 100);
reg.register(Entity {
id: EntityId::new("alice-a"),
canonical_name: "Alice".into(),
kind: EntityKind::Person,
scope: scope_a.clone(),
aliases: vec!["ally".into()],
metadata: None,
})
.unwrap();
let result = reg.resolve("ally", &scope_a);
assert_eq!(result.quality, MatchQuality::ExactAlias);
let result = reg.resolve("ally", &scope_b);
assert_eq!(result.quality, MatchQuality::Unresolved);
assert_eq!(result.alternatives.len(), 1);
assert_eq!(result.alternatives[0].id.as_str(), "alice-a");
}
#[test]
fn scoped_fallback_to_namespace_only() {
let ns_scope = ScopeKey::namespace_only("ns");
let narrow_scope = ScopeKey {
namespace: "ns".into(),
domain: Some("code".into()),
workspace_id: None,
repo_id: None,
};
let mut reg = EntityRegistry::new(16, 100);
reg.register(Entity {
id: EntityId::new("alice-ns"),
canonical_name: "Alice".into(),
kind: EntityKind::Person,
scope: ns_scope.clone(),
aliases: vec![],
metadata: None,
})
.unwrap();
let result = reg.resolve("Alice", &narrow_scope);
assert_eq!(result.quality, MatchQuality::ScopedFallback);
assert!(result.entity.is_some());
}
#[test]
fn registry_full_rejects() {
let scope = ScopeKey::namespace_only("test");
let mut reg = EntityRegistry::new(16, 1);
reg.register(make_entity("Alice", scope.clone(), vec![]))
.unwrap();
let err = reg.register(make_entity("Bob", scope, vec![])).unwrap_err();
assert_eq!(err.kind(), "registry_full");
}
#[test]
fn update_replaces_old_indexes() {
let scope = ScopeKey::namespace_only("test");
let mut reg = EntityRegistry::new(16, 100);
reg.register(make_entity("Alice", scope.clone(), vec!["ally"]))
.unwrap();
let updated = Entity {
id: EntityId::new(format!("alice@{}", scope)),
canonical_name: "Alice".into(),
kind: EntityKind::Person,
scope: scope.clone(),
aliases: vec!["al".into()],
metadata: None,
};
reg.register(updated).unwrap();
assert_eq!(
reg.resolve("ally", &scope).quality,
MatchQuality::Unresolved
);
assert_eq!(reg.resolve("al", &scope).quality, MatchQuality::ExactAlias);
}
#[test]
fn clear_scope_removes_only_targeted_scope() {
let scope_a = ScopeKey::namespace_only("a");
let scope_b = ScopeKey::namespace_only("b");
let mut reg = EntityRegistry::new(16, 100);
reg.register(make_entity("Alice", scope_a.clone(), vec![]))
.unwrap();
reg.register(make_entity("Bob", scope_b.clone(), vec![]))
.unwrap();
reg.clear_scope(&scope_a);
assert_eq!(reg.len(), 1);
assert_eq!(
reg.resolve("Alice", &scope_a).quality,
MatchQuality::Unresolved
);
assert_eq!(
reg.resolve("Bob", &scope_b).quality,
MatchQuality::ExactCanonical
);
}
#[test]
fn clear_all_removes_everything() {
let scope_a = ScopeKey::namespace_only("a");
let scope_b = ScopeKey::namespace_only("b");
let mut reg = EntityRegistry::new(16, 100);
reg.register(make_entity("Alice", scope_a.clone(), vec!["ally"]))
.unwrap();
reg.register(make_entity("Bob", scope_b.clone(), vec!["bobby"]))
.unwrap();
assert_eq!(reg.len(), 2);
reg.clear_all();
assert_eq!(reg.len(), 0);
assert!(reg.is_empty());
assert_eq!(
reg.resolve("Alice", &scope_a).quality,
MatchQuality::Unresolved
);
assert_eq!(
reg.resolve("ally", &scope_a).quality,
MatchQuality::Unresolved
);
assert_eq!(
reg.resolve("Bob", &scope_b).quality,
MatchQuality::Unresolved
);
assert_eq!(
reg.resolve("bobby", &scope_b).quality,
MatchQuality::Unresolved
);
}
#[test]
fn clear_all_allows_re_registration() {
let scope = ScopeKey::namespace_only("test");
let mut reg = EntityRegistry::new(16, 2);
reg.register(make_entity("Alice", scope.clone(), vec![]))
.unwrap();
reg.register(make_entity("Bob", scope.clone(), vec![]))
.unwrap();
assert!(reg
.register(make_entity("Charlie", scope.clone(), vec![]))
.is_err());
reg.clear_all();
reg.register(make_entity("Charlie", scope.clone(), vec![]))
.unwrap();
assert_eq!(reg.len(), 1);
assert_eq!(
reg.resolve("Charlie", &scope).quality,
MatchQuality::ExactCanonical
);
}
}