1use std::collections::BTreeMap;
5use std::fmt::Display;
6
7use schemars::JsonSchema;
8use serde::{Deserialize, Serialize};
9
10use crate::provenance::Exactness;
11
12#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
17pub struct Annotations {
18 #[serde(default, skip_serializing_if = "Vec::is_empty")]
20 pub streams: Vec<String>,
21 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
23 pub provenance: BTreeMap<String, Provenance>,
24 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
26 pub exactness: BTreeMap<String, ExactnessNote>,
27}
28
29#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
31pub struct Provenance {
32 pub stream: u32,
34 pub offset: u64,
36 #[serde(default, skip_serializing_if = "Option::is_none")]
38 pub tag: Option<String>,
39}
40
41#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
43pub struct ExactnessNote {
44 pub entity: Exactness,
46 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
48 pub fields: BTreeMap<String, Exactness>,
49}
50
51impl Default for ExactnessNote {
52 fn default() -> Self {
53 Self {
54 entity: Exactness::ByteExact,
55 fields: BTreeMap::new(),
56 }
57 }
58}
59
60#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
62pub struct StreamHandle(u32);
63
64#[derive(Debug, Default, Clone)]
66pub struct AnnotationBuilder {
67 annotations: Annotations,
68}
69
70impl AnnotationBuilder {
71 pub fn new() -> Self {
73 Self::default()
74 }
75
76 pub fn stream(&mut self, stream: impl Into<String>) -> StreamHandle {
78 let stream = stream.into();
79 if let Some(index) = self
80 .annotations
81 .streams
82 .iter()
83 .position(|existing| existing == &stream)
84 {
85 return StreamHandle(index as u32);
86 }
87
88 let index = u32::try_from(self.annotations.streams.len())
89 .expect("annotation stream count exceeds u32::MAX");
90 self.annotations.streams.push(stream);
91 StreamHandle(index)
92 }
93
94 pub fn note(
99 &mut self,
100 id: impl Display,
101 stream: StreamHandle,
102 offset: u64,
103 ) -> ProvenanceNote<'_> {
104 let id = id.to_string();
105 self.annotations.provenance.insert(
106 id.clone(),
107 Provenance {
108 stream: stream.0,
109 offset,
110 tag: None,
111 },
112 );
113 ProvenanceNote {
114 provenance: self
115 .annotations
116 .provenance
117 .get_mut(&id)
118 .expect("provenance was just inserted"),
119 }
120 }
121
122 pub fn exactness(&mut self, id: impl Display, exactness: Exactness) -> &mut Self {
125 let id = id.to_string();
126 let note = self.annotations.exactness.entry(id.clone()).or_default();
127 note.entity = exactness;
128 note.fields.retain(|_, value| *value != exactness);
129 if exactness == Exactness::ByteExact && note.fields.is_empty() {
130 self.annotations.exactness.remove(&id);
131 }
132 self
133 }
134
135 pub fn derived(&mut self, id: impl Display, field: impl Into<String>) -> &mut Self {
137 self.field_exactness(id, field, Exactness::Derived)
138 }
139
140 pub fn field_exactness(
145 &mut self,
146 id: impl Display,
147 field: impl Into<String>,
148 exactness: Exactness,
149 ) -> &mut Self {
150 let id = id.to_string();
151 let field = field.into();
152 if exactness == Exactness::ByteExact {
153 if let Some(note) = self.annotations.exactness.get_mut(&id) {
154 if note.entity == Exactness::ByteExact {
155 note.fields.remove(&field);
156 if note.fields.is_empty() {
157 self.annotations.exactness.remove(&id);
158 }
159 } else {
160 note.fields.insert(field, Exactness::ByteExact);
161 }
162 }
163 } else {
164 self.annotations
165 .exactness
166 .entry(id)
167 .or_default()
168 .fields
169 .insert(field, exactness);
170 }
171 self
172 }
173
174 pub fn remove_entity(&mut self, id: impl Display) {
176 let id = id.to_string();
177 self.annotations.provenance.remove(&id);
178 self.annotations.exactness.remove(&id);
179 }
180
181 pub fn annotations(&self) -> &Annotations {
183 &self.annotations
184 }
185
186 pub fn build(self) -> Annotations {
188 self.annotations
189 }
190}
191
192pub struct ProvenanceNote<'a> {
194 provenance: &'a mut Provenance,
195}
196
197impl ProvenanceNote<'_> {
198 pub fn tag(self, tag: impl Into<String>) {
200 self.provenance.tag = Some(tag.into());
201 }
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207
208 #[test]
209 fn builder_interns_streams_and_records_provenance() {
210 let mut builder = AnnotationBuilder::new();
211 let first = builder.stream("f3d:Breps.BlobParts/body.smbh");
212 let second = builder.stream("f3d:Breps.BlobParts/body.smbh");
213
214 assert_eq!(first, second);
215 builder.note("f3d:body#0", first, 42).tag("body");
216
217 let annotations = builder.build();
218 assert_eq!(annotations.streams.len(), 1);
219 assert_eq!(
220 annotations.provenance["f3d:body#0"],
221 Provenance {
222 stream: 0,
223 offset: 42,
224 tag: Some("body".to_string()),
225 }
226 );
227 }
228
229 #[test]
230 fn exactness_table_stays_sparse() {
231 let mut builder = AnnotationBuilder::new();
232
233 builder
234 .derived("f3d:edge#0", "param_range")
235 .exactness("f3d:edge#0", Exactness::Inferred);
236 builder.field_exactness("f3d:edge#0", "param_range", Exactness::ByteExact);
237
238 let expected_fields = BTreeMap::from([("param_range".to_string(), Exactness::ByteExact)]);
239 assert_eq!(
240 builder.annotations().exactness["f3d:edge#0"],
241 ExactnessNote {
242 entity: Exactness::Inferred,
243 fields: expected_fields,
244 }
245 );
246
247 builder.exactness("f3d:edge#0", Exactness::ByteExact);
248 assert!(builder.annotations().exactness.is_empty());
249 }
250
251 #[test]
252 fn removing_an_entity_removes_provenance_and_exactness() {
253 let mut builder = AnnotationBuilder::new();
254 let stream = builder.stream("catia:e5_0d_03");
255 builder.note("catia:e5:curve#0", stream, 42).tag("circle");
256 builder.derived("catia:e5:curve#0", "geometry");
257
258 builder.remove_entity("catia:e5:curve#0");
259
260 assert!(!builder
261 .annotations()
262 .provenance
263 .contains_key("catia:e5:curve#0"));
264 assert!(!builder
265 .annotations()
266 .exactness
267 .contains_key("catia:e5:curve#0"));
268 }
269}