1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
//! A datastructure for holding key map entries
use std::fmt::Debug;
#[derive(Debug, Clone)]
struct Node<Value: Debug> {
label: u8,
children: Vec<Node<Value>>,
value: Option<Value>,
}
impl<Value: Debug> Node<Value> {
fn new(label: u8) -> Self {
Self {
label,
children: Vec::new(),
value: None,
}
}
fn insert(&mut self, key: &[u8], value: Value) {
if key.is_empty() {
// We've reached the leaf
self.value = Some(value);
return;
}
match self
.children
.binary_search_by(|node| node.label.cmp(&key[0]))
{
Ok(idx) => {
self.children[idx].insert(&key[1..], value);
}
Err(idx) => {
self.children.insert(idx, Node::new(key[0]));
self.children[idx].insert(&key[1..], value);
}
}
}
fn lookup(&self, key: &[u8], depth: usize) -> NodeFind<&Value> {
if key.is_empty() {
// We've matched the maximum extent of the input key.
if self.children.is_empty() {
match self.value.as_ref() {
Some(value) => {
// An unambiguous match for the entire input
return NodeFind::Exact(depth, value);
}
None => panic!("Node has no children and no value!?"),
}
}
return match self.value.as_ref() {
Some(value) => NodeFind::AmbiguousMatch(depth, value),
None => NodeFind::AmbiguousBackTrack,
};
}
match self
.children
.binary_search_by(|node| node.label.cmp(&key[0]))
{
Ok(idx) => {
match self.children[idx].lookup(&key[1..], depth + 1) {
NodeFind::AmbiguousBackTrack => {
// The child didn't have an exact match, so check
// whether we do
match self.value.as_ref() {
Some(value) => NodeFind::AmbiguousMatch(depth, value),
None => NodeFind::AmbiguousBackTrack,
}
}
result => result,
}
}
Err(_) => {
if depth == 0 {
NodeFind::None
} else {
match self.value.as_ref() {
Some(value) => NodeFind::Exact(depth, value),
None => NodeFind::AmbiguousBackTrack,
}
}
}
}
}
}
/// Internal lookup disposition
enum NodeFind<Value> {
/// No possible matches
None,
/// Found an exact match. (match-len, value)
Exact(usize, Value),
/// Didn't find an exact match at the full extent of the key,
/// so ask the upper layer to back track to find a partial.
AmbiguousBackTrack,
/// After backtracking, found a prefix match, but we know that
/// there might be a more specific match given more data. (match-len,
/// value).
AmbiguousMatch(usize, Value),
}
/// Holds the result of a lookup operation
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Found<Value> {
/// There are definitively no possible matches
None,
/// We found an unambiguous match.
/// The data is (length-of-match, value)
Exact(usize, Value),
/// We found a match, but there are other longer matches
/// that are possible. Ideally we'd accumulate more data
/// to know for sure.
/// The data is (length-of-shortest-match, value)
Ambiguous(usize, Value),
/// If we had more data, we might match something
NeedData,
}
/// The `KeyMap` struct is intended to hold the text sequences
/// generated by unix terminal programs. Those sequences have
/// overlapping/ambiguous meaning which requires having more
/// data to correctly interpret the sequence.
/// The `lookup` operation returns an enum describing the
/// confidence of a match rather than a simple map lookup.
#[derive(Debug, Clone)]
pub struct KeyMap<Value: Debug + Clone> {
root: Node<Value>,
}
impl<Value: Debug + Clone> Default for KeyMap<Value> {
fn default() -> Self {
Self::new()
}
}
impl<Value: Debug + Clone> KeyMap<Value> {
pub fn new() -> Self {
Self { root: Node::new(0) }
}
/// Insert a value into the keymap
pub fn insert<K: AsRef<[u8]>>(&mut self, key: K, value: Value) {
self.root.insert(key.as_ref(), value)
}
/// Perform a lookup for `key`.
/// `key` can be a string consisting of a sequence of bytes.
/// The `lookup` operation considers the prefix of `key` and searches
/// for a match.
///
/// If `Found::None` is returned then the prefix for key has no matching
/// keymap entry.
///
/// If `Found::Exact` is returned then the returned value informs the caller
/// of the length of the key that was matched; the remainder of the key
/// was not considered and is something that should be considered again
/// in a subsequent lookup operation.
///
/// If `Found::Ambiguous` is returned then the key matched a valid
/// entry (which is returned as the value), but there is at least one
/// other entry that could match if more data were available. If the caller
/// knows that no more data is available immediately then it may be valid
/// to treat this result as equivalent to `Found::Exact`. The intended
/// use for this variant is to handle the case where a sequence straddles
/// a buffer boundary (eg: fixed size buffer receives a partial sequence,
/// and the remainder is immediately available on the next read) without
/// misinterpreting the read data.
///
/// If `Found::NeedData` is returned it indicates that `key` is too short
/// to determine a match. The purpose is similar to the `Found::Ambiguous`
/// case; if the caller knows that no more data is available this can be
/// treated as `Found::None`, but otherwise it would be best to read more
/// data from the stream and retry with a longer input.
pub fn lookup<S: AsRef<[u8]>>(&self, key: S) -> Found<Value> {
match self.root.lookup(key.as_ref(), 0) {
NodeFind::None => Found::None,
NodeFind::AmbiguousBackTrack => Found::NeedData,
NodeFind::Exact(depth, value) => Found::Exact(depth, value.clone()),
NodeFind::AmbiguousMatch(depth, value) => Found::Ambiguous(depth, value.clone()),
}
}
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn lookup_empty() {
let km: KeyMap<bool> = KeyMap::new();
assert_eq!(km.lookup("boo"), Found::None);
}
#[test]
fn lookup() {
let mut km = KeyMap::new();
km.insert("boa", true);
km.insert("boo", true);
km.insert("boom", false);
assert_eq!(km.lookup("b"), Found::NeedData);
assert_eq!(km.lookup("bo"), Found::NeedData);
assert_eq!(km.lookup("boa"), Found::Exact(3, true),);
assert_eq!(km.lookup("boo"), Found::Ambiguous(3, true),);
assert_eq!(km.lookup("boom"), Found::Exact(4, false),);
assert_eq!(km.lookup("boom!"), Found::Exact(4, false),);
}
#[test]
fn sequence() {
let mut km = KeyMap::new();
km.insert("\x03", true);
km.insert("\x27", true);
km.insert("\x03XYZ", true);
assert_eq!(km.lookup("\x03"), Found::Ambiguous(1, true),);
assert_eq!(km.lookup("\x03foo"), Found::Exact(1, true),);
assert_eq!(km.lookup("\x03X"), Found::Ambiguous(1, true),);
}
}