Combining ZRange options in SortedSetObject (#1021)

* Combining ZRange options

* SortedSetRange ele fix

* Fixed test

* Simplified rangeOpts creation and avoided some scratchBufferManager

---------

Co-authored-by: Tal Zaccai <talzacc@microsoft.com>

libs/server/Objects/SortedSet/SortedSetObject.cs

@@ -27,18 +27,12 @@ public enum SortedSetOperation : byte
         ZINCRBY,
         ZRANK,
         ZRANGE,
-        ZRANGEBYLEX,
-        ZRANGEBYSCORE,
-        ZRANGESTORE,
         GEOADD,
         GEOHASH,
         GEODIST,
         GEOPOS,
         GEOSEARCH,
         GEOSEARCHSTORE,
-        ZREVRANGE,
-        ZREVRANGEBYLEX,
-        ZREVRANGEBYSCORE,
         ZREVRANK,
         ZREMRANGEBYLEX,
         ZREMRANGEBYRANK,
@@ -51,6 +45,38 @@ public enum SortedSetOperation : byte
         ZMSCORE
     }
 
+    /// <summary>
+    /// Options for specifying the range in sorted set operations.
+    /// </summary>
+    [Flags]
+    public enum SortedSetRangeOpts : byte
+    {
+        /// <summary>
+        /// No options specified.
+        /// </summary>
+        None = 0,
+        /// <summary>
+        /// Range by score.
+        /// </summary>
+        ByScore = 1,
+        /// <summary>
+        /// Range by lexicographical order.
+        /// </summary>
+        ByLex = 1 << 1,
+        /// <summary>
+        /// Reverse the range order.
+        /// </summary>
+        Reverse = 1 << 2,
+        /// <summary>
+        /// Store the result.
+        /// </summary>
+        Store = 1 << 3,
+        /// <summary>
+        /// Include scores in the result.
+        /// </summary>
+        WithScores = 1 << 4
+    }
+
     [Flags]
     public enum SortedSetAddOption
     {
@@ -259,14 +285,6 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore
                     case SortedSetOperation.ZRANK:
                         SortedSetRank(ref input, ref output.SpanByteAndMemory);
                         break;
-                    case SortedSetOperation.ZRANGE:
-                    case SortedSetOperation.ZRANGESTORE:
-                        SortedSetRange(ref input, ref output.SpanByteAndMemory);
-                        break;
-                    case SortedSetOperation.ZRANGEBYLEX:
-                    case SortedSetOperation.ZRANGEBYSCORE:
-                        SortedSetRange(ref input, ref output.SpanByteAndMemory);
-                        break;
                     case SortedSetOperation.GEOADD:
                         GeoAdd(ref input, ref output.SpanByteAndMemory);
                         break;
@@ -283,11 +301,7 @@ public override unsafe bool Operate(ref ObjectInput input, ref GarnetObjectStore
                     case SortedSetOperation.GEOSEARCHSTORE:
                         GeoSearch(ref input, ref output.SpanByteAndMemory);
                         break;
-                    case SortedSetOperation.ZREVRANGE:
-                        SortedSetRange(ref input, ref output.SpanByteAndMemory);
-                        break;
-                    case SortedSetOperation.ZREVRANGEBYLEX:
-                    case SortedSetOperation.ZREVRANGEBYSCORE:
+                    case SortedSetOperation.ZRANGE:
                         SortedSetRange(ref input, ref output.SpanByteAndMemory);
                         break;
                     case SortedSetOperation.ZREVRANK:

libs/server/Objects/SortedSet/SortedSetObjectImpl.cs

@@ -416,6 +416,7 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output)
             //ZRANGE key min max [BYSCORE|BYLEX] [REV] [LIMIT offset count] [WITHSCORES]
             //ZRANGEBYSCORE key min max [WITHSCORES] [LIMIT offset count]
             //ZREVRANGEBYSCORE key max min [WITHSCORES] [LIMIT offset count]
+            var rangeOpts = (SortedSetRangeOpts)input.arg2;
             var count = input.parseState.Count;
             var respProtocolVersion = input.arg1;
 
@@ -436,30 +437,13 @@ private void SortedSetRange(ref ObjectInput input, ref SpanByteAndMemory output)
                 var maxSpan = input.parseState.GetArgSliceByRef(currIdx++).ReadOnlySpan;
 
                 // read the rest of the arguments
-                ZRangeOptions options = new();
-                switch (input.header.SortedSetOp)
+                ZRangeOptions options = new()
                 {
-                    case SortedSetOperation.ZRANGEBYLEX:
-                        options.ByLex = true;
-                        break;
-                    case SortedSetOperation.ZRANGESTORE:
-                        options.WithScores = true;
-                        break;
-                    case SortedSetOperation.ZRANGEBYSCORE:
-                        options.ByScore = true;
-                        break;
-                    case SortedSetOperation.ZREVRANGE:
-                        options.Reverse = true;
-                        break;
-                    case SortedSetOperation.ZREVRANGEBYSCORE:
-                        options.ByScore = true;
-                        options.Reverse = true;
-                        break;
-                    case SortedSetOperation.ZREVRANGEBYLEX:
-                        options.ByLex = true;
-                        options.Reverse = true;
-                        break;
-                }
+                    ByScore = (rangeOpts & SortedSetRangeOpts.ByScore) != 0,
+                    ByLex = (rangeOpts & SortedSetRangeOpts.ByLex) != 0,
+                    Reverse = (rangeOpts & SortedSetRangeOpts.Reverse) != 0,
+                    WithScores = (rangeOpts & SortedSetRangeOpts.WithScores) != 0 || (rangeOpts & SortedSetRangeOpts.Store) != 0
+                };
 
                 if (count > 2)
                 {

libs/server/Resp/Objects/SortedSetCommands.cs

@@ -159,20 +159,34 @@ private unsafe bool SortedSetRange<TGarnetApi>(RespCommand command, ref TGarnetA
             var sbKey = parseState.GetArgSliceByRef(0).SpanByte;
             var keyBytes = sbKey.ToByteArray();
 
-            var op =
-                command switch
-                {
-                    RespCommand.ZRANGE => SortedSetOperation.ZRANGE,
-                    RespCommand.ZREVRANGE => SortedSetOperation.ZREVRANGE,
-                    RespCommand.ZRANGEBYLEX => SortedSetOperation.ZRANGEBYLEX,
-                    RespCommand.ZRANGEBYSCORE => SortedSetOperation.ZRANGEBYSCORE,
-                    RespCommand.ZREVRANGEBYLEX => SortedSetOperation.ZREVRANGEBYLEX,
-                    RespCommand.ZREVRANGEBYSCORE => SortedSetOperation.ZREVRANGEBYSCORE,
-                    _ => throw new Exception($"Unexpected {nameof(SortedSetOperation)}: {command}")
-                };
+            var rangeOpts = SortedSetRangeOpts.None;
 
-            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = op };
-            var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion);
+            switch (command)
+            {
+                case RespCommand.ZRANGE:
+                    break;
+                case RespCommand.ZREVRANGE:
+                    rangeOpts = SortedSetRangeOpts.Reverse;
+                    break;
+                case RespCommand.ZRANGEBYLEX:
+                    rangeOpts = SortedSetRangeOpts.ByLex;
+                    break;
+                case RespCommand.ZRANGEBYSCORE:
+                    rangeOpts = SortedSetRangeOpts.ByScore;
+                    break;
+                case RespCommand.ZREVRANGEBYLEX:
+                    rangeOpts = SortedSetRangeOpts.ByLex | SortedSetRangeOpts.Reverse;
+                    break;
+                case RespCommand.ZREVRANGEBYSCORE:
+                    rangeOpts = SortedSetRangeOpts.ByScore | SortedSetRangeOpts.Reverse;
+                    break;
+                case RespCommand.ZRANGESTORE:
+                    rangeOpts = SortedSetRangeOpts.Store;
+                    break;
+            }
+
+            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE };
+            var input = new ObjectInput(header, ref parseState, startIdx: 1, arg1: respProtocolVersion, arg2: (int)rangeOpts);
 
             var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(dcurr, (int)(dend - dcurr)) };
 
