UnownedUserDataKey: Refactor lambda -> loops

Noticed this in binary size review as having a lot of nested classes.
Also removes most explicit null checks in favor of implicit ones.

Bug: None
Change-Id: I9777788512371b497c4236f16073a2b342e1acb5
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2419121Reviewed-by: Tommy Nyquist <nyquist@chromium.org>
Commit-Queue: Andrew Grieve <agrieve@chromium.org>
Cr-Commit-Position: refs/heads/master@{#811247}

base/android/java/src/org/chromium/base/UnownedUserDataKey.java

@@ -9,10 +9,11 @@ import androidx.annotation.Nullable;
 import androidx.annotation.VisibleForTesting;
 
 import java.lang.ref.WeakReference;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.HashSet;
+import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.atomic.AtomicInteger;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * UnownedUserDataKey is used in conjunction with a particular {@link UnownedUserData} as the key
@@ -60,14 +61,6 @@ import java.util.concurrent.atomic.AtomicReference;
  * @see UnownedUserData for the marker interface used for this type of data.
  */
 public final class UnownedUserDataKey<T extends UnownedUserData> {
-    private interface Predicate<T> {
-        boolean test(T t);
-    }
-
-    private interface Action<T> {
-        void execute(T t, WeakReference<T> weakT);
-    }
-
     @NonNull
     private final Class<T> mClazz;
     private final Set<WeakReference<UnownedUserDataHost>> mHostAttachments = new HashSet<>();
@@ -94,15 +87,14 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      * @param object The object to attach.
      */
     public final void attachToHost(@NonNull UnownedUserDataHost host, @NonNull T object) {
-        checkNotNull(host);
-        checkNotNull(object);
+        Objects.requireNonNull(object);
         // Setting a new value might lead to detachment of previously attached data, including
         // re-entry to this key, to happen before we update the {@link #mHostAttachments}.
         host.set(this, object);
 
-        AtomicReference<UnownedUserDataHost> retrievedHost = new AtomicReference<>();
-        executeHostAction(host::equals, (attachedHost, unused) -> retrievedHost.set(attachedHost));
-        if (retrievedHost.get() == null) mHostAttachments.add(new WeakReference<>(host));
+        if (!isAttachedToHost(host)) {
+            mHostAttachments.add(new WeakReference<>(host));
+        }
     }
 
     /**
@@ -116,12 +108,12 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      */
     @Nullable
     public final T retrieveDataFromHost(@NonNull UnownedUserDataHost host) {
-        checkNotNull(host);
-
-        AtomicReference<T> retrievedValue = new AtomicReference<>();
-        executeHostAction(
-                host::equals, (attachedHost, unused) -> retrievedValue.set(host.get(this)));
-        return retrievedValue.get();
+        for (UnownedUserDataHost attachedHost : getStrongRefs()) {
+            if (host.equals(attachedHost)) {
+                return host.get(this);
+            }
+        }
+        return null;
     }
 
     /**
@@ -131,11 +123,11 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      * @param host The host to detach from.
      */
     public final void detachFromHost(@NonNull UnownedUserDataHost host) {
-        checkNotNull(host);
-        executeHostAction(host::equals, (attachedHost, weakAttachedHost) -> {
-            attachedHost.remove(this);
-            mHostAttachments.remove(weakAttachedHost);
-        });
+        for (UnownedUserDataHost attachedHost : getStrongRefs()) {
+            if (host.equals(attachedHost)) {
+                removeHostAttachment(attachedHost);
+            }
+        }
     }
 
     /**
@@ -143,13 +135,11 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      * this key. It is OK to call this for already detached objects.
      */
     public final void detachFromAllHosts(@NonNull T object) {
-        checkNotNull(object);
-        executeHostAction((attachedHost)
-                                  -> object.equals(attachedHost.get(this)),
-                (attachedHost, weakAttachedHost) -> {
-                    attachedHost.remove(this);
-                    mHostAttachments.remove(weakAttachedHost);
-                });
+        for (UnownedUserDataHost attachedHost : getStrongRefs()) {
+            if (object.equals(attachedHost.get(this))) {
+                removeHostAttachment(attachedHost);
+            }
+        }
     }
 
     /**
@@ -159,7 +149,6 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      * @return true if currently attached, false otherwise.
      */
     public final boolean isAttachedToHost(@NonNull UnownedUserDataHost host) {
-        checkNotNull(host);
         T t = retrieveDataFromHost(host);
         return t != null;
     }
@@ -168,21 +157,35 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
      * @return Whether the {@link UnownedUserData} is currently attached to any hosts with this key.
      */
     public final boolean isAttachedToAnyHost(@NonNull T object) {
-        checkNotNull(object);
         return getHostAttachmentCount(object) > 0;
     }
 
     @VisibleForTesting(otherwise = VisibleForTesting.PRIVATE)
     /* package */ int getHostAttachmentCount(@NonNull T object) {
-        AtomicInteger count = new AtomicInteger(0);
-        executeHostAction((attachedHost)
-                                  -> object.equals(attachedHost.get(this)),
-                (attachedHost, unused) -> count.incrementAndGet());
-        return count.get();
+        int ret = 0;
+        for (UnownedUserDataHost attachedHost : getStrongRefs()) {
+            if (object.equals(attachedHost.get(this))) {
+                ret++;
+            }
+        }
+        return ret;
     }
 
-    private void executeHostAction(@NonNull Predicate<UnownedUserDataHost> predicate,
-            @NonNull Action<UnownedUserDataHost> action) {
+    private void removeHostAttachment(UnownedUserDataHost host) {
+        host.remove(this);
+        for (WeakReference<UnownedUserDataHost> hostWeakReference : mHostAttachments) {
+            if (host.equals(hostWeakReference.get())) {
+                // Modifying mHostAttachments while iterating over it is okay here because we
+                // break out of the loop right away.
+                mHostAttachments.remove(hostWeakReference);
+                return;
+            }
+        }
+    }
+
+    // TODO(https://crbug.com/1131047): Make a //base helper for this.
+    private Collection<UnownedUserDataHost> getStrongRefs() {
+        ArrayList<UnownedUserDataHost> ret = new ArrayList<>();
         Set<WeakReference<UnownedUserDataHost>> hosts = new HashSet<>(mHostAttachments);
         for (WeakReference<UnownedUserDataHost> hostWeakReference : hosts) {
             UnownedUserDataHost hostStrongReference = hostWeakReference.get();
@@ -191,23 +194,11 @@ public final class UnownedUserDataKey<T extends UnownedUserData> {
                 continue;
             }
             if (hostStrongReference.isDestroyed()) {
-                throw new IllegalStateException("Host should have been removed already.");
+                assert false : "Host should have been removed already.";
+                throw new IllegalStateException();
             }
-            if (!predicate.test(hostStrongReference)) continue;
-
-            action.execute(hostStrongReference, hostWeakReference);
-        }
-    }
-
-    private void checkNotNull(T unownedUserData) {
-        if (unownedUserData == null) {
-            throw new IllegalArgumentException("UnownedUserData can not be null.");
-        }
-    }
-
-    private static void checkNotNull(UnownedUserDataHost host) {
-        if (host == null) {
-            throw new IllegalArgumentException("UnownedUserDataHost can not be null.");
+            ret.add(hostStrongReference);
         }
+        return ret;
     }
 }

base/android/junit/src/org/chromium/base/UnownedUserDataKeyTest.java

@@ -843,14 +843,17 @@ public class UnownedUserDataKeyTest {
 
     @Test
     public void testNullKeyOrDataShouldBeDisallowed() {
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.attachToHost(null, null));
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.attachToHost(mHost1, null));
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.attachToHost(null, mFoo));
+        assertThrows(NullPointerException.class, () -> Foo.KEY.attachToHost(null, null));
+        assertThrows(NullPointerException.class, () -> Foo.KEY.attachToHost(mHost1, null));
+        assertThrows(NullPointerException.class, () -> Foo.KEY.attachToHost(null, mFoo));
 
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.retrieveDataFromHost(null));
+        // Need a non-empty registry to avoid no-op.
+        Foo.KEY.attachToHost(mHost1, mFoo);
+        assertThrows(NullPointerException.class, () -> Foo.KEY.retrieveDataFromHost(null));
 
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.detachFromHost(null));
-        assertThrows(IllegalArgumentException.class, () -> Foo.KEY.detachFromAllHosts(null));
+        assertThrows(NullPointerException.class, () -> Foo.KEY.detachFromHost(null));
+        assertThrows(NullPointerException.class, () -> Foo.KEY.detachFromAllHosts(null));
+        Foo.KEY.detachFromAllHosts(mFoo);
     }
 
     @Test