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	#[must_use]
23	pub const fn new(b: u8) -> Self {
24		Self(b.to_ascii_lowercase())
25	}
26}
27
28impl Key for IgnoreAsciiCase {
29	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream) {
30		let uppercase = self.0.to_ascii_uppercase();
31		if self.0 == uppercase {
32			return self.0.append_eq(out, lhs);
33		}
34
35		let lowercase_literal = Literal::byte_character(self.0);
36		let uppercase_literal = Literal::byte_character(uppercase);
37		out.append_all(quote! { (#lhs == #lowercase_literal || #lhs == #uppercase_literal) });
38	}
39}
40
41pub struct Builder<K, V> {
42	arms: BTreeMap<usize, Node<K, V>>,
43}
44
45impl<K, V> Default for Builder<K, V> {
46	fn default() -> Self {
47		Self::new()
48	}
49}
50
51impl<K, V> Builder<K, V> {
52	#[must_use]
53	pub const fn new() -> Self {
54		Self {
55			arms: BTreeMap::new(),
56		}
57	}
58}
59
60impl<K: Key, V> Builder<K, V> {
61	pub fn insert<IK: IntoIterator<Item = K>>(&mut self, key: IK, value: V) -> Option<V>
62	where
63		IK::IntoIter: ExactSizeIterator,
64	{
65		let key = key.into_iter();
66		let arm = self.arms.entry(key.len()).or_default();
67		arm.insert(key, value)
68	}
69}
70
71fn lifetime(name: &str, span: Span) -> TokenStream {
72	let mut ts = TokenStream::new();
73	let mut p = Punct::new('\'', Spacing::Joint);
74	p.set_span(span);
75	ts.append(p);
76	ts.append(Ident::new(name, span));
77	ts
78}
79
80impl<K: Key, V: ToTokens> Builder<K, V> {
81	pub fn build<I: ToTokens, F: ToTokens>(&self, input: I, fallback: F) -> TokenStream {
82		let block_label = lifetime("block", Span::mixed_site());
83		let input_assign = Ident::new("b", Span::mixed_site());
84
85		let mut arms = TokenStream::new();
86		for (len, node) in &self.arms {
87			let len_literal = Literal::usize_unsuffixed(*len);
88			arms.append_all(quote! { #len_literal => });
89			node.if_chain(&mut arms, &input_assign, &block_label, 0);
90			arms.append_all(quote! { , });
91		}
92
93		quote! {{
94			let block = #block_label : {
95				let #input_assign = #input;
96				match #input_assign.len() {
97					#arms
98					_ => {}
99				};
100				#fallback
101			};
102			block
103		}}
104	}
105}
106
107pub struct Node<K, V> {
108	data: Option<V>,
109	children: BTreeMap<K, Node<K, V>>,
110}
111
112impl<K, V> Default for Node<K, V> {
113	fn default() -> Self {
114		Self {
115			data: None,
116			children: BTreeMap::new(),
117		}
118	}
119}
120
121impl<K: Key, V> Node<K, V> {
122	fn insert(&mut self, key: impl Iterator<Item = K>, value: V) -> Option<V> {
123		let mut node = self;
124		for key in key {
125			node = node.children.entry(key).or_default();
126		}
127		node.data.replace(value)
128	}
129}
130
131impl<K: Key, V: ToTokens> Node<K, V> {
132	fn if_chain(&self, out: &mut TokenStream, key: &Ident, label: &TokenStream, mut index: usize) {
133		let mut node = self;
134		let mut same = false;
135		while node.children.len() == 1 {
136			out.append_all(if same {
137				quote! { && }
138			} else {
139				quote! { if }
140			});
141			same = true;
142
143			let val;
144			(val, node) = node.children.iter().next().unwrap();
145
146			let index_literal = Literal::usize_unsuffixed(index);
147			val.append_eq(out, quote! { #key[#index_literal] });
148
149			index += 1;
150		}
151
152		let node = node;
153		let res = if node.children.is_empty() {
154			let data = node.data.as_ref().unwrap();
155			quote! { break #label #data; }
156		} else {
157			let mut res = TokenStream::new();
158			for (val, child) in &node.children {
159				if !res.is_empty() {
160					res.append_all(quote! { else });
161				}
162
163				let index_literal = Literal::usize_unsuffixed(index);
164				res.append_all(quote! { if });
165				val.append_eq(&mut res, quote! { #key[#index_literal] });
166
167				let mut if_body = TokenStream::new();
168				child.if_chain(&mut if_body, key, label, index + 1);
169				res.append_all(quote! { { #if_body } });
170			}
171			res
172		};
173
174		if same || node.children.is_empty() {
175			out.append_all(quote! { { #res } });
176		} else {
177			out.append_all(res);
178		}
179	}
180}
forestrie_builder/lib.rs

forestrie_builder/
lib.rs
