ferron/config/lookup/
tree.rs

1use std::collections::{BTreeMap, VecDeque};
2
3use ferron_common::config::Conditional;
4
5use crate::config::lookup::conditionals::{match_conditional, ConditionMatchData};
6
7/// A single key for the configuration tree lookup
8#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
9pub enum ConfigFilterTreeSingleKey {
10  /// The configuration is a host configuration
11  IsHostConfiguration,
12  /// The port number
13  Port(u16),
14  /// An octet of an IPv4 address
15  IPv4Octet(u8),
16  /// An octet of an IPv6 address
17  IPv6Octet(u8),
18  /// The configuration is for localhost
19  IsLocalhost,
20  /// A hostname domain level
21  HostDomainLevel(String),
22  /// A hostname domain level with wildcard
23  HostDomainLevelWildcard,
24  /// A location path segment
25  LocationSegment(String),
26  /// A conditional
27  Conditional(Conditional),
28  // Note how error handler status isn't included in the tree key,
29  // because error handlers are stored separately and don't affect the tree structure
30}
31
32impl ConfigFilterTreeSingleKey {
33  /// Checks whether the single tree key contains a predicate
34  pub fn is_predicate(&self) -> bool {
35    matches!(self, Self::HostDomainLevelWildcard | Self::Conditional(_))
36  }
37}
38
39/// A multi-key for the configuration tree lookup
40#[derive(Clone, Debug, Eq, PartialEq)]
41struct ConfigFilterTreeMultiKey(Vec<ConfigFilterTreeSingleKey>);
42
43#[allow(clippy::non_canonical_partial_ord_impl)]
44impl PartialOrd for ConfigFilterTreeMultiKey {
45  fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
46    for i in 0..self.0.len().max(other.0.len()) {
47      let self_element = self.0.get(i);
48      let other_element = other.0.get(i);
49      match (self_element, other_element) {
50        (Some(a), Some(b)) => {
51          let cmp = a.cmp(b);
52          if cmp != std::cmp::Ordering::Equal {
53            return Some(cmp);
54          }
55        }
56        _ => return None,
57      }
58    }
59    Some(std::cmp::Ordering::Equal)
60  }
61}
62
63impl Ord for ConfigFilterTreeMultiKey {
64  fn cmp(&self, other: &Self) -> std::cmp::Ordering {
65    // The `partial_cmp` method returns `None` if the keys are of different lengths,
66    // but for the `Ord` trait we need to define a total order.
67    // In this case, we can consider keys of different lengths as equal for ordering purposes,
68    // since they will be handled separately in the tree structure.
69    //
70    // This is especially important in radix trees with BTreeMaps as children.
71    self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
72  }
73}
74
75/// A node in the configuration filter tree
76#[derive(Debug)]
77struct ConfigFilterTreeNode<T> {
78  /// The configuration value associated with this node, if any
79  pub value: Option<T>,
80  /// The child nodes of this node, indexed by their fixed value multi-key
81  pub children_fixed: BTreeMap<ConfigFilterTreeMultiKey, ConfigFilterTreeNode<T>>,
82  /// The child nodes of this node, indexed by their predicate single key
83  pub children_predicate: BTreeMap<ConfigFilterTreeSingleKey, ConfigFilterTreeNode<T>>,
84}
85
86/// The configuration filter tree, used for efficient lookup of configurations based on multiple criteria.
87/// The tree is a hybrid of a radix tree and a trie, where each node can have multiple fixed value children (like a radix tree)
88/// and multiple predicate children (like a trie).
89#[derive(Debug)]
90pub struct ConfigFilterTree<T> {
91  /// The root node of the configuration tree
92  root: ConfigFilterTreeNode<T>,
93}
94
95impl<T> ConfigFilterTree<T> {
96  /// Creates a new empty configuration filter tree
97  pub fn new() -> Self {
98    Self {
99      root: ConfigFilterTreeNode {
100        value: None,
101        children_fixed: BTreeMap::new(),
102        children_predicate: BTreeMap::new(),
103      },
104    }
105  }
106
107  /// Inserts a configuration value into the tree based on the provided multi-key
108  #[allow(dead_code)]
109  pub fn insert(&mut self, key: Vec<ConfigFilterTreeSingleKey>, value: T) {
110    self.insert_node(key).replace(value);
111  }
112
113  /// Inserts a node into the tree based on the provided multi-key and returns a mutable reference to the value at that node.
114  pub fn insert_node(&mut self, key: Vec<ConfigFilterTreeSingleKey>) -> &mut Option<T> {
115    let mut current_node = &mut self.root;
116    let mut key_iter = key.into_iter();
117    let mut key_option = key_iter.next();
118    while let Some(key) = key_option.take() {
119      if key.is_predicate() {
120        // This key is a predicate, so we need to look for it in the children_predicate map
121        match current_node.children_predicate.entry(key) {
122          std::collections::btree_map::Entry::Occupied(entry) => {
123            current_node = entry.into_mut();
124          }
125          std::collections::btree_map::Entry::Vacant(entry) => {
126            current_node = entry.insert(ConfigFilterTreeNode {
127              value: None,
128              children_fixed: BTreeMap::new(),
129              children_predicate: BTreeMap::new(),
130            });
131          }
132        }
133        key_option = key_iter.next();
134      } else {
135        // This key is not a predicate, so we can try to insert it into the children_fixed map as part of a multi-key
136        let mut multi_key = ConfigFilterTreeMultiKey(vec![key]);
137        match current_node.children_fixed.entry(multi_key) {
138          std::collections::btree_map::Entry::Occupied(mut entry) => {
139            let entry_key = entry.key();
140            for i in 1..=entry_key.0.len() {
141              if i == entry_key.0.len() {
142                // All keys match, continue with the existing multi-key
143
144                key_option = key_iter.next();
145                // Safety: after the current node is changed, we immediately break the inner loop and continue
146                // the outer loop. Without "unsafe", the borrow checker would not allow us to borrow the entry
147                // as mutable again in the next iteration of the outer loop.
148                current_node = unsafe {
149                  std::mem::transmute::<&mut ConfigFilterTreeNode<T>, &mut ConfigFilterTreeNode<T>>(entry.get_mut())
150                };
151                break;
152              }
153              key_option = key_iter.next();
154              let mut break_multi_key = false;
155              if let Some(key) = &key_option {
156                if key != &entry_key.0[i] {
157                  // Keys differ, break the multi-key and insert a new node
158                  break_multi_key = true;
159                }
160              } else {
161                // No more keys, break the multi-key
162                break_multi_key = true;
163              }
164              if break_multi_key {
165                // Break the multi-key at index i and insert a new node for the remaining keys
166                let (mut entry_key, entry_value) = entry.remove_entry();
167                let entry_key_right = ConfigFilterTreeMultiKey(entry_key.0.split_off(i));
168                #[allow(clippy::mutable_key_type)]
169                let mut new_children_fixed = BTreeMap::new();
170                new_children_fixed.insert(entry_key_right, entry_value);
171                match current_node.children_fixed.entry(entry_key) {
172                  std::collections::btree_map::Entry::Occupied(entry) => {
173                    current_node = entry.into_mut();
174                  }
175                  std::collections::btree_map::Entry::Vacant(entry) => {
176                    current_node = entry.insert(ConfigFilterTreeNode {
177                      value: None,
178                      children_fixed: new_children_fixed,
179                      children_predicate: BTreeMap::new(),
180                    });
181                  }
182                }
183                break;
184              }
185            }
186          }
187          std::collections::btree_map::Entry::Vacant(entry) => {
188            multi_key = entry.into_key();
189
190            key_option = key_iter.next();
191            while let Some(key) = &key_option {
192              if !key.is_predicate() {
193                let key = key_option.take().expect("key_option should be Some here");
194                multi_key.0.push(key);
195                key_option = key_iter.next();
196              } else {
197                break;
198              }
199            }
200
201            match current_node.children_fixed.entry(multi_key) {
202              std::collections::btree_map::Entry::Occupied(entry) => {
203                current_node = entry.into_mut();
204              }
205              std::collections::btree_map::Entry::Vacant(entry) => {
206                current_node = entry.insert(ConfigFilterTreeNode {
207                  value: None,
208                  children_fixed: BTreeMap::new(),
209                  children_predicate: BTreeMap::new(),
210                });
211              }
212            };
213          }
214        }
215      }
216    }
217
218    &mut current_node.value
219  }
220
221  /// Obtains a reference to the configuration value associated with the provided multi-key, if it exists
222  pub fn get<'a, 'b>(
223    &'a self,
224    key: Vec<ConfigFilterTreeSingleKey>,
225    condition_match_data: Option<ConditionMatchData<'b>>,
226  ) -> Result<Option<&'a T>, Box<dyn std::error::Error + Send + Sync>> {
227    let mut current_nodes = VecDeque::new();
228    current_nodes.push_back((&self.root, ConfigFilterTreeMultiKey(key)));
229    let mut value = current_nodes[0].0.value.as_ref();
230    while let Some((mut current_node, mut key)) = current_nodes.pop_front() {
231      while !key.0.is_empty() {
232        let key_end = key.0.split_off(1);
233        let mut partial_key = ConfigFilterTreeMultiKey(key.0);
234        key.0 = key_end;
235        if let Some((child_key, child)) = current_node.children_fixed.get_key_value(&partial_key) {
236          // If a fixed multi-key matches, continue with the child node and the remaining key
237          current_node = child;
238          let mut index = 0;
239          let mut secondary_index = 1;
240          let mut is_matching = true;
241          while let (Some(key_single), Some(child_key_single)) = (key.0.get(index), child_key.0.get(secondary_index)) {
242            if std::mem::discriminant(key_single) == std::mem::discriminant(child_key_single) {
243              // The keys have the same variant, so we can compare them
244              if key_single != child_key_single {
245                // The keys differ, so the fixed multi-key does not match
246                is_matching = false;
247                break;
248              } else {
249                // The keys match, continue with the next key
250                secondary_index += 1;
251              }
252            }
253            index += 1;
254          }
255          if !is_matching || (index >= key.0.len() && secondary_index < child_key.0.len()) {
256            // The keys differ or are out of bounds, so the fixed multi-key does not match
257            break;
258          }
259          key.0 = key.0.split_off(index);
260          if current_node.value.is_some() {
261            value = current_node.value.as_ref();
262          }
263          continue;
264        }
265
266        let partial_key = partial_key.0.remove(0);
267        if partial_key.is_predicate() {
268          // The next key is a predicate, so try to match the predicate single keys
269          if let Some(child) = current_node.children_predicate.get(&partial_key) {
270            // If a predicate key matches, continue with the child node and the remaining key
271            current_node = child;
272            if current_node.value.is_some() {
273              value = current_node.value.as_ref();
274            }
275            continue;
276          }
277        }
278
279        // If no fixed multi-key matches, try to match the predicate single keys
280        let mut conditional_predicate_matched = false;
281        for (predicate_key, child) in &current_node.children_predicate {
282          if predicate_key.is_predicate() {
283            // This is a predicate key, so we need to check if it matches the condition data
284            if predicate_key == &ConfigFilterTreeSingleKey::HostDomainLevelWildcard
285              && matches!(partial_key, ConfigFilterTreeSingleKey::HostDomainLevel(_))
286            {
287              // Host domain level wildcard matches any host domain level
288              current_node = child;
289              let mut index = 0;
290              while let Some(ConfigFilterTreeSingleKey::HostDomainLevel(_)) = key.0.get(index) {
291                index += 1;
292              }
293              key.0 = key.0.split_off(index);
294              if current_node.value.is_some() {
295                value = current_node.value.as_ref();
296              }
297              // There can be only one variation of HostDomainLevelWildcard at the same level of the tree,
298              // so we can break after the first match. This is in contrast to Conditional predicates,
299              // where there can be multiple different conditionals at the same level of the tree,
300              // and we need to check all of them.
301              break;
302            } else if let ConfigFilterTreeSingleKey::Conditional(conditional) = predicate_key {
303              // Conditional key, check if the condition matches
304              if let Some(condition_match_data) = condition_match_data.as_ref() {
305                if match_conditional(conditional, condition_match_data)? {
306                  let current_node = child;
307                  if current_node.value.is_some() {
308                    value = current_node.value.as_ref();
309                  }
310                  // Push to current nodes queue. At the end, a value with the longest matching key will be returned,
311                  // so we want to continue searching for other matches after this one.
312                  current_nodes.push_back((current_node, key.clone()));
313                  conditional_predicate_matched = true;
314                }
315              }
316            }
317          }
318        }
319        if conditional_predicate_matched {
320          break;
321        }
322
323        // If nothing happened, probably no matching fixed multi-key or predicate single key found...
324      }
325    }
326
327    Ok(value)
328  }
329}
330
331#[cfg(test)]
332mod tests {
333  use super::*;
334
335  #[test]
336  fn test_basic_get() {
337    let mut tree = ConfigFilterTree::new();
338    tree.insert(
339      vec![
340        ConfigFilterTreeSingleKey::Port(80),
341        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
342        ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
343        ConfigFilterTreeSingleKey::HostDomainLevelWildcard,
344      ],
345      "Example",
346    );
347    tree.insert(
348      vec![
349        ConfigFilterTreeSingleKey::Port(80),
350        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
351        ConfigFilterTreeSingleKey::HostDomainLevel("example2".to_string()),
352        ConfigFilterTreeSingleKey::HostDomainLevel("www".to_string()),
353      ],
354      "Example 2",
355    );
356
357    assert_eq!(
358      tree
359        .get(
360          vec![
361            ConfigFilterTreeSingleKey::Port(80),
362            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
363            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
364            ConfigFilterTreeSingleKey::HostDomainLevel("www".to_string()),
365          ],
366          None
367        )
368        .unwrap(),
369      Some(&"Example")
370    );
371
372    assert_eq!(
373      tree
374        .get(
375          vec![
376            ConfigFilterTreeSingleKey::Port(80),
377            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
378            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
379            ConfigFilterTreeSingleKey::HostDomainLevel("subsite".to_string()),
380          ],
381          None
382        )
383        .unwrap(),
384      Some(&"Example")
385    );
386
387    assert_eq!(
388      tree
389        .get(
390          vec![
391            ConfigFilterTreeSingleKey::Port(80),
392            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
393            ConfigFilterTreeSingleKey::HostDomainLevel("example2".to_string()),
394            ConfigFilterTreeSingleKey::HostDomainLevel("www".to_string()),
395          ],
396          None
397        )
398        .unwrap(),
399      Some(&"Example 2")
400    );
401
402    assert_eq!(
403      tree
404        .get(
405          vec![
406            ConfigFilterTreeSingleKey::Port(80),
407            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
408            ConfigFilterTreeSingleKey::HostDomainLevel("example3".to_string()),
409            ConfigFilterTreeSingleKey::HostDomainLevel("www".to_string()),
410          ],
411          None
412        )
413        .unwrap(),
414      None
415    );
416
417    assert_eq!(
418      tree
419        .get(
420          vec![
421            ConfigFilterTreeSingleKey::Port(80),
422            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
423            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
424          ],
425          None
426        )
427        .unwrap(),
428      None
429    );
430  }
431
432  #[test]
433  fn test_empty_tree() {
434    let tree: ConfigFilterTree<&str> = ConfigFilterTree::new();
435    assert_eq!(tree.get(vec![], None).unwrap(), None);
436  }
437
438  #[test]
439  fn test_insert_and_get_single_key() {
440    let mut tree = ConfigFilterTree::new();
441    tree.insert(vec![ConfigFilterTreeSingleKey::Port(80)], "Port 80");
442    assert_eq!(
443      tree.get(vec![ConfigFilterTreeSingleKey::Port(80)], None).unwrap(),
444      Some(&"Port 80")
445    );
446  }
447
448  #[test]
449  fn test_insert_and_get_multi_key() {
450    let mut tree = ConfigFilterTree::new();
451    tree.insert(
452      vec![
453        ConfigFilterTreeSingleKey::Port(80),
454        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
455      ],
456      "Port 80, com",
457    );
458    assert_eq!(
459      tree
460        .get(
461          vec![
462            ConfigFilterTreeSingleKey::Port(80),
463            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
464          ],
465          None
466        )
467        .unwrap(),
468      Some(&"Port 80, com")
469    );
470  }
471
472  #[test]
473  fn test_wildcard_matching() {
474    let mut tree = ConfigFilterTree::new();
475    tree.insert(
476      vec![
477        ConfigFilterTreeSingleKey::Port(80),
478        ConfigFilterTreeSingleKey::HostDomainLevelWildcard,
479      ],
480      "Wildcard",
481    );
482    assert_eq!(
483      tree
484        .get(
485          vec![
486            ConfigFilterTreeSingleKey::Port(80),
487            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
488          ],
489          None
490        )
491        .unwrap(),
492      Some(&"Wildcard")
493    );
494  }
495
496  #[test]
497  fn test_partial_key_matching() {
498    let mut tree = ConfigFilterTree::new();
499    tree.insert(
500      vec![
501        ConfigFilterTreeSingleKey::Port(80),
502        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
503        ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
504      ],
505      "Partial",
506    );
507    assert_eq!(
508      tree
509        .get(
510          vec![
511            ConfigFilterTreeSingleKey::Port(80),
512            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
513          ],
514          None
515        )
516        .unwrap(),
517      None
518    );
519  }
520
521  #[test]
522  fn test_overlapping_keys() {
523    let mut tree = ConfigFilterTree::new();
524    tree.insert(
525      vec![
526        ConfigFilterTreeSingleKey::Port(80),
527        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
528      ],
529      "First",
530    );
531    tree.insert(
532      vec![
533        ConfigFilterTreeSingleKey::Port(80),
534        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
535        ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
536      ],
537      "Second",
538    );
539    assert_eq!(
540      tree
541        .get(
542          vec![
543            ConfigFilterTreeSingleKey::Port(80),
544            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
545          ],
546          None
547        )
548        .unwrap(),
549      Some(&"First")
550    );
551    assert_eq!(
552      tree
553        .get(
554          vec![
555            ConfigFilterTreeSingleKey::Port(80),
556            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
557            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
558          ],
559          None
560        )
561        .unwrap(),
562      Some(&"Second")
563    );
564  }
565
566  #[test]
567  fn test_mixed_predicate_and_fixed_keys() {
568    let mut tree = ConfigFilterTree::new();
569    tree.insert(
570      vec![
571        ConfigFilterTreeSingleKey::Port(80),
572        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
573        ConfigFilterTreeSingleKey::HostDomainLevelWildcard,
574      ],
575      "Mixed",
576    );
577    assert_eq!(
578      tree
579        .get(
580          vec![
581            ConfigFilterTreeSingleKey::Port(80),
582            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
583            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
584          ],
585          None
586        )
587        .unwrap(),
588      Some(&"Mixed")
589    );
590  }
591
592  #[test]
593  fn test_keys_with_redundant_in_between() {
594    let mut tree = ConfigFilterTree::new();
595    tree.insert(
596      vec![
597        ConfigFilterTreeSingleKey::Port(80),
598        ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
599        ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
600      ],
601      "Value",
602    );
603    assert_eq!(
604      tree
605        .get(
606          vec![
607            ConfigFilterTreeSingleKey::Port(80),
608            ConfigFilterTreeSingleKey::IPv4Octet(127),
609            ConfigFilterTreeSingleKey::IPv4Octet(0),
610            ConfigFilterTreeSingleKey::IPv4Octet(0),
611            ConfigFilterTreeSingleKey::IPv4Octet(1),
612            ConfigFilterTreeSingleKey::HostDomainLevel("com".to_string()),
613            ConfigFilterTreeSingleKey::HostDomainLevel("example".to_string()),
614          ],
615          None
616        )
617        .unwrap(),
618      Some(&"Value")
619    );
620  }
621}