Implement runTransaction and batch support

Sources/FirebaseFirestore/Firestore+Swift.swift

@@ -34,6 +34,63 @@ extension Firestore {
   public func collection(_ collectionPath: String) -> CollectionReference {
     swift_firebase.swift_cxx_shims.firebase.firestore.firestore_collection(self, std.string(collectionPath))
   }
+
+  public func runTransaction(_ updateBlock: @escaping (Transaction, UnsafePointer<NSError?>?) -> Any?, completion: @escaping (Any?, Error?) -> Void) {
+    runTransaction(with: nil, block: updateBlock, completion: completion)
+  }
+
+  public func runTransaction(
+    with options: TransactionOptions?,
+    block updateBlock: @escaping (Transaction, UnsafePointer<NSError?>?) -> Any?,
+    completion: @escaping (Any?, Error?) -> Void
+  ) {
+    let context = TransactionContext(updateBlock: updateBlock)
+    let boxed = Unmanaged.passRetained(context as AnyObject)
+    let future = swift_firebase.swift_cxx_shims.firebase.firestore.firestore_run_transaction(
+      self, options ?? .init(), { transaction, pErrorMessage, pvUpdateBlock in
+        let context = Unmanaged<AnyObject>.fromOpaque(pvUpdateBlock!).takeUnretainedValue() as! TransactionContext
+
+        // Instead of trying to relay the generated `NSError` through firebase's `Error` type
+        // and error message string, just store the `NSError` on `context` and access it later.
+        // We then only need to tell firebase if the update block succeeded or failed and can
+        // just not bother setting `pErrorMessage`.
+
+        // It is expected to run `updateBlock` on whatever thread this happens to be. This is
+        // consistent with the behavior of the ObjC API as well.
+
+        // Since we could run `updateBlock` multiple times, we need to take care to reset any
+        // residue from previous runs. That means clearing out this error field.
+        context.error = nil
+
+        withUnsafePointer(to: context.error) { pError in
+          context.result = context.updateBlock(transaction!.pointee, pError)
+        }
+
+        return context.error != nil ? firebase.firestore.kErrorNone : firebase.firestore.kErrorCancelled
+      },
+      boxed.toOpaque()
+    )
+    future.setCompletion({
+      completion(context.result, context.error)
+      boxed.release()
+    })
+  }
+
+  public func batch() -> WriteBatch {
+    swift_firebase.swift_cxx_shims.firebase.firestore.firestore_batch(self)
+  }
+
+  private class TransactionContext {
+    typealias UpdateBlock = (Transaction, UnsafePointer<NSError?>?) -> Any?
+
+    let updateBlock: UpdateBlock
+    var result: Any?
+    var error: NSError?
+
+    init(updateBlock: @escaping UpdateBlock) {
+      self.updateBlock = updateBlock
+    }
+  }
 }
 
 // An extension that adds the encoder and decoder functions required

Sources/FirebaseFirestore/NSError+FirestoreError.swift

@@ -8,11 +8,18 @@ extension NSError {
   internal static func firestore(_ error: firebase.firestore.Error?, errorMessage: UnsafePointer<CChar>? = nil) -> NSError? {
     guard let actualError = error, actualError.rawValue != 0 else { return nil }
 
-    var userInfo = [String: Any]()
+    var errorMessageString: String?
     if let errorMessage {
-      userInfo[NSLocalizedDescriptionKey] = String(cString: errorMessage)
+      errorMessageString = .init(cString: errorMessage)
     }
+    return firestore(actualError, errorMessage: errorMessageString)
+  }
 
-    return NSError(domain: "firebase.firestore", code: Int(actualError.rawValue), userInfo: userInfo)
+  internal static func firestore(_ error: firebase.firestore.Error, errorMessage: String?) -> NSError {
+    var userInfo = [String: Any]()
+    if let errorMessage {
+      userInfo[NSLocalizedDescriptionKey] = errorMessage
+    }
+    return NSError(domain: "firebase.firestore", code: Int(error.rawValue), userInfo: userInfo)
   }
 }

Sources/FirebaseFirestore/Transaction+Swift.swift

