use std::collections::HashSet;
use serde_json::Value;
use crate::schema::HaysonRecord;
use super::exact_index::{value_key, value_keys, ExactIndex};
use super::sorted_index::{ArraySortedIndex, RangeBounds, SortValue};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RegistryIndexKind {
Exact,
Sorted,
Fulltext,
}
#[derive(Debug, Clone)]
pub struct RegistryIndex {
pub name: String,
pub fields: Vec<String>,
pub kind: RegistryIndexKind,
pub unique: bool,
}
pub fn composite_key(parts: &[String]) -> String {
serde_json::to_string(parts).expect("a vector of strings always serializes")
}
pub fn lookup_key_of(values: &[Value]) -> String {
let parts: Vec<String> = values.iter().map(value_key).collect();
composite_key(&parts)
}
pub fn unique_key_of(record: &HaysonRecord, fields: &[String]) -> Option<String> {
let mut parts = Vec::with_capacity(fields.len());
for field in fields {
match record.get(field) {
None | Some(Value::Null) => return None,
Some(value) => parts.push(value_key(value)),
}
}
Some(composite_key(&parts))
}
pub fn sort_vector_of(record: &HaysonRecord, fields: &[String]) -> Option<Vec<SortValue>> {
let mut vector = Vec::new();
for field in fields {
let Some(value) = record.get(field) else { break };
let Some(sortable) = SortValue::from_json(value) else {
break;
};
vector.push(sortable);
}
if vector.is_empty() {
None
} else {
Some(vector)
}
}
pub struct CompositeExactIndex {
fields: Vec<String>,
prefixes: Vec<ExactIndex>,
}
impl CompositeExactIndex {
pub fn new(fields: &[String]) -> Self {
let prefixes = fields.iter().map(|_| ExactIndex::new()).collect();
CompositeExactIndex {
fields: fields.to_vec(),
prefixes,
}
}
fn keys_of(&self, record: &HaysonRecord) -> Vec<Vec<String>> {
let mut per_prefix = Vec::new();
let mut combinations: Vec<Vec<String>> = vec![Vec::new()];
for field in &self.fields {
let encoded = value_keys(record.get(field));
if encoded.is_empty() {
break;
}
let mut wider = Vec::with_capacity(combinations.len() * encoded.len());
for combination in &combinations {
for value in &encoded {
let mut next = combination.clone();
next.push(value.clone());
wider.push(next);
}
}
combinations = wider;
per_prefix.push(combinations.iter().map(|c| composite_key(c)).collect());
}
per_prefix
}
pub fn add(&mut self, record: &HaysonRecord, key: &str) {
for (at, encoded) in self.keys_of(record).into_iter().enumerate() {
let Some(prefix) = self.prefixes.get_mut(at) else {
continue;
};
for value in encoded {
prefix.add(value, key.to_string());
}
}
}
pub fn remove(&mut self, record: &HaysonRecord, key: &str) {
for (at, encoded) in self.keys_of(record).into_iter().enumerate() {
let Some(prefix) = self.prefixes.get_mut(at) else {
continue;
};
for value in encoded {
prefix.remove(&value, key);
}
}
}
pub fn lookup(&self, values: &[Value]) -> Option<&HashSet<String>> {
let at = values.len().checked_sub(1)?;
let prefix = self.prefixes.get(at)?;
prefix.lookup(&lookup_key_of(values))
}
pub fn holder_of(&self, record: &HaysonRecord) -> Option<String> {
let key = unique_key_of(record, &self.fields)?;
let widest = self.prefixes.last()?;
widest.holder_of(&key).map(str::to_string)
}
}
pub struct CompositeSortedIndex {
fields: Vec<String>,
sorted: ArraySortedIndex,
}
impl CompositeSortedIndex {
pub fn new(fields: &[String]) -> Self {
CompositeSortedIndex {
fields: fields.to_vec(),
sorted: ArraySortedIndex::new(),
}
}
pub fn insert(&mut self, record: &HaysonRecord, key: &str) {
if let Some(vector) = sort_vector_of(record, &self.fields) {
self.sorted.insert(key, vector);
}
}
pub fn remove(&mut self, record: &HaysonRecord, key: &str) {
if let Some(vector) = sort_vector_of(record, &self.fields) {
self.sorted.remove(key, &vector);
}
}
pub fn range(&self, bounds: &RangeBounds) -> Vec<String> {
self.sorted.range(bounds)
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use super::*;
fn record(pairs: &[(&str, Value)]) -> HaysonRecord {
pairs.iter().map(|(k, v)| (k.to_string(), v.clone())).collect()
}
#[test]
fn reads_its_fields_in_the_order_they_were_declared() {
let fields = vec!["userId".to_string(), "orderId".to_string()];
let mut index = CompositeExactIndex::new(&fields);
index.add(&record(&[("userId", json!("u1")), ("orderId", json!(2))]), "o2");
assert!(index.lookup(&[json!("u1"), json!(2)]).unwrap().contains("o2"));
assert!(index.lookup(&[json!(2), json!("u1")]).is_none());
}
#[test]
fn answers_a_lookup_on_the_leading_field_alone_from_the_prefix_of_the_key() {
let fields = vec!["userId".to_string(), "orderId".to_string()];
let mut index = CompositeExactIndex::new(&fields);
index.add(&record(&[("userId", json!("u1")), ("orderId", json!(1))]), "o1");
index.add(&record(&[("userId", json!("u1")), ("orderId", json!(2))]), "o2");
index.add(&record(&[("userId", json!("u2")), ("orderId", json!(1))]), "o4");
let mut u1: Vec<&String> = index.lookup(&[json!("u1")]).unwrap().iter().collect();
u1.sort();
assert_eq!(u1, vec!["o1", "o2"]);
}
#[test]
fn holds_an_array_field_of_the_key_per_element() {
let fields = vec!["labels".to_string(), "userId".to_string()];
let mut index = CompositeExactIndex::new(&fields);
let o1 = record(&[("labels", json!(["a", "b"])), ("userId", json!("u1"))]);
index.add(&o1, "o1");
assert!(index.lookup(&[json!("b"), json!("u1")]).unwrap().contains("o1"));
index.remove(&o1, "o1");
assert!(index.lookup(&[json!("b"), json!("u1")]).is_none());
}
#[test]
fn holds_a_unique_key_only_when_every_field_of_it_repeats() {
let fields = vec!["userId".to_string(), "orderId".to_string()];
let mut index = CompositeExactIndex::new(&fields);
index.add(&record(&[("userId", json!("u1")), ("orderId", json!(2))]), "o2");
assert_eq!(
index.holder_of(&record(&[("userId", json!("u1")), ("orderId", json!(2))])),
Some("o2".to_string())
);
assert_eq!(
index.holder_of(&record(&[("userId", json!("u1")), ("orderId", json!(9))])),
None
);
assert_eq!(index.holder_of(&record(&[("userId", json!("u1"))])), None);
}
#[test]
fn answers_a_range_bounded_on_the_whole_vector() {
let fields = vec!["userId".to_string(), "total".to_string()];
let mut index = CompositeSortedIndex::new(&fields);
index.insert(&record(&[("userId", json!("u1")), ("total", json!(10))]), "o1");
index.insert(&record(&[("userId", json!("u1")), ("total", json!(20))]), "o2");
index.insert(&record(&[("userId", json!("u1")), ("total", json!(30))]), "o3");
index.insert(&record(&[("userId", json!("u2")), ("total", json!(40))]), "o4");
let bounds = RangeBounds {
gte: Some(vec![SortValue::Str("u1".into()), SortValue::Num(20.0)]),
lte: Some(vec![SortValue::Str("u1".into()), SortValue::Num(30.0)]),
..Default::default()
};
assert_eq!(index.range(&bounds), vec!["o2", "o3"]);
}
}