@@ -208,8 +222,8 @@ private unsafe bool SortedSetRangeStore<TGarnetApi>(ref TGarnetApi storageApi)
             var dstKey = parseState.GetArgSliceByRef(0);
             var srcKey = parseState.GetArgSliceByRef(1);
 
-            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGESTORE };
-            var input = new ObjectInput(header, ref parseState, startIdx: 2, arg1: respProtocolVersion);
+            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE };
+            var input = new ObjectInput(header, ref parseState, startIdx: 2, arg1: respProtocolVersion, arg2: (int)SortedSetRangeOpts.Store);
 
             var status = storageApi.SortedSetRangeStore(dstKey, srcKey, ref input, out int result);
 

libs/server/Storage/Session/ObjectStore/SortedSetOps.cs

@@ -466,37 +466,18 @@ public unsafe GarnetStatus SortedSetRange<TObjectContext>(ArgSlice key, ArgSlice
                 return GarnetStatus.NOTFOUND;
             }
 
-            ReadOnlySpan<byte> operation = default;
-            var sortedOperation = SortedSetOperation.ZRANGE;
-            switch (sortedSetOrderOperation)
+            var rangeOpts = sortedSetOrderOperation switch
             {
-                case SortedSetOrderOperation.ByScore:
-                    sortedOperation = SortedSetOperation.ZRANGEBYSCORE;
-                    operation = "BYSCORE"u8;
-                    break;
-                case SortedSetOrderOperation.ByLex:
-                    sortedOperation = SortedSetOperation.ZRANGE;
-                    operation = "BYLEX"u8;
-                    break;
-                case SortedSetOrderOperation.ByRank:
-                    if (reverse)
-                        sortedOperation = SortedSetOperation.ZREVRANGE;
-                    operation = default;
-                    break;
-            }
+                SortedSetOrderOperation.ByScore => SortedSetRangeOpts.ByScore,
+                SortedSetOrderOperation.ByLex => SortedSetRangeOpts.ByLex,
+                _ => SortedSetRangeOpts.None
+            };
 
             var arguments = new List<ArgSlice> { min, max };
 
