1use std::collections::HashMap;
2
3use serde::{Deserialize, Serialize};
4
5use crate::error::{Error, Result};
6use crate::iri::IriId;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
10pub struct EntityId(pub u32);
12
13impl EntityId {
14 #[must_use]
16 pub fn index(self) -> u32 {
17 self.0
18 }
19}
20
21#[non_exhaustive]
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
24pub enum EntityKind {
25 Class,
27 Individual,
29 ObjectProperty,
31 DataProperty,
33 AnnotationProperty,
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
39pub struct EntityRecord {
40 pub iri: IriId,
42 pub kind: EntityKind,
44}
45
46#[derive(Debug, Default, Clone, PartialEq, Eq)]
48pub struct EntityRegistry {
49 entities: Vec<EntityRecord>,
50 by_iri: HashMap<IriId, EntityId>,
51}
52
53impl EntityRegistry {
54 #[must_use]
56 pub fn new() -> Self {
57 Self::default()
58 }
59
60 #[must_use]
62 pub fn len(&self) -> usize {
63 self.entities.len()
64 }
65
66 #[must_use]
68 pub fn is_empty(&self) -> bool {
69 self.entities.is_empty()
70 }
71
72 #[must_use]
74 pub fn entity_by_iri(&self, iri: IriId) -> Option<EntityId> {
75 self.by_iri.get(&iri).copied()
76 }
77
78 pub fn entity(&self, id: EntityId) -> Result<&EntityRecord> {
80 self.entities
81 .get(id.0 as usize)
82 .ok_or(Error::UnknownEntity(id))
83 }
84
85 pub fn get_or_register(
87 &mut self,
88 iri: IriId,
89 iri_str: &str,
90 kind: EntityKind,
91 ) -> Result<EntityId> {
92 if let Some(&existing) = self.by_iri.get(&iri) {
93 let record = &self.entities[existing.0 as usize];
94 if record.kind != kind {
95 return Err(Error::EntityKindMismatch {
96 iri: iri_str.to_owned(),
97 expected: kind,
98 found: record.kind,
99 });
100 }
101 return Ok(existing);
102 }
103
104 let id = EntityId(
105 u32::try_from(self.entities.len())
106 .map_err(|_| Error::InvalidAxiom("entity registry capacity exceeded".into()))?,
107 );
108 self.by_iri.insert(iri, id);
109 self.entities.push(EntityRecord { iri, kind });
110 Ok(id)
111 }
112
113 pub fn iter(&self) -> impl Iterator<Item = (EntityId, &EntityRecord)> {
115 self.entities
116 .iter()
117 .enumerate()
118 .map(|(i, record)| (EntityId(i as u32), record))
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use super::*;
125 use crate::iri::InternPool;
126
127 #[test]
128 fn register_and_lookup() {
129 let mut pool = InternPool::new();
130 let iri = pool.intern("http://example.org/A").expect("intern");
131 let mut registry = EntityRegistry::new();
132 let id = registry
133 .get_or_register(iri, "http://example.org/A", EntityKind::Class)
134 .expect("register");
135 assert_eq!(registry.entity_by_iri(iri), Some(id));
136 assert_eq!(registry.entity(id).expect("entity").kind, EntityKind::Class);
137 }
138
139 #[test]
140 fn kind_mismatch_includes_iri_string() {
141 let mut pool = InternPool::new();
142 let iri = pool.intern("http://example.org/A").expect("intern");
143 let mut registry = EntityRegistry::new();
144 registry
145 .get_or_register(iri, "http://example.org/A", EntityKind::Class)
146 .expect("register");
147 let err = registry
148 .get_or_register(iri, "http://example.org/A", EntityKind::Individual)
149 .expect_err("mismatch");
150 if let Error::EntityKindMismatch { iri, .. } = err {
151 assert_eq!(iri, "http://example.org/A");
152 } else {
153 panic!("expected EntityKindMismatch");
154 }
155 }
156
157 #[test]
158 fn kind_mismatch_errors() {
159 let mut pool = InternPool::new();
160 let iri = pool.intern("http://example.org/A").expect("intern");
161 let mut registry = EntityRegistry::new();
162 registry
163 .get_or_register(iri, "http://example.org/A", EntityKind::Class)
164 .expect("register");
165 let err = registry
166 .get_or_register(iri, "http://example.org/A", EntityKind::Individual)
167 .expect_err("mismatch");
168 assert!(matches!(
169 err,
170 Error::EntityKindMismatch {
171 expected: EntityKind::Individual,
172 found: EntityKind::Class,
173 ..
174 }
175 ));
176 }
177
178 #[test]
179 fn unknown_entity_errors() {
180 let registry = EntityRegistry::new();
181 let err = registry.entity(EntityId(0)).expect_err("unknown");
182 assert_eq!(err, Error::UnknownEntity(EntityId(0)));
183 }
184
185 #[test]
186 fn reregister_same_kind_returns_stable_id() {
187 let mut pool = InternPool::new();
188 let iri = pool.intern("http://example.org/A").expect("intern");
189 let mut registry = EntityRegistry::new();
190 let first = registry
191 .get_or_register(iri, "http://example.org/A", EntityKind::Class)
192 .expect("register");
193 let second = registry
194 .get_or_register(iri, "http://example.org/A", EntityKind::Class)
195 .expect("register");
196 assert_eq!(first, second);
197 }
198}