Added convenience overloads for Connection.createAggregation.

kasei · kasei · commit 1e6bfbf2c046 · 2019-02-21T21:58:34.000-08:00
diff --git a/Sources/SQLite/Core/Connection.swift b/Sources/SQLite/Core/Connection.swift
@@ -626,7 +626,15 @@ public final class Connection {
     ///   - state: A block of code to run to produce a fresh state variable for
     ///     each aggregation group. The block should return an
     ///     UnsafeMutablePointer to the fresh state variable.
-    public func createAggregation<T>(_ aggregate: String, argumentCount: UInt? = nil, deterministic: Bool = false, step: @escaping ([Binding?], UnsafeMutablePointer<T>) -> (), final: @escaping (UnsafeMutablePointer<T>) -> Binding?, state: @escaping () -> UnsafeMutablePointer<T>) {
+    public func createAggregation<T>(
+        _ aggregate: String,
+        argumentCount: UInt? = nil,
+        deterministic: Bool = false,
+        step: @escaping ([Binding?], UnsafeMutablePointer<T>) -> (),
+        final: @escaping (UnsafeMutablePointer<T>) -> Binding?,
+        state: @escaping () -> UnsafeMutablePointer<T>) {
+        
+        
         let argc = argumentCount.map { Int($0) } ?? -1
         let box : Aggregate = { (stepFlag: Int, context: OpaquePointer?, argc: Int32, argv: UnsafeMutablePointer<OpaquePointer?>?) in
             let ptr = sqlite3_aggregate_context(context, 64)! // needs to be at least as large as uintptr_t; better way to do this?
@@ -690,17 +698,17 @@ public final class Connection {
             { context, argc, value in
                 let function = unsafeBitCast(sqlite3_user_data(context), to: Aggregate.self)
                 function(1, context, argc, value)
-            },
+        },
             { context in
                 let function = unsafeBitCast(sqlite3_user_data(context), to: Aggregate.self)
                 function(0, context, 0, nil)
-            },
+        },
             nil
         )
         if aggregations[aggregate] == nil { self.aggregations[aggregate] = [:] }
         aggregations[aggregate]?[argc] = box
     }
-    
+
     fileprivate typealias Aggregate = @convention(block) (Int, OpaquePointer?, Int32, UnsafeMutablePointer<OpaquePointer?>?) -> Void
     fileprivate typealias Function = @convention(block) (OpaquePointer?, Int32, UnsafeMutablePointer<OpaquePointer?>?) -> Void
     fileprivate var functions = [String: [Int: Function]]()
diff --git a/Sources/SQLite/Typed/CustomFunctions.swift b/Sources/SQLite/Typed/CustomFunctions.swift
@@ -133,4 +133,69 @@ public extension Connection {
         }
     }
 
+    // MARK: -
+    
+    public func createAggregation<T: AnyObject>(
+        _ aggregate: String,
+        argumentCount: UInt? = nil,
+        deterministic: Bool = false,
+        initialValue: T,
+        reduce: @escaping (T, [Binding?]) -> T,
+        result: @escaping (T) -> Binding?
+        ) {
+        
+        let step: ([Binding?], UnsafeMutablePointer<UnsafeMutableRawPointer>) -> () = { (bindings, ptr) in
+            let p = ptr.pointee.assumingMemoryBound(to: T.self)
+            let current = Unmanaged<T>.fromOpaque(p).takeRetainedValue()
+            let next = reduce(current, bindings)
+            ptr.pointee = Unmanaged.passRetained(next).toOpaque()
+        }
+        
+        let final: (UnsafeMutablePointer<UnsafeMutableRawPointer>) -> Binding? = { (ptr) in
+            let p = ptr.pointee.assumingMemoryBound(to: T.self)
+            let obj = Unmanaged<T>.fromOpaque(p).takeRetainedValue()
+            let value = result(obj)
+            ptr.deallocate()
+            return value
+        }
+        
+        let state: () -> UnsafeMutablePointer<UnsafeMutableRawPointer> = {
+            let p = UnsafeMutablePointer<UnsafeMutableRawPointer>.allocate(capacity: 1)
+            p.pointee = Unmanaged.passRetained(initialValue).toOpaque()
+            return p
+        }
+        
+        createAggregation(aggregate, step: step, final: final, state: state)
+    }
+    
+    public func createAggregation<T>(
+        _ aggregate: String,
+        argumentCount: UInt? = nil,
+        deterministic: Bool = false,
+        initialValue: T,
+        reduce: @escaping (T, [Binding?]) -> T,
+        result: @escaping (T) -> Binding?
+        ) {
+        
+        let step: ([Binding?], UnsafeMutablePointer<T>) -> () = { (bindings, p) in
+            let current = p.pointee
+            let next = reduce(current, bindings)
+            p.pointee = next
+        }
+        
+        let final: (UnsafeMutablePointer<T>) -> Binding? = { (p) in
+            let v = result(p.pointee)
+            p.deallocate()
+            return v
+        }
+        
+        let state: () -> UnsafeMutablePointer<T> = {
+            let p = UnsafeMutablePointer<T>.allocate(capacity: 1)
+            p.pointee = initialValue
+            return p
+        }
+        
+        createAggregation(aggregate, step: step, final: final, state: state)
+    }
+
 }
