[GR-17457] Array sample fix

PullRequest: truffleruby/3356

Author: aardvark179

spec/ruby/core/array/fixtures/classes.rb

@@ -40,6 +40,68 @@ def self.empty_recursive_array
     a
   end
 
+  # Chi squared critical values for tests with n degrees of freedom at 99% confidence.
+  # Values obtained from NIST Engineering Statistic Handbook at
+  # https://www.itl.nist.gov/div898/handbook/eda/section3/eda3674.htm
+
+  CHI_SQUARED_CRITICAL_VALUES = [
+    0,
+    6.635, 9.210, 11.345, 13.277, 15.086, 16.812, 18.475, 20.090, 21.666, 23.209,
+    24.725, 26.217, 27.688, 29.141, 30.578, 32.000, 33.409, 34.805, 36.191, 37.566,
+    38.932, 40.289, 41.638, 42.980, 44.314, 45.642, 46.963, 48.278, 49.588, 50.892,
+    52.191, 53.486, 54.776, 56.061, 57.342, 58.619, 59.893, 61.162, 62.428, 63.691,
+    64.950, 66.206, 67.459, 68.710, 69.957, 71.201, 72.443, 73.683, 74.919, 76.154,
+    77.386, 78.616, 79.843, 81.069, 82.292, 83.513, 84.733, 85.950, 87.166, 88.379,
+    89.591, 90.802, 92.010, 93.217, 94.422, 95.626, 96.828, 98.028, 99.228, 100.425,
+    101.621, 102.816, 104.010, 105.202, 106.393, 107.583, 108.771, 109.958, 111.144, 112.329,
+    113.512, 114.695, 115.876, 117.057, 118.236, 119.414, 120.591, 121.767, 122.942, 124.116,
+    125.289, 126.462, 127.633, 128.803, 129.973, 131.141, 132.309, 133.476, 134.642, 135.807,
+      ]
+
+  def self.measure_sample_fairness(size, samples, iters)
+    ary = Array.new(size) { |x| x }
+    (samples).times do |i|
+      chi_results = []
+      3.times do
+        counts = Array.new(size) { 0 }
+        expected = iters / size
+        iters.times do
+          x = ary.sample(samples)[i]
+          counts[x] += 1
+        end
+        chi_squared = 0.0
+        counts.each do |count|
+          chi_squared += (((count - expected) ** 2) * 1.0 / expected)
+        end
+        chi_results << chi_squared
+        break if chi_squared <= CHI_SQUARED_CRITICAL_VALUES[size]
+      end
+
+      chi_results.min.should <= CHI_SQUARED_CRITICAL_VALUES[size]
+    end
+  end
+
+  def self.measure_sample_fairness_large_sample_size(size, samples, iters)
+    ary = Array.new(size) { |x| x }
+    counts = Array.new(size) { 0 }
+    expected = iters * samples / size
+    iters.times do
+      ary.sample(samples).each do |sample|
+        counts[sample] += 1
+      end
+    end
+    chi_squared = 0.0
+    counts.each do |count|
+      chi_squared += (((count - expected) ** 2) * 1.0 / expected)
+    end
+
+    # Chi squared critical values for tests with 4 degrees of freedom
+    # Values obtained from NIST Engineering Statistic Handbook at
+    # https://www.itl.nist.gov/div898/handbook/eda/section3/eda3674.htm
+
+    chi_squared.should <= CHI_SQUARED_CRITICAL_VALUES[size]
+  end
+
   class MyArray < Array
     # The #initialize method has a different signature than Array to help
     # catch places in the specs that do not assert the #initialize is not

spec/ruby/core/array/sample_spec.rb

@@ -3,16 +3,14 @@
 
 describe "Array#sample" do
   it "samples evenly" do
-    ary = [0, 1, 2, 3]
-    3.times do |i|
-      counts = [0, 0, 0, 0]
-      4000.times do
-        counts[ary.sample(3)[i]] += 1
-      end
-      counts.each do |count|
-        (800..1200).should include(count)
-      end
-    end
+    ArraySpecs.measure_sample_fairness(4, 1, 4000)
+    ArraySpecs.measure_sample_fairness(4, 2, 4000)
+    ArraySpecs.measure_sample_fairness(4, 3, 4000)
+    ArraySpecs.measure_sample_fairness(40, 3, 4000)
+    ArraySpecs.measure_sample_fairness(40, 4, 4000)
+    ArraySpecs.measure_sample_fairness(40, 8, 4000)
+    ArraySpecs.measure_sample_fairness(40, 16, 4000)
+    ArraySpecs.measure_sample_fairness_large_sample_size(100, 80, 40000)
   end
 
   it "returns nil for an empty Array" do

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

@@ -1073,11 +1073,12 @@ def sample(count = undefined, random: nil)
   end
 
   private def sample_many_swap(count, rng)
-    # linear dependence on array size, therefore very slow for small count / size
+    # linear dependence on count,
     result = Array.new(self)
 
-    count.times do |c|
-      result.__send__ :swap, c, rng.rand(size)
+    count.times do |i|
+      r = i + rng.rand(result.size - i)
+      result.__send__ :swap, i, r
     end
 
     count == size ? result : result[0, count]