namespace/
tree.rs

1// Copyright 2023 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use crate::NamespaceError;
6use cm_types::{IterablePath, Name, NamespacePath};
7use std::collections::BTreeMap;
8
9/// A tree representation of a namespace.
10///
11/// Each leaf node may hold an arbitrary payload `T`. Usually this is a directory client endpoint
12/// or client proxy.
13#[derive(Debug)]
14pub struct Tree<T> {
15    root: Node<T>,
16}
17
18impl<T> Tree<T> {
19    pub fn new() -> Self {
20        Tree { root: Node::Internal(BTreeMap::new()) }
21    }
22
23    pub fn add(self: &mut Self, path: &NamespacePath, thing: T) -> Result<&mut T, NamespaceError> {
24        let names = path.split();
25        self.root.add(names.into_iter(), thing).map_err(|e| match e {
26            AddError::Shadow => NamespaceError::Shadow(path.clone()),
27            AddError::Duplicate => NamespaceError::Duplicate(path.clone()),
28        })
29    }
30
31    pub fn get_mut(&mut self, path: &NamespacePath) -> Option<&mut T> {
32        self.root.get_mut(path.iter_segments())
33    }
34
35    pub fn get(&self, path: &NamespacePath) -> Option<&T> {
36        self.root.get(path.iter_segments())
37    }
38
39    pub fn remove(&mut self, path: &NamespacePath) -> Option<T> {
40        self.root.remove(path.iter_segments().peekable())
41    }
42
43    pub fn flatten(self) -> Vec<(NamespacePath, T)> {
44        self.root
45            .flatten()
46            .into_iter()
47            .map(|(path, thing)| (NamespacePath::new(path).unwrap(), thing))
48            .collect()
49    }
50
51    pub fn map_ref<R>(&self, mut f: impl FnMut(&T) -> R) -> Tree<R> {
52        Tree { root: self.root.map_ref(&mut f) }
53    }
54}
55
56impl<T> Default for Tree<T> {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl<T: Clone> Clone for Tree<T> {
63    fn clone(&self) -> Self {
64        Self { root: self.root.map_ref(&mut <T as Clone>::clone) }
65    }
66}
67
68#[derive(Debug)]
69enum Node<T> {
70    Internal(BTreeMap<Name, Node<T>>),
71    Leaf(T),
72}
73
74enum AddError {
75    Shadow,
76    Duplicate,
77}
78
79impl<T> Node<T> {
80    fn add(&mut self, mut path: std::vec::IntoIter<Name>, thing: T) -> Result<&mut T, AddError> {
81        match path.next() {
82            Some(name) => match self {
83                Node::Leaf(_) => Err(AddError::Shadow),
84                Node::Internal(children) => {
85                    let entry =
86                        children.entry(name).or_insert_with(|| Node::Internal(BTreeMap::new()));
87                    entry.add(path, thing)
88                }
89            },
90            None => {
91                match self {
92                    Node::Internal(children) => {
93                        if !children.is_empty() {
94                            return Err(AddError::Shadow);
95                        }
96                    }
97                    Node::Leaf(_) => {
98                        return Err(AddError::Duplicate);
99                    }
100                }
101                *self = Node::Leaf(thing);
102                match self {
103                    Node::Internal(_) => unreachable!(),
104                    Node::Leaf(ref mut t) => Ok(t),
105                }
106            }
107        }
108    }
109
110    fn get_mut<'a, 'b, I>(&'a mut self, mut path: I) -> Option<&'a mut T>
111    where
112        I: Iterator<Item = &'b Name>,
113    {
114        match path.next() {
115            Some(name) => match self {
116                Node::Leaf(_) => None,
117                Node::Internal(children) => match children.get_mut(name) {
118                    Some(node) => node.get_mut(path),
119                    None => None,
120                },
121            },
122            None => match self {
123                Node::Internal(_) => None,
124                Node::Leaf(ref mut n) => Some(n),
125            },
126        }
127    }
128
129    fn get<'a, 'b, I>(&'a self, mut path: I) -> Option<&'a T>
130    where
131        I: Iterator<Item = &'b Name>,
132    {
133        match path.next() {
134            Some(name) => match self {
135                Node::Leaf(_) => None,
136                Node::Internal(children) => match children.get(name) {
137                    Some(node) => node.get(path),
138                    None => None,
139                },
140            },
141            None => match self {
142                Node::Internal(_) => None,
143                Node::Leaf(ref n) => Some(n),
144            },
145        }
146    }
147
148    fn remove<'a, I>(&mut self, mut path: std::iter::Peekable<I>) -> Option<T>
149    where
150        I: Iterator<Item = &'a Name>,
151    {
152        match path.next() {
153            Some(name) => match self {
154                Node::Leaf(_) => None,
155                Node::Internal(children) => {
156                    if path.peek().is_none() {
157                        match children.remove(name) {
158                            Some(Node::Leaf(n)) => Some(n),
159                            Some(Node::Internal(c)) => {
160                                children.insert(name.clone(), Node::Internal(c));
161                                return None;
162                            }
163                            None => None,
164                        }
165                    } else {
166                        match children.get_mut(name) {
167                            Some(node) => node.remove(path),
168                            None => None,
169                        }
170                    }
171                }
172            },
173            None => None,
174        }
175    }
176
177    pub fn flatten(self) -> Vec<(String, T)> {
178        fn recurse<T>(path: &str, node: Node<T>, result: &mut Vec<(String, T)>) {
179            match node {
180                Node::Internal(map) => {
181                    for (key, value) in map {
182                        let path = format!("{}/{}", path, key);
183                        recurse(&path, value, result);
184                    }
185                }
186                Node::Leaf(leaf) => {
187                    // The single empty slash is a special case.
188                    // When there are intermediate nodes, we can append "/{path}" to the previous
189                    // path segment. But if the root is a leaf node, there will be no slash unless
190                    // we add one here.
191                    let path = if path.is_empty() { "/" } else { path };
192                    result.push((path.to_string(), leaf));
193                }
194            }
195        }
196
197        let mut result = Vec::new();
198        recurse("", self, &mut result);
199        result
200    }
201
202    pub fn map_ref<R>(&self, f: &mut impl FnMut(&T) -> R) -> Node<R> {
203        match self {
204            Node::Internal(map) => Node::Internal(BTreeMap::from_iter(
205                map.iter().map(|(k, v)| (k.clone(), v.map_ref(f))),
206            )),
207            Node::Leaf(t) => Node::Leaf(f(t)),
208        }
209    }
210}
211
212#[cfg(test)]
213mod tests {
214    use super::*;
215    use assert_matches::assert_matches;
216
217    fn ns_path(str: &str) -> NamespacePath {
218        str.parse().unwrap()
219    }
220
221    #[test]
222    fn test_add() {
223        let mut tree = Tree::new();
224        tree.add(&ns_path("/foo"), 1).unwrap();
225        tree.add(&ns_path("/bar"), 2).unwrap();
226        tree.add(&ns_path("/x/y/z"), 3).unwrap();
227        tree.add(&ns_path("/x/x/z"), 4).unwrap();
228        tree.add(&ns_path("/x/y/w"), 5).unwrap();
229        assert_matches!(tree.add(&ns_path("/foo"), 6), Err(NamespaceError::Duplicate(_)));
230        assert_matches!(tree.add(&ns_path("/bar"), 7), Err(NamespaceError::Duplicate(_)));
231
232        tree.add(&ns_path("/a/b/c"), 0).unwrap();
233        assert_matches!(tree.add(&ns_path("/"), 0), Err(NamespaceError::Shadow(_)));
234        assert_matches!(tree.add(&ns_path("/a"), 0), Err(NamespaceError::Shadow(_)));
235        assert_matches!(tree.add(&ns_path("/a/b"), 0), Err(NamespaceError::Shadow(_)));
236        assert_matches!(tree.add(&ns_path("/a/b/c"), 0), Err(NamespaceError::Duplicate(_)));
237        assert_matches!(tree.add(&ns_path("/a/b/c/d"), 0), Err(NamespaceError::Shadow(_)));
238        assert_matches!(tree.add(&ns_path("/a/b/c/d/e"), 0), Err(NamespaceError::Shadow(_)));
239    }
240
241    #[test]
242    fn test_root() {
243        let mut tree = Tree::new();
244        tree.add(&ns_path("/"), 1).unwrap();
245        assert_matches!(
246            tree.add(&ns_path("/"), 2),
247            Err(NamespaceError::Duplicate(path)) if path.to_string() == "/"
248        );
249        assert_matches!(
250            tree.add(&ns_path("/a"), 3),
251            Err(NamespaceError::Shadow(path)) if path.to_string() == "/a"
252        );
253    }
254
255    #[test]
256    fn test_get_mut() {
257        let mut tree = Tree::new();
258        tree.add(&ns_path("/foo"), 1).unwrap();
259        assert_eq!(*tree.get_mut(&ns_path("/foo")).unwrap(), 1);
260        *tree.get_mut(&ns_path("/foo")).unwrap() = 2;
261        assert_eq!(*tree.get_mut(&ns_path("/foo")).unwrap(), 2);
262        assert_matches!(tree.get_mut(&ns_path("/bar")), None);
263
264        tree.add(&ns_path("/a/b"), 1).unwrap();
265        assert_matches!(tree.get_mut(&ns_path("/a")), None);
266        assert_matches!(tree.get_mut(&ns_path("/a/b/c")), None);
267        assert_eq!(*tree.get_mut(&ns_path("/a/b")).unwrap(), 1);
268        *tree.get_mut(&ns_path("/a/b")).unwrap() = 2;
269        assert_eq!(*tree.get_mut(&ns_path("/a/b")).unwrap(), 2);
270    }
271
272    #[test]
273    fn test_get() {
274        let mut tree = Tree::new();
275        assert_eq!(tree.get(&ns_path("/foo")), None);
276        tree.add(&ns_path("/foo"), 1).unwrap();
277        assert_eq!(*tree.get(&ns_path("/foo")).unwrap(), 1);
278
279        tree.add(&ns_path("/a/b"), 1).unwrap();
280        assert_matches!(tree.get_mut(&ns_path("/a")), None);
281        assert_eq!(*tree.get(&ns_path("/a/b")).unwrap(), 1);
282        assert_matches!(tree.get_mut(&ns_path("/a/b/c")), None);
283    }
284
285    #[test]
286    fn test_remove() {
287        let mut tree = Tree::new();
288        assert_matches!(tree.remove(&ns_path("/foo")), None);
289        tree.add(&ns_path("/foo"), 1).unwrap();
290        assert_matches!(tree.remove(&ns_path("/foo")), Some(1));
291        assert_matches!(tree.remove(&ns_path("/foo")), None);
292
293        tree.add(&ns_path("/foo/bar"), 1).unwrap();
294        assert_matches!(tree.remove(&ns_path("/foo")), None);
295        assert_matches!(tree.remove(&ns_path("/foo/bar/baz")), None);
296        assert_matches!(tree.remove(&ns_path("/foo/bar")), Some(1));
297        assert_matches!(tree.remove(&ns_path("/foo/bar")), None);
298    }
299
300    #[test]
301    fn test_flatten() {
302        let mut tree = Tree::new();
303        tree.add(&ns_path("/a/b/c"), 1).unwrap();
304        tree.add(&ns_path("/b/c/d/e"), 2).unwrap();
305        tree.add(&ns_path("/b/c/e/f"), 3).unwrap();
306        let mut entries = tree.flatten();
307        entries.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
308        assert_eq!(entries.len(), 3);
309        assert_eq!(entries[0].0.to_string(), "/a/b/c");
310        assert_eq!(entries[0].1, 1);
311        assert_eq!(entries[1].0.to_string(), "/b/c/d/e");
312        assert_eq!(entries[1].1, 2);
313        assert_eq!(entries[2].0.to_string(), "/b/c/e/f");
314        assert_eq!(entries[2].1, 3);
315    }
316
317    #[test]
318    fn test_flatten_root() {
319        let mut tree = Tree::new();
320        tree.add(&ns_path("/"), 1).unwrap();
321        let entries = tree.flatten();
322        assert_eq!(entries.len(), 1);
323        assert_eq!(entries[0].0.to_string(), "/");
324        assert_eq!(entries[0].1, 1);
325    }
326
327    #[test]
328    fn test_clone() {
329        let mut tree = Tree::new();
330        tree.add(&ns_path("/foo"), 1).unwrap();
331        tree.add(&ns_path("/bar/baz"), 2).unwrap();
332
333        let tree_clone = tree.clone();
334
335        let mut entries = tree.flatten();
336        entries.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
337
338        let mut entries_clone = tree_clone.flatten();
339        entries_clone.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
340
341        assert_eq!(entries, entries_clone);
342    }
343}
namespace/tree.rs

namespace/
tree.rs