@@ -0,0 +1,53 @@
+// SPDX-License-Identifier: BSD-3-Clause
+
+@_exported
+import firebase
+
+import CxxShim
+import Foundation
+
+public typealias Transaction = swift_firebase.swift_cxx_shims.firebase.firestore.TransactionWeakReference
+
+extension Transaction {
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference) -> Transaction {
+    setData(data, forDocument: document, merge: false)
+  }
+
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference, merge: Bool) -> Transaction {
+    assert(is_valid())
+    self.Set(document, FirestoreDataConverter.firestoreValue(document: data), merge ? .Merge() : .init())
+    return self
+  }
+
+  /* TODO: implement
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference, mergeFields: [Any]) -> Transaction {
+  }
+  */
+
+  public mutating func updateData(_ fields: [String : Any], forDocument document: DocumentReference) -> Transaction {
+    assert(is_valid())
+    self.Update(document, FirestoreDataConverter.firestoreValue(document: fields))
+    return self
+  }
+
+  public mutating func deleteDocument(_ document: DocumentReference) -> Transaction {
+    assert(is_valid())
+    Delete(document)
+    return self
+  }
+
+  public mutating func getDocument(_ document: DocumentReference) throws -> DocumentSnapshot {
+    assert(is_valid())
+
+    var error = firebase.firestore.kErrorNone
+    var errorMessage = std.string()
+
+    let snapshot = Get(document, &error, &errorMessage)
+
+    if error != firebase.firestore.kErrorNone {
+      throw NSError.firestore(error, errorMessage: String(errorMessage))
+    }
+
+    return snapshot
+  }
+}

Sources/FirebaseFirestore/TransactionOptions+Swift.swift

@@ -0,0 +1,17 @@
+// SPDX-License-Identifier: BSD-3-Clause
+
+@_exported
+import firebase
+
+public typealias TransactionOptions = firebase.firestore.TransactionOptions
+
+extension TransactionOptions {
+  public var maxAttempts: Int {
+    get {
+      Int(max_attempts())
+    }
+    set {
+      set_max_attempts(Int32(newValue))
+    }
+  }
+}

Sources/FirebaseFirestore/WriteBatch+Swift.swift

@@ -0,0 +1,67 @@
+// SPDX-License-Identifier: BSD-3-Clause
+
+@_exported
+import firebase
+@_spi(FirebaseInternal)
+import FirebaseCore
+
+import CxxShim
+import Foundation
+
+public typealias WriteBatch = firebase.firestore.WriteBatch
+
+extension WriteBatch {
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference) -> WriteBatch {
+    setData(data, forDocument: document, merge: false)
+  }
+
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference, merge: Bool) -> WriteBatch {
+    _ = swift_firebase.swift_cxx_shims.firebase.firestore.write_batch_set(
+      self, document, FirestoreDataConverter.firestoreValue(document: data), merge ? .Merge() : .init()
+    )
+    return self
+  }
+
+  /* TODO: implement
+  public mutating func setData(_ data: [String : Any], forDocument document: DocumentReference, mergeFields: [Any]) -> WriteBatch {
+  }
+  */
+
+  public mutating func updateData(_ fields: [String : Any], forDocument document: DocumentReference) -> WriteBatch {
+    _ = swift_firebase.swift_cxx_shims.firebase.firestore.write_batch_update(
+      self, document, FirestoreDataConverter.firestoreValue(document: fields)
+    )
+    return self
+  }
+
+  public mutating func deleteDocument(_ document: DocumentReference) -> WriteBatch {
+    _ = swift_firebase.swift_cxx_shims.firebase.firestore.write_batch_delete(
+      self, document
+    )
+    return self
+  }
+
+  public mutating func commit(completion: @escaping ((Error?) -> Void) = { _ in }) {
+    let future = swift_firebase.swift_cxx_shims.firebase.firestore.write_batch_commit(self)
+    future.setCompletion({
+      let (_, error) = future.resultAndError
+      DispatchQueue.main.async {
+        completion(error)
+      }
+    })
+  }
+
+  public mutating func commit() async throws {
+    try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, any Error>) in
+      let future = swift_firebase.swift_cxx_shims.firebase.firestore.write_batch_commit(self)
+      future.setCompletion({
+        let (_, error) = future.resultAndError
+        if let error {
+          continuation.resume(throwing: error)
+        } else {
+          continuation.resume()
+        }
+      })
+    }
+  }
+}

Sources/firebase/include/FirebaseCore.hh

@@ -21,6 +21,12 @@ class __attribute__((swift_attr("conforms_to:FirebaseCore.FutureProtocol")))
 
   Future(const ::firebase::Future<R>& rhs) : ::firebase::Future<R>(rhs) {}
 
+  // Allow explicit conversion from `Future<void>` in support of `VoidFuture`.
+  static Future From(const ::firebase::Future<void>& other) {
+    static_assert(sizeof(::firebase::Future<void>) == sizeof(::firebase::Future<R>));
+    return Future(*reinterpret_cast<const ::firebase::Future<R>*>(&other));
+  }
+
   void OnCompletion(
       _Nonnull FutureCompletionType completion,
       _Nullable void* user_data) const {
@@ -31,6 +37,10 @@ class __attribute__((swift_attr("conforms_to:FirebaseCore.FutureProtocol")))
   }
 };
 
+// As a workaround, use `int` here instead of `void` for futures with no
+// result. Swift is not able to handle a `ResultType` of `void`.
+typedef Future<int> VoidFuture;
+
 } // namespace swift_firebase::swift_cxx_shims::firebase
 
 #endif

Sources/firebase/include/FirebaseFirestore.hh

@@ -3,14 +3,19 @@
 #ifndef firebase_include_FirebaseFirestore_hh
 #define firebase_include_FirebaseFirestore_hh
 
+#include <string>
+#include <utility>
+
 #include <firebase/firestore.h>
 
 #include "FirebaseCore.hh"
+#include "TransactionWeakReference.hh"
 
 // Functions defined in this namespace are used to get around the lack of
 // virtual function support currently in Swift. As that support changes
 // these functions will go away whenever possible.
 namespace swift_firebase::swift_cxx_shims::firebase::firestore {
+
 inline ::firebase::firestore::Settings
 firestore_settings(::firebase::firestore::Firestore *firestore) {
   return firestore->settings();
@@ -28,6 +33,34 @@ firestore_collection(::firebase::firestore::Firestore *firestore,
   return firestore->Collection(collection_path);
 }
 
+typedef ::firebase::firestore::Error (*FirebaseRunTransactionUpdateCallback)(
+    TransactionWeakReference *transaction,
+    std::string& error_message,
+    void *user_data);
+inline VoidFuture
+firestore_run_transaction(
+    ::firebase::firestore::Firestore *firestore,
+    ::firebase::firestore::TransactionOptions options,
+    FirebaseRunTransactionUpdateCallback update_callback,
+    void *user_data) {
+  return VoidFuture::From(
+      firestore->RunTransaction(options, [update_callback, user_data](
+          ::firebase::firestore::Transaction& transaction,
+          std::string& error_message
+      ) -> ::firebase::firestore::Error {
+        TransactionWeakReference transaction_ref(&transaction);
+        ::firebase::firestore::Error error =
+            update_callback(&transaction_ref, error_message, user_data);
+        transaction_ref.reset();
+        return error;
+      }));
+}
+
+inline ::firebase::firestore::WriteBatch
+firestore_batch(::firebase::firestore::Firestore *firestore) {
+  return firestore->batch();
+}
+
 // MARK: - DocumentReference
 
 inline ::firebase::firestore::Firestore *
@@ -279,6 +312,46 @@ query_snapshot_size(const ::firebase::firestore::QuerySnapshot& snapshot) {
   return snapshot.size();
 }
 
+// MARK: WriteBatch
+
+inline ::firebase::firestore::WriteBatch&
+write_batch_set(
+    ::firebase::firestore::WriteBatch write_batch,
+    const ::firebase::firestore::DocumentReference& document,
+    const ::firebase::firestore::MapFieldValue& data,
+    const ::firebase::firestore::SetOptions& options =
+        ::firebase::firestore::SetOptions()) {
+  return write_batch.Set(document, data, options);
+}
+
+inline ::firebase::firestore::WriteBatch&
+write_batch_update(
+    ::firebase::firestore::WriteBatch write_batch,
+    const ::firebase::firestore::DocumentReference& document,
+    const ::firebase::firestore::MapFieldValue& data) {
+  return write_batch.Update(document, data);
+}
+
+inline ::firebase::firestore::WriteBatch&
+write_batch_update(
+    ::firebase::firestore::WriteBatch write_batch,
+    const ::firebase::firestore::DocumentReference& document,
+    const ::firebase::firestore::MapFieldPathValue& data) {
+  return write_batch.Update(document, data);
+}
+
+inline ::firebase::firestore::WriteBatch&
+write_batch_delete(
+    ::firebase::firestore::WriteBatch write_batch,
+    const ::firebase::firestore::DocumentReference& document) {
+  return write_batch.Delete(document);
+}
+
+inline VoidFuture
+write_batch_commit(::firebase::firestore::WriteBatch write_batch) {
+  return VoidFuture::From(write_batch.Commit());
+}
+
 } // namespace swift_firebase::swift_cxx_shims::firebase::firestore
 
 #endif