wedb_embed 0.1.0

Embedded Kvrocks-compatible storage engine for WeDb
Documentation
use crate::error::{Error, Result};
use std::str;

/// ZADD 选项标志(对标 Apache Kvrocks ZSetFlags)
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ZAdd {
    Nx,
    Xx,
    Gt,
    Lt,
    Ch,
    Incr,
}

/// 聚合函数类型(对标 Apache Kvrocks AggregateMethod)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Aggregate {
    #[default]
    Sum,
    Min,
    Max,
}

impl Aggregate {
    #[inline]
    pub fn parse(s: &str) -> Self {
        match s.to_ascii_uppercase().as_str() {
            "MIN" => Self::Min,
            "MAX" => Self::Max,
            _ => Self::Sum,
        }
    }

    #[inline]
    pub fn apply(&self, current: f64, new_val: f64) -> f64 {
        let res = match self {
            Self::Sum => current + new_val,
            Self::Min => current.min(new_val),
            Self::Max => current.max(new_val),
        };
        if res.is_nan() { 0.0 } else { res }
    }
}

/// 分数范围规格(对标 Apache Kvrocks RangeScoreSpec)
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct RangeScoreSpec {
    pub min: f64,
    pub max: f64,
    pub minex: bool,
    pub maxex: bool,
    pub offset: usize,
    pub count: Option<usize>,
}

impl Default for RangeScoreSpec {
    #[inline]
    fn default() -> Self {
        Self {
            min: f64::NEG_INFINITY,
            max: f64::INFINITY,
            minex: false,
            maxex: false,
            offset: 0,
            count: None,
        }
    }
}

impl RangeScoreSpec {
    #[inline]
    pub fn new(min: f64, max: f64) -> Self {
        Self {
            min,
            max,
            ..Default::default()
        }
    }

    #[inline]
    pub fn with_limit(min: f64, max: f64, offset: usize, count: usize) -> Self {
        Self {
            min,
            max,
            minex: false,
            maxex: false,
            offset,
            count: Some(count),
        }
    }

    #[inline]
    pub fn check(&self, score: f64) -> bool {
        let min_ok = if self.minex {
            score > self.min
        } else {
            score >= self.min
        };
        let max_ok = if self.maxex {
            score < self.max
        } else {
            score <= self.max
        };
        min_ok && max_ok
    }

    /// 解析 Redis 风格的分数边界字符串(例如 "(1.5", "[10", "10", "-inf", "+inf", "(-inf" 等)
    pub fn parse_bound(s: &str) -> Result<(f64, bool)> {
        let s = s.trim();
        if s.is_empty() {
            return Err(Error::invalid_data("ERR min or max is not a float"));
        }

        let (val_str, is_exclusive) = if let Some(rest) = s.strip_prefix('(') {
            (rest.trim(), true)
        } else if let Some(rest) = s.strip_prefix('[') {
            (rest.trim(), false)
        } else {
            (s, false)
        };

        if val_str.eq_ignore_ascii_case("-inf") || val_str.eq_ignore_ascii_case("-infinity") {
            return Ok((f64::NEG_INFINITY, is_exclusive));
        }
        if val_str.eq_ignore_ascii_case("+inf")
            || val_str.eq_ignore_ascii_case("+infinity")
            || val_str.eq_ignore_ascii_case("inf")
            || val_str.eq_ignore_ascii_case("infinity")
        {
            return Ok((f64::INFINITY, is_exclusive));
        }

        let val = val_str
            .parse::<f64>()
            .map_err(|_| Error::invalid_data("ERR min or max is not a float"))?;
        if val.is_nan() {
            return Err(Error::invalid_data("ERR min or max is not a float"));
        }
        Ok((val, is_exclusive))
    }
}

/// 字典序范围规格(对标 Apache Kvrocks RangeLexSpec)
#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct RangeLexSpec {
    pub min: Vec<u8>,
    pub max: Vec<u8>,
    pub minex: bool,
    pub maxex: bool,
    pub min_infinite: bool,
    pub max_infinite: bool,
    pub offset: usize,
    pub count: Option<usize>,
}

impl RangeLexSpec {
    #[inline]
    pub fn new(min: impl Into<Vec<u8>>, max: impl Into<Vec<u8>>) -> Self {
        Self {
            min: min.into(),
            max: max.into(),
            minex: false,
            maxex: false,
            min_infinite: false,
            max_infinite: false,
            offset: 0,
            count: None,
        }
    }

    #[inline]
    pub fn unbounded() -> Self {
        Self {
            min_infinite: true,
            max_infinite: true,
            ..Default::default()
        }
    }

