Skip to content

Commit 7880dd8

Browse files
x-sheepx-sheep
authored
Add Frequency generators
1 parent bae7879 commit 7880dd8

File tree

4 files changed

+294
-1
lines changed

4 files changed

+294
-1
lines changed

‎Sources/PropertyBased/Documentation.docc/Gen.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,11 @@ You can generate individual characters, and use ``/Generator/string(of:)`` to fo
5353
- ``/Gen/year``
5454
- ``/Gen/year(in:)``
5555

56+
### Generating enums with associated values
57+
58+
- ``/Gen/oneOf(_:)``
59+
- ``/Gen/frequency(_:)``
60+
5661
### Handling immutable collections
5762

5863
- ``/Gen/case``

‎Sources/PropertyBased/Gen+Frequency.swift‎

Lines changed: 176 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,176 @@
1+
//
2+
// Gen+Frequency.swift
3+
// PropertyBased
4+
//
5+
// Created by Lennard Sprong on 20/06/2025.
6+
//
7+
8+
#if swift(>=6.2)
9+
extension Gen {
10+
/// Produces a generator that randomly selects one of the provided generators.
11+
///
12+
/// Every generator has an equal chance of being run.
13+
///
14+
/// ## Precondition
15+
/// At least one generator must be provided.
16+
///
17+
/// ## Example
18+
/// ```swift
19+
/// enum Choice {
20+
/// case plain
21+
/// case number(Int)
22+
/// case text(String)
23+
/// }
24+
///
25+
/// Gen.oneOf(
26+
/// Gen.always(Choice.plain),
27+
/// Gen.int().map(Choice.number),
28+
/// Gen.lowercaseLetter.string(of: 8).map(Choice.text),
29+
/// )
30+
/// ```
31+
///
32+
/// - Parameter generators: The generators to choose from.
33+
/// - Returns: A new generator.
34+
@_documentation(visibility: internal)
35+
public static func oneOf<Input>(_ generators: Generator<Value, AnySequence<Input>>...)
36+
-> Generator<Value, AnySequence<(index: Int, value: Input)>>
37+
{
38+
return frequency(
39+
generators.map { gen in (weight: 1.0, gen) }
40+
)
41+
}
42+
43+
/// Produces a generator that randomly selects one of the provided generators, with a weighted distribution.
44+
///
45+
/// ## Precondition
46+
/// At least one generator with a weight above zero must be provided.
47+
///
48+
/// ## Example
49+
/// ```swift
50+
/// enum Choice {
51+
/// case plain
52+
/// case number(Int)
53+
/// case text(String)
54+
/// }
55+
///
56+
/// Gen.frequency(
57+
/// (0.5, Gen.always(Choice.plain)),
58+
/// (1, Gen.int().map(Choice.number)),
59+
/// (2, Gen.lowercaseLetter.string(of: 8).map(Choice.text)),
60+
/// )
61+
/// ```
62+
///
63+
/// - Parameter generators: The generators to choose from, with a weight for each.
64+
/// - Returns: A new generator.
65+
@_documentation(visibility: internal)
66+
public static func frequency<Input>(
67+
_ generators: [(weight: FloatLiteralType, gen: Generator<Value, AnySequence<Input>>)]
68+
)
69+
-> Generator<Value, AnySequence<(index: Int, value: Input)>>
70+
{
71+
precondition(!generators.isEmpty, "At least one generator must be specified.")
72+
73+
var total: FloatLiteralType = 0
74+
let options = generators.map { weight, gen in
75+
precondition(weight >= 0, "Weight must be non-negative, found a weight of \(weight)")
76+
77+
total += weight
78+
return (limit: total, gen: gen)
79+
}
80+
81+
precondition(total > 0, "At least one generator with a weight above 0 must be specified.")
82+
83+
return Generator(
84+
run: { [total] rng in
85+
let pick = FloatLiteralType.random(in: 0..<total, using: &rng)
86+
let index = options.firstIndex { $0.limit > pick }! as Int
87+
88+
return (index: index, value: options[index].gen._runIntermediate(&rng))
89+
},
90+
shrink: { pair in
91+
let opt = options[pair.index]
92+
let shrunk = opt.gen._shrinker(pair.value).lazy.map { (index: pair.index, value: $0) }
93+
return AnySequence(shrunk)
94+
},
95+
finalResult: { pair in
96+
let opt = options[pair.index]
97+
return opt.gen._mapFilter(pair.value)
98+
}
99+
)
100+
}
101+
102+
/// Produces a generator that randomly selects one of the provided generators.
103+
///
104+
/// Every generator has an equal chance of being run.
105+
///
106+
/// ## Precondition
107+
/// At least one generator must be provided.
108+
///
109+
/// ## Example
110+
/// ```swift
111+
/// enum Choice {
112+
/// case plain
113+
/// case number(Int)
114+
/// case text(String)
115+
/// }
116+
///
117+
/// Gen.oneOf(
118+
/// Gen.always(Choice.plain),
119+
/// Gen.int().map(Choice.number),
120+
/// Gen.lowercaseLetter.string(of: 8).map(Choice.text),
121+
/// )
122+
/// ```
123+
///
124+
/// - Parameter generators: The generators to choose from.
125+
/// - Returns: A new generator.
126+
@_disfavoredOverload
127+
public static func oneOf<each Seq: Sequence>(_ generators: repeat Generator<Value, each Seq>)
128+
-> Generator<Value, AnySequence<(index: Int, value: Any)>>
129+
{
130+
var gens: [(weight: FloatLiteralType, gen: Generator<Value, AnySequence<Any>>)] = []
131+
132+
for gen in repeat each generators {
133+
gens.append((weight: 1.0, gen.eraseToAny()))
134+
}
135+
136+
return frequency(gens)
137+
}
138+
139+
/// Produces a generator that randomly selects one of the provided generators, with a weighted distribution.
140+
///
141+
/// ## Precondition
142+
/// At least one generator with a weight above zero must be provided.
143+
///
144+
/// ## Example
145+
/// ```swift
146+
/// enum Choice {
147+
/// case plain
148+
/// case number(Int)
149+
/// case text(String)
150+
/// }
151+
///
152+
/// Gen.frequency(
153+
/// (0.5, Gen.always(Choice.plain)),
154+
/// (1, Gen.int().map(Choice.number)),
155+
/// (2, Gen.lowercaseLetter.string(of: 8).map(Choice.text)),
156+
/// )
157+
/// ```
158+
///
159+
/// - Parameter generators: The generators to choose from, with a weight for each.
160+
/// - Returns: A new generator.
161+
@_disfavoredOverload
162+
public static func frequency<each Seq: Sequence>(
163+
_ generators: repeat (weight: FloatLiteralType, gen: Generator<Value, each Seq>)
164+
)
165+
-> Generator<Value, AnySequence<(index: Int, value: Any)>>
166+
{
167+
var gens: [(weight: FloatLiteralType, gen: Generator<Value, AnySequence<Any>>)] = []
168+
169+
for (weight, gen) in repeat each generators {
170+
gens.append((weight: weight, gen.eraseToAny()))
171+
}
172+
173+
return frequency(gens)
174+
}
175+
}
176+
#endif // swift(>=6.2)

