diff --git a/core/java/android/service/autofill/FillContext.java b/core/java/android/service/autofill/FillContext.java
index d6616bf..251d346 100644
--- a/core/java/android/service/autofill/FillContext.java
+++ b/core/java/android/service/autofill/FillContext.java
@@ -19,11 +19,19 @@
 import static android.view.autofill.Helper.sDebug;
 
 import android.annotation.NonNull;
+import android.annotation.Nullable;
 import android.app.assist.AssistStructure;
+import android.app.assist.AssistStructure.ViewNode;
 import android.os.Bundle;
 import android.os.CancellationSignal;
 import android.os.Parcel;
 import android.os.Parcelable;
+import android.util.ArrayMap;
+import android.util.SparseIntArray;
+import android.view.autofill.AutofillId;
+
+import java.util.ArrayList;
+import java.util.LinkedList;
 
 /**
  * This class represents a context for each fill request made via {@link
@@ -43,6 +51,13 @@
     private final int mRequestId;
     private final @NonNull AssistStructure mStructure;
 
+    /**
+     * Lookup table AutofillId->ViewNode to speed up {@link #findViewNodesByAutofillIds}
+     * This is purely a cache and can be deleted at any time
+     */
+    @Nullable private ArrayMap<AutofillId, AssistStructure.ViewNode> mViewNodeLookupTable;
+
+
     /** @hide */
     public FillContext(int requestId, @NonNull AssistStructure structure) {
         mRequestId = requestId;
@@ -90,6 +105,79 @@
         parcel.writeParcelable(mStructure, flags);
     }
 
+    /**
+     * Finds {@link ViewNode}s that have the requested ids.
+     *
+     * @param ids The ids of the node to find
+     *
+     * @return The nodes indexed in the same way as the ids
+     *
+     * @hide
+     */
+    @NonNull public ViewNode[] findViewNodesByAutofillIds(@NonNull AutofillId[] ids) {
+        final LinkedList<ViewNode> nodesToProcess = new LinkedList<>();
+        final ViewNode[] foundNodes = new AssistStructure.ViewNode[ids.length];
+
+        // Indexes of foundNodes that are not found yet
+        final SparseIntArray missingNodeIndexes = new SparseIntArray(ids.length);
+
+        for (int i = 0; i < ids.length; i++) {
+            if (mViewNodeLookupTable != null) {
+                int lookupTableIndex = mViewNodeLookupTable.indexOfKey(ids[i]);
+
+                if (lookupTableIndex >= 0) {
+                    foundNodes[i] = mViewNodeLookupTable.valueAt(lookupTableIndex);
+                } else {
+                    missingNodeIndexes.put(i, /* ignored */ 0);
+                }
+            } else {
+                missingNodeIndexes.put(i, /* ignored */ 0);
+            }
+        }
+
+        final int numWindowNodes = mStructure.getWindowNodeCount();
+        for (int i = 0; i < numWindowNodes; i++) {
+            nodesToProcess.add(mStructure.getWindowNodeAt(i).getRootViewNode());
+        }
+
+        while (missingNodeIndexes.size() > 0 && !nodesToProcess.isEmpty()) {
+            final ViewNode node = nodesToProcess.removeFirst();
+
+            for (int i = 0; i < missingNodeIndexes.size(); i++) {
+                final int index = missingNodeIndexes.keyAt(i);
+                final AutofillId id = ids[index];
+
+                if (id.equals(node.getAutofillId())) {
+                    foundNodes[index] = node;
+
+                    if (mViewNodeLookupTable == null) {
+                        mViewNodeLookupTable = new ArrayMap<>(ids.length);
+                    }
+
+                    mViewNodeLookupTable.put(id, node);
+
+                    missingNodeIndexes.removeAt(i);
+                    break;
+                }
+            }
+
+            for (int i = 0; i < node.getChildCount(); i++) {
+                nodesToProcess.addLast(node.getChildAt(i));
+            }
+        }
+
+        // Remember which ids could not be resolved to not search for them again the next time
+        for (int i = 0; i < missingNodeIndexes.size(); i++) {
+            if (mViewNodeLookupTable == null) {
+                mViewNodeLookupTable = new ArrayMap<>(missingNodeIndexes.size());
+            }
+
+            mViewNodeLookupTable.put(ids[missingNodeIndexes.keyAt(i)], null);
+        }
+
+        return foundNodes;
+    }
+
     public static final Parcelable.Creator<FillContext> CREATOR =
             new Parcelable.Creator<FillContext>() {
         @Override
diff --git a/services/autofill/java/com/android/server/autofill/Helper.java b/services/autofill/java/com/android/server/autofill/Helper.java
index 8d947b9..ffcde8d 100644
--- a/services/autofill/java/com/android/server/autofill/Helper.java
+++ b/services/autofill/java/com/android/server/autofill/Helper.java
@@ -16,11 +16,7 @@
 
 package com.android.server.autofill;
 
-import android.annotation.NonNull;
-import android.app.assist.AssistStructure;
-import android.app.assist.AssistStructure.ViewNode;
 import android.os.Bundle;
-import android.view.autofill.AutofillId;
 
 import java.util.Arrays;
 import java.util.Objects;
@@ -65,37 +61,4 @@
         append(builder, bundle);
         return builder.toString();
     }
