Skip to main content

asdf_yaml/
document.rs

1//! The document: a node arena plus the directives that introduced it.
2
3use crate::node::{CollectionStyle, Entry, Node, NodeData, NodeId, ScalarStyle};
4use crate::tag::{Tag, TagHandle};
5
6/// The `%YAML` version directive.
7#[derive(Clone, Copy, PartialEq, Eq, Debug)]
8pub struct YamlVersion {
9    /// Major version.
10    pub major: u32,
11    /// Minor version.
12    pub minor: u32,
13}
14
15impl YamlVersion {
16    /// The version ASDF mandates in the tree.
17    pub const V1_1: YamlVersion = YamlVersion { major: 1, minor: 1 };
18}
19
20impl Default for YamlVersion {
21    fn default() -> Self {
22        Self::V1_1
23    }
24}
25
26/// A parsed YAML document: an arena of nodes reachable from a single root.
27#[derive(Clone, Debug)]
28pub struct Document {
29    nodes: Vec<Node>,
30    root: Option<NodeId>,
31    /// The `%YAML` directive, if the document carried one.
32    pub version: Option<YamlVersion>,
33    /// The `%TAG` directives, in the order they appeared.
34    pub tag_handles: Vec<TagHandle>,
35}
36
37impl Default for Document {
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43impl Document {
44    /// An empty document with no root.
45    pub fn new() -> Self {
46        Self { nodes: Vec::new(), root: None, version: None, tag_handles: Vec::new() }
47    }
48
49    /// An empty document carrying the directives ASDF conventionally writes.
50    pub fn new_asdf() -> Self {
51        Self {
52            nodes: Vec::new(),
53            root: None,
54            version: Some(YamlVersion::V1_1),
55            tag_handles: vec![TagHandle::asdf_default()],
56        }
57    }
58
59    /// The number of nodes in the arena.
60    ///
61    /// This counts every node including those only reachable through an
62    /// alias, so it is not the same as the number of distinct tree positions.
63    pub fn node_count(&self) -> usize {
64        self.nodes.len()
65    }
66
67    /// The document's root node, if it has one.
68    pub fn root(&self) -> Option<NodeId> {
69        self.root
70    }
71
72    /// Set the document's root node.
73    pub fn set_root(&mut self, id: NodeId) {
74        self.root = Some(id);
75    }
76
77    /// Add a node to the arena and return its id.
78    pub fn add(&mut self, node: Node) -> NodeId {
79        let id = NodeId(u32::try_from(self.nodes.len()).expect("node arena overflow"));
80        self.nodes.push(node);
81        id
82    }
83
84    /// Borrow a node.
85    ///
86    /// # Panics
87    /// Panics if `id` did not come from this document.
88    pub fn node(&self, id: NodeId) -> &Node {
89        &self.nodes[id.index()]
90    }
91
92    /// Mutably borrow a node.
93    ///
94    /// # Panics
95    /// Panics if `id` did not come from this document.
96    pub fn node_mut(&mut self, id: NodeId) -> &mut Node {
97        &mut self.nodes[id.index()]
98    }
99
100    /// Borrow a node if the id is in range.
101    pub fn get(&self, id: NodeId) -> Option<&Node> {
102        self.nodes.get(id.index())
103    }
104
105    /// Follow aliases until reaching a non-alias node.
106    ///
107    /// Alias chains are bounded by the arena size, so a cycle -- which a
108    /// conforming YAML document cannot contain, but a hand-edited one might --
109    /// terminates rather than looping forever.
110    pub fn resolve(&self, mut id: NodeId) -> NodeId {
111        for _ in 0..=self.nodes.len() {
112            match self.nodes.get(id.index()).map(|n| &n.data) {
113                Some(NodeData::Alias(target)) => id = *target,
114                _ => return id,
115            }
116        }
117        id
118    }
119
120    /// Borrow a node, following aliases first.
121    pub fn resolved(&self, id: NodeId) -> &Node {
122        self.node(self.resolve(id))
123    }
124
125    /// The effective tag of a node, following aliases.
126    ///
127    /// An alias node carries no tag of its own, so the tag of its target is
128    /// what callers mean when they ask.
129    pub fn tag_of(&self, id: NodeId) -> Option<&Tag> {
130        self.resolved(id).tag.as_ref()
131    }
132
133    // ---- convenience constructors -------------------------------------
134
135    /// Add a plain scalar node.
136    pub fn add_scalar(&mut self, value: impl Into<String>) -> NodeId {
137        self.add(Node::scalar(value))
138    }
139
140    /// Add a scalar node with an explicit style.
141    pub fn add_scalar_styled(&mut self, value: impl Into<String>, style: ScalarStyle) -> NodeId {
142        self.add(Node::scalar_styled(value, style))
143    }
144
145    /// Add a sequence node built from existing nodes.
146    pub fn add_sequence(&mut self, items: Vec<NodeId>) -> NodeId {
147        self.add(Node::new(NodeData::Sequence { items, style: CollectionStyle::Auto }))
148    }
149
150    /// Add a mapping node built from existing key/value node pairs.
151    pub fn add_mapping(&mut self, pairs: Vec<(NodeId, NodeId)>) -> NodeId {
152        let entries = pairs.into_iter().map(|(key, value)| Entry { key, value }).collect();
153        self.add(Node::new(NodeData::Mapping { entries, style: CollectionStyle::Auto }))
154    }
155
156    // ---- accessors ----------------------------------------------------
157
158    /// The items of a sequence node, following aliases.
159    pub fn sequence_items(&self, id: NodeId) -> Option<&[NodeId]> {
160        match &self.resolved(id).data {
161            NodeData::Sequence { items, .. } => Some(items),
162            _ => None,
163        }
164    }
165
166    /// The entries of a mapping node, following aliases.
167    pub fn mapping_entries(&self, id: NodeId) -> Option<&[Entry]> {
168        match &self.resolved(id).data {
169            NodeData::Mapping { entries, .. } => Some(entries),
170            _ => None,
171        }
172    }
173
174    /// Look up a mapping value by string key, following aliases.
175    ///
176    /// Where a key appears more than once -- which YAML permits and ASDF
177    /// files occasionally contain -- the first occurrence wins, matching how
178    /// a streaming reader would see it.
179    pub fn mapping_get(&self, id: NodeId, key: &str) -> Option<NodeId> {
180        let entries = self.mapping_entries(id)?;
181        entries.iter().find(|e| self.resolved(e.key).as_str() == Some(key)).map(|e| e.value)
182    }
183
184    /// Insert or replace a mapping entry by string key.
185    ///
186    /// Returns the previous value node when the key was already present.
187    pub fn mapping_set(&mut self, id: NodeId, key: &str, value: NodeId) -> Option<NodeId> {
188        let target = self.resolve(id);
189
190        let existing = self.mapping_entries(target).and_then(|entries| {
191            entries.iter().position(|e| self.resolved(e.key).as_str() == Some(key))
192        });
193
194        match existing {
195            Some(pos) => {
196                let NodeData::Mapping { entries, .. } = &mut self.node_mut(target).data else {
197                    return None;
198                };
199                Some(core::mem::replace(&mut entries[pos].value, value))
200            }
201            None => {
202                let key_id = self.add_scalar(key);
203                let NodeData::Mapping { entries, .. } = &mut self.node_mut(target).data else {
204                    return None;
205                };
206                entries.push(Entry { key: key_id, value });
207                None
208            }
209        }
210    }
211
212    /// Remove a mapping entry by string key, returning the value node.
213    pub fn mapping_remove(&mut self, id: NodeId, key: &str) -> Option<NodeId> {
214        let target = self.resolve(id);
215        let pos = self
216            .mapping_entries(target)?
217            .iter()
218            .position(|e| self.resolved(e.key).as_str() == Some(key))?;
219        let NodeData::Mapping { entries, .. } = &mut self.node_mut(target).data else {
220            return None;
221        };
222        Some(entries.remove(pos).value)
223    }
224
225    /// Index into a sequence, following aliases and accepting negative
226    /// indices that count back from the end.
227    pub fn sequence_get(&self, id: NodeId, index: i64) -> Option<NodeId> {
228        let items = self.sequence_items(id)?;
229        let len = i64::try_from(items.len()).ok()?;
230        let idx = if index < 0 { len + index } else { index };
231        if idx < 0 || idx >= len {
232            return None;
233        }
234        items.get(usize::try_from(idx).ok()?).copied()
235    }
236
237    /// The number of children of a container node, or `None` for scalars.
238    pub fn container_len(&self, id: NodeId) -> Option<usize> {
239        match &self.resolved(id).data {
240            NodeData::Sequence { items, .. } => Some(items.len()),
241            NodeData::Mapping { entries, .. } => Some(entries.len()),
242            _ => None,
243        }
244    }
245}
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    fn tiny() -> (Document, NodeId) {
252        let mut doc = Document::new();
253        let v1 = doc.add_scalar("1");
254        let v2 = doc.add_scalar("2");
255        let k1 = doc.add_scalar("a");
256        let k2 = doc.add_scalar("b");
257        let map = doc.add_mapping(vec![(k1, v1), (k2, v2)]);
258        doc.set_root(map);
259        (doc, map)
260    }
261
262    #[test]
263    fn mapping_lookup_and_order() {
264        let (doc, map) = tiny();
265        assert_eq!(doc.container_len(map), Some(2));
266        let a = doc.mapping_get(map, "a").unwrap();
267        assert_eq!(doc.node(a).as_str(), Some("1"));
268        assert!(doc.mapping_get(map, "missing").is_none());
269
270        let keys: Vec<_> = doc
271            .mapping_entries(map)
272            .unwrap()
273            .iter()
274            .map(|e| doc.node(e.key).as_str().unwrap())
275            .collect();
276        assert_eq!(keys, ["a", "b"], "insertion order must be preserved");
277    }
278
279    #[test]
280    fn mapping_set_replaces_in_place() {
281        let (mut doc, map) = tiny();
282        let three = doc.add_scalar("3");
283        let old = doc.mapping_set(map, "a", three);
284        assert!(old.is_some());
285        assert_eq!(doc.node(doc.mapping_get(map, "a").unwrap()).as_str(), Some("3"));
286        // replacing must not reorder or grow the mapping
287        assert_eq!(doc.container_len(map), Some(2));
288        let keys: Vec<_> = doc
289            .mapping_entries(map)
290            .unwrap()
291            .iter()
292            .map(|e| doc.node(e.key).as_str().unwrap())
293            .collect();
294        assert_eq!(keys, ["a", "b"]);
295    }
296
297    #[test]
298    fn mapping_set_appends_new_key() {
299        let (mut doc, map) = tiny();
300        let v = doc.add_scalar("9");
301        assert!(doc.mapping_set(map, "c", v).is_none());
302        assert_eq!(doc.container_len(map), Some(3));
303        assert_eq!(doc.node(doc.mapping_get(map, "c").unwrap()).as_str(), Some("9"));
304    }
305
306    #[test]
307    fn mapping_remove_works() {
308        let (mut doc, map) = tiny();
309        let removed = doc.mapping_remove(map, "a").unwrap();
310        assert_eq!(doc.node(removed).as_str(), Some("1"));
311        assert_eq!(doc.container_len(map), Some(1));
312        assert!(doc.mapping_get(map, "a").is_none());
313    }
314
315    #[test]
316    fn negative_sequence_indices_count_from_end() {
317        let mut doc = Document::new();
318        let b = doc.add_scalar("b");
319        let c = doc.add_scalar("c");
320        let seq = doc.add_sequence(vec![b, c]);
321
322        assert_eq!(doc.node(doc.sequence_get(seq, 0).unwrap()).as_str(), Some("b"));
323        assert_eq!(doc.node(doc.sequence_get(seq, -1).unwrap()).as_str(), Some("c"));
324        assert_eq!(doc.node(doc.sequence_get(seq, -2).unwrap()).as_str(), Some("b"));
325        assert!(doc.sequence_get(seq, 2).is_none());
326        assert!(doc.sequence_get(seq, -3).is_none());
327    }
328
329    #[test]
330    fn aliases_resolve_through() {
331        let mut doc = Document::new();
332        let target = doc.add_scalar("shared");
333        doc.node_mut(target).anchor = Some("anc".into());
334        let alias = doc.add(Node::new(NodeData::Alias(target)));
335
336        assert!(doc.node(alias).is_alias());
337        assert_eq!(doc.resolve(alias), target);
338        assert_eq!(doc.resolved(alias).as_str(), Some("shared"));
339    }
340
341    #[test]
342    fn alias_cycle_terminates() {
343        // A conforming document cannot express this, but a corrupt one might;
344        // resolution must not hang.
345        let mut doc = Document::new();
346        let a = doc.add(Node::new(NodeData::Alias(NodeId(1))));
347        let b = doc.add(Node::new(NodeData::Alias(NodeId(0))));
348        let _ = doc.resolve(a);
349        let _ = doc.resolve(b);
350    }
351}