From e22a2051a59906775765a0d40bf824c0f3e3f09f Mon Sep 17 00:00:00 2001
From: Bruce Merry <bmerry@ska.ac.za>
Date: Mon, 7 Jan 2019 11:45:49 +0200
Subject: [PATCH] Minor optimisations to zset

Instead of storing a separate SortedDict sorted lexically, it just
stores a regular dict to look up scores. The *BYLEX commands now take
advantage of the assumption that all scores are equal by searching in
the _byscore list, using the score of the first key.

This makes a very small improvement in performance, and also reduces
memory usage.
---
 fakeredis/_server.py |  8 ++++----
 fakeredis/_zset.py   | 40 ++++++++++++++++++++++++++--------------
 2 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/fakeredis/_server.py b/fakeredis/_server.py
index e568776..c4d56b9 100644
--- a/fakeredis/_server.py
+++ b/fakeredis/_server.py
@@ -1889,6 +1889,7 @@ def _apply_withscores(items, withscores):
     @command((Key(ZSet), bytes, bytes), (bytes,))
     def zadd(self, key, *args):
         # TODO: handle NX, XX, CH, INCR
+        zset = key.value
         if len(args) % 2 != 0:
             raise redis.ResponseError(SYNTAX_ERROR_MSG)
         items = []
@@ -1896,12 +1897,11 @@ def zadd(self, key, *args):
         for i in range(0, len(args), 2):
             score = Float.decode(args[i])
             items.append((score, args[i + 1]))
-        old_len = len(key.value)
+        old_len = len(zset)
         for item in items:
-            if item[1] not in key.value or key.value[item[1]] != item[0]:
-                key.value[item[1]] = item[0]
+            if zset.add(item[1], item[0]):
                 key.updated()
-        return len(key.value) - old_len
+        return len(zset) - old_len
 
     @command((Key(ZSet),))
     def zcard(self, key):
diff --git a/fakeredis/_zset.py b/fakeredis/_zset.py
index b3cf611..95097b4 100644
--- a/fakeredis/_zset.py
+++ b/fakeredis/_zset.py
@@ -3,21 +3,25 @@
 
 class ZSet(object):
     def __init__(self):
-        self._bylex = sortedcontainers.SortedDict()
-        self._byscore = sortedcontainers.SortedSet()
+        self._bylex = {}     # Maps value to score
+        self._byscore = sortedcontainers.SortedList()
 
     def __contains__(self, value):
         return value in self._bylex
 
-    def __setitem__(self, value, score):
-        try:
-            old_score = self._bylex[value]
-        except KeyError:
-            pass
-        else:
-            self._byscore.discard((old_score, value))
+    def add(self, value, score):
+        """Update the item and return whether it modified the zset"""
+        old_score = self._bylex.get(value, None)
+        if old_score is not None:
+            if score == old_score:
+                return False
+            self._byscore.remove((old_score, value))
         self._bylex[value] = score
         self._byscore.add((score, value))
+        return True
+
+    def __setitem__(self, value, score):
+        self.add(value, score)
 
     def __getitem__(self, key):
         return self._bylex[key]
@@ -49,21 +53,29 @@ def zcount(self, min_, max_):
         return max(0, pos2 - pos1)
 
     def zlexcount(self, min_value, min_exclusive, max_value, max_exclusive):
+        if not self._byscore:
+            return 0
+        score = self._byscore[0][0]
         if min_exclusive:
-            pos1 = self._bylex.bisect_right(min_value)
+            pos1 = self._byscore.bisect_right((score, min_value))
         else:
-            pos1 = self._bylex.bisect_left(min_value)
+            pos1 = self._byscore.bisect_left((score, min_value))
         if max_exclusive:
-            pos2 = self._bylex.bisect_left(max_value)
+            pos2 = self._byscore.bisect_left((score, max_value))
         else:
-            pos2 = self._bylex.bisect_right(max_value)
+            pos2 = self._byscore.bisect_right((score, max_value))
         return max(0, pos2 - pos1)
 
     def islice_score(self, start, stop, reverse=False):
         return self._byscore.islice(start, stop, reverse)
 
     def irange_lex(self, start, stop, inclusive=(True, True), reverse=False):
-        return self._bylex.irange(start, stop, inclusive=inclusive, reverse=reverse)
+        if not self._byscore:
+            return iter([])
+        score = self._byscore[0][0]
+        it = self._byscore.irange((score, start), (score, stop),
+                                  inclusive=inclusive, reverse=reverse)
+        return (item[1] for item in it)
 
     def irange_score(self, start, stop, reverse=False):
         return self._byscore.irange(start, stop, reverse=reverse)