forestrie_builder/
lib.rs

1use std::collections::BTreeMap;
2
3use proc_macro2::{Ident, Literal, Span, TokenStream};
4use quote::{ToTokens, TokenStreamExt, quote};
5
6pub trait Key: Ord {
7	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream);
8}
9
10impl Key for u8 {
11	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream) {
12		out.append_all(lhs);
13		out.append_all(quote! { == });
14		out.append(Literal::byte_character(*self));
15	}
16}
17
18#[derive(PartialEq, Eq, PartialOrd, Ord)]
19pub struct AsciiIgnoreCase(u8);
20
21impl AsciiIgnoreCase {
22	pub fn new(b: u8) -> Self {
23		Self(b.to_ascii_lowercase())
24	}
25}
26
27impl Key for AsciiIgnoreCase {
28	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream) {
29		let uppercase = self.0.to_ascii_uppercase();
30		if self.0 == uppercase {
31			return self.0.append_eq(out, lhs);
32		}
33
34		let lowercase_literal = Literal::byte_character(self.0);
35		let uppercase_literal = Literal::byte_character(uppercase);
36		out.append_all(quote! { (#lhs == #lowercase_literal || #lhs == #uppercase_literal) });
37	}
38}
39
40pub struct Builder<K, V> {
41	arms: BTreeMap<usize, Node<K, V>>,
42}
43
44impl<K, V> Default for Builder<K, V> {
45	fn default() -> Self {
46		Self::new()
47	}
48}
49
50impl<K, V> Builder<K, V> {
51	pub const fn new() -> Self {
52		Self {
53			arms: BTreeMap::new(),
54		}
55	}
56}
57
58impl<K: Key, V> Builder<K, V> {
59	pub fn insert<IK: IntoIterator<Item = K>>(&mut self, key: IK, value: V) -> Option<V>
60	where
61		IK::IntoIter: ExactSizeIterator,
62	{
63		let key = key.into_iter();
64		let arm = self.arms.entry(key.len()).or_default();
65		arm.insert(key, value)
66	}
67}
68
69impl<K: Key, V: ToTokens> Builder<K, V> {
70	pub fn build<I: ToTokens, F: ToTokens>(&self, input: I, fallback: F) -> TokenStream {
71		let block_label = quote! { 'ret };
72		let input_assign = Ident::new("b", Span::call_site());
73
74		let mut arms = TokenStream::new();
75		for (len, node) in &self.arms {
76			arms.append_all(quote! { #len => });
77			node.if_chain(&mut arms, &input_assign, &block_label, 0);
78			arms.append_all(quote! { , });
79		}
80
81		quote! { #block_label : {
82			let #input_assign = #input;
83			match #input_assign.len() {
84				#arms
85				_ => {}
86			};
87			#fallback
88		}}
89	}
90}
91
92pub struct Node<K, V> {
93	data: Option<V>,
94	children: BTreeMap<K, Node<K, V>>,
95}
96
97impl<K, V> Default for Node<K, V> {
98	fn default() -> Self {
99		Self {
100			data: None,
101			children: BTreeMap::new(),
102		}
103	}
104}
105
106impl<K: Key, V> Node<K, V> {
107	fn insert(&mut self, key: impl Iterator<Item = K>, value: V) -> Option<V> {
108		let mut node = self;
109		for key in key {
110			node = node.children.entry(key).or_default();
111		}
112		node.data.replace(value)
113	}
114}
115
116impl<K: Key, V: ToTokens> Node<K, V> {
117	fn if_chain(&self, out: &mut TokenStream, key: &Ident, label: &TokenStream, mut index: usize) {
118		let mut node = self;
119
120		let mut same = false;
121		while node.children.len() == 1 {
122			out.append_all(if same {
123				quote! { && }
124			} else {
125				quote! { if }
126			});
127			same = true;
128
129			let val;
130			(val, node) = node.children.iter().next().unwrap();
131
132			let index_literal = Literal::usize_unsuffixed(index);
133			val.append_eq(out, quote! { #key[#index_literal] });
134
135			index += 1;
136		}
137
138		let res = if node.children.is_empty() {
139			let data = node.data.as_ref().unwrap();
140			quote! { break #label #data; }
141		} else {
142			let mut res = TokenStream::new();
143			for (val, child) in &node.children {
144				if !res.is_empty() {
145					res.append_all(quote! { else });
146				}
147
148				let index_literal = Literal::usize_unsuffixed(index);
149				res.append_all(quote! { if });
150				val.append_eq(&mut res, quote! { #key[#index_literal] });
151
152				let mut if_body = TokenStream::new();
153				child.if_chain(&mut if_body, key, label, index + 1);
154				res.append_all(quote! { { #if_body } });
155			}
156			res
157		};
158
159		if same {
160			out.append_all(quote! { { #res } });
161		} else {
162			out.append_all(res);
163		}
164	}
165}
forestrie_builder/lib.rs

forestrie_builder/
lib.rs
