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::{BorrowedName, 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(
81        &mut self,
82        mut path: std::vec::IntoIter<&BorrowedName>,
83        thing: T,
84    ) -> Result<&mut T, AddError> {
85        match path.next() {
86            Some(name) => match self {
87                Node::Leaf(_) => Err(AddError::Shadow),
88                Node::Internal(children) => {
89                    let entry = children
90                        .entry(name.into())
91                        .or_insert_with(|| Node::Internal(BTreeMap::new()));
92                    entry.add(path, thing)
93                }
94            },
95            None => {
96                match self {
97                    Node::Internal(children) => {
98                        if !children.is_empty() {
99                            return Err(AddError::Shadow);
100                        }
101                    }
102                    Node::Leaf(_) => {
103                        return Err(AddError::Duplicate);
104                    }
105                }
106                *self = Node::Leaf(thing);
107                match self {
108                    Node::Internal(_) => unreachable!(),
109                    Node::Leaf(ref mut t) => Ok(t),
110                }
111            }
112        }
113    }
114
115    fn get_mut<'a, 'b, I>(&'a mut self, mut path: I) -> Option<&'a mut T>
116    where
117        I: Iterator<Item = &'b BorrowedName>,
118    {
119        match path.next() {
120            Some(name) => match self {
121                Node::Leaf(_) => None,
122                Node::Internal(children) => match children.get_mut(name) {
123                    Some(node) => node.get_mut(path),
124                    None => None,
125                },
126            },
127            None => match self {
128                Node::Internal(_) => None,
129                Node::Leaf(ref mut n) => Some(n),
130            },
131        }
132    }
133
134    fn get<'a, 'b, I>(&'a self, mut path: I) -> Option<&'a T>
135    where
136        I: Iterator<Item = &'b BorrowedName>,
137    {
138        match path.next() {
139            Some(name) => match self {
140                Node::Leaf(_) => None,
141                Node::Internal(children) => match children.get(name) {
142                    Some(node) => node.get(path),
143                    None => None,
144                },
145            },
146            None => match self {
147                Node::Internal(_) => None,
148                Node::Leaf(ref n) => Some(n),
149            },
150        }
151    }
152
153    fn remove<'a, I>(&mut self, mut path: std::iter::Peekable<I>) -> Option<T>
154    where
155        I: Iterator<Item = &'a BorrowedName>,
156    {
157        match path.next() {
158            Some(name) => match self {
159                Node::Leaf(_) => None,
160                Node::Internal(children) => {
161                    if path.peek().is_none() {
162                        match children.remove(name) {
163                            Some(Node::Leaf(n)) => Some(n),
164                            Some(Node::Internal(c)) => {
165                                children.insert(name.into(), Node::Internal(c));
166                                return None;
167                            }
168                            None => None,
169                        }
170                    } else {
171                        match children.get_mut(name) {
172                            Some(node) => node.remove(path),
173                            None => None,
174                        }
175                    }
176                }
177            },
178            None => None,
179        }
180    }
181
182    pub fn flatten(self) -> Vec<(String, T)> {
183        fn recurse<T>(path: &str, node: Node<T>, result: &mut Vec<(String, T)>) {
184            match node {
185                Node::Internal(map) => {
186                    for (key, value) in map {
187                        let path = format!("{}/{}", path, key);
188                        recurse(&path, value, result);
189                    }
190                }
191                Node::Leaf(leaf) => {
192                    // The single empty slash is a special case.
193                    // When there are intermediate nodes, we can append "/{path}" to the previous
194                    // path segment. But if the root is a leaf node, there will be no slash unless
195                    // we add one here.
196                    let path = if path.is_empty() { "/" } else { path };
197                    result.push((path.to_string(), leaf));
198                }
199            }
200        }
201
202        let mut result = Vec::new();
203        recurse("", self, &mut result);
204        result
205    }
206
207    pub fn map_ref<R>(&self, f: &mut impl FnMut(&T) -> R) -> Node<R> {
208        match self {
209            Node::Internal(map) => Node::Internal(BTreeMap::from_iter(
210                map.iter().map(|(k, v)| (k.clone(), v.map_ref(f))),
211            )),
212            Node::Leaf(t) => Node::Leaf(f(t)),
213        }
214    }
215}
216
217#[cfg(test)]
218mod tests {
219    use super::*;
220    use assert_matches::assert_matches;
221
222    fn ns_path(str: &str) -> NamespacePath {
223        str.parse().unwrap()
224    }
225
226    #[test]
227    fn test_add() {
228        let mut tree = Tree::new();
229        tree.add(&ns_path("/foo"), 1).unwrap();
230        tree.add(&ns_path("/bar"), 2).unwrap();
231        tree.add(&ns_path("/x/y/z"), 3).unwrap();
232        tree.add(&ns_path("/x/x/z"), 4).unwrap();
233        tree.add(&ns_path("/x/y/w"), 5).unwrap();
234        assert_matches!(tree.add(&ns_path("/foo"), 6), Err(NamespaceError::Duplicate(_)));
235        assert_matches!(tree.add(&ns_path("/bar"), 7), Err(NamespaceError::Duplicate(_)));
236
237        tree.add(&ns_path("/a/b/c"), 0).unwrap();
238        assert_matches!(tree.add(&ns_path("/"), 0), Err(NamespaceError::Shadow(_)));
239        assert_matches!(tree.add(&ns_path("/a"), 0), Err(NamespaceError::Shadow(_)));
240        assert_matches!(tree.add(&ns_path("/a/b"), 0), Err(NamespaceError::Shadow(_)));
241        assert_matches!(tree.add(&ns_path("/a/b/c"), 0), Err(NamespaceError::Duplicate(_)));
242        assert_matches!(tree.add(&ns_path("/a/b/c/d"), 0), Err(NamespaceError::Shadow(_)));
243        assert_matches!(tree.add(&ns_path("/a/b/c/d/e"), 0), Err(NamespaceError::Shadow(_)));
244    }
245
246    #[test]
247    fn test_root() {
248        let mut tree = Tree::new();
249        tree.add(&ns_path("/"), 1).unwrap();
250        assert_matches!(
251            tree.add(&ns_path("/"), 2),
252            Err(NamespaceError::Duplicate(path)) if path.to_string() == "/"
253        );
254        assert_matches!(
255            tree.add(&ns_path("/a"), 3),
256            Err(NamespaceError::Shadow(path)) if path.to_string() == "/a"
257        );
258    }
259
260    #[test]
261    fn test_get_mut() {
262        let mut tree = Tree::new();
263        tree.add(&ns_path("/foo"), 1).unwrap();
264        assert_eq!(*tree.get_mut(&ns_path("/foo")).unwrap(), 1);
265        *tree.get_mut(&ns_path("/foo")).unwrap() = 2;
266        assert_eq!(*tree.get_mut(&ns_path("/foo")).unwrap(), 2);
267        assert_matches!(tree.get_mut(&ns_path("/bar")), None);
268
269        tree.add(&ns_path("/a/b"), 1).unwrap();
270        assert_matches!(tree.get_mut(&ns_path("/a")), None);
271        assert_matches!(tree.get_mut(&ns_path("/a/b/c")), None);
272        assert_eq!(*tree.get_mut(&ns_path("/a/b")).unwrap(), 1);
273        *tree.get_mut(&ns_path("/a/b")).unwrap() = 2;
274        assert_eq!(*tree.get_mut(&ns_path("/a/b")).unwrap(), 2);
275    }
276
277    #[test]
278    fn test_get() {
279        let mut tree = Tree::new();
280        assert_eq!(tree.get(&ns_path("/foo")), None);
281        tree.add(&ns_path("/foo"), 1).unwrap();
282        assert_eq!(*tree.get(&ns_path("/foo")).unwrap(), 1);
283
284        tree.add(&ns_path("/a/b"), 1).unwrap();
285        assert_matches!(tree.get_mut(&ns_path("/a")), None);
286        assert_eq!(*tree.get(&ns_path("/a/b")).unwrap(), 1);
287        assert_matches!(tree.get_mut(&ns_path("/a/b/c")), None);
288    }
289
290    #[test]
291    fn test_remove() {
292        let mut tree = Tree::new();
293        assert_matches!(tree.remove(&ns_path("/foo")), None);
294        tree.add(&ns_path("/foo"), 1).unwrap();
295        assert_matches!(tree.remove(&ns_path("/foo")), Some(1));
296        assert_matches!(tree.remove(&ns_path("/foo")), None);
297
298        tree.add(&ns_path("/foo/bar"), 1).unwrap();
299        assert_matches!(tree.remove(&ns_path("/foo")), None);
300        assert_matches!(tree.remove(&ns_path("/foo/bar/baz")), None);
301        assert_matches!(tree.remove(&ns_path("/foo/bar")), Some(1));
302        assert_matches!(tree.remove(&ns_path("/foo/bar")), None);
303    }
304
305    #[test]
306    fn test_flatten() {
307        let mut tree = Tree::new();
308        tree.add(&ns_path("/a/b/c"), 1).unwrap();
309        tree.add(&ns_path("/b/c/d/e"), 2).unwrap();
310        tree.add(&ns_path("/b/c/e/f"), 3).unwrap();
311        let mut entries = tree.flatten();
312        entries.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
313        assert_eq!(entries.len(), 3);
314        assert_eq!(entries[0].0.to_string(), "/a/b/c");
315        assert_eq!(entries[0].1, 1);
316        assert_eq!(entries[1].0.to_string(), "/b/c/d/e");
317        assert_eq!(entries[1].1, 2);
318        assert_eq!(entries[2].0.to_string(), "/b/c/e/f");
319        assert_eq!(entries[2].1, 3);
320    }
321
322    #[test]
323    fn test_flatten_root() {
324        let mut tree = Tree::new();
325        tree.add(&ns_path("/"), 1).unwrap();
326        let entries = tree.flatten();
327        assert_eq!(entries.len(), 1);
328        assert_eq!(entries[0].0.to_string(), "/");
329        assert_eq!(entries[0].1, 1);
330    }
331
332    #[test]
333    fn test_clone() {
334        let mut tree = Tree::new();
335        tree.add(&ns_path("/foo"), 1).unwrap();
336        tree.add(&ns_path("/bar/baz"), 2).unwrap();
337
338        let tree_clone = tree.clone();
339
340        let mut entries = tree.flatten();
341        entries.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
342
343        let mut entries_clone = tree_clone.flatten();
344        entries_clone.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
345
346        assert_eq!(entries, entries_clone);
347    }
348}
namespace/tree.rs

namespace/
tree.rs
