Aggregate functions support (#262)

While this landed in StructuredQueries proper, we need to expose the
functionality to GRDB.

Author: stephencelis

Package.resolved

@@ -36,7 +36,7 @@ let package = Package(
     .package(url: "https://github.com/pointfreeco/swift-snapshot-testing", from: "1.18.4"),
     .package(
       url: "https://github.com/pointfreeco/swift-structured-queries",
-      from: "0.22.3",
+      from: "0.24.0",
       traits: [
         .trait(name: "StructuredQueriesTagged", condition: .when(traits: ["SQLiteDataTagged"]))
       ]

Package.swift

@@ -28,7 +28,7 @@ let package = Package(
     .package(url: "https://github.com/pointfreeco/swift-dependencies", from: "1.9.0"),
     .package(url: "https://github.com/pointfreeco/swift-sharing", from: "2.3.0"),
     .package(url: "https://github.com/pointfreeco/swift-snapshot-testing", from: "1.18.4"),
-    .package(url: "https://github.com/pointfreeco/swift-structured-queries", from: "0.22.3"),
+    .package(url: "https://github.com/pointfreeco/swift-structured-queries", from: "0.24.0"),
     .package(url: "https://github.com/pointfreeco/xctest-dynamic-overlay", from: "1.5.0"),
   ],
   targets: [

[email protected]

@@ -2,20 +2,20 @@ import Foundation
 import GRDBSQLite
 
 extension Database {
-  /// Adds a user-defined `@DatabaseFunction` to a connection.
+  /// Adds a user-defined scalar `@DatabaseFunction` to a connection.
   ///
-  /// - Parameter function: A database function to add.
+  /// - Parameter function: A scalar database function to add.
   public func add(function: some ScalarDatabaseFunction) {
     sqlite3_create_function_v2(
       sqliteConnection,
       function.name,
       function.argumentCount,
       function.textEncoding,
-      Unmanaged.passRetained(ScalarDatabaseFunctionBox(function)).toOpaque(),
+      Unmanaged.passRetained(ScalarDatabaseFunctionDefinition(function)).toOpaque(),
       { context, argumentCount, arguments in
         do {
           var decoder = SQLiteFunctionDecoder(argumentCount: argumentCount, arguments: arguments)
-          try Unmanaged<ScalarDatabaseFunctionBox>
+          try Unmanaged<ScalarDatabaseFunctionDefinition>
             .fromOpaque(sqlite3_user_data(context))
             .takeUnretainedValue()
             .function
@@ -27,17 +27,56 @@ extension Database {
       },
       nil,
       nil,
-      { box in
-        guard let box else { return }
-        Unmanaged<ScalarDatabaseFunctionBox>.fromOpaque(box).release()
+      { context in
+        guard let context else { return }
+        Unmanaged<ScalarDatabaseFunctionDefinition>.fromOpaque(context).release()
+      }
+    )
+  }
+
+  /// Adds a user-defined aggregate `@DatabaseFunction` to a connection.
+  ///
+  /// - Parameter function: An aggregate database function to add.
+  public func add(function: some AggregateDatabaseFunction) {
+    let body = Unmanaged.passRetained(AggregateDatabaseFunctionDefinition(function)).toOpaque()
+    sqlite3_create_function_v2(
+      sqliteConnection,
+      function.name,
+      function.argumentCount,
+      function.textEncoding,
+      body,
+      nil,
+      { context, argumentCount, arguments in
+        var decoder = SQLiteFunctionDecoder(argumentCount: argumentCount, arguments: arguments)
+        let function = AggregateDatabaseFunctionContext[context].takeUnretainedValue()
+        do {
+          try function.iterator.step(&decoder)
+        } catch {
+          sqlite3_result_error(context, error.localizedDescription, -1)
+        }
+      },
+      { context in
+        let unmanagedFunction = AggregateDatabaseFunctionContext[context]
+        let function = unmanagedFunction.takeUnretainedValue()
+        unmanagedFunction.release()
+        function.iterator.finish()
+        do {
+          try function.iterator.result.result(db: context)
+        } catch {
+          sqlite3_result_error(context, error.localizedDescription, -1)
+        }
+      },
+      { context in
+        guard let context else { return }
+        Unmanaged<AggregateDatabaseFunctionContext>.fromOpaque(context).release()
       }
     )
   }
 
   /// Deletes a user-defined `@DatabaseFunction` from a connection.
   ///
   /// - Parameter function: A database function to delete.
-  public func remove(function: some ScalarDatabaseFunction) {
+  public func remove(function: some DatabaseFunction) {
     sqlite3_create_function_v2(
       sqliteConnection,
       function.name,
@@ -52,7 +91,7 @@ extension Database {
   }
 }
 
-extension ScalarDatabaseFunction {
+extension DatabaseFunction {
   fileprivate var argumentCount: Int32 {
     Int32(argumentCount ?? -1)
   }
@@ -62,41 +101,132 @@ extension ScalarDatabaseFunction {
   }
 }
 
-private final class ScalarDatabaseFunctionBox {
+private final class ScalarDatabaseFunctionDefinition {
   let function: any ScalarDatabaseFunction
   init(_ function: some ScalarDatabaseFunction) {
     self.function = function
   }
 }
 
-extension [QueryBinding] {
-  fileprivate init(argumentCount: Int32, arguments: UnsafeMutablePointer<OpaquePointer?>?) {
-    self = (0..<argumentCount).map { offset in
-      let value = arguments?[Int(offset)]
-      switch sqlite3_value_type(value) {
-      case SQLITE_BLOB:
-        if let blob = sqlite3_value_blob(value) {
-          let count = Int(sqlite3_value_bytes(value))
-          let buffer = UnsafeRawBufferPointer(start: blob, count: count)
-          return .blob([UInt8](buffer))
-        } else {
-          return .blob([])
+private final class AggregateDatabaseFunctionDefinition {
+  let function: any AggregateDatabaseFunction
+  init(_ function: some AggregateDatabaseFunction) {
+    self.function = function
+  }
+}
+
+private final class AggregateDatabaseFunctionContext {
+  static subscript(context: OpaquePointer?) -> Unmanaged<AggregateDatabaseFunctionContext> {
+    let size = MemoryLayout<Unmanaged<AggregateDatabaseFunctionContext>>.size
+    let pointer = sqlite3_aggregate_context(context, Int32(size))!
+    if pointer.load(as: Int.self) == 0 {
+      let definition = Unmanaged<AggregateDatabaseFunctionDefinition>
+        .fromOpaque(sqlite3_user_data(context))
+        .takeUnretainedValue()
+      let context = AggregateDatabaseFunctionContext(definition.function)
+      let unmanagedContext = Unmanaged.passRetained(context)
+      pointer
+        .assumingMemoryBound(to: Unmanaged<AggregateDatabaseFunctionContext>.self)
+        .pointee = unmanagedContext
+      return unmanagedContext
+    } else {
+      return
+        pointer
+        .assumingMemoryBound(to: Unmanaged<AggregateDatabaseFunctionContext>.self)
+        .pointee
+    }
+  }
+  let iterator: any AggregateDatabaseFunctionIteratorProtocol
+  init(_ body: some AggregateDatabaseFunction) {
+    self.iterator = AggregateDatabaseFunctionIterator(body)
+  }
+}
+
+private protocol AggregateDatabaseFunctionIteratorProtocol<Body> {
+  associatedtype Body: AggregateDatabaseFunction
+
+  var body: Body { get }
+  var stream: Stream<Body.Element> { get }
+  func start()
+  func step(_ decoder: inout some QueryDecoder) throws
+  func finish()
+  var result: QueryBinding { get throws }
+}
+
+private final class AggregateDatabaseFunctionIterator<
+  Body: AggregateDatabaseFunction
+>: AggregateDatabaseFunctionIteratorProtocol {
+  let body: Body
+  let stream = Stream<Body.Element>()
+  let queue: DispatchQueue
+  var _result: QueryBinding?
+  init(_ body: Body) {
+    self.body = body
+    self.queue = DispatchQueue(
+      label: "co.pointfree.StructuredQueriesSQLite.AggregateDatabaseFunction.\(body.name)"
+    )
+    nonisolated(unsafe) let iterator: any AggregateDatabaseFunctionIteratorProtocol = self
+    queue.async {
+      iterator.start()
+    }
+  }
+  func start() {
+    do {
+      _result = try body.invoke(stream)
+    } catch {
+      _result = .invalid(error)
+    }
+  }
+  func step(_ decoder: inout some QueryDecoder) throws {
+    try stream.send(body.step(&decoder))
+  }
+  func finish() {
+    stream.finish()
+  }
+  var result: QueryBinding {
+    get throws {
+      while true {
+        if let result = queue.sync(execute: { _result }) {
+          return result
         }
-      case SQLITE_FLOAT:
-        return .double(sqlite3_value_double(value))
-      case SQLITE_INTEGER:
-        return .int(sqlite3_value_int64(value))
-      case SQLITE_NULL:
-        return .null
-      case SQLITE_TEXT:
-        return .text(String(cString: UnsafePointer(sqlite3_value_text(value))))
-      default:
-        return .invalid(UnknownType())
       }
     }
   }
+}
+
+private final class Stream<Element>: Sequence {
+  let condition = NSCondition()
+  private var buffer: [Element] = []
+  private var isFinished = false
 
-  private struct UnknownType: Error {}
+  func send(_ element: Element) {
+    condition.withLock {
+      buffer.append(element)
+      condition.signal()
+    }
+  }
+
+  func finish() {
+    condition.withLock {
+      isFinished = true
+      condition.broadcast()
+    }
+  }
+
+  func makeIterator() -> Iterator { Iterator(base: self) }
+
+  struct Iterator: IteratorProtocol {
+    fileprivate let base: Stream
+    mutating func next() -> Element? {
+      base.condition.withLock {
+        while base.buffer.isEmpty && !base.isFinished {
+          base.condition.wait()
+        }
+        guard !base.buffer.isEmpty else { return nil }
+        return base.buffer.removeFirst()
+      }
+    }
+  }
 }
 
 extension QueryBinding {

Sources/SQLiteData/StructuredQueries+GRDB/CustomFunctions.swift

@@ -0,0 +1,70 @@
+import DependenciesTestSupport
+import Foundation
+import SQLiteData
+import SQLiteDataTestSupport
+import Testing
+
+struct DatabaseFunctionTests {
+  @DatabaseFunction
+  func exclaim(_ text: String) -> String {
+    text + "!"
+  }
+  @Test func scalarFunction() async throws {
+    var configuration = Configuration()
+    configuration.prepareDatabase { db in
+      db.add(function: $exclaim)
+    }
+    let database = try DatabaseQueue(configuration: configuration)
+    assertQuery(Values($exclaim("Blob")), database: database) {
+      """
+      ┌─────────┐
+      │ "Blob!" │
+      └─────────┘
+      """
+    }
+  }
+
+  @Test(.dependency(\.defaultDatabase, try .database())) func aggregateFunction() async throws {
+    assertQuery(Record.select { $sum($0.id) }) {
+      """
+      ┌───┐
+      │ 6 │
+      └───┘
+      """
+    }
+  }
+}
+
+@Table
+private struct Record: Equatable {
+  let id: Int
+}
+
+@DatabaseFunction
+func sum(_ xs: some Sequence<Int>) -> Int {
+  xs.reduce(0, +)
+}
+
+extension DatabaseWriter where Self == DatabaseQueue {
+  fileprivate static func database() throws -> DatabaseQueue {
+    var configuration = Configuration()
+    configuration.prepareDatabase { db in
+      db.add(function: $sum)
+    }
+    let database = try DatabaseQueue(configuration: configuration)
+    try database.write { db in
+      try #sql(
+        """
+        CREATE TABLE "records" (
+          "id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT
+        )
+        """
+      )
+      .execute(db)
+      for _ in 1...3 {
+        _ = try Record.insert { Record.Draft() }.execute(db)
+      }
+    }
+    return database
+  }
+}