#![allow(clippy::cast_precision_loss)]
use super::helpers::{make_label_prop_key, safe_bitmap_id, PostcardPersistence};
use roaring::RoaringBitmap;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::ops::Bound;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum OrderedValue {
Null,
Integer(i64),
Float(OrderedFloat),
String(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct OrderedFloat(pub f64);
impl PartialOrd for OrderedFloat {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Eq for OrderedFloat {}
impl Ord for OrderedFloat {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.0.partial_cmp(&other.0).unwrap_or_else(|| {
match (self.0.is_nan(), other.0.is_nan()) {
(true, true) => std::cmp::Ordering::Equal,
(true, false) => std::cmp::Ordering::Greater,
(false, true) => std::cmp::Ordering::Less,
(false, false) => unreachable!(),
}
})
}
}
impl PartialOrd for OrderedValue {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl Ord for OrderedValue {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
match (self, other) {
(Self::Null, Self::Null) => std::cmp::Ordering::Equal,
(Self::Null, _) | (Self::Integer(_) | Self::Float(_), Self::String(_)) => {
std::cmp::Ordering::Less
}
(_, Self::Null) | (Self::String(_), Self::Integer(_) | Self::Float(_)) => {
std::cmp::Ordering::Greater
}
(Self::Integer(a), Self::Integer(b)) => a.cmp(b),
(Self::Float(a), Self::Float(b)) => a.cmp(b),
(Self::Integer(a), Self::Float(b)) => OrderedFloat(*a as f64).cmp(b),
(Self::Float(a), Self::Integer(b)) => a.cmp(&OrderedFloat(*b as f64)),
(Self::String(a), Self::String(b)) => a.cmp(b),
}
}
}
impl OrderedValue {
#[must_use]
pub fn from_json(value: &serde_json::Value) -> Option<Self> {
match value {
serde_json::Value::Null => Some(Self::Null),
serde_json::Value::Number(n) => {
if let Some(i) = n.as_i64() {
Some(Self::Integer(i))
} else {
n.as_f64().map(|f| Self::Float(OrderedFloat(f)))
}
}
serde_json::Value::String(s) => Some(Self::String(s.clone())),
_ => None, }
}
}
pub const RANGE_INDEX_VERSION: u32 = 1;
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct RangeIndex {
#[serde(default = "default_range_version")]
version: u32,
indexes: HashMap<(String, String), BTreeMap<OrderedValue, RoaringBitmap>>,
}
fn default_range_version() -> u32 {
RANGE_INDEX_VERSION
}
impl RangeIndex {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn create_index(&mut self, label: &str, property: &str) {
self.indexes
.entry(make_label_prop_key(label, property))
.or_default();
}
#[must_use]
pub fn has_index(&self, label: &str, property: &str) -> bool {
self.indexes
.contains_key(&make_label_prop_key(label, property))
}
pub fn insert(
&mut self,
label: &str,
property: &str,
value: &serde_json::Value,
node_id: u64,
) -> bool {
let Some(safe_id) = safe_bitmap_id(node_id) else {
return false;
};
let key = make_label_prop_key(label, property);
if let Some(btree) = self.indexes.get_mut(&key) {
if let Some(ordered) = OrderedValue::from_json(value) {
btree.entry(ordered).or_default().insert(safe_id);
return true;
}
}
false
}
pub fn remove(
&mut self,
label: &str,
property: &str,
value: &serde_json::Value,
node_id: u64,
) -> bool {
let Some(safe_id) = safe_bitmap_id(node_id) else {
return false;
};
let key = make_label_prop_key(label, property);
if let Some(btree) = self.indexes.get_mut(&key) {
if let Some(ordered) = OrderedValue::from_json(value) {
if let Some(bitmap) = btree.get_mut(&ordered) {
let removed = bitmap.remove(safe_id);
if bitmap.is_empty() {
btree.remove(&ordered);
}
return removed;
}
}
}
false
}
#[must_use]
pub fn range_greater_than(
&self,
label: &str,
property: &str,
threshold: &serde_json::Value,
) -> RoaringBitmap {
self.range_query(
label,
property,
Bound::Excluded(threshold),
Bound::Unbounded,
)
}
#[must_use]
pub fn range_greater_or_equal(
&self,
label: &str,
property: &str,
threshold: &serde_json::Value,
) -> RoaringBitmap {
self.range_query(
label,
property,
Bound::Included(threshold),
Bound::Unbounded,
)
}
#[must_use]
pub fn range_less_than(
&self,
label: &str,
property: &str,
threshold: &serde_json::Value,
) -> RoaringBitmap {
self.range_query(
label,
property,
Bound::Unbounded,
Bound::Excluded(threshold),
)
}
#[must_use]
pub fn range_less_or_equal(
&self,
label: &str,
property: &str,
threshold: &serde_json::Value,
) -> RoaringBitmap {
self.range_query(
label,
property,
Bound::Unbounded,
Bound::Included(threshold),
)
}
#[must_use]
pub fn range_between(
&self,
label: &str,
property: &str,
start: &serde_json::Value,
end: &serde_json::Value,
) -> RoaringBitmap {
self.range_query(
label,
property,
Bound::Included(start),
Bound::Included(end),
)
}
fn convert_bound(bound: Bound<&serde_json::Value>) -> Option<Bound<OrderedValue>> {
match bound {
Bound::Included(v) => OrderedValue::from_json(v).map(Bound::Included),
Bound::Excluded(v) => OrderedValue::from_json(v).map(Bound::Excluded),
Bound::Unbounded => Some(Bound::Unbounded),
}
}
fn range_query(
&self,
label: &str,
property: &str,
start: Bound<&serde_json::Value>,
end: Bound<&serde_json::Value>,
) -> RoaringBitmap {
let mut result = RoaringBitmap::new();
let key = make_label_prop_key(label, property);
let Some(btree) = self.indexes.get(&key) else {
return result;
};
let (Some(start_ord), Some(end_ord)) =
(Self::convert_bound(start), Self::convert_bound(end))
else {
return result;
};
for bitmap in btree.range((start_ord, end_ord)).map(|(_k, v)| v) {
result |= bitmap;
}
result
}
pub fn drop_index(&mut self, label: &str, property: &str) -> bool {
self.indexes
.remove(&make_label_prop_key(label, property))
.is_some()
}
pub fn clear(&mut self) {
self.indexes.clear();
}
#[must_use]
pub fn memory_usage(&self) -> usize {
let mut total = std::mem::size_of::<Self>();
for ((label, prop), btree) in &self.indexes {
total += label.len() + prop.len();
for bitmap in btree.values() {
total += std::mem::size_of::<OrderedValue>();
total += bitmap.serialized_size();
}
}
total
}
#[must_use]
pub fn indexed_properties(&self) -> Vec<(String, String)> {
self.indexes.keys().cloned().collect()
}
}
impl PostcardPersistence for RangeIndex {}
impl RangeIndex {
pub fn to_bytes(&self) -> Result<Vec<u8>, postcard::Error> {
<Self as PostcardPersistence>::to_bytes(self)
}
pub fn from_bytes(bytes: &[u8]) -> Result<Self, postcard::Error> {
<Self as PostcardPersistence>::from_bytes(bytes)
}
pub fn save_to_file(&self, path: &std::path::Path) -> std::io::Result<()> {
<Self as PostcardPersistence>::save_to_file(self, path)
}
pub fn load_from_file(path: &std::path::Path) -> std::io::Result<Self> {
<Self as PostcardPersistence>::load_from_file(path)
}
}