Fix RandomSample crash when 0 was out from random. (#70)

Author: sidepelican

Sources/Algorithms/RandomSample.swift

@@ -94,7 +94,8 @@ internal func nextW<G: RandomNumberGenerator>(
 internal func nextOffset<G: RandomNumberGenerator>(
   w: Double, using rng: inout G
 ) -> Int {
-  Int(Double.log(.random(in: 0..<1, using: &rng)) / .log(1 - w))
+  let offset = Double.log(.random(in: 0..<1, using: &rng)) / .log(onePlus: -w)
+  return offset < Double(Int.max) ? Int(offset) : Int.max
 }
 
 extension Collection {
@@ -201,10 +202,10 @@ extension Sequence {
       w *= nextW(k: k, using: &rng)
 
       // Find the offset of the next element to swap into the reservoir.
-      var offset = nextOffset(w: w, using: &rng) + 1
+      var offset = nextOffset(w: w, using: &rng)
 
-      // Skip over `offset - 1` elements to find the selected element.
-      while offset > 1, let _ = iterator.next() {
+      // Skip over `offset` elements to find the selected element.
+      while offset > 0, let _ = iterator.next() {
         offset -= 1
       }
       guard let nextElement = iterator.next() else { break }

Tests/SwiftAlgorithmsTests/RandomSampleTests.swift

@@ -10,7 +10,7 @@
 //===----------------------------------------------------------------------===//
 
 import XCTest
-import Algorithms
+@testable import Algorithms
 
 func validateRandomSamples<S: Sequence>(
   _ samples: [Int: Int],
@@ -95,4 +95,32 @@ final class RandomSampleTests: XCTestCase {
     let sample2c = c.randomStableSample(count: k, using: &generator)
     XCTAssertEqual(sample1c, sample2c)
   }
+
+  func testRandomSampleRandomEdgeCasesInternal() {
+    struct ZeroGenerator: RandomNumberGenerator {
+      mutating func next() -> UInt64 { 0 }
+    }
+    var zero = ZeroGenerator()
+    _ = nextOffset(w: 1, using: &zero) // must not crash
+
+    struct AlmostAllZeroGenerator: RandomNumberGenerator {
+      private var forward: SplitMix64
+      private var count: Int = 0
+
+      init(seed: UInt64) {
+        forward = SplitMix64(seed: seed)
+      }
+
+      mutating func next() -> UInt64 {
+        defer { count &+= 1 }
+        if count % 1000 == 0 { return forward.next() }
+        return 0
+      }
+    }
+
+    var almostAllZero = AlmostAllZeroGenerator(seed: 0)
+    _ = s.randomSample(count: k, using: &almostAllZero) // must not crash
+    almostAllZero = AlmostAllZeroGenerator(seed: 0)
+    _ = c.randomSample(count: k, using: &almostAllZero) // must not crash
+  }
 }