forestrie-builder 0.1.0

use std::collections::BTreeMap;

use proc_macro2::{Ident, Literal, Span, TokenStream};
use quote::{ToTokens, TokenStreamExt, quote};

pub trait Key: Ord {
	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream);
}

impl Key for u8 {
	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream) {
		out.append_all(lhs);
		out.append_all(quote! { == });
		out.append(Literal::byte_character(*self));
	}
}

#[derive(PartialEq, Eq, PartialOrd, Ord)]
pub struct AsciiIgnoreCase(u8);

impl AsciiIgnoreCase {
	pub fn new(b: u8) -> Self {
		Self(b.to_ascii_lowercase())
	}
}

impl Key for AsciiIgnoreCase {
	fn append_eq(&self, out: &mut TokenStream, lhs: TokenStream) {
		let uppercase = self.0.to_ascii_uppercase();
		if self.0 == uppercase {
			return self.0.append_eq(out, lhs);
		}

		let lowercase_literal = Literal::byte_character(self.0);
		let uppercase_literal = Literal::byte_character(uppercase);
		out.append_all(quote! { (#lhs == #lowercase_literal || #lhs == #uppercase_literal) });
	}
}

pub struct Builder<K, V> {
	arms: BTreeMap<usize, Node<K, V>>,
}

impl<K, V> Default for Builder<K, V> {
	fn default() -> Self {
		Self::new()
	}
}

impl<K, V> Builder<K, V> {
	pub const fn new() -> Self {
		Self {
			arms: BTreeMap::new(),
		}
	}
}

impl<K: Key, V> Builder<K, V> {
	pub fn insert<IK: IntoIterator<Item = K>>(&mut self, key: IK, value: V) -> Option<V>
	where
		IK::IntoIter: ExactSizeIterator,
	{
		let key = key.into_iter();
		let arm = self.arms.entry(key.len()).or_default();
		arm.insert(key, value)
	}
}

impl<K: Key, V: ToTokens> Builder<K, V> {
	pub fn build<I: ToTokens, F: ToTokens>(&self, input: I, fallback: F) -> TokenStream {
		let block_label = quote! { 'ret };
		let input_assign = Ident::new("b", Span::call_site());

		let mut arms = TokenStream::new();
		for (len, node) in &self.arms {
			arms.append_all(quote! { #len => });
			node.if_chain(&mut arms, &input_assign, &block_label, 0);
			arms.append_all(quote! { , });
		}

		quote! { #block_label : {
			let #input_assign = #input;
			match #input_assign.len() {
				#arms
				_ => {}
			};
			#fallback
		}}
	}
}

pub struct Node<K, V> {
	data: Option<V>,
	children: BTreeMap<K, Node<K, V>>,
}

impl<K, V> Default for Node<K, V> {
	fn default() -> Self {
		Self {
			data: None,
			children: BTreeMap::new(),
		}
	}
}

impl<K: Key, V> Node<K, V> {
	fn insert(&mut self, key: impl Iterator<Item = K>, value: V) -> Option<V> {
		let mut node = self;
		for key in key {
			node = node.children.entry(key).or_default();
		}
		node.data.replace(value)
	}
}

impl<K: Key, V: ToTokens> Node<K, V> {
	fn if_chain(&self, out: &mut TokenStream, key: &Ident, label: &TokenStream, mut index: usize) {
		let mut node = self;

		let mut same = false;
		while node.children.len() == 1 {
			out.append_all(if same {
				quote! { && }
			} else {
				quote! { if }
			});
			same = true;

			let val;
			(val, node) = node.children.iter().next().unwrap();

			let index_literal = Literal::usize_unsuffixed(index);
			val.append_eq(out, quote! { #key[#index_literal] });

			index += 1;
		}

		let res = if node.children.is_empty() {
			let data = node.data.as_ref().unwrap();
			quote! { break #label #data; }
		} else {
			let mut res = TokenStream::new();
			for (val, child) in &node.children {
				if !res.is_empty() {
					res.append_all(quote! { else });
				}

				let index_literal = Literal::usize_unsuffixed(index);
				res.append_all(quote! { if });
				val.append_eq(&mut res, quote! { #key[#index_literal] });

				let mut if_body = TokenStream::new();
				child.if_chain(&mut if_body, key, label, index + 1);
				res.append_all(quote! { { #if_body } });
			}
			res
		};

		if same {
			out.append_all(quote! { { #res } });
		} else {
			out.append_all(res);
		}
	}
}