Repository: google/swift-benchmark (swift)

Issue:
run setup code before the benchmark
I would like a way to set up some data before the timed portion of the benchmark runs.

Concrete use case: In https://github.com/borglab/SwiftFusion/pull/61, I add a benchmark that requires a dataset from the internet. I put the code that downloads the dataset outside of the benchmark, but this means that the dataset download happens even when the benchmark that requires it is filtered out.

Pull request description:
This change adds support for lower-overhead version of writing explicit measurement loop within the benchmark itself:

```swift
benchmark("my bench") { state in
    // set-up, not measured

    state.measure {
        // body of the benchmark, runs #iterations times
    }

    // tear-down, not measured
}
```

By using measure we get less overhead of measurements by around 10%:

```
    $ swift run -c release BenchmarkMinimalExample --min-time 10
    --max-iterations 100000000
    Updating https://github.com/apple/swift-argument-parser
    Resolving https://github.com/apple/swift-argument-parser at 0.0.6
    [4/4] Linking BenchmarkMinimalExample
    running noop... done! (6095.58 ms)
    running measure noop... done! (5549.44 ms)

    name            time    std         iterations
    ----------------------------------------------
    noop            23.0 ns ±  83.88 %   100000000
    measure noop    21.0 ns ±  88.80 %   100000000
```

`state.measure` has the same precision as if the clock is used directly in a tight loop without any closures. To improve precision further we would need to provide a higher precision clock than the one provided by DispatchTime.

Fixes #29

Patch:
diff --git a/Sources/Benchmark/Benchmark.swift b/Sources/Benchmark/Benchmark.swift
index c929922..539e9df 100644
--- a/Sources/Benchmark/Benchmark.swift
+++ b/Sources/Benchmark/Benchmark.swift
@@ -16,7 +16,7 @@ public protocol AnyBenchmark {
     var name: String { get }
     var settings: [BenchmarkSetting] { get }
     func setUp()
-    func run()
+    func run(_ state: inout BenchmarkState) throws
     func tearDown()
 }
 
