Improvements to Array#sample

Co-authored with Lillian

Author: chrisseaton

src/main/ruby/truffleruby/core/array.rb

@@ -975,7 +975,7 @@ def sample(count=undefined, options=undefined)
     end
 
     if count and count < 0
-      raise ArgumentError, 'count must be greater than 0'
+      raise ArgumentError, 'count must be >= 0'
     end
 
     rng = options[:random] if options
@@ -996,9 +996,9 @@ def sample(count=undefined, options=undefined)
       return [at(rng.rand(size))]
     when 2
       i = rng.rand(size)
-      j = rng.rand(size)
-      if i == j
-        j = i == 0 ? i + 1 : i - 1
+      j = rng.rand(size - 1)
+      if j >= i
+        j += 1
       end
       return [at(i), at(j)]
     else
@@ -1007,69 +1007,107 @@ def sample(count=undefined, options=undefined)
   end
 
   def sample_many(count, rng)
-    if size / count > 3
-      abandon = false
-
-      result = Array.new count
-      i = 1
+    if count <= 70 # three implementations; choice determined experimentally
+      if 2.0 * size / count  <= count  + 13
+        sample_many_swap(count, rng)
+      else
+        sample_many_quad(count, rng)
+      end
+    else
+      if size <= -1100.0 + 59.5 * count
+        sample_many_swap(count, rng)
+      else
+        sample_many_hash(count,rng)
+      end
+    end
+  end
+  private :sample_many
 
-      result[0] = rng.rand(size)
-      while i < count
-        k = rng.rand(size)
+  def sample_many_swap(count, rng)
+    #linear dependence on array size, therefore very slow for small count / size
+    result = Array.new(self)
 
-        spin = false
-        spin_count = 0
+    count.times do |c|
+      result.__send__ :swap, c, rng.rand(size)
+    end
+  end
 
-        while true # rubocop:disable Lint/LiteralAsCondition
-          j = 0
-          while j < i
-            if k == result[j]
-              spin = true
-              break
-            end
+  def sample_many_quad(count, rng)
+    #quadratic time due to linear time collision check but low overhead
+    result = Array.new count
+    i = 1
 
-            j += 1
-          end
+    result[0] = rng.rand(size)
 
-          if spin
-            if (spin_count += 1) > 100
-              abandon = true
-              break
-            end
+    while i < count
+      k = rng.rand(size)
+      spin = false
 
-            k = rng.rand(size)
-          else
+      while true # rubocop:disable Lint/LiteralAsCondition
+        j = 0
+        while j < i
+          if k == result[j]
+            spin = true
             break
           end
+
+          j += 1
         end
 
-        break if abandon
+        if spin
+          k = rng.rand(size)
+          spin = false
+        else
+          break
+        end
+      end
 
-        result[i] = k
+      result[i] = k
+      i += 1
+    end
 
-        i += 1
-      end
+    i = 0
+    while i < count
+      result[i] = at result[i]
+      i += 1
+    end
 
-      unless abandon
-        i = 0
-        while i < count
-          result[i] = at result[i]
-          i += 1
-        end
+    result
+  end
+
+  def sample_many_hash(count, rng)
+    # use hash for constant time collision check but higher overhead
+    result = Array.new count
+    i = 1
+
+    result[0] = rng.rand(size)
+    result_set = { result[0] => 0 }
 
-        return result
+    while i < count
+      k = rng.rand(size)
+
+      while true # rubocop:disable Lint/LiteralAsCondition
+        if result_set.include?(k)
+          k = rng.rand(size)
+        else
+          break
+        end
       end
-    end
 
-    result = Array.new(self)
+      result[i] = k
+      result_set[i] = k
 
-    count.times do |c|
-      result.__send__ :swap, c, rng.rand(size)
+      i += 1
     end
 
-    count == size ? result : result[0, count]
+    i = 0
+    while i < count
+      result[i] = at result[i]
+      i += 1
+    end
+
+    result
   end
-  private :sample_many
 
   def select!(&block)
     return to_enum(:select!) { size } unless block_given?