‎Sources/PropertyBased/Generator.swift‎

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,7 @@ extension Generator {
261261
}
262262

263263
extension Generator {
264-
/// Wrap the shrinking sequence into a type-erased `AnySequence` struct.
264+
/// Wrap the shrinking sequence into an `AnySequence` struct.
265265
///
266266
/// This can be used if multiple generators must have the exact same type.
267267
/// - Returns: A copy of this generator.
@@ -272,4 +272,22 @@ extension Generator {
272272
finalResult: _mapFilter
273273
)
274274
}
275+
276+
/// Wrap the shrinking sequence into a type-erased `AnySequence` struct.
277+
///
278+
/// This can be used if multiple generators must have the exact same type, and the underlying input value must also be hidden.
279+
/// - Returns: A copy of this generator.
280+
@inlinable public func eraseToAny() -> Generator<ResultValue, AnySequence<Any>> {
281+
return .init(
282+
run: { rng in
283+
self._runIntermediate(&rng) as Any
284+
},
285+
shrink: {
286+
AnySequence(_shrinker($0 as! InputValue).lazy.map { $0 as Any })
287+
},
288+
finalResult: {
289+
self._mapFilter($0 as! InputValue)
290+
}
291+
)
292+
}
275293
}

‎Tests/PropertyBasedTests/GenTests+Frequency.swift‎

