Add MaximumAttemptsTrait to control Generator behavior (#25)

Fixes #24

Author: x-sheep

README.md

@@ -30,7 +30,7 @@ Property-Based Testing can be used as an alternative for (or in addition to) tes
 
 Add the following line to the dependencies array in your `Package.swift` file:
 ```swift
-.package(url: "https://github.com/x-sheep/swift-property-based.git", from: "1.0.0")
+.package(url: "https://github.com/x-sheep/swift-property-based.git", from: "2.0.0")
 ```
 
 Then add the following to your test target:

Sources/PropertyBased/Documentation.docc/Generator.md

@@ -10,6 +10,7 @@
 ### Testing a generator
 
 - ``run(using:)``
+- ``run(using:limit:)``
 
 ### Grouping generated values
 

Sources/PropertyBased/Gen+Collection.swift

@@ -79,14 +79,14 @@ extension Generator {
     @inlinable
     public func array(of count: ClosedRange<Int>) -> Generator<[ResultValue], ArrayShrink> {
         return .init(
-            run: { rng in
+            run: { rng, limit in
                 let itemCount = Int.random(in: count, using: &rng)
 
                 var collection: [InputValue] = []
 
                 collection.reserveCapacity(itemCount)
                 for _ in 0..<itemCount {
-                    collection.append(self.runFull(&rng).input)
+                    collection.append(try self.runFull(&rng, limit).input)
                 }
                 return collection
             },

Sources/PropertyBased/Gen+Date.swift

@@ -123,7 +123,7 @@ extension Gen where Value == Date {
         let end = Date().timeIntervalSinceReferenceDate
 
         return .init(
-            run: { rng in TimeInterval.random(in: interval, using: &rng) },
+            run: { rng, _ in TimeInterval.random(in: interval, using: &rng) },
             shrink: {
                 let seq =
                     abs($0.distance(to: end)) > preferredDistanceFromNow
@@ -240,7 +240,7 @@ extension Gen where Value == Date {
         let end = Int(Date().timeIntervalSinceReferenceDate / secondsPerDay)
 
         return .init(
-            run: { rng in Int.random(in: interval, using: &rng) },
+            run: { rng, _ in Int.random(in: interval, using: &rng) },
             shrink: { $0.shrink(within: interval, towards: end) },
             finalResult: { Date(timeIntervalSinceReferenceDate: TimeInterval($0) * secondsPerDay) }
         )

Sources/PropertyBased/Gen+Frequency.swift

@@ -83,11 +83,11 @@ extension Gen {
         precondition(total > 0, "At least one generator with a weight above 0 must be specified.")
 
         return Generator(
-            run: { [total] rng in
+            run: { [total] rng, limit 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.runFull(&rng).input)
+                return try (index: index, value: options[index].gen.runFull(&rng, limit).input)
             },
             shrink: { pair in
                 let opt = options[pair.index]

Sources/PropertyBased/Generator.swift

@@ -23,7 +23,7 @@ public struct Generator<ResultValue, ShrinkSequence: SendableSequenceType>: Send
 
     /// Generate a single result, before mapping or filtering.
     @usableFromInline
-    internal var _runIntermediate: @Sendable (inout any SeededRandomNumberGenerator) -> sending InputValue
+    internal var _runIntermediate: @Sendable (inout any SeededRandomNumberGenerator, Int) throws -> sending InputValue
 
     /// Map an intermediate result to its final value, or return `nil` if the value should be filtered.
     @usableFromInline
@@ -33,22 +33,27 @@ public struct Generator<ResultValue, ShrinkSequence: SendableSequenceType>: Send
     internal var _shrinker: @Sendable (InputValue) -> ShrinkSequence
 
     /// Run the generator until a single unfiltered value is found.
-    @inlinable
-    internal func runFull<G: SeededRandomNumberGenerator>(_ rng: inout G)
+    @usableFromInline
+    internal func runFull<G: SeededRandomNumberGenerator>(_ rng: inout G, _ limit: Int)
+        throws
         -> sending (
             input: InputValue, result: ResultValue
         )
     {
         var arng: any SeededRandomNumberGenerator = rng
         defer { rng = arng as! G }
 
-        while true {
-            let run = _runIntermediate(&arng)
+        var attempts = 0
+
+        while attempts <= limit {
+            let run = try _runIntermediate(&arng, limit)
 
             if let ret = _mapFilter(run) {
                 return (run, ret)
             }
+            attempts += 1
         }
+        throw GeneratorError.runLimitExceeded(limit)
     }
 }
 
@@ -57,7 +62,22 @@ extension Generator {
     /// - Parameter rng: The random number generator to use.
     /// - Returns: A randomly generated value.
     public func run<G: SeededRandomNumberGenerator>(using rng: inout G) -> sending ResultValue {
-        runFull(&rng).result
+        try! runFull(&rng, Int.max).result
+    }
+
+    /// Generate a single value within a certain amount of attempts.
+    /// - Parameter rng: The random number generator to use.
+    /// - Parameter limit: The maximum amount of attempts before the generator stops.
+    /// - Returns: A randomly generated value.
+    /// - Throws: When the limit is reached.
+    public func run<G: SeededRandomNumberGenerator>(using rng: inout G, limit: Int)
+        throws(GeneratorError) -> sending ResultValue
+    {
+        do {
+            return try runFull(&rng, limit).result
+        } catch {
+            throw error as! GeneratorError
+        }
     }
 
     /// Remove the shrinker for this generator.
@@ -100,7 +120,7 @@ extension Generator where InputValue == ResultValue {
         run: @Sendable @escaping (inout any SeededRandomNumberGenerator) -> sending ResultValue,
         shrink: @Sendable @escaping (InputValue) -> sending ShrinkSequence,
     ) {
-        self._runIntermediate = run
+        _runIntermediate = { rng, _ in run(&rng) }
         self._shrinker = shrink
         self._mapFilter = { $0 }
     }
@@ -109,7 +129,7 @@ extension Generator where InputValue == ResultValue {
 extension Generator {
     @inlinable
     internal init(
-        run: @Sendable @escaping (inout any SeededRandomNumberGenerator) -> sending InputValue,
+        run: @Sendable @escaping (inout any SeededRandomNumberGenerator, Int) throws -> sending InputValue,
         shrink: @Sendable @escaping (InputValue) -> ShrinkSequence,
         finalResult: @Sendable @escaping (InputValue) -> ResultValue?
     ) {
@@ -128,7 +148,7 @@ extension Generator where ShrinkSequence == Shrink.None<ResultValue> {
     public init(
         run: @Sendable @escaping (inout any SeededRandomNumberGenerator) -> sending ResultValue
     ) {
-        _runIntermediate = run
+        _runIntermediate = { rng, _ in run(&rng) }
         _shrinker = { _ in .init() }
         _mapFilter = { $0 }
     }
@@ -220,11 +240,11 @@ extension Generator {
     /// - Returns: A generator of optional values.
     public func optional(valueRate: Float = 0.75) -> Generator<ResultValue?, Shrink.WithNil<ShrinkSequence>> {
         return .init(
-            run: { rng in
+            run: { rng, limit in
                 if Float.random(in: 0..<1, using: &rng) >= valueRate {
                     return nil as InputValue?
                 }
-                return self._runIntermediate(&rng)
+                return try self._runIntermediate(&rng, limit)
             },
             shrink: { value in
                 if let value {
@@ -283,8 +303,8 @@ extension Generator {
     /// - Returns: A copy of this generator.
     @inlinable public func eraseToAny() -> Generator<ResultValue, AnySequence<Any>> {
         return .init(
-            run: { rng in
-                self._runIntermediate(&rng) as Any
+            run: { rng, limit in
+                try self._runIntermediate(&rng, limit) as Any
             },
             shrink: {
                 AnySequence(_shrinker($0 as! InputValue).lazy.map { $0 as Any })

Sources/PropertyBased/GeneratorError.swift

@@ -0,0 +1,29 @@
+//
+//  GeneratorError.swift
+//  PropertyBased
+//
+//  Created by Lennard Sprong on 13/08/2026.
+//
+
+#if canImport(Foundation)
+import Foundation
+#endif
+
+/// Errors that may be thrown by a generator.
+public enum GeneratorError: Equatable, Error, CustomStringConvertible {
+    /// A generator failed to generate a valid value within the specified amount of attempts.
+    case runLimitExceeded(Int)
+
+    public var description: String {
+        switch self {
+        case .runLimitExceeded(let count):
+            "Failed to generate a valid input after \(count) attempts. Check if the Generator is filtering too many values."
+        }
+    }
+}
+
+#if canImport(Foundation)
+extension GeneratorError: LocalizedError {
+    public var localizedDescription: String { description }
+}
+#endif

Sources/PropertyBased/MaximumAttemptsTrait.swift

@@ -0,0 +1,39 @@
+//
+//  MaximumAttemptsTrait.swift
+//  PropertyBased
+//
+//  Created by Lennard Sprong on 10/08/2026.
+//
+
+import Testing
+
+/// A trait that changes how often a generator can reject values before stopping.
+///
+/// Use ``maximumAttempts(_:)`` to construct an instance of this trait.
+public struct MaximumAttemptsTrait: TestTrait, SuiteTrait, TestScoping {
+    @_documentation(visibility: internal)
+    public var isRecursive: Bool { false }
+
+    public func provideScope(
+        for test: Test, testCase: Test.Case?, performing function: @Sendable () async throws -> Void
+    ) async throws {
+        try await Self.$_maxAttempts.withValue(value) {
+            try await function()
+        }
+    }
+
+    @TaskLocal static var _maxAttempts: Int?
+
+    var value: Int
+}
+
+extension Trait where Self == MaximumAttemptsTrait {
+    /// Change how often a generator can run before stopping.
+    /// - Parameter limit: The new limit.
+    /// - Returns: An instance of ``MaximumAttemptsTrait``.
+    /// - Precondition: Limit must have a positive value.
+    public static func maximumAttempts(_ limit: Int) -> Self {
+        precondition(limit >= 0, "Limit must have a positive value.")
+        return Self(value: limit)
+    }
+}

Sources/PropertyBased/PropertyCheck.swift

@@ -118,14 +118,32 @@ public func propertyCheck<InputValue, ResultValue>(
     var rngWithIssues: (rng: Xoshiro, value: InputValue, isError: Bool)?
 
     let actualCount = fixedRng != nil ? 1 : count
+    let runLimit = MaximumAttemptsTrait._maxAttempts
 
     for _ in 0..<actualCount {
         guard !Task.isCancelled else { return }
 
         var rng = fixedRng?.rng ?? Xoshiro()
         let rngCopy = rng
 
-        let (inputValue, resultValue) = input.runFull(&rng)
+        let inputValue: InputValue
+        let resultValue: ResultValue
+        do {
+            (inputValue, resultValue) = try input.runFull(&rng, runLimit ?? 10000)
+        } catch {
+            var failureMessage = String(describing: error)
+            if runLimit == nil, let genError = error as? GeneratorError, case .runLimitExceeded = genError {
+                failureMessage += "\n\nYou can add `.maximumAttempts()` to the Test or Suite to increase the limit."
+            }
+
+            if fixedRng == nil {
+                let seed = rngCopy.traitHint
+                failureMessage += "\n\nAdd `.fixedSeed\(seed)` to the Test to reproduce this issue."
+            }
+
+            Issue.record("\(failureMessage)", sourceLocation: sourceLocation)
+            return
+        }
 
         let foundIssues = await countIssues(isolation: isolation, suppress: EnableShrinkTrait.isEnabled) {
             try await body(resultValue)