Add Frequency generators

Author: x-sheep

Sources/PropertyBased/Gen+Frequency.swift

@@ -0,0 +1,78 @@
+//
+//  Gen+Frequency.swift
+//  PropertyBased
+//
+//  Created by Lennard Sprong on 20/06/2025.
+//
+
+#if swift(>=6.2)
+extension Gen {
+    public static func oneOf<Input>(_ generators: Generator<Value, AnySequence<Input>>...)
+        -> Generator<Value, AnySequence<(index: Int, value: Input)>>
+    {
+        return frequency(
+            generators.map { gen in (rate: 1.0, gen) }
+        )
+    }
+
+    public static func frequency<Input>(
+        _ generators: [(rate: FloatLiteralType, gen: Generator<Value, AnySequence<Input>>)]
+    )
+        -> Generator<Value, AnySequence<(index: Int, value: Input)>>
+    {
+        var total: FloatLiteralType = 0
+        let options = generators.map { rate, gen in
+            precondition(rate >= 0, "Rate must be non-negative, found a rate of \(rate)")
+
+            total += rate
+            return (limit: total, gen: gen)
+        }
+
+        return Generator(
+            run: { [total] rng in
+                let pick = FloatLiteralType.random(in: 0..<total, using: &rng)
+                let index = options.firstIndex { $0.limit > pick }! as Int
+
+                return (index: index, value: options[index].gen._runIntermediate(&rng))
+            },
+            shrink: { pair in
+                let opt = options[pair.index]
+                let shrunk = opt.gen._shrinker(pair.value).lazy.map { (index: pair.index, value: $0) }
+                return AnySequence(shrunk)
+            },
+            finalResult: { pair in
+                let opt = options[pair.index]
+                return opt.gen._mapFilter(pair.value)
+            }
+        )
+    }
+
+    @_disfavoredOverload
+    public static func oneOf<each Seq: Sequence>(_ generators: repeat Generator<Value, each Seq>)
+        -> Generator<Value, AnySequence<(index: Int, value: Any)>>
+    {
+        var gens: [(rate: FloatLiteralType, gen: Generator<Value, AnySequence<Any>>)] = []
+
+        for gen in repeat each generators {
+            gens.append((rate: 1.0, gen.eraseToAny()))
+        }
+
+        return frequency(gens)
+    }
+
+    @_disfavoredOverload
+    public static func frequency<each Seq: Sequence>(
+        _ generators: repeat (rate: FloatLiteralType, gen: Generator<Value, each Seq>)
+    )
+        -> Generator<Value, AnySequence<(index: Int, value: Any)>>
+    {
+        var gens: [(rate: FloatLiteralType, gen: Generator<Value, AnySequence<Any>>)] = []
+
+        for (rate, gen) in repeat each generators {
+            gens.append((rate: rate, gen.eraseToAny()))
+        }
+
+        return frequency(gens)
+    }
+}
+#endif  // swift(>=6.2)

Sources/PropertyBased/Generator.swift

@@ -261,7 +261,7 @@ extension Generator {
 }
 
 extension Generator {
-    /// Wrap the shrinking sequence into a type-erased `AnySequence` struct.
+    /// Wrap the shrinking sequence into an `AnySequence` struct.
     ///
     /// This can be used if multiple generators must have the exact same type.
     /// - Returns: A copy of this generator.
@@ -272,4 +272,22 @@ extension Generator {
             finalResult: _mapFilter
         )
     }
+
+    /// Wrap the shrinking sequence into a type-erased `AnySequence` struct.
+    ///
+    /// This can be used if multiple generators must have the exact same type, and the underlying input value must also be hidden.
+    /// - Returns: A copy of this generator.
+    @inlinable public func eraseToAny() -> Generator<ResultValue, AnySequence<Any>> {
+        return .init(
+            run: { rng in
+                self._runIntermediate(&rng) as Any
+            },
+            shrink: {
+                AnySequence(_shrinker($0 as! InputValue).lazy.map { $0 as Any })
+            },
+            finalResult: {
+                self._mapFilter($0 as! InputValue)
+            }
+        )
+    }
 }

Tests/PropertyBasedTests/GenTests+Frequency.swift

@@ -0,0 +1,94 @@
+//
+//  GenTests+Frequency.swift
+//  PropertyBased
+//
+//  Created by Lennard Sprong on 08/07/2025.
+//
+
+import Testing
+
+@testable import PropertyBased
+
+#if swift(>=6.2)
+@Suite struct GenFrequencyTests {
+    enum Choice: Hashable {
+        case plain
+        case number(Int)
+        case text(String)
+    }
+
+    // MARK: oneOf
+
+    static let unpackedGen = Gen.oneOf(
+        Gen.always(Choice.plain).eraseToAny(),
+        Gen.int().map(Choice.number).eraseToAny(),
+        Gen.lowercaseLetter.string(of: 8).map(Choice.text).eraseToAny(),
+    )
+
+    static let packedGen = Gen.oneOf(
+        Gen.always(Choice.plain),
+        Gen.int().map(Choice.number),
+        Gen.lowercaseLetter.string(of: 8).map(Choice.text),
+    )
+
+    static let generators = [(0, unpackedGen), (1, packedGen)]
+
+    @Test(arguments: generators)
+    func testGenerateEnum(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async {
+        let gen = pair.1
+        await testGen(gen)
+
+        await confirmation(expectedCount: 1...) { confirm1 in
+            await confirmation(expectedCount: 1...) { confirm2 in
+                await confirmation(expectedCount: 1...) { confirm3 in
+                    await propertyCheck(count: 200, input: gen) { item in
+                        switch item {
+                        case .plain:
+                            confirm1()
+                        case .number:
+                            confirm2()
+                        case .text:
+                            confirm3()
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    @Test(arguments: generators)
+    func testShrinkChoice(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async throws {
+        let gen = pair.1
+        let value = (index: 1, value: 500 as Any)
+        let results = gen._shrinker(value).compactMap(gen._mapFilter)
+        try #require(results.count > 1)
+        #expect(results.first == .number(0))
+        #expect(!results.contains(.number(500)))
+    }
+
+    // MARK: frequency
+
+    static let unpackedFreqGen = Gen<Choice>.frequency([
+        (1, Gen.int().map(Choice.number).eraseToAny()),
+        (2.0, Gen.lowercaseLetter.string(of: 8).map(Choice.text).eraseToAny()),
+        (0, Gen.always(Choice.plain).eraseToAny()),
+    ])
+
+    static let packedFreqGen = Gen.frequency(
+        (1, Gen.int().map(Choice.number)),
+        (2, Gen.lowercaseLetter.string(of: 8).map(Choice.text)),
+        (0, Gen.always(Choice.plain)),
+    )
+    static let freqGenerators = [(0, unpackedFreqGen), (1, packedFreqGen)]
+
+    @Test(arguments: freqGenerators)
+    func testGenerateWithFrequency(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async {
+        let gen = pair.1
+        await testGen(gen)
+
+        await propertyCheck(count: 200, input: gen) { item in
+            #expect(item != .plain)
+        }
+    }
+}
+#endif