forestrie_builder/
lib.rs

1use std::collections::BTreeMap;
2
3use proc_macro2::{Ident, Literal, Punct, Spacing, 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 IgnoreAsciiCase(u8);
20
21impl IgnoreAsciiCase {
22	pub fn new(b: u8) -> Self {
23		Self(b.to_ascii_lowercase())
24	}
25}
26
27impl Key for IgnoreAsciiCase {
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
69fn lifetime(name: &str, span: Span) -> TokenStream {
70	let mut ts = TokenStream::new();
71	let mut p = Punct::new('\'', Spacing::Joint);
72	p.set_span(span);
73	ts.append(p);
74	ts.append(Ident::new(name, span));
75	ts
76}
77
78impl<K: Key, V: ToTokens> Builder<K, V> {
79	pub fn build<I: ToTokens, F: ToTokens>(&self, input: I, fallback: F) -> TokenStream {
80		let block_label = lifetime("block", Span::mixed_site());
81		let input_assign = Ident::new("b", Span::mixed_site());
82
83		let mut arms = TokenStream::new();
84		for (len, node) in &self.arms {
85			let len_literal = Literal::usize_unsuffixed(*len);
86			arms.append_all(quote! { #len_literal => });
87			node.if_chain(&mut arms, &input_assign, &block_label, 0);
88			arms.append_all(quote! { , });
89		}
90
91		quote! {{
92			let block = #block_label : {
93				let #input_assign = #input;
94				match #input_assign.len() {
95					#arms
96					_ => {}
97				};
98				#fallback
99			};
100			block
101		}}
102	}
103}
104
105pub struct Node<K, V> {
106	data: Option<V>,
107	children: BTreeMap<K, Node<K, V>>,
108}
109
110impl<K, V> Default for Node<K, V> {
111	fn default() -> Self {
112		Self {
113			data: None,
114			children: BTreeMap::new(),
115		}
116	}
117}
118
119impl<K: Key, V> Node<K, V> {
120	fn insert(&mut self, key: impl Iterator<Item = K>, value: V) -> Option<V> {
121		let mut node = self;
122		for key in key {
123			node = node.children.entry(key).or_default();
124		}
125		node.data.replace(value)
126	}
127}
128
129impl<K: Key, V: ToTokens> Node<K, V> {
130	fn if_chain(&self, out: &mut TokenStream, key: &Ident, label: &TokenStream, mut index: usize) {
131		let mut node = self;
132		let mut same = false;
133		while node.children.len() == 1 {
134			out.append_all(if same {
135				quote! { && }
136			} else {
137				quote! { if }
138			});
139			same = true;
140
141			let val;
142			(val, node) = node.children.iter().next().unwrap();
143
144			let index_literal = Literal::usize_unsuffixed(index);
145			val.append_eq(out, quote! { #key[#index_literal] });
146
147			index += 1;
148		}
149
150		let node = node;
151		let res = if node.children.is_empty() {
152			let data = node.data.as_ref().unwrap();
153			quote! { break #label #data; }
154		} else {
155			let mut res = TokenStream::new();
156			for (val, child) in &node.children {
157				if !res.is_empty() {
158					res.append_all(quote! { else });
159				}
160
161				let index_literal = Literal::usize_unsuffixed(index);
162				res.append_all(quote! { if });
163				val.append_eq(&mut res, quote! { #key[#index_literal] });
164
165				let mut if_body = TokenStream::new();
166				child.if_chain(&mut if_body, key, label, index + 1);
167				res.append_all(quote! { { #if_body } });
168			}
169			res
170		};
171
172		if same || node.children.is_empty() {
173			out.append_all(quote! { { #res } });
174		} else {
175			out.append_all(res);
176		}
177	}
178}
forestrie_builder/lib.rs

forestrie_builder/
lib.rs