diff --git a/Tests/SQLiteTests/CustomAggregationTests.swift b/Tests/SQLiteTests/CustomAggregationTests.swift
@@ -22,33 +22,33 @@ class CustomAggregationTests : SQLiteTestCase {
         try! InsertUser("Eve", age: 28, admin: false)
     }
     
-    func testCustomSum() {
+    func testUnsafeCustomSum() {
         let step = { (bindings: [Binding?], state: UnsafeMutablePointer<Int64>) in
             if let v = bindings[0] as? Int64 {
                 state.pointee += v
             }
         }
-
+        
         let final = { (state: UnsafeMutablePointer<Int64>) -> Binding? in
             let v = state.pointee
             let p = UnsafeMutableBufferPointer(start: state, count: 1)
             p.deallocate()
             return v
         }
-        let _ = db.createAggregation("mySUM", step: step, final: final) {
+        let _ = db.createAggregation("mySUM1", step: step, final: final) {
             let v = UnsafeMutableBufferPointer<Int64>.allocate(capacity: 1)
             v[0] = 0
             return v.baseAddress!
         }
-        let result = try! db.prepare("SELECT mySUM(age) AS s FROM users")
+        let result = try! db.prepare("SELECT mySUM1(age) AS s FROM users")
         let i = result.columnNames.index(of: "s")!
         for row in result {
             let value = row[i] as? Int64
             XCTAssertEqual(83, value)
         }
     }
     
-    func testCustomSumGrouping() {
+    func testUnsafeCustomSumGrouping() {
         let step = { (bindings: [Binding?], state: UnsafeMutablePointer<Int64>) in
             if let v = bindings[0] as? Int64 {
                 state.pointee += v
@@ -60,14 +60,71 @@ class CustomAggregationTests : SQLiteTestCase {
             p.deallocate()
             return v
         }
-        let _ = db.createAggregation("mySUM", step: step, final: final) {
+        let _ = db.createAggregation("mySUM2", step: step, final: final) {
             let v = UnsafeMutableBufferPointer<Int64>.allocate(capacity: 1)
             v[0] = 0
             return v.baseAddress!
         }
-        let result = try! db.prepare("SELECT mySUM(age) AS s FROM users GROUP BY admin ORDER BY s")
+        let result = try! db.prepare("SELECT mySUM2(age) AS s FROM users GROUP BY admin ORDER BY s")
         let i = result.columnNames.index(of: "s")!
         let values = result.compactMap { $0[i] as? Int64 }
         XCTAssertTrue(values.elementsEqual([28, 55]))
     }
+    
+    func testCustomSum() {
+        let reduce : (Int64, [Binding?]) -> Int64 = { (last, bindings) in
+            let v = (bindings[0] as? Int64) ?? 0
+            return last + v
+        }
+        let _ = db.createAggregation("myReduceSUM1", initialValue: Int64(2000), reduce: reduce, result: { $0 })
+        let result = try! db.prepare("SELECT myReduceSUM1(age) AS s FROM users")
+        let i = result.columnNames.index(of: "s")!
+        for row in result {
+            let value = row[i] as? Int64
+            XCTAssertEqual(2083, value)
+        }
+    }
+
+    func testCustomSumGrouping() {
+        let reduce : (Int64, [Binding?]) -> Int64 = { (last, bindings) in
+            let v = (bindings[0] as? Int64) ?? 0
+            return last + v
+        }
+        let _ = db.createAggregation("myReduceSUM2", initialValue: Int64(3000), reduce: reduce, result: { $0 })
+        let result = try! db.prepare("SELECT myReduceSUM2(age) AS s FROM users GROUP BY admin ORDER BY s")
+        let i = result.columnNames.index(of: "s")!
+        let values = result.compactMap { $0[i] as? Int64 }
+        XCTAssertTrue(values.elementsEqual([3028, 3055]))
+    }
+    
+    func testCustomObjectSum() {
+        {
+            let initial = TestObject(value: 1000)
+            let reduce : (TestObject, [Binding?]) -> TestObject = { (last, bindings) in
+                let v = (bindings[0] as? Int64) ?? 0
+                return TestObject(value: last.value + v)
+            }
+            let _ = db.createAggregation("myReduceSUMX", initialValue: initial, reduce: reduce, result: { $0.value })
+            // end this scope to ensure that the initial value is retained
+            // by the createAggregation call.
+        }()
+        let result = try! db.prepare("SELECT myReduceSUMX(age) AS s FROM users")
+        let i = result.columnNames.index(of: "s")!
+        for row in result {
+            let value = row[i] as? Int64
+            XCTAssertEqual(1083, value)
+        }
+    }
+}
+
+/// This class is used to test that aggregation state variables
+/// can be reference types and are properly memory managed when
+/// crossing the Swift<->C boundary multiple times.
+class TestObject {
+    var value: Int64
+    init(value: Int64) {
+        self.value = value
+    }
+    deinit {
+    }
 }