-
-    static ViewNode findViewNodeById(@NonNull AssistStructure structure, @NonNull AutofillId id) {
-        final int size = structure.getWindowNodeCount();
-        for (int i = 0; i < size; i++) {
-            final AssistStructure.WindowNode window = structure.getWindowNodeAt(i);
-            final ViewNode root = window.getRootViewNode();
-            if (id.equals(root.getAutofillId())) {
-                return root;
-            }
-            final ViewNode child = findViewNodeById(root, id);
-            if (child != null) {
-                return child;
-            }
-        }
-        return null;
-    }
-
-    static ViewNode findViewNodeById(@NonNull ViewNode parent, @NonNull AutofillId id) {
-        final int childrenSize = parent.getChildCount();
-        if (childrenSize > 0) {
-            for (int i = 0; i < childrenSize; i++) {
-                final ViewNode child = parent.getChildAt(i);
-                if (id.equals(child.getAutofillId())) {
-                    return child;
-                }
-                final ViewNode grandChild = findViewNodeById(child, id);
-                if (grandChild != null && id.equals(grandChild.getAutofillId())) {
-                    return grandChild;
-                }
-            }
-        }
-        return null;
-    }
 }
diff --git a/services/autofill/java/com/android/server/autofill/Session.java b/services/autofill/java/com/android/server/autofill/Session.java
index c455eda..ef5cdd1 100644
--- a/services/autofill/java/com/android/server/autofill/Session.java
+++ b/services/autofill/java/com/android/server/autofill/Session.java
@@ -25,7 +25,6 @@
 import static android.view.autofill.AutofillManager.ACTION_VIEW_ENTERED;
 import static android.view.autofill.AutofillManager.ACTION_VIEW_EXITED;
 
-import static com.android.server.autofill.Helper.findViewNodeById;
 import static com.android.server.autofill.Helper.sDebug;
 import static com.android.server.autofill.Helper.sVerbose;
 import static com.android.server.autofill.ViewState.STATE_AUTOFILLED;
@@ -79,7 +78,6 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.Map.Entry;
 import java.util.concurrent.atomic.AtomicInteger;
 
 /**
@@ -212,7 +210,7 @@
 
                 final int numContexts = mContexts.size();
                 for (int i = 0; i < numContexts; i++) {
-                    fillStructureWithAllowedValues(mContexts.get(i).getStructure(), flags);
+                    fillContextWithAllowedValues(mContexts.get(i), flags);
                 }
 
                 request = new FillRequest(requestId, mContexts, mClientState, flags);
@@ -223,20 +221,35 @@
     };
 
     /**
-     * Updates values of the nodes in the structure so that:
+     * Returns the ids of all entries in {@link #mViewStates} in the same order.
+     */
+    private AutofillId[] getIdsOfAllViewStates() {
+        final int numViewState = mViewStates.size();
+        final AutofillId[] ids = new AutofillId[numViewState];
+        for (int i = 0; i < numViewState; i++) {
+            ids[i] = mViewStates.valueAt(i).id;
+        }
+
+        return ids;
+    }
+
+    /**
+     * Updates values of the nodes in the context's structure so that:
      * - proper node is focused
      * - autofillValue is sent back to service when it was previously autofilled
      * - autofillValue is sent in the view used to force a request
      *
-     * @param structure The structure to be filled
+     * @param fillContext The context to be filled
      * @param flags The flags that started the session
      */
-    private void fillStructureWithAllowedValues(@NonNull AssistStructure structure, int flags) {
-        final int numViewStates = mViewStates.size();
-        for (int i = 0; i < numViewStates; i++) {
+    private void fillContextWithAllowedValues(@NonNull FillContext fillContext, int flags) {
+        final ViewNode[] nodes = fillContext.findViewNodesByAutofillIds(getIdsOfAllViewStates());
+
+        final int numViewState = mViewStates.size();
+        for (int i = 0; i < numViewState; i++) {
             final ViewState viewState = mViewStates.valueAt(i);
 
-            final ViewNode node = findViewNodeById(structure, viewState.id);
+            final ViewNode node = nodes[i];
             if (node == null) {
                 Slog.w(TAG, "fillStructureWithAllowedValues(): no node for " + viewState.id);
                 continue;
@@ -841,19 +854,23 @@
 
         final int numContexts = mContexts.size();
 
-        for (int i = 0; i < numContexts; i++) {
-            final FillContext context = mContexts.get(i);
+        for (int contextNum = 0; contextNum < numContexts; contextNum++) {
+            final FillContext context = mContexts.get(contextNum);
+
+            final ViewNode[] nodes = context.findViewNodesByAutofillIds(getIdsOfAllViewStates());
 
             if (sVerbose) Slog.v(TAG, "callSaveLocked(): updating " + context);
 
-            for (Entry<AutofillId, ViewState> entry : mViewStates.entrySet()) {
-                final AutofillValue value = entry.getValue().getCurrentValue();
+            for (int viewStateNum = 0; viewStateNum < mViewStates.size(); viewStateNum++) {
+                final ViewState state = mViewStates.valueAt(viewStateNum);
+
+                final AutofillId id = state.id;
+                final AutofillValue value = state.getCurrentValue();
                 if (value == null) {
-                    if (sVerbose) Slog.v(TAG, "callSaveLocked(): skipping " + entry.getKey());
+                    if (sVerbose) Slog.v(TAG, "callSaveLocked(): skipping " + id);
                     continue;
                 }
-                final AutofillId id = entry.getKey();
-                final ViewNode node = findViewNodeById(context.getStructure(), id);
+                final ViewNode node = nodes[viewStateNum];
                 if (node == null) {
                     Slog.w(TAG, "callSaveLocked(): did not find node with id " + id);
                     continue;