-            // Operation order
-            if (!operation.IsEmpty)
-            {
-                arguments.Add(scratchBufferManager.CreateArgSlice(operation));
-            }
-
-            // Reverse
-            if (sortedOperation != SortedSetOperation.ZREVRANGE && reverse)
+            if (reverse)
             {
-                arguments.Add(scratchBufferManager.CreateArgSlice("REV"u8));
+                rangeOpts |= SortedSetRangeOpts.Reverse;
             }
 
             // Limit parameter
@@ -517,9 +498,9 @@ public unsafe GarnetStatus SortedSetRange<TObjectContext>(ArgSlice key, ArgSlice
             parseState.InitializeWithArguments([.. arguments]);
 
             // Prepare the input
-            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = sortedOperation };
+            var header = new RespInputHeader(GarnetObjectType.SortedSet) { SortedSetOp = SortedSetOperation.ZRANGE };
             var inputArg = 2; // Default RESP server protocol version
-            var input = new ObjectInput(header, ref parseState, arg1: inputArg);
+            var input = new ObjectInput(header, ref parseState, arg1: inputArg, arg2: (int)rangeOpts);
 
             var outputFooter = new GarnetObjectStoreOutput { SpanByteAndMemory = new SpanByteAndMemory(null) };
             var status = ReadObjectStoreOperationWithOutput(key.ToArray(), ref input, ref objectContext, ref outputFooter);

libs/server/Transaction/TxnKeyManager.cs

@@ -90,8 +90,8 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan<byte>
                 RespCommand.ZINCRBY => SortedSetObjectKeys(SortedSetOperation.ZINCRBY, inputCount),
                 RespCommand.ZRANK => SortedSetObjectKeys(SortedSetOperation.ZRANK, inputCount),
                 RespCommand.ZRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
-                RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYLEX, inputCount),
-                RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGEBYSCORE, inputCount),
+                RespCommand.ZRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
+                RespCommand.ZRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
                 RespCommand.ZREVRANK => SortedSetObjectKeys(SortedSetOperation.ZREVRANK, inputCount),
                 RespCommand.ZREMRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYLEX, inputCount),
                 RespCommand.ZREMRANGEBYRANK => SortedSetObjectKeys(SortedSetOperation.ZREMRANGEBYRANK, inputCount),
@@ -105,9 +105,9 @@ internal int GetKeys(RespCommand command, int inputCount, out ReadOnlySpan<byte>
                 RespCommand.GEODIST => SortedSetObjectKeys(SortedSetOperation.GEODIST, inputCount),
                 RespCommand.GEOPOS => SortedSetObjectKeys(SortedSetOperation.GEOPOS, inputCount),
                 RespCommand.GEOSEARCH => SortedSetObjectKeys(SortedSetOperation.GEOSEARCH, inputCount),
-                RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGE, inputCount),
-                RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYLEX, inputCount),
-                RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZREVRANGEBYSCORE, inputCount),
+                RespCommand.ZREVRANGE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
+                RespCommand.ZREVRANGEBYLEX => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
+                RespCommand.ZREVRANGEBYSCORE => SortedSetObjectKeys(SortedSetOperation.ZRANGE, inputCount),
                 RespCommand.LINDEX => ListObjectKeys((byte)ListOperation.LINDEX),
                 RespCommand.LINSERT => ListObjectKeys((byte)ListOperation.LINSERT),
                 RespCommand.LLEN => ListObjectKeys((byte)ListOperation.LLEN),
@@ -210,7 +210,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount)
                 SortedSetOperation.ZINCRBY => SingleKey(1, true, LockType.Exclusive),
                 SortedSetOperation.ZRANK => SingleKey(1, true, LockType.Exclusive),
                 SortedSetOperation.ZRANGE => SingleKey(1, true, LockType.Shared),
-                SortedSetOperation.ZRANGEBYSCORE => SingleKey(1, true, LockType.Shared),
                 SortedSetOperation.ZREVRANK => SingleKey(1, true, LockType.Exclusive),
                 SortedSetOperation.ZREMRANGEBYLEX => SingleKey(1, true, LockType.Exclusive),
                 SortedSetOperation.ZREMRANGEBYRANK => SingleKey(1, true, LockType.Exclusive),
@@ -224,9 +223,6 @@ private int SortedSetObjectKeys(SortedSetOperation command, int inputCount)
                 SortedSetOperation.GEODIST => SingleKey(1, true, LockType.Shared),
                 SortedSetOperation.GEOPOS => SingleKey(1, true, LockType.Shared),
                 SortedSetOperation.GEOSEARCH => SingleKey(1, true, LockType.Shared),
-                SortedSetOperation.ZREVRANGE => SingleKey(1, true, LockType.Shared),
-                SortedSetOperation.ZREVRANGEBYLEX => SingleKey(1, true, LockType.Shared),
-                SortedSetOperation.ZREVRANGEBYSCORE => SingleKey(1, true, LockType.Shared),
                 _ => -1
             };
         }