use crate::error::{Error, Result};
use std::str;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ZAdd {
Nx,
Xx,
Gt,
Lt,
Ch,
Incr,
}
#[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 }
}
}
#[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
}
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))
}
}
#[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
}
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",
))
}
}
}
#[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,
}
}
}
#[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());
}
}