use std::collections::{BTreeMap, HashMap, HashSet};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use crate::tile::TileId;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TileMeta {
pub id: TileId,
pub valence: f64,
pub created: DateTime<Utc>,
pub accessed: DateTime<Utc>,
pub constraints: HashSet<String>,
}
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct CrystalIndex {
tiles: HashMap<TileId, TileMeta>,
by_valence: BTreeMap<OrderedValence, Vec<TileId>>,
by_time: BTreeMap<i64, TileId>,
constraint_index: HashMap<String, HashSet<TileId>>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct OrderedValence(pub f64);
impl PartialEq for OrderedValence {
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
}
}
impl Eq for OrderedValence {}
impl PartialOrd for OrderedValence {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
self.0.partial_cmp(&other.0)
}
}
impl Ord for OrderedValence {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.partial_cmp(other).unwrap_or(std::cmp::Ordering::Equal)
}
}
impl CrystalIndex {
pub fn new() -> Self {
Self::default()
}
pub fn insert(&mut self, meta: TileMeta) {
let id = meta.id.clone();
let valence = meta.valence;
let created_ts = meta.created.timestamp();
let constraints = meta.constraints.clone();
self.tiles.insert(id.clone(), meta);
self.by_valence
.entry(OrderedValence(valence))
.or_default()
.push(id.clone());
self.by_time.insert(created_ts, id.clone());
for key in &constraints {
self.constraint_index
.entry(key.clone())
.or_default()
.insert(id.clone());
}
}
pub fn remove(&mut self, id: &TileId) -> Option<TileMeta> {
let meta = self.tiles.remove(id)?;
if let Some(ids) = self.by_valence.get_mut(&OrderedValence(meta.valence)) {
ids.retain(|i| i != id);
if ids.is_empty() {
self.by_valence.remove(&OrderedValence(meta.valence));
}
}
let ts = meta.created.timestamp();
self.by_time.remove(&ts);
for key in &meta.constraints {
if let Some(set) = self.constraint_index.get_mut(key) {
set.remove(id);
if set.is_empty() {
self.constraint_index.remove(key);
}
}
}
Some(meta)
}
pub fn query(&self, constraints: &[&str], limit: usize) -> Vec<TileId> {
let mut candidates: HashMap<TileId, usize> = HashMap::new();
for key in constraints {
if let Some(ids) = self.constraint_index.get(*key) {
for id in ids {
*candidates.entry(id.clone()).or_default() += 1;
}
}
}
let mut ranked: Vec<_> = candidates.into_iter().collect();
ranked.sort_by(|a, b| b.1.cmp(&a.1));
ranked.into_iter().take(limit).map(|(id, _)| id).collect()
}
pub fn top_by_valence(&self, limit: usize) -> Vec<TileId> {
self.by_valence
.iter()
.rev() .flat_map(|(_, ids)| ids.iter().cloned())
.take(limit)
.collect()
}
pub fn recent(&self, limit: usize) -> Vec<TileId> {
self.by_time
.iter()
.rev() .map(|(_, id)| id.clone())
.take(limit)
.collect()
}
pub fn get(&self, id: &TileId) -> Option<&TileMeta> {
self.tiles.get(id)
}
pub fn len(&self) -> usize {
self.tiles.len()
}
pub fn is_empty(&self) -> bool {
self.tiles.is_empty()
}
pub fn all_ids(&self) -> Vec<TileId> {
self.tiles.keys().cloned().collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn meta(valence: f64, constraints: &[&str]) -> TileMeta {
TileMeta {
id: TileId::new(),
valence,
created: Utc::now(),
accessed: Utc::now(),
constraints: constraints.iter().map(|s| s.to_string()).collect(),
}
}
#[test]
fn insert_and_query() {
let mut index = CrystalIndex::new();
let m = meta(0.5, &["rust", "programming"]);
let id = m.id.clone();
index.insert(m);
let results = index.query(&["rust"], 10);
assert_eq!(results.len(), 1);
assert_eq!(results[0], id);
}
#[test]
fn top_by_valence() {
let mut index = CrystalIndex::new();
let m1 = meta(0.3, &[]);
let m2 = meta(0.9, &[]);
let m3 = meta(0.6, &[]);
let high_id = m2.id.clone();
index.insert(m1);
index.insert(m2);
index.insert(m3);
let top = index.top_by_valence(1);
assert_eq!(top.len(), 1);
assert_eq!(top[0], high_id);
}
#[test]
fn recent() {
let mut index = CrystalIndex::new();
let mut m1 = meta(0.5, &[]);
m1.created = Utc::now() - chrono::Duration::days(1);
let m2 = meta(0.5, &[]);
let recent_id = m2.id.clone();
index.insert(m1);
index.insert(m2);
let recent = index.recent(1);
assert_eq!(recent.len(), 1);
assert_eq!(recent[0], recent_id);
}
#[test]
fn remove_tile() {
let mut index = CrystalIndex::new();
let m = meta(0.5, &["test"]);
let id = m.id.clone();
index.insert(m);
assert_eq!(index.len(), 1);
index.remove(&id);
assert!(index.is_empty());
assert!(index.query(&["test"], 10).is_empty());
}
#[test]
fn multi_constraint_query() {
let mut index = CrystalIndex::new();
let m1 = meta(0.5, &["rust", "systems"]);
let m2 = meta(0.5, &["rust", "web"]);
index.insert(m1);
index.insert(m2);
let results = index.query(&["rust"], 10);
assert_eq!(results.len(), 2);
let results = index.query(&["systems"], 10);
assert_eq!(results.len(), 1);
}
#[test]
fn empty_query() {
let index = CrystalIndex::new();
let results = index.query(&["nonexistent"], 10);
assert!(results.is_empty());
}
}