Add ZREMRANGEBYLEX command

chowc · Jul 14, 2015 · 2db650c · 2db650c
1 parent 48bbdcb
commit 2db650c
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 1 deletion.
diff --git a/TODO.md b/TODO.md
@@ -53,7 +53,7 @@
     - [x] zrem
     - [x] zremrangebyscore
     - [x] zremrangebyrank
-    - [ ] zremrangebylex
+    - [x] zremrangebylex
     - [x] zunionstore
     - [x] zinterstore
     - [x] zrange

diff --git a/command/src/lib.rs b/command/src/lib.rs
@@ -1152,6 +1152,33 @@ fn zremrangebyscore(parser: &Parser, db: &mut Database, dbindex: usize) -> Respo
     }
 }
 
+fn zremrangebylex(parser: &Parser, db: &mut Database, dbindex: usize) -> Response {
+    validate!(parser.argv.len() == 4, "Wrong number of parameters");
+    let key = try_validate!(parser.get_vec(1), "Invalid key");
+    let min = {
+        let m = try_validate!(parser.get_vec(2), "Invalid min");
+        match get_vec_bound(m) {
+            Ok(v) => v,
+            Err(e) => return e,
+        }
+    };
+    let max = {
+        let m = try_validate!(parser.get_vec(3), "Invalid max");
+        match get_vec_bound(m) {
+            Ok(v) => v,
+            Err(e) => return e,
+        }
+    };
+    let el = match db.get_mut(dbindex, &key) {
+        Some(e) => e,
+        None => return Response::Integer(0),
+    };
+    match el.zremrangebylex(min, max) {
+        Ok(c) => Response::Integer(c as i64),
+        Err(err) => Response::Error(err.to_string()),
+    }
+}
+
 fn zremrangebyrank(parser: &Parser, db: &mut Database, dbindex: usize) -> Response {
     validate!(parser.argv.len() == 4, "Wrong number of parameters");
     let key = try_validate!(parser.get_vec(1), "Invalid key");
@@ -1645,6 +1672,7 @@ pub fn command(
         "zscore" => zscore(parser, db, dbindex),
         "zincrby" => zincrby(parser, db, dbindex),
         "zrem" => zrem(parser, db, dbindex),
+        "zremrangebylex" => zremrangebylex(parser, db, dbindex),
         "zremrangebyscore" => zremrangebyscore(parser, db, dbindex),
         "zremrangebyrank" => zremrangebyrank(parser, db, dbindex),
         "zcount" => zcount(parser, db, dbindex),
@@ -2578,6 +2606,16 @@ mod test_command {
         assert_eq!(command(&parser!(b"zremrangebyscore key -inf inf"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(1));
     }
 
+    #[test]
+    fn zremrangebylex_command() {
+        let mut db = Database::new(Config::new(Logger::new(Level::Warning)));
+        assert_eq!(command(&parser!(b"zadd key 0 a 0 b 0 c 0 d"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(4));
+        assert_eq!(command(&parser!(b"zremrangebylex key [b (d"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(2));
+        assert_eq!(command(&parser!(b"zremrangebylex key [b (d"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(0));
+        assert_eq!(command(&parser!(b"zremrangebylex key (b [d"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(1));
+        assert_eq!(command(&parser!(b"zremrangebylex key - +"), &mut db, &mut 0, &mut true, None, None, None).unwrap(), Response::Integer(1));
+    }
+
     #[test]
     fn zremrangebyrank_command() {
         let mut db = Database::new(Config::new(Logger::new(Level::Warning)));

diff --git a/database/src/lib.rs b/database/src/lib.rs
@@ -1271,6 +1271,31 @@ impl Value {
         }
     }
 
+    /// Removes all elements within a lex range. Returns the number of removed elements
+    ///
+    /// # Examples
+    /// ```
+    /// #![feature(collections_bound)]
+    /// use database::Value;
+    /// use std::collections::Bound;
+    ///
+    /// let mut val = Value::Nil;
+    /// assert_eq!(val.zremrangebylex(Bound::Unbounded, Bound::Unbounded).unwrap(), 0);
+    /// val.zadd(0.0, vec![1], false, false, false, false).unwrap();
+    /// val.zadd(0.0, vec![2], false, false, false, false).unwrap();
+    /// val.zadd(0.0, vec![3], false, false, false, false).unwrap();
+    /// assert_eq!(val.zremrangebylex(Bound::Included(vec![2]), Bound::Excluded(vec![3])).unwrap(), 1);
+    /// assert_eq!(val.zremrangebylex(Bound::Included(vec![2]), Bound::Excluded(vec![3])).unwrap(), 0);
+    /// assert_eq!(val.zcard().unwrap(), 2);
+    /// ```
+    pub fn zremrangebylex(&mut self, min: Bound<Vec<u8>>, max: Bound<Vec<u8>>) -> Result<usize, OperationError> {
+        match *self {
+            Value::Nil => Ok(0),
+            Value::SortedSet(ref mut value) => Ok(value.zremrangebylex(min, max)),
+            _ => Err(OperationError::WrongTypeError),
+        }
+    }
+
     /// Removes all elements within a rank range. Returns the number of removed elements
     ///
     /// # Examples

diff --git a/database/src/zset.rs b/database/src/zset.rs
@@ -192,6 +192,10 @@ impl ValueSortedSet {
             ValueSortedSet::Data(ref skiplist, _) => skiplist,
         };
 
+        if skiplist.len() == 0 {
+            return vec![];
+        }
+
         let f = skiplist.front().unwrap().get_f64();
         let mut f1 = SortedSetMember::new(f.clone(), vec![]);
         let mut f2 = SortedSetMember::new(f.clone(), vec![]);
@@ -244,6 +248,24 @@ impl ValueSortedSet {
         count
     }
 
+    pub fn zremrangebylex(&mut self, min: Bound<Vec<u8>>, max: Bound<Vec<u8>>) -> usize {
+        let pos = match min {
+            Bound::Included(ref s) => self.zlexcount(Bound::Unbounded, Bound::Excluded(s.clone())),
+            Bound::Excluded(ref s) => self.zlexcount(Bound::Unbounded, Bound::Included(s.clone())),
+            Bound::Unbounded => 0,
+        };
+        let count = self.zlexcount(min, max);
+        let (skiplist, hmap) = match *self {
+            ValueSortedSet::Data(ref mut skiplist, ref mut hmap) => (skiplist, hmap),
+        };
+
+        for _ in 0..count {
+            let el = skiplist.remove_index(&pos);
+            hmap.remove(&el.s);
+        };
+        count
+    }
+
     fn normalize_range(&self, start: i64, stop: i64, rev: bool) -> (usize, usize) {
         let skiplist = match *self {
             ValueSortedSet::Data(ref skiplist, _) => skiplist,
@@ -568,6 +590,21 @@ fn zremrangebyscore() {
     assert_eq!(zset.zremrangebyscore(Bound::Unbounded, Bound::Unbounded), 0);
 }
 
+#[test]
+fn zremrangebylex() {
+    let mut zset = ValueSortedSet::new();
+    zset.zadd(0.0, vec![1], false, false, false, false);
+    zset.zadd(0.0, vec![2], false, false, false, false);
+    zset.zadd(0.0, vec![3], false, false, false, false);
+    zset.zadd(0.0, vec![4], false, false, false, false);
+    assert_eq!(zset.zremrangebylex(Bound::Included(vec![2]), Bound::Excluded(vec![4])), 2);
+    assert_eq!(zset.zrank(vec![1]).unwrap(), 0);
+    assert_eq!(zset.zrank(vec![2]), None);
+    assert_eq!(zset.zremrangebylex(Bound::Unbounded, Bound::Unbounded), 2);
+    assert_eq!(zset.zrank(vec![1]), None);
+    assert_eq!(zset.zremrangebylex(Bound::Unbounded, Bound::Unbounded), 0);
+}
+
 #[test]
 fn zremrangebyrank() {
     let mut zset = ValueSortedSet::new();