Lines changed: 94 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
//
2+
// GenTests+Frequency.swift
3+
// PropertyBased
4+
//
5+
// Created by Lennard Sprong on 08/07/2025.
6+
//
7+
8+
import Testing
9+
10+
@testable import PropertyBased
11+
12+
#if swift(>=6.2)
13+
@Suite struct GenFrequencyTests {
14+
enum Choice: Hashable {
15+
case plain
16+
case number(Int)
17+
case text(String)
18+
}
19+
20+
// MARK: oneOf
21+
22+
static let unpackedGen = Gen.oneOf(
23+
Gen.always(Choice.plain).eraseToAny(),
24+
Gen.int().map(Choice.number).eraseToAny(),
25+
Gen.lowercaseLetter.string(of: 8).map(Choice.text).eraseToAny(),
26+
)
27+
28+
static let packedGen = Gen.oneOf(
29+
Gen.always(Choice.plain),
30+
Gen.int().map(Choice.number),
31+
Gen.lowercaseLetter.string(of: 8).map(Choice.text),
32+
)
33+
34+
static let generators = [(0, unpackedGen), (1, packedGen)]
35+
36+
@Test(arguments: generators)
37+
func testGenerateEnum(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async {
38+
let gen = pair.1
39+
await testGen(gen)
40+
41+
await confirmation(expectedCount: 1...) { confirm1 in
42+
await confirmation(expectedCount: 1...) { confirm2 in
43+
await confirmation(expectedCount: 1...) { confirm3 in
44+
await propertyCheck(count: 200, input: gen) { item in
45+
switch item {
46+
case .plain:
47+
confirm1()
48+
case .number:
49+
confirm2()
50+
case .text:
51+
confirm3()
52+
}
53+
}
54+
}
55+
}
56+
}
57+
}
58+
59+
@Test(arguments: generators)
60+
func testShrinkChoice(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async throws {
61+
let gen = pair.1
62+
let value = (index: 1, value: 500 as Any)
63+
let results = gen._shrinker(value).compactMap(gen._mapFilter)
64+
try #require(results.count > 1)
65+
#expect(results.first == .number(0))
66+
#expect(!results.contains(.number(500)))
67+
}
68+
69+
// MARK: frequency
70+
71+
static let unpackedFreqGen = Gen<Choice>.frequency([
72+
(1, Gen.int().map(Choice.number).eraseToAny()),
73+
(2.0, Gen.lowercaseLetter.string(of: 8).map(Choice.text).eraseToAny()),
74+
(0, Gen.always(Choice.plain).eraseToAny()),
75+
])
76+
77+
static let packedFreqGen = Gen.frequency(
78+
(1, Gen.int().map(Choice.number)),
79+
(2, Gen.lowercaseLetter.string(of: 8).map(Choice.text)),
80+
(0, Gen.always(Choice.plain)),
81+
)
82+
static let freqGenerators = [(0, unpackedFreqGen), (1, packedFreqGen)]
83+
84+
@Test(arguments: freqGenerators)
85+
func testGenerateWithFrequency(_ pair: (Int, Generator<Choice, AnySequence<(index: Int, value: Any)>>)) async {
86+
let gen = pair.1
87+
await testGen(gen)
88+
89+
await propertyCheck(count: 200, input: gen) { item in
90+
#expect(item != .plain)
91+
}
92+
}
93+
}
94+
#endif

0 commit comments

Comments
 (0)