Added SortedSet#between

Author: alexdowad

lib/hamster/sorted_set.rb

@@ -560,6 +560,14 @@ def up_to(item, &block)
       end
     end
 
+    def between(from, to, &block)
+      if block_given?
+        @node.each_between(from, to, @comparator, &block)
+      else
+        self.class.alloc(@node.between(from, to, @comparator), @comparator)
+      end
+    end
+
     # Return a randomly chosen item from this set. If the set is empty, return `nil`.
     #
     # @return [Object]
@@ -778,6 +786,18 @@ def suffix(item, comparator, inclusive)
         end
       end
 
+      def between(from, to, comparator)
+        if direction(from, comparator) > 0 # all on the right
+          @right.between(from, to, comparator)
+        elsif direction(to, comparator) < 0 # all on the left
+          @left.between(from, to, comparator)
+        else
+          left = @left.suffix(from, comparator, true)
+          right = @right.prefix(to, comparator, true)
+          rebalance(left, right)
+        end
+      end
+
       def each_less(item, comparator, inclusive, &block)
         dir = direction(item, comparator)
         if dir > 0 || (inclusive && dir == 0)
@@ -800,6 +820,18 @@ def each_greater(item, comparator, inclusive, &block)
         end
       end
 
+      def each_between(from, to, comparator, &block)
+        if direction(from, comparator) > 0 # all on the right
+          @right.each_between(from, to, comparator, &block)
+        elsif direction(to, comparator) < 0 # all on the left
+          @left.each_between(from, to, comparator, &block)
+        else
+          @left.each_greater(from, comparator, true, &block)
+          yield @item
+          @right.each_less(to, comparator, true, &block)
+        end
+      end
+
       def each(&block)
         @left.each(&block)
         yield @item
@@ -942,8 +974,10 @@ def e.delete(item, comparator); self; end
       def e.include?(item, comparator); false; end
       def e.prefix(item, comparator, inclusive); self; end
       def e.suffix(item, comparator, inclusive); self; end
+      def e.between(from, to, comparator); self; end
       def e.each_greater(item, comparator, inclusive); end
       def e.each_less(item, comparator, inclusive); end
+      def e.each_between(item, comparator, inclusive); end
       def e.empty?; true; end
       def e.slice(from, length); self; end
     end.freeze

spec/lib/hamster/sorted_set/between_spec.rb

@@ -0,0 +1,52 @@
+require "spec_helper"
+require "hamster/sorted_set"
+
+describe Hamster::SortedSet do
+  describe "#between" do
+    context "when called without a block" do
+      it "returns a sorted set of all items from the first argument to the second" do
+        100.times do
+          items = rand(100).times.collect { rand(1000) }
+          set   = Hamster::SortedSet.new(items)
+          from,to = [rand(1000),rand(1000)].sort
+          result = set.between(from, to)
+          array  = items.select { |x| x >= from && x <= to }.sort
+          result.class.should be(Hamster::SortedSet)
+          result.size.should == array.size
+          result.to_a.should == array
+        end
+      end
+    end
+
+    context "when called with a block" do
+      it "yields all the items lower than the argument" do
+        100.times do
+          items = rand(100).times.collect { rand(1000) }
+          set   = Hamster::SortedSet.new(items)
+          from,to = [rand(1000),rand(1000)].sort
+          result = []
+          set.between(from, to) { |x| result << x }
+          array  = items.select { |x| x >= from && x <= to }.sort
+          result.size.should == array.size
+          result.should == array
+        end
+      end
+    end
+
+    context "on an empty set" do
+      it "returns an empty set" do
+        Hamster::SortedSet.empty.between(1, 2).should be_empty
+        Hamster::SortedSet.empty.between('abc', 'def').should be_empty
+        Hamster::SortedSet.empty.between(:symbol, :another).should be_empty
+      end
+    end
+
+    context "with a 'to' argument lower than the 'from' argument" do
+      it "returns an empty set" do
+        result = Hamster.sorted_set(*(1..100)).between(6, 5)
+        result.class.should be(Hamster::SortedSet)
+        result.should be_empty
+      end
+    end
+  end
+end