    #[inline]
    pub fn check(&self, member: &[u8]) -> bool {
        let min_ok = if self.min_infinite || self.min.is_empty() || self.min.as_slice() == b"-" {
            true
        } else if self.minex {
            member > self.min.as_slice()
        } else {
            member >= self.min.as_slice()
        };

        let max_ok = if self.max_infinite || self.max.is_empty() || self.max.as_slice() == b"+" {
            true
        } else if self.maxex {
            member < self.max.as_slice()
        } else {
            member <= self.max.as_slice()
        };

        min_ok && max_ok
    }

    /// 解析 Redis 字典序边界(例如 "-", "+", "(abc", "[abc")
    pub fn parse_bound(bound: &[u8]) -> Result<(Vec<u8>, bool, bool)> {
        if bound == b"-" {
            return Ok((Vec::new(), false, true));
        }
        if bound == b"+" {
            return Ok((Vec::new(), false, true));
        }
        if let Some(rest) = bound.strip_prefix(b"(") {
            Ok((rest.to_vec(), true, false))
        } else if let Some(rest) = bound.strip_prefix(b"[") {
            Ok((rest.to_vec(), false, false))
        } else {
            Err(Error::invalid_data(
                "ERR min or max not valid string range item",
            ))
        }
    }
}

/// 排名范围规格(对标 Apache Kvrocks RangeRankSpec)
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct RangeRankSpec {
    pub start: i64,
    pub stop: i64,
    pub reversed: bool,
}

impl RangeRankSpec {
    #[inline]
    pub fn new(start: i64, stop: i64) -> Self {
        Self {
            start,
            stop,
            reversed: false,
        }
    }

    #[inline]
    pub fn rev(start: i64, stop: i64) -> Self {
        Self {
            start,
            stop,
            reversed: true,
        }
    }
}

/// 统一 ZRANGE 选项配置(对标 Redis 6.2+ ZRANGE 全参选项)
#[derive(Debug, Clone, PartialEq, Default)]
pub struct ZRangeSpec {
    pub by_score: bool,
    pub by_lex: bool,
    pub rev: bool,
    pub with_scores: bool,
    pub offset: usize,
    pub count: Option<usize>,
}

pub type ZRange = ZRangeSpec;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_aggregate_apply_and_parse() {
        assert_eq!(Aggregate::parse("min"), Aggregate::Min);
        assert_eq!(Aggregate::parse("MAX"), Aggregate::Max);
        assert_eq!(Aggregate::parse("sum"), Aggregate::Sum);
        assert_eq!(Aggregate::parse("unknown"), Aggregate::Sum);

        assert_eq!(Aggregate::Sum.apply(10.0, 20.0), 30.0);
        assert_eq!(Aggregate::Min.apply(10.0, 20.0), 10.0);
        assert_eq!(Aggregate::Max.apply(10.0, 20.0), 20.0);

        assert_eq!(Aggregate::Sum.apply(10.0, f64::NAN), 0.0);
        assert_eq!(Aggregate::Min.apply(-f64::INFINITY, 10.0), -f64::INFINITY);
        assert_eq!(Aggregate::Max.apply(f64::INFINITY, 10.0), f64::INFINITY);
    }

    #[test]
    fn test_range_score_spec_parse_bound() {
        let (v, ex) = RangeScoreSpec::parse_bound("-inf").unwrap();
        assert_eq!(v, f64::NEG_INFINITY);
        assert!(!ex);

        let (v, ex) = RangeScoreSpec::parse_bound("+inf").unwrap();
        assert_eq!(v, f64::INFINITY);
        assert!(!ex);

        let (v, ex) = RangeScoreSpec::parse_bound("(1.5").unwrap();
        assert_eq!(v, 1.5);
        assert!(ex);

        let (v, ex) = RangeScoreSpec::parse_bound("[2.5").unwrap();
        assert_eq!(v, 2.5);
        assert!(!ex);

        let (v, ex) = RangeScoreSpec::parse_bound("10").unwrap();
        assert_eq!(v, 10.0);
        assert!(!ex);

        assert!(RangeScoreSpec::parse_bound("nan").is_err());
        assert!(RangeScoreSpec::parse_bound("").is_err());
    }

    #[test]
    fn test_range_lex_spec_parse_bound() {
        let (val, ex, inf) = RangeLexSpec::parse_bound(b"-").unwrap();
        assert!(val.is_empty());
        assert!(!ex);
        assert!(inf);

        let (val, ex, inf) = RangeLexSpec::parse_bound(b"+").unwrap();
        assert!(val.is_empty());
        assert!(!ex);
        assert!(inf);

        let (val, ex, inf) = RangeLexSpec::parse_bound(b"(hello").unwrap();
        assert_eq!(val, b"hello");
        assert!(ex);
        assert!(!inf);

        let (val, ex, inf) = RangeLexSpec::parse_bound(b"[world").unwrap();
        assert_eq!(val, b"world");
        assert!(!ex);
        assert!(!inf);

        assert!(RangeLexSpec::parse_bound(b"invalid").is_err());
    }
}