use crate::Key;
use indexmap::IndexMap;
use proc_macro::TokenStream;
use quote::quote;
use std::collections::HashSet;
use syn::Expr;
#[derive(Debug, Clone)]
pub(crate) enum Algorithm {
Position { idx: usize },
Xor { idx1: usize, idx2: usize },
}
#[derive(Clone)]
pub(crate) struct Table<'x> {
pub algorithm: Algorithm,
pub positions: Vec<(u8, &'x Key, Value<'x>)>,
pub default: Value<'x>,
pub ignore_case: bool,
}
#[derive(Clone)]
pub(crate) enum Value<'x> {
Some(&'x Expr),
None,
True,
False,
Expr(&'x Expr),
Table(Box<Table<'x>>),
}
pub fn build_tiny_map(
name: Expr,
options: Vec<(&Key, Value)>,
default: Value,
ignore_case: bool,
) -> TokenStream {
let mut map: IndexMap<usize, Vec<(&Key, Value)>> = IndexMap::new();
let mut min_key_size = usize::MAX;
let mut max_key_size = 0;
for (key, value) in &options {
let key_size = key.len();
min_key_size = min_key_size.min(key_size);
max_key_size = max_key_size.max(key_size);
map.entry(key_size).or_default().push((key, value.clone()));
}
let (header, footer) = if matches!(default, Value::Expr(_)) {
(
quote! {
let __key = #name;
},
quote! {
if !__matched {
#default
}
},
)
} else {
(
quote! {
let __key = #name;
},
quote! {},
)
};
if let Some(table) = try_hash(&options, &default, min_key_size, false, ignore_case) {
let else_cond = if !matches!(default, Value::Expr(_)) {
quote! {
else {
#default
}
}
} else {
quote! {
#footer
}
};
TokenStream::from(quote! {{
#header
let mut __matched = (#min_key_size..=#max_key_size).contains(&__key.len());
if __matched {
#table
} #else_cond
}})
} else {
let match_default = if matches!(default, Value::Expr(_)) {
quote! {
{ __matched = false; }
}
} else {
quote! {
#default
}
};
let match_arms = map.iter().map(|(size, keys)| {
if keys.len() == 1 {
let (key, value) = keys.iter().next().unwrap();
if ignore_case {
quote! { #size if __key.eq_ignore_ascii_case(#key) => #value, }
} else {
quote! { #size if __key == #key => #value, }
}
} else {
let table =
try_hash(keys, &default, *size, true, ignore_case).unwrap_or_else(|| {
panic!(
"Failed to build lookup table for {} keys: {:?}",
keys.len(),
keys.iter().map(|(k, _)| k).collect::<Vec<_>>()
)
});
quote! { #size => { #table } }
}
});
TokenStream::from(quote! {{
#header
let mut __matched = true;
match __key.len() {
#(#match_arms)*
_ => #match_default,
}
#footer
}})
}
}
impl Algorithm {
pub fn hash(&self, value: &[u8]) -> u8 {
match self {
Algorithm::Position { idx } => value[*idx],
Algorithm::Xor { idx1, idx2 } => value[*idx1] ^ value[*idx2],
}
}
}
pub(crate) fn try_hash<'x>(
keys: &Vec<(&'x Key, Value<'x>)>,
default: &Value<'x>,
size: usize,
is_final_pass: bool,
ignore_case: bool,
) -> Option<Table<'x>> {
if size == 1 && is_final_pass {
return Some(Table {
algorithm: Algorithm::Position { idx: 0 },
positions: keys
.iter()
.map(|(key, value)| (key.as_bytes()[0], *key, (*value).clone()))
.collect(),
ignore_case,
default: default.clone(),
});
}
let mut best_match_count = 0;
let mut best_match_algo = Algorithm::Position { idx: 0 };
for idx in 0..size {
let mut byte_set = HashSet::new();
let algorithm = Algorithm::Position { idx };
for (key, _) in keys {
byte_set.insert(algorithm.hash(key.as_bytes()));
}
if byte_set.len() == keys.len() {
return Some(Table {
positions: keys
.iter()
.map(|(key, value)| (algorithm.hash(key.as_bytes()), *key, (*value).clone()))
.collect(),
algorithm,
ignore_case,
default: default.clone(),
});
} else if byte_set.len() > best_match_count {
best_match_count = byte_set.len();
best_match_algo = Algorithm::Position { idx };
}
}
for i in 0..size {
for j in i + 1..size {
let mut byte_set = HashSet::new();
let algorithm = Algorithm::Xor { idx1: i, idx2: j };
for (key, _) in keys {
byte_set.insert(algorithm.hash(key.as_bytes()));
}
if byte_set.len() == keys.len() {
return Some(Table {
positions: keys
.iter()
.map(|(key, value)| (algorithm.hash(key.as_bytes()), *key, value.clone()))
.collect(),
algorithm,
ignore_case,
default: default.clone(),
});
} else if byte_set.len() > best_match_count {
best_match_count = byte_set.len();
best_match_algo = algorithm;
}
}
}
if is_final_pass {
let mut key_groups = IndexMap::new();
for (key, value) in keys {
key_groups
.entry(best_match_algo.hash(key.as_bytes()))
.or_insert_with(Vec::new)
.push((*key, value.clone()));
}
let mut table = Table {
algorithm: best_match_algo,
positions: Vec::with_capacity(keys.len()),
default: default.clone(),
ignore_case,
};
for (hash, keys) in key_groups {
if keys.len() > 1 {
let sub_table = try_hash(&keys, default, size, true, ignore_case).unwrap();
table
.positions
.push((hash, keys[0].0, Value::Table(Box::new(sub_table))));
} else {
let (key, value) = keys.into_iter().next().unwrap();
table.positions.push((hash, key, value));
}
}
Some(table)
} else {
None
}
}
impl quote::ToTokens for Table<'_> {
fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
let algorithm = match &self.algorithm {
Algorithm::Position { idx } => {
if self.ignore_case {
quote! { let hash = __key[#idx].to_ascii_lowercase(); }
} else {
quote! { let hash = __key[#idx]; }
}
}
Algorithm::Xor { idx1, idx2 } => {
if self.ignore_case {
quote! { let hash = __key[#idx1].to_ascii_lowercase() ^ __key[#idx2].to_ascii_lowercase(); }
} else {
quote! { let hash = __key[#idx1] ^ __key[#idx2]; }
}
}
};
let match_default = if matches!(self.default, Value::Expr(_)) {
quote! {
{ __matched = false; }
}
} else {
let default = &self.default;
quote! {
#default
}
};
let match_arms = self.positions.iter().map(|(hash, key, value)| {
if key.len() > 1 && !matches!(value, Value::Table(_)) {
if self.ignore_case {
quote! { #hash if __key.eq_ignore_ascii_case(#key) => #value, }
} else {
quote! { #hash if __key == #key => #value, }
}
} else {
quote! { #hash => #value, }
}
});
tokens.extend(quote! {
#algorithm
match hash {
#(#match_arms)*
_ => #match_default,
}
});
}
}
impl quote::ToTokens for Value<'_> {
fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
tokens.extend(match self {
Value::Some(expr) => quote! { Some(#expr) },
Value::None => quote! { None },
Value::True => quote! { true },
Value::False => quote! { false },
Value::Expr(expr) => quote! { #expr },
Value::Table(table) => quote! {{ #table }},
});
}
}