@@ -28,25 +28,55 @@ extension AnyBenchmark {
 internal class ClosureBenchmark: AnyBenchmark {
     let name: String
     let settings: [BenchmarkSetting]
-    let closure: () -> Void
+    let closure: () throws -> Void
 
-    init(_ name: String, settings: [BenchmarkSetting], closure: @escaping () -> Void) {
+    init(_ name: String, settings: [BenchmarkSetting], closure: @escaping () throws -> Void) {
         self.name = name
         self.settings = settings
         self.closure = closure
     }
 
-    func run() {
-        return self.closure()
+    func run(_ state: inout BenchmarkState) throws {
+        return try self.closure()
     }
 }
 
-public func benchmark(_ name: String, function: @escaping () -> Void) {
+internal class InoutClosureBenchmark: AnyBenchmark {
+    let name: String
+    let settings: [BenchmarkSetting]
+    let closure: (inout BenchmarkState) throws -> Void
+
+    init(
+        _ name: String, settings: [BenchmarkSetting],
+        closure: @escaping (inout BenchmarkState) throws -> Void
+    ) {
+        self.name = name
+        self.settings = settings
+        self.closure = closure
+    }
+
+    func run(_ state: inout BenchmarkState) throws {
+        return try self.closure(&state)
+    }
+}
+
+public func benchmark(_ name: String, function: @escaping () throws -> Void) {
     defaultBenchmarkSuite.benchmark(name, function: function)
 }
 
+public func benchmark(_ name: String, function: @escaping (inout BenchmarkState) throws -> Void) {
+    defaultBenchmarkSuite.benchmark(name, function: function)
+}
+
+public func benchmark(
+    _ name: String, settings: BenchmarkSetting..., function: @escaping () throws -> Void
+) {
+    defaultBenchmarkSuite.benchmark(name, settings: settings, function: function)
+}
+
 public func benchmark(
-    _ name: String, settings: BenchmarkSetting..., function: @escaping () -> Void
+    _ name: String, settings: BenchmarkSetting...,
+    function: @escaping (inout BenchmarkState) throws -> Void
 ) {
     defaultBenchmarkSuite.benchmark(name, settings: settings, function: function)
 }
diff --git a/Sources/Benchmark/BenchmarkRunner.swift b/Sources/Benchmark/BenchmarkRunner.swift
index 2f3f1aa..af35fab 100644
--- a/Sources/Benchmark/BenchmarkRunner.swift
+++ b/Sources/Benchmark/BenchmarkRunner.swift
@@ -51,24 +51,25 @@ public struct BenchmarkRunner {
         }
 
         reporter.report(running: benchmark.name, suite: suite.name)
-        var totalTime = BenchmarkClock()
-        totalTime.recordStart()
+        let totalStart = now()
 
         if let n = settings.warmupIterations {
-            let _ = doNIterations(n, benchmark: benchmark, suite: suite)
+            let _ = doNIterations(n, benchmark: benchmark, suite: suite, settings: settings)
         }
 
         var measurements: [Double] = []
         if let n = settings.iterations {
-            measurements = doNIterations(n, benchmark: benchmark, suite: suite)
+            measurements = doNIterations(n, benchmark: benchmark, suite: suite, settings: settings)
         } else {
             measurements = doAdaptiveIterations(
                 benchmark: benchmark, suite: suite, settings: settings)
         }
 
-        totalTime.recordEnd()
+        let totalEnd = now()
+        let totalElapsed = totalEnd - totalStart
+
         reporter.report(
-            finishedRunning: benchmark.name, suite: suite.name, nanosTaken: totalTime.elapsed)
+            finishedRunning: benchmark.name, suite: suite.name, nanosTaken: totalElapsed)
 
         let result = BenchmarkResult(
             benchmarkName: benchmark.name,
@@ -124,7 +125,7 @@ public struct BenchmarkRunner {
         var n: Int = 1
 
         while true {
-            measurements = doNIterations(n, benchmark: benchmark, suite: suite)
+            measurements = doNIterations(n, benchmark: benchmark, suite: suite, settings: settings)
             if n != 1 && hasCollectedEnoughData(measurements, settings: settings) { break }
             n = predictNumberOfIterationsNeeded(measurements, settings: settings)
             assert(n > measurements.count, "Number of iterations should increase with every retry.")
@@ -133,20 +134,16 @@ public struct BenchmarkRunner {
         return measurements
     }
 
-    func doNIterations(_ n: Int, benchmark: AnyBenchmark, suite: BenchmarkSuite) -> [Double] {
-        var clock = BenchmarkClock()
-        var measurements: [Double] = []
-        measurements.reserveCapacity(n)
-
-        for _ in 1...n {
-            benchmark.setUp()
-            clock.recordStart()
-            benchmark.run()
-            clock.recordEnd()
-            benchmark.tearDown()
-            measurements.append(Double(clock.elapsed))
+    func doNIterations(
+        _ n: Int, benchmark: AnyBenchmark, suite: BenchmarkSuite, settings: BenchmarkSettings
+    ) -> [Double] {
+        var state = BenchmarkState(iterations: n, settings: settings)
+        do {
+            try state.loop(benchmark)
+        } catch is BenchmarkTermination {
+        } catch {
+            fatalError("Unexpected error: \(error).")
         }
-
-        return measurements
+        return state.measurements
     }
 }
diff --git a/Sources/Benchmark/BenchmarkState.swift b/Sources/Benchmark/BenchmarkState.swift
new file mode 100644
index 0000000..3e18eff
--- /dev/null
+++ b/Sources/Benchmark/BenchmarkState.swift
@@ -0,0 +1,93 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// Benchmark state is used to collect the 
+/// benchmark measurements and view the settings it 
+/// was configured with.
+/// 
+/// Apart from the standard benchmark loop, you can also use
+/// it with customized benchmark measurement sections via either
+/// `start`/`end` or `measure` functions.
+public struct BenchmarkState {
+    let iterations: Int
+    var startTime: UInt64
+    var endTime: UInt64
+    var measurements: [Double]
+
+    /// Aggregated settings for the current benchmark run. 
+    public let settings: BenchmarkSettings
+
+    @inline(__always)
+    init(iterations: Int, settings: BenchmarkSettings) {
+        self.iterations = iterations
+        self.startTime = 0
+        self.endTime = 0
+        self.measurements = []
+        self.measurements.reserveCapacity(iterations)
+        self.settings = settings
+    }
+
+    /// Explicitly marks the start of the measurement section.
+    @inline(__always)
+    public mutating func start() {
+        self.endTime = 0
+        self.startTime = now()
+    }
+
+    /// Explicitly marks the end of the measurement section
+    /// and records the time since start of the benchmark.
+    @inline(__always)
+    public mutating func end() throws {
+        let value = now()
+        if self.endTime == 0 {
+            self.endTime = value
+            try record()
+        }
+    }
+
+    @inline(__always)
+    mutating func record() throws {
+        if measurements.count < iterations {
+            measurements.append(self.duration)
+        } else {
+            throw BenchmarkTermination()
+        }
+    }
+
+    @inline(__always)
+    var duration: Double {
+        return Double(self.endTime - self.startTime)
+    }
+
+    @inline(__always)
+    mutating func loop(_ benchmark: AnyBenchmark) throws {
+        while true {
+            benchmark.setUp()
+            start()
+            try benchmark.run(&self)
+            try end()
+            benchmark.tearDown()
+        }
+    }
+
+    /// Run the closure within within benchmark measurement section.
+    /// It may throw errors to propagate early termination to 
+    /// the outer benchmark loop.
+    @inline(__always)
+    public mutating func measure(f: () -> Void) throws {
+        start()
+        f()
+        try end()
+    }
+}
diff --git a/Sources/Benchmark/BenchmarkSuite.swift b/Sources/Benchmark/BenchmarkSuite.swift
index 7d211fd..dda724b 100644
--- a/Sources/Benchmark/BenchmarkSuite.swift
+++ b/Sources/Benchmark/BenchmarkSuite.swift
@@ -44,24 +44,47 @@ public class BenchmarkSuite {
         benchmarks.append(benchmark)
     }
 
-    public func benchmark(_ name: String, function f: @escaping () -> Void) {
+    public func benchmark(_ name: String, function f: @escaping () throws -> Void) {
         let benchmark = ClosureBenchmark(name, settings: [], closure: f)
         register(benchmark: benchmark)
     }
 
     public func benchmark(
-        _ name: String, settings: BenchmarkSetting..., function f: @escaping () -> Void
+        _ name: String, function f: @escaping (inout BenchmarkState) throws -> Void
+    ) {
+        let benchmark = InoutClosureBenchmark(name, settings: [], closure: f)
+        register(benchmark: benchmark)
+    }
+
+    public func benchmark(
+        _ name: String, settings: BenchmarkSetting..., function f: @escaping () throws -> Void
     ) {
         let benchmark = ClosureBenchmark(name, settings: settings, closure: f)
         register(benchmark: benchmark)
     }
 
     public func benchmark(
-        _ name: String, settings: [BenchmarkSetting], function f: @escaping () -> Void
+        _ name: String, settings: BenchmarkSetting...,
+        function f: @escaping (inout BenchmarkState) throws -> Void
+    ) {
+        let benchmark = InoutClosureBenchmark(name, settings: settings, closure: f)
+        register(benchmark: benchmark)
+    }
+
+    public func benchmark(
+        _ name: String, settings: [BenchmarkSetting], function f: @escaping () throws -> Void
     ) {
         let benchmark = ClosureBenchmark(name, settings: settings, closure: f)
         register(benchmark: benchmark)
     }
+
+    public func benchmark(
+        _ name: String, settings: [BenchmarkSetting],
+        function f: @escaping (inout BenchmarkState) throws -> Void
+    ) {
+        let benchmark = InoutClosureBenchmark(name, settings: settings, closure: f)
+        register(benchmark: benchmark)
+    }
 }
 
 var defaultBenchmarkSuite = BenchmarkSuite(name: "")
diff --git a/Sources/Benchmark/BenchmarkTermination.swift b/Sources/Benchmark/BenchmarkTermination.swift
new file mode 100644
index 0000000..6c5baca
--- /dev/null
+++ b/Sources/Benchmark/BenchmarkTermination.swift
@@ -0,0 +1,15 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+struct BenchmarkTermination: Error {}
diff --git a/Sources/Benchmark/BenchmarkClock.swift b/Sources/Benchmark/BenchmarkTime.swift
similarity index 58%
rename from Sources/Benchmark/BenchmarkClock.swift
rename to Sources/Benchmark/BenchmarkTime.swift
index 18f404c..f8ab693 100644
--- a/Sources/Benchmark/BenchmarkClock.swift
+++ b/Sources/Benchmark/BenchmarkTime.swift
@@ -14,26 +14,7 @@
 
 import Dispatch
 
-struct BenchmarkClock {
-    var start: UInt64 = 0
-    var end: UInt64 = 0
-
-    private static func now() -> UInt64 {
-        return DispatchTime.now().uptimeNanoseconds
-    }
-
-    @inlinable
-    mutating func recordStart() {
-        start = BenchmarkClock.now()
-    }
-
-    @inlinable
-    mutating func recordEnd() {
-        end = BenchmarkClock.now()
-    }
-
-    @inlinable
-    var elapsed: UInt64 {
-        return end - start
-    }
+@inline(__always)
+func now() -> UInt64 {
+    return DispatchTime.now().uptimeNanoseconds
 }
diff --git a/Sources/Benchmark/CMakeLists.txt b/Sources/Benchmark/CMakeLists.txt
index bf6818f..7a0f121 100644
--- a/Sources/Benchmark/CMakeLists.txt
+++ b/Sources/Benchmark/CMakeLists.txt
@@ -1,7 +1,6 @@
 add_library(Benchmark
   Benchmark.swift
   BenchmarkArguments.swift
-  BenchmarkClock.swift
   BenchmarkCommand.swift
   BenchmarkFilter.swift
   BenchmarkMain.swift
@@ -9,9 +8,12 @@ add_library(Benchmark
   BenchmarkResult.swift
   BenchmarkRunner.swift
   BenchmarkSetting.swift
+  BenchmarkState.swift
   BenchmarkSuite.swift
-  MakeTests.swift
-  Stats.swift)
+  BenchmarkTermination.swift
+  BenchmarkTime.swift
+  Stats.swift
+  TestIntegration.swift)
 # NOTE: workaround for CMake <3.17 which does not propagate the property
 set_target_properties(Benchmark PROPERTIES
   INTERFACE_INCLUDE_DIRECTORIES ${CMAKE_Swift_MODULE_DIRECTORY})


Test patch:
diff --git a/Sources/Benchmark/TestIntegration.swift b/Sources/Benchmark/TestIntegration.swift
index 30c3104..119b207 100644
--- a/Sources/Benchmark/TestIntegration.swift
+++ b/Sources/Benchmark/TestIntegration.swift
@@ -16,7 +16,13 @@
 public func runTests(suites: [BenchmarkSuite]) {
     for suite in suites {
         for benchmark in suite.benchmarks {
-            benchmark.run()
+            var state = BenchmarkState(iterations: 1, settings: BenchmarkSettings())
+            do {
+                try benchmark.run(&state)
+            } catch is BenchmarkTermination {
+            } catch {
+                fatalError("Unexpected error: \(error).")
+            }
         }
     }
 }
@@ -31,7 +37,13 @@ public func makeTests<T>(_ type: T.Type, suites: [BenchmarkSuite]) -> [(String,
             let name = "\(suite.name): \(benchmark.name)"
             let closure: (T) -> () -> Void = { _ in
                 return {
-                    benchmark.run()
+                    var state = BenchmarkState(iterations: 1, settings: BenchmarkSettings())
+                    do {
+                        try benchmark.run(&state)
+                    } catch is BenchmarkTermination {
+                    } catch {
+                        fatalError("Unexpected error: \(error).")
+                    }
                     return
                 }
             }
diff --git a/Tests/BenchmarkTests/BenchmarkRunnerTests.swift b/Tests/BenchmarkTests/BenchmarkRunnerTests.swift
index 81bc4d8..b98347d 100644
--- a/Tests/BenchmarkTests/BenchmarkRunnerTests.swift
+++ b/Tests/BenchmarkTests/BenchmarkRunnerTests.swift
@@ -39,11 +39,57 @@ final class BenchmarkRunnerTests: XCTestCase {
             runBenchmarks(settings: settings))
     }
 
+    func testCustomMeasurements() throws {
+        let suite = BenchmarkSuite(name: "Suite")
+
+        suite.benchmark("noop") {
+        }
+
+        suite.benchmark("start/end noop") { state in
+            state.start()
+            try state.end()
+        }
+
+        suite.benchmark("measure noop") { state in
+            try state.measure {
+            }
+        }
+
+        suite.benchmark("measure/while noop") { state in
+            while true {
+                try state.measure {
+                }
+            }
+        }
+
+        suite.benchmark("measure uneven iterations noop") { state in
+            for _ in 1...1337 {
+                try state.measure {
+                }
+            }
+        }
+
+        var runner = BenchmarkRunner(
+            suites: [suite],
+            settings: [Iterations(100000)],
+            reporter: BlackHoleReporter())
+        try runner.run()
+
+        let noopResults = runner.results[0].measurements
+        let customResults = runner.results.map { $0.measurements }.dropFirst()
+
+        XCTAssertEqual(noopResults.count, 100000)
+        for customResult in customResults {
+            XCTAssertEqual(customResult.count, 100000)
+        }
+    }
+
     static var allTests = [
         ("testFilterBenchmarksSuffix", testFilterBenchmarksSuffix),
         ("testFilterBenchmarksSuiteName", testFilterBenchmarksSuiteName),
         ("testFilterBenchmarksFullName", testFilterBenchmarksFullName),
         ("testAutomaticallyDetectIterations", testAutomaticallyDetectIterations),
+        ("testCustomMeasurements", testCustomMeasurements),
     ]
 }
 
diff --git a/Tests/BenchmarkTests/CustomBenchmarkTests.swift b/Tests/BenchmarkTests/CustomBenchmarkTests.swift
index 7cf974c..24f0ef1 100644
--- a/Tests/BenchmarkTests/CustomBenchmarkTests.swift
+++ b/Tests/BenchmarkTests/CustomBenchmarkTests.swift
@@ -50,7 +50,7 @@ fileprivate class MockBenchmark: AnyBenchmark {
         didSetUp = true
     }
 
-    func run() {
+    func run(_ state: inout BenchmarkState) {
         didRun = true
     }