abbrev_tree/
lib.rs

1#[cfg(feature = "serde")]
2use serde::{Serialize, Deserialize};
3use std::fmt;
4use std::mem;
5
6#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
7#[derive(Clone, Default, Hash)]
8pub struct AbbrevTree<T> {
9    v: Vec<(String, AbbrevTree<T>)>,
10    pub data: T,
11}
12
13impl<T: Default> AbbrevTree<T> {
14    pub fn new() -> Self {
15        Default::default()
16    }
17
18    // TODO: Recursion is probably bad but oh well.
19    pub fn add(&mut self, item: &str, data: T) {
20        // Find match.
21        for (chunk, subtree) in &mut self.v {
22            let prefix_len = common_prefix_length(chunk, item);
23            if prefix_len > 0 {
24                if prefix_len == chunk.len() {
25                    // Full match. Recurse.
26                    if subtree.v.len() == 0 {
27                        let d = mem::replace(
28                            &mut subtree.data, Default::default()
29                        );
30                        subtree.v.push((
31                            "".to_string(),
32                            AbbrevTree { v: Vec::new(), data: d },
33                        ))
34                    }
35                    return subtree.add(&item[prefix_len..], data);
36                } else {
37                    // Partial match. Split and then add.
38                    let chunk_suffix = chunk.split_off(prefix_len);
39                    let d = mem::replace(&mut subtree.data, Default::default());
40                    let v: Vec<_> = subtree.v.drain(..).collect();
41                    subtree.v.push((
42                        chunk_suffix,
43                        AbbrevTree { v, data: d },
44                    ));
45                    return subtree.add(&item[prefix_len..], data);
46                }
47            }
48        }
49
50        // Else add new subtree.
51        self.v.push((
52            item.to_string(),
53            AbbrevTree { v: Vec::new(), data },
54        ));
55    }
56
57    pub fn complete<'d>(&'d self, item: &str) -> Vec<(String, &'d T)> {
58        let mut v = Vec::new();
59        self._complete("", item, &mut v);
60        v
61    }
62
63    fn _complete<'d>(
64        &'d self, left: &str, item: &str, v: &mut Vec<(String, &'d T)>
65    ) {
66        if self.v.len() == 0 && item == "" {
67            v.push((left.to_string(), &self.data));
68        }
69
70        for (chunk, subtree) in &self.v {
71            let prefix_len = common_prefix_length(chunk, item);
72            // TODO: Make sure this makes sense.
73            if item == "" || item.len() == prefix_len
74                    || chunk.len() == prefix_len {
75                let mut s = left.to_string();
76                s.push_str(chunk);
77                subtree._complete(&s, &item[prefix_len..], v);
78            }
79        }
80    }
81
82    pub fn get_mut<'d>(&'d mut self, item: &str) -> Option<&'d mut T> {
83        self._get_mut("", item)
84    }
85
86    fn _get_mut<'d>(&'d mut self, left: &str, item: &str) -> Option<&'d mut T> {
87        if self.v.len() == 0 && item == "" {
88            // We're a leaf and item is exhausted.
89            return Some(&mut self.data);
90        }
91
92        for (chunk, subtree) in &mut self.v {
93            let prefix_len = common_prefix_length(chunk, item);
94            if prefix_len == chunk.len() {
95                let mut s = left.to_string();
96                s.push_str(chunk);
97                match subtree._get_mut(&s, &item[prefix_len..]) {
98                    Some(d) => return Some(d),
99                    None => (),
100                }
101            }
102        }
103
104        None
105    }
106}
107
108impl<T> fmt::Debug for AbbrevTree<T> {
109    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
110        // FIXME: We should check for f.alternate().
111        let mut stack = vec![self.v.iter()];
112        let mut first = true;
113        while stack.len() > 0 {
114            match stack.last_mut().unwrap().next() {
115                Some(x) => {
116                    if !first {
117                        write!(f, "\n")?;
118                    }
119                    write!(
120                        f,
121                        "{}{:?}",
122                        " ".repeat(2 * (stack.len()-1)),
123                        x.0,
124                    )?;
125                    stack.push((x.1).v.iter());
126                },
127                None => { stack.pop(); },
128            }
129            first = false;
130        }
131        Ok(())
132    }
133}
134
135#[cfg(test)]
136#[test]
137fn test_abbrev_tree() {
138    let mut t = AbbrevTree::new();
139    println!("{:?}", t);
140    assert_eq!(t.v.len(), 0);
141
142    t.add("cat", ());
143    println!("{:?}", t);
144    assert_eq!(t.v.len(), 1);
145    assert_eq!(t.v[0].0, "cat");
146    assert_eq!((t.v[0].1).v.len(), 0);
147
148    t.add("cargo", ());
149    println!("{:?}", t);
150    assert_eq!(t.v.len(), 1);
151    assert_eq!(t.v[0].0, "ca");
152    assert_eq!((t.v[0].1).v.len(), 2);
153    assert_eq!((t.v[0].1).v[0].0, "t");
154    assert_eq!(((t.v[0].1).v[0].1).v.len(), 0);
155    assert_eq!((t.v[0].1).v[1].0, "rgo");
156    assert_eq!(((t.v[0].1).v[1].1).v.len(), 0);
157
158    t.add("chmod", ());
159    println!("{:?}", t);
160    assert_eq!(t.v.len(), 1);
161    assert_eq!(t.v[0].0, "c");
162
163    t.add("chown", ());
164    println!("{:?}", t);
165    assert_eq!(t.v.len(), 1);
166    assert_eq!(t.v[0].0, "c");
167
168    t.add("ls", ());
169    println!("{:?}", t);
170    assert_eq!(t.v.len(), 2);
171    assert_eq!(t.v[0].0, "c");
172    assert_eq!(t.v[1].0, "ls");
173
174    t.add("lshw", ());
175    println!("{:?}", t);
176
177    fn first<A, B, I: IntoIterator<Item = (A, B)>>(i: I) -> Vec<A> {
178        i.into_iter().map(|x: (_, _)| x.0).collect()
179    }
180
181    assert_eq!(first(t.complete("c")), vec![
182        "cat".to_string(),
183        "cargo".to_string(),
184        "chmod".to_string(),
185        "chown".to_string(),
186    ]);
187    assert_eq!(first(t.complete("ca")), vec![
188        "cat".to_string(),
189        "cargo".to_string(),
190    ]);
191    assert_eq!(first(t.complete("cat")), vec!["cat".to_string()]);
192    assert_eq!(first(t.complete("ch")), vec![
193        "chmod".to_string(),
194        "chown".to_string(),
195    ]);
196    assert_eq!(first(t.complete("cho")), vec!["chown".to_string()]);
197    assert_eq!(first(t.complete("chow")), vec!["chown".to_string()]);
198    assert_eq!(first(t.complete("chown")), vec!["chown".to_string()]);
199    assert_eq!(first(t.complete("l")), vec![
200        "ls".to_string(),
201        "lshw".to_string(),
202    ]);
203    assert_eq!(first(t.complete("ls")), vec![
204        "ls".to_string(),
205        "lshw".to_string(),
206    ]);
207    assert_eq!(first(t.complete("lsh")), vec!["lshw".to_string()]);
208    assert_eq!(first(t.complete("lshw")), vec!["lshw".to_string()]);
209    assert_eq!(first(t.complete("x")), Vec::<String>::new());
210    assert_eq!(first(t.complete("xyz")), Vec::<String>::new());
211
212    assert!(t.get_mut("c").is_none());
213    assert!(t.get_mut("ca").is_none());
214    t.get_mut("cat").unwrap();
215    t.get_mut("cargo").unwrap();
216    t.get_mut("chmod").unwrap();
217    t.get_mut("chown").unwrap();
218    t.get_mut("ls").unwrap();
219    t.get_mut("lshw").unwrap();
220    assert!(t.get_mut("xyz").is_none());
221}
222
223fn common_prefix_length(a: &str, b: &str) -> usize {
224    let mut aa = a.char_indices();
225    let mut bb = b.char_indices();
226
227    loop {
228        match (aa.next(), bb.next()) {
229            (Some((ai, ac)), Some((_, bc))) =>
230                if ac != bc {
231                    return ai;
232                },
233            (None, Some((bi, _))) => return bi,
234            (Some((ai, _)), None) => return ai,
235            (None, None) => return a.len(),
236        }
237    }
238}
239
240#[cfg(test)]
241#[test]
242fn test_common_prefix_length() {
243    assert_eq!(common_prefix_length("", "foo"), 0);
244    assert_eq!(common_prefix_length("foo", "foo"), 3);
245    assert_eq!(common_prefix_length("foo", "foobar"), 3);
246    assert_eq!(common_prefix_length("foobar", "foo"), 3);
247    assert_eq!(common_prefix_length("foobar", "bar"), 0);
248    assert_eq!(common_prefix_length("foobar", "foofuzz"), 3);
249}