From f41adf7c88d784faf78df5512bac2b53780cb56d Mon Sep 17 00:00:00 2001
From: Tom Andersen <tom-andersen@users.noreply.github.com>
Date: Fri, 19 Jan 2024 16:32:41 -0500
Subject: [PATCH] refactor: Optimize FieldMask instantiation (#1536)

* Optimize FieldMask instantiation

* Pretty
---
 .../com/google/cloud/firestore/FieldMask.java   | 16 ++++++++--------
 .../google/cloud/firestore/UpdateBuilder.java   | 17 ++++++++++++-----
 2 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java
index 0277029da..bec096abb 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/FieldMask.java
@@ -16,10 +16,8 @@
 
 package com.google.cloud.firestore;
 
-import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.Collection;
-import java.util.List;
+import java.util.Collections;
 import java.util.Map;
 import java.util.SortedSet;
 import java.util.TreeSet;
@@ -47,7 +45,7 @@ private FieldMask(SortedSet<FieldPath> fieldPaths) {
    */
   @Nonnull
   public static FieldMask of(String... fieldPaths) {
-    List<FieldPath> paths = new ArrayList<>();
+    TreeSet<FieldPath> paths = new TreeSet<>();
     for (String fieldPath : fieldPaths) {
       paths.add(FieldPath.fromDotSeparatedString(fieldPath));
     }
@@ -62,16 +60,18 @@ public static FieldMask of(String... fieldPaths) {
    */
   @Nonnull
   public static FieldMask of(FieldPath... fieldPaths) {
-    return new FieldMask(Arrays.asList(fieldPaths));
+    TreeSet<FieldPath> paths = new TreeSet<>();
+    Collections.addAll(paths, fieldPaths);
+    return new FieldMask(paths);
   }
 
   static FieldMask fromObject(Map<String, Object> values) {
-    List<FieldPath> fieldPaths = extractFromMap(values, FieldPath.empty());
+    TreeSet<FieldPath> fieldPaths = extractFromMap(values, FieldPath.empty());
     return new FieldMask(fieldPaths);
   }
 
-  private static List<FieldPath> extractFromMap(Map<String, Object> values, FieldPath path) {
-    List<FieldPath> fieldPaths = new ArrayList<>();
+  private static TreeSet<FieldPath> extractFromMap(Map<String, Object> values, FieldPath path) {
+    TreeSet<FieldPath> fieldPaths = new TreeSet<>();
 
     for (Map.Entry<String, Object> entry : values.entrySet()) {
       Object value = entry.getValue();
diff --git a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
index 87cf83307..4336e255c 100644
--- a/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
+++ b/google-cloud-firestore/src/main/java/com/google/cloud/firestore/UpdateBuilder.java
@@ -16,6 +16,9 @@
 
 package com.google.cloud.firestore;
 
+import static com.google.common.base.Predicates.not;
+import static java.util.stream.Collectors.toCollection;
+
 import com.google.api.core.ApiFuture;
 import com.google.api.core.ApiFutures;
 import com.google.api.core.InternalExtensionOnly;
@@ -272,9 +275,11 @@ private T performSet(
         DocumentTransform.fromFieldPathMap(documentReference, documentData);
 
     if (options.getFieldMask() != null) {
-      List<FieldPath> fieldMask = new ArrayList<>(options.getFieldMask());
-      fieldMask.removeAll(documentTransform.getFields());
-      documentMask = new FieldMask(fieldMask);
+      TreeSet<FieldPath> fieldPaths =
+          options.getFieldMask().stream()
+              .filter(not(documentTransform.getFields()::contains))
+              .collect(toCollection(TreeSet::new));
+      documentMask = new FieldMask(fieldPaths);
     } else if (options.isMerge()) {
       documentMask = FieldMask.fromObject(fields);
     }
@@ -544,10 +549,12 @@ public boolean allowTransform() {
                 return true;
               }
             });
-    List<FieldPath> fieldPaths = new ArrayList<>(fields.keySet());
     DocumentTransform documentTransform =
         DocumentTransform.fromFieldPathMap(documentReference, fields);
-    fieldPaths.removeAll(documentTransform.getFields());
+    TreeSet<FieldPath> fieldPaths =
+        fields.keySet().stream()
+            .filter(not(documentTransform.getFields()::contains))
+            .collect(toCollection(TreeSet::new));
     FieldMask fieldMask = new FieldMask(fieldPaths);
 
     Write.Builder write = documentSnapshot.toPb();