Reland "Introduce a member rebinding lens prior to writing"

This reverts commit c98e5037cad0c58bb1ce0fd8805680647c205597.

Change-Id: If366f4e47bf73fc57262b2ea13a8b4a30ae93096
diff --git a/src/main/java/com/android/tools/r8/R8.java b/src/main/java/com/android/tools/r8/R8.java
index f61edb9..60306fc 100644
--- a/src/main/java/com/android/tools/r8/R8.java
+++ b/src/main/java/com/android/tools/r8/R8.java
@@ -74,6 +74,8 @@
 import com.android.tools.r8.optimize.BridgeHoisting;
 import com.android.tools.r8.optimize.ClassAndMemberPublicizer;
 import com.android.tools.r8.optimize.MemberRebindingAnalysis;
+import com.android.tools.r8.optimize.MemberRebindingIdentityLens;
+import com.android.tools.r8.optimize.MemberRebindingIdentityLensFactory;
 import com.android.tools.r8.optimize.VisibilityBridgeRemover;
 import com.android.tools.r8.origin.CommandLineOrigin;
 import com.android.tools.r8.repackaging.Repackaging;
@@ -795,6 +797,10 @@
         }
       }
 
+      MemberRebindingIdentityLens memberRebindingLens =
+          MemberRebindingIdentityLensFactory.create(appView, executorService);
+      appView.setGraphLens(memberRebindingLens);
+
       // Perform repackaging.
       // TODO(b/165783399): Consider making repacking available without minification.
       if (options.isMinifying() && options.testing.enableExperimentalRepackaging) {
@@ -805,7 +811,13 @@
         RepackagingLens lens =
             new Repackaging(appView.withLiveness()).run(appBuilder, executorService, timing);
         if (lens != null) {
-          appView.rewriteWithLensAndApplication(lens, appBuilder.build());
+          // Specify to use the member rebinding lens as the parent lens during the rewriting. This
+          // is needed to ensure that the rebound references are available during lens lookups.
+          // TODO(b/168282032): This call-site should not have to think about the parent lens that
+          //  is used for the rewriting. Once the new member rebinding lens replaces the old member
+          //  rebinding analysis it should be possible to clean this up.
+          appView.rewriteWithLensAndApplication(
+              lens, appBuilder.build(), memberRebindingLens.getPrevious());
         }
       }
 
diff --git a/src/main/java/com/android/tools/r8/graph/AbstractAccessContexts.java b/src/main/java/com/android/tools/r8/graph/AbstractAccessContexts.java
index e078248..4615f1f 100644
--- a/src/main/java/com/android/tools/r8/graph/AbstractAccessContexts.java
+++ b/src/main/java/com/android/tools/r8/graph/AbstractAccessContexts.java
@@ -182,7 +182,7 @@
       }
     }
 
-    Map<DexField, ProgramMethodSet> getAccessesWithContexts() {
+    public Map<DexField, ProgramMethodSet> getAccessesWithContexts() {
       return accessesWithContexts;
     }
 
diff --git a/src/main/java/com/android/tools/r8/graph/AppInfo.java b/src/main/java/com/android/tools/r8/graph/AppInfo.java
index 6da4e08..d2d137b 100644
--- a/src/main/java/com/android/tools/r8/graph/AppInfo.java
+++ b/src/main/java/com/android/tools/r8/graph/AppInfo.java
@@ -13,6 +13,7 @@
 import com.android.tools.r8.utils.BooleanBox;
 import com.android.tools.r8.utils.InternalOptions;
 import java.util.Collection;
+import java.util.function.Consumer;
 
 public class AppInfo implements DexDefinitionSupplier {
 
@@ -125,6 +126,12 @@
     return app.classesWithDeterministicOrder();
   }
 
+  public void forEachMethod(Consumer<ProgramMethod> consumer) {
+    for (DexProgramClass clazz : classes()) {
+      clazz.forEachProgramMethod(consumer);
+    }
+  }
+
   @Override
   public DexClass definitionFor(DexType type) {
     return definitionForWithoutExistenceAssert(type);
diff --git a/src/main/java/com/android/tools/r8/graph/AppView.java b/src/main/java/com/android/tools/r8/graph/AppView.java
index f89f265..a24eb49 100644
--- a/src/main/java/com/android/tools/r8/graph/AppView.java
+++ b/src/main/java/com/android/tools/r8/graph/AppView.java
@@ -6,7 +6,7 @@
 
 import com.android.tools.r8.features.ClassToFeatureSplitMap;
 import com.android.tools.r8.graph.DexValue.DexValueString;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.graph.analysis.InitializedClassesInInstanceMethodsAnalysis.InitializedClassesInInstanceMethods;
 import com.android.tools.r8.graph.classmerging.HorizontallyMergedLambdaClasses;
 import com.android.tools.r8.graph.classmerging.MergedClasses;
@@ -520,36 +520,54 @@
     return !cfByteCodePassThrough.isEmpty();
   }
 
-  public void rewriteWithLens(NestedGraphLens lens) {
+  public void rewriteWithLens(NonIdentityGraphLens lens) {
     if (lens != null) {
-      rewriteWithLens(lens, appInfo().app().asDirect(), withLiveness());
+      rewriteWithLens(lens, appInfo().app().asDirect(), withLiveness(), lens.getPrevious());
     }
   }
 
-  public void rewriteWithApplication(DirectMappedDexApplication application) {
-    assert application != null;
-    rewriteWithLens(null, application, withLiveness());
-  }
-
   public void rewriteWithLensAndApplication(
-      NestedGraphLens lens, DirectMappedDexApplication application) {
+      NonIdentityGraphLens lens, DirectMappedDexApplication application) {
+    rewriteWithLensAndApplication(lens, application, lens.getPrevious());
+  }
+
+  public void rewriteWithLensAndApplication(
+      NonIdentityGraphLens lens, DirectMappedDexApplication application, GraphLens appliedLens) {
     assert lens != null;
     assert application != null;
-    rewriteWithLens(lens, application, withLiveness());
+    rewriteWithLens(lens, application, withLiveness(), appliedLens);
   }
 
   private static void rewriteWithLens(
-      NestedGraphLens lens,
+      NonIdentityGraphLens lens,
       DirectMappedDexApplication application,
-      AppView<AppInfoWithLiveness> appView) {
-    if (lens != null) {
-      boolean changed = appView.setGraphLens(lens);
-      assert changed;
-      assert application.verifyWithLens(lens);
+      AppView<AppInfoWithLiveness> appView,
+      GraphLens appliedLens) {
+    if (lens == null) {
+      return;
     }
-    appView.setAppInfo(appView.appInfo().rewrittenWithLens(application, lens));
-    if (appView.hasInitClassLens()) {
-      appView.setInitClassLens(appView.initClassLens().rewrittenWithLens(lens));
+
+    boolean changed = appView.setGraphLens(lens);
+    assert changed;
+    assert application.verifyWithLens(lens);
+
+    // The application has already been rewritten with the given applied lens. Therefore, we
+    // temporarily replace that lens with the identity lens to avoid the overhead of traversing
+    // the entire lens chain upon each lookup during the rewriting.
+    NonIdentityGraphLens temporaryRootLens = lens;
+    while (temporaryRootLens.getPrevious() != appliedLens) {
+      GraphLens previousLens = temporaryRootLens.getPrevious();
+      assert previousLens.isNonIdentityLens();
+      temporaryRootLens = previousLens.asNonIdentityLens();
     }
+
+    temporaryRootLens.withAlternativeParentLens(
+        GraphLens.getIdentityLens(),
+        () -> {
+          appView.setAppInfo(appView.appInfo().rewrittenWithLens(application, lens));
+          if (appView.hasInitClassLens()) {
+            appView.setInitClassLens(appView.initClassLens().rewrittenWithLens(lens));
+          }
+        });
   }
 }
diff --git a/src/main/java/com/android/tools/r8/graph/AppliedGraphLens.java b/src/main/java/com/android/tools/r8/graph/AppliedGraphLens.java
index 0a03d50..f3515e9 100644
--- a/src/main/java/com/android/tools/r8/graph/AppliedGraphLens.java
+++ b/src/main/java/com/android/tools/r8/graph/AppliedGraphLens.java
@@ -19,7 +19,7 @@
  *
  * <p>The mappings from the original program to the generated program are kept, though.
  */
-public class AppliedGraphLens extends NonIdentityGraphLens {
+public final class AppliedGraphLens extends NonIdentityGraphLens {
 
   private final AppView<?> appView;
 
@@ -85,6 +85,11 @@
   }
 
   @Override
+  public boolean isAppliedLens() {
+    return true;
+  }
+
+  @Override
   public DexType getOriginalType(DexType type) {
     return originalTypeNames.getOrDefault(type, type);
   }
diff --git a/src/main/java/com/android/tools/r8/graph/DexField.java b/src/main/java/com/android/tools/r8/graph/DexField.java
index 0997619..fcb79fe 100644
--- a/src/main/java/com/android/tools/r8/graph/DexField.java
+++ b/src/main/java/com/android/tools/r8/graph/DexField.java
@@ -23,6 +23,10 @@
     }
   }
 
+  public DexType getHolderType() {
+    return holder;
+  }
+
   public DexType getType() {
     return type;
   }
@@ -153,4 +157,8 @@
   public String toSourceString() {
     return type.toSourceString() + " " + holder.toSourceString() + "." + name.toSourceString();
   }
+
+  public DexField withHolder(DexType holder, DexItemFactory dexItemFactory) {
+    return dexItemFactory.createField(holder, type, name);
+  }
 }
diff --git a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoCollectionImpl.java b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoCollectionImpl.java
index c82fcf0..88869df 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoCollectionImpl.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoCollectionImpl.java
@@ -9,11 +9,20 @@
 import java.util.Map;
 import java.util.function.BiPredicate;
 import java.util.function.Consumer;
+import java.util.function.Function;
 
 public class FieldAccessInfoCollectionImpl
     implements FieldAccessInfoCollection<FieldAccessInfoImpl> {
 
-  private Map<DexField, FieldAccessInfoImpl> infos = new IdentityHashMap<>();
+  private final Map<DexField, FieldAccessInfoImpl> infos;
+
+  public FieldAccessInfoCollectionImpl() {
+    this(new IdentityHashMap<>());
+  }
+
+  public FieldAccessInfoCollectionImpl(Map<DexField, FieldAccessInfoImpl> infos) {
+    this.infos = infos;
+  }
 
   @Override
   public void destroyAccessContexts() {
@@ -25,6 +34,11 @@
     infos.values().forEach(FieldAccessInfoImpl::flattenAccessContexts);
   }
 
+  public FieldAccessInfoImpl computeIfAbsent(
+      DexField field, Function<DexField, FieldAccessInfoImpl> fn) {
+    return infos.computeIfAbsent(field, fn);
+  }
+
   @Override
   public boolean contains(DexField field) {
     return infos.containsKey(field);
diff --git a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
index 023e423..9dae761 100644
--- a/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
+++ b/src/main/java/com/android/tools/r8/graph/FieldAccessInfoImpl.java
@@ -71,6 +71,14 @@
     return field;
   }
 
+  public AbstractAccessContexts getReadsWithContexts() {
+    return readsWithContexts;
+  }
+
+  public void setReadsWithContexts(AbstractAccessContexts readsWithContexts) {
+    this.readsWithContexts = readsWithContexts;
+  }
+
   @Override
   public int getNumberOfReadContexts() {
     return readsWithContexts.getNumberOfAccessContexts();
diff --git a/src/main/java/com/android/tools/r8/graph/GraphLens.java b/src/main/java/com/android/tools/r8/graph/GraphLens.java
index f3c96d5..ad7cd8b 100644
--- a/src/main/java/com/android/tools/r8/graph/GraphLens.java
+++ b/src/main/java/com/android/tools/r8/graph/GraphLens.java
@@ -7,6 +7,7 @@
 import com.android.tools.r8.ir.code.Invoke.Type;
 import com.android.tools.r8.ir.desugar.InterfaceProcessor.InterfaceProcessorNestedGraphLens;
 import com.android.tools.r8.shaking.KeepInfoCollection;
+import com.android.tools.r8.utils.Action;
 import com.android.tools.r8.utils.SetUtils;
 import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
@@ -21,7 +22,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.function.Supplier;
 
 /**
  * A GraphLens implements a virtual view on top of the graph, used to delay global rewrites until
@@ -63,10 +63,18 @@
       return rewritings.getOrDefault(reference, reference);
     }
 
+    public boolean hasReboundReference() {
+      return reboundReference != null;
+    }
+
     public DexField getReboundReference() {
       return reboundReference;
     }
 
+    public DexField getRewrittenReboundReference(Map<DexField, DexField> rewritings) {
+      return rewritings.getOrDefault(reboundReference, reboundReference);
+    }
+
     public static Builder builder(GraphLens lens) {
       return new Builder(lens);
     }
@@ -92,8 +100,7 @@
       }
 
       public FieldLookupResult build() {
-        // All non-identity graph lenses should set the rebound reference.
-        assert lens.isIdentityLens() || reboundReference != null;
+        // TODO(b/168282032): All non-identity graph lenses should set the rebound reference.
         return new FieldLookupResult(reference, reboundReference);
       }
     }
@@ -199,10 +206,6 @@
     }
   }
 
-  public static Builder builder() {
-    return new Builder();
-  }
-
   /**
    * Intentionally private. All graph lenses except for {@link IdentityGraphLens} should inherit
    * from {@link NonIdentityGraphLens}.
@@ -331,6 +334,10 @@
     return true;
   }
 
+  public boolean isAppliedLens() {
+    return false;
+  }
+
   public abstract boolean isIdentityLens();
 
   public abstract boolean isNonIdentityLens();
@@ -517,12 +524,11 @@
       return previousLens;
     }
 
-    public final <T> T withAlternativeParentLens(GraphLens lens, Supplier<T> action) {
+    public final void withAlternativeParentLens(GraphLens lens, Action action) {
       GraphLens oldParent = getPrevious();
       previousLens = lens;
-      T result = action.get();
+      action.execute();
       previousLens = oldParent;
-      return result;
     }
 
     @Override
@@ -744,6 +750,10 @@
       this.dexItemFactory = dexItemFactory;
     }
 
+    public static Builder builder() {
+      return new Builder();
+    }
+
     @Override
     public DexType getOriginalType(DexType type) {
       return getPrevious().getOriginalType(type);
@@ -805,7 +815,10 @@
           return result;
         }
       }
-      DexType previous = getPrevious().lookupType(type);
+      return internalDescribeLookupType(getPrevious().lookupType(type));
+    }
+
+    private DexType internalDescribeLookupType(DexType previous) {
       return typeMap != null ? typeMap.getOrDefault(previous, previous) : previous;
     }
 
@@ -901,12 +914,22 @@
 
     @Override
     protected FieldLookupResult internalDescribeLookupField(FieldLookupResult previous) {
-      DexField rewrittenReference = previous.getRewrittenReference(fieldMap);
-      return FieldLookupResult.builder(this)
-          .setReference(rewrittenReference)
-          // TODO(b/168282032): This should set the rebound reference.
-          .setReboundReference(rewrittenReference)
-          .build();
+      if (previous.hasReboundReference()) {
+        // Rewrite the rebound reference and then "fixup" the non-rebound reference.
+        DexField rewrittenReboundReference = previous.getRewrittenReboundReference(fieldMap);
+        return FieldLookupResult.builder(this)
+            .setReboundReference(rewrittenReboundReference)
+            .setReference(
+                rewrittenReboundReference.withHolder(
+                    internalDescribeLookupType(previous.getReference().getHolderType()),
+                    dexItemFactory))
+            .build();
+      } else {
+        // TODO(b/168282032): We should always have the rebound reference, so this should become
+        //  unreachable.
+        DexField rewrittenReference = previous.getRewrittenReference(fieldMap);
+        return FieldLookupResult.builder(this).setReference(rewrittenReference).build();
+      }
     }
 
     @Override
diff --git a/src/main/java/com/android/tools/r8/graph/UseRegistry.java b/src/main/java/com/android/tools/r8/graph/UseRegistry.java
index 75e6565..5830357 100644
--- a/src/main/java/com/android/tools/r8/graph/UseRegistry.java
+++ b/src/main/java/com/android/tools/r8/graph/UseRegistry.java
@@ -17,6 +17,10 @@
     this.factory = factory;
   }
 
+  public final void accept(ProgramMethod method) {
+    method.registerCodeReferences(this);
+  }
+
   public abstract void registerInitClass(DexType type);
 
   public abstract void registerInvokeVirtual(DexMethod method);
diff --git a/src/main/java/com/android/tools/r8/optimize/BridgeHoisting.java b/src/main/java/com/android/tools/r8/optimize/BridgeHoisting.java
index c19bcd2..038d078 100644
--- a/src/main/java/com/android/tools/r8/optimize/BridgeHoisting.java
+++ b/src/main/java/com/android/tools/r8/optimize/BridgeHoisting.java
@@ -379,6 +379,11 @@
     }
 
     @Override
+    public boolean hasCodeRewritings() {
+      return getPrevious().hasCodeRewritings();
+    }
+
+    @Override
     public boolean isLegitimateToHaveEmptyMappings() {
       return true;
     }
diff --git a/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLens.java b/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLens.java
new file mode 100644
index 0000000..70b14cd
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLens.java
@@ -0,0 +1,119 @@
+// Copyright (c) 2018, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.optimize;
+
+import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.FieldAccessInfo;
+import com.android.tools.r8.graph.GraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
+import com.android.tools.r8.graph.RewrittenPrototypeDescription;
+import com.android.tools.r8.ir.code.Invoke.Type;
+import java.util.IdentityHashMap;
+import java.util.Map;
+
+/**
+ * This lens is used to populate the rebound field and method reference during lookup, such that
+ * both the non-rebound and rebound field and method references are available to all descendants of
+ * this lens.
+ */
+public class MemberRebindingIdentityLens extends NonIdentityGraphLens {
+
+  private final Map<DexField, DexField> nonReboundFieldReferenceToDefinitionMap;
+
+  private MemberRebindingIdentityLens(
+      Map<DexField, DexField> nonReboundFieldReferenceToDefinitionMap, GraphLens previousLens) {
+    super(previousLens);
+    assert !previousLens.hasCodeRewritings();
+    this.nonReboundFieldReferenceToDefinitionMap = nonReboundFieldReferenceToDefinitionMap;
+  }
+
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  @Override
+  protected FieldLookupResult internalDescribeLookupField(FieldLookupResult previous) {
+    assert previous.getReboundReference() == null;
+    return FieldLookupResult.builder(this)
+        .setReference(previous.getReference())
+        .setReboundReference(getReboundFieldReference(previous.getReference()))
+        .build();
+  }
+
+  private DexField getReboundFieldReference(DexField field) {
+    return nonReboundFieldReferenceToDefinitionMap.getOrDefault(field, field);
+  }
+
+  @Override
+  public DexType getOriginalType(DexType type) {
+    return getPrevious().getOriginalType(type);
+  }
+
+  @Override
+  public DexField getOriginalFieldSignature(DexField field) {
+    return getPrevious().getOriginalFieldSignature(field);
+  }
+
+  @Override
+  public DexMethod getOriginalMethodSignature(DexMethod method) {
+    return getPrevious().getOriginalMethodSignature(method);
+  }
+
+  @Override
+  public DexField getRenamedFieldSignature(DexField originalField) {
+    return getPrevious().getRenamedFieldSignature(originalField);
+  }
+
+  @Override
+  public DexMethod getRenamedMethodSignature(DexMethod originalMethod, GraphLens applied) {
+    return getPrevious().getRenamedMethodSignature(originalMethod, applied);
+  }
+
+  @Override
+  public DexType lookupType(DexType type) {
+    return getPrevious().lookupType(type);
+  }
+
+  @Override
+  public GraphLensLookupResult lookupMethod(DexMethod method, DexMethod context, Type type) {
+    return getPrevious().lookupMethod(method, context, type);
+  }
+
+  @Override
+  public RewrittenPrototypeDescription lookupPrototypeChangesForMethodDefinition(DexMethod method) {
+    return getPrevious().lookupPrototypeChangesForMethodDefinition(method);
+  }
+
+  @Override
+  public boolean isContextFreeForMethods() {
+    return getPrevious().isContextFreeForMethods();
+  }
+
+  public static class Builder {
+
+    private final Map<DexField, DexField> nonReboundFieldReferenceToDefinitionMap =
+        new IdentityHashMap<>();
+
+    void recordNonReboundFieldAccesses(FieldAccessInfo fieldAccessInfo) {
+      fieldAccessInfo.forEachIndirectAccess(
+          nonReboundFieldReference ->
+              recordNonReboundFieldAccess(nonReboundFieldReference, fieldAccessInfo.getField()));
+    }
+
+    private void recordNonReboundFieldAccess(
+        DexField nonReboundFieldReference, DexField reboundFieldReference) {
+      nonReboundFieldReferenceToDefinitionMap.put(nonReboundFieldReference, reboundFieldReference);
+    }
+
+    MemberRebindingIdentityLens build(GraphLens previousLens) {
+      // This intentionally does not return null when the maps are empty. In this case there are no
+      // non-rebound field or method references, but the member rebinding lens is still needed to
+      // populate the rebound reference during field and method lookup.
+      return new MemberRebindingIdentityLens(nonReboundFieldReferenceToDefinitionMap, previousLens);
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLensFactory.java b/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLensFactory.java
new file mode 100644
index 0000000..48a01f8
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/optimize/MemberRebindingIdentityLensFactory.java
@@ -0,0 +1,181 @@
+// Copyright (c) 2020, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.optimize;
+
+import com.android.tools.r8.graph.AbstractAccessContexts.ConcreteAccessContexts;
+import com.android.tools.r8.graph.AppInfoWithClassHierarchy;
+import com.android.tools.r8.graph.AppView;
+import com.android.tools.r8.graph.DexField;
+import com.android.tools.r8.graph.DexMethod;
+import com.android.tools.r8.graph.DexType;
+import com.android.tools.r8.graph.FieldAccessInfoCollection;
+import com.android.tools.r8.graph.FieldAccessInfoCollectionImpl;
+import com.android.tools.r8.graph.FieldAccessInfoImpl;
+import com.android.tools.r8.graph.FieldResolutionResult.SuccessfulFieldResolutionResult;
+import com.android.tools.r8.graph.GraphLens;
+import com.android.tools.r8.graph.UseRegistry;
+import com.android.tools.r8.shaking.AppInfoWithLiveness;
+import com.android.tools.r8.utils.InternalOptions.TestingOptions;
+import com.android.tools.r8.utils.ThreadUtils;
+import com.android.tools.r8.utils.collections.ProgramMethodSet;
+import com.google.common.collect.Sets;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+
+public class MemberRebindingIdentityLensFactory {
+
+  /**
+   * In order to construct an instance of {@link MemberRebindingIdentityLens} we need a mapping from
+   * non-rebound field and method references to their definitions.
+   *
+   * <p>If shrinking or minification is enabled, we retrieve these from {@link AppInfoWithLiveness}.
+   * Otherwise we apply the {@link NonReboundMemberReferencesRegistry} below to all code objects to
+   * compute the mapping.
+   */
+  public static MemberRebindingIdentityLens create(
+      AppView<? extends AppInfoWithClassHierarchy> appView, ExecutorService executorService)
+      throws ExecutionException {
+    TestingOptions testingOptions = appView.options().testing;
+    FieldAccessInfoCollection<?> fieldAccessInfoCollection =
+        appView.appInfo().hasLiveness()
+                && testingOptions.alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding
+            ? appView.appInfo().withLiveness().getFieldAccessInfoCollection()
+            : createFieldAccessInfoCollectionForMemberRebinding(appView, executorService);
+    return create(fieldAccessInfoCollection, appView.graphLens());
+  }
+
+  public static MemberRebindingIdentityLens create(
+      FieldAccessInfoCollection<?> fieldAccessInfoCollection, GraphLens previousLens) {
+    MemberRebindingIdentityLens.Builder builder = MemberRebindingIdentityLens.builder();
+    fieldAccessInfoCollection.forEach(builder::recordNonReboundFieldAccesses);
+    return builder.build(previousLens);
+  }
+
+  /**
+   * Applies {@link NonReboundMemberReferencesRegistry} to all code objects to construct a mapping
+   * from non-rebound field references to their definition.
+   */
+  private static FieldAccessInfoCollection<?> createFieldAccessInfoCollectionForMemberRebinding(
+      AppView<? extends AppInfoWithClassHierarchy> appView, ExecutorService executorService)
+      throws ExecutionException {
+    NonReboundMemberReferencesRegistry registry = new NonReboundMemberReferencesRegistry(appView);
+    ThreadUtils.processItems(appView.appInfo()::forEachMethod, registry::accept, executorService);
+    return registry.getFieldAccessInfoCollection();
+  }
+
+  private static class NonReboundMemberReferencesRegistry extends UseRegistry {
+
+    private final AppInfoWithClassHierarchy appInfo;
+    private final FieldAccessInfoCollectionImpl fieldAccessInfoCollection =
+        new FieldAccessInfoCollectionImpl(new ConcurrentHashMap<>());
+    private final Set<DexField> seenFieldReferences = Sets.newConcurrentHashSet();
+
+    FieldAccessInfoCollection<?> getFieldAccessInfoCollection() {
+      return fieldAccessInfoCollection;
+    }
+
+    public NonReboundMemberReferencesRegistry(
+        AppView<? extends AppInfoWithClassHierarchy> appView) {
+      super(appView.dexItemFactory());
+      this.appInfo = appView.appInfo();
+    }
+
+    @Override
+    public void registerInstanceFieldRead(DexField field) {
+      registerFieldAccess(field);
+    }
+
+    @Override
+    public void registerInstanceFieldWrite(DexField field) {
+      registerFieldAccess(field);
+    }
+
+    @Override
+    public void registerStaticFieldRead(DexField field) {
+      registerFieldAccess(field);
+    }
+
+    @Override
+    public void registerStaticFieldWrite(DexField field) {
+      registerFieldAccess(field);
+    }
+
+    private void registerFieldAccess(DexField field) {
+      if (!seenFieldReferences.add(field)) {
+        return;
+      }
+      SuccessfulFieldResolutionResult resolutionResult =
+          appInfo.resolveField(field).asSuccessfulResolution();
+      if (resolutionResult == null) {
+        return;
+      }
+      DexField reboundReference = resolutionResult.getResolvedField().toReference();
+      if (field == reboundReference) {
+        // For the purpose of member rebinding, we don't care about already rebound references.
+        return;
+      }
+      FieldAccessInfoImpl fieldAccessInfo =
+          fieldAccessInfoCollection.computeIfAbsent(reboundReference, FieldAccessInfoImpl::new);
+      synchronized (fieldAccessInfo) {
+        // Record the fact that there is a non-rebound access to the given field. We don't
+        // distinguish between non-rebound reads and writes, so we just record it as a read.
+        ConcreteAccessContexts accessContexts =
+            fieldAccessInfo.getReadsWithContexts().isConcrete()
+                ? fieldAccessInfo.getReadsWithContexts().asConcrete()
+                : new ConcreteAccessContexts();
+        // For the purpose of member rebinding, we don't care about the access contexts, so we
+        // simply use the empty set.
+        accessContexts.getAccessesWithContexts().put(field, ProgramMethodSet.empty());
+      }
+    }
+
+    @Override
+    public void registerInvokeVirtual(DexMethod method) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInvokeDirect(DexMethod method) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInvokeStatic(DexMethod method) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInvokeInterface(DexMethod method) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInvokeSuper(DexMethod method) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInitClass(DexType type) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerNewInstance(DexType type) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerTypeReference(DexType type) {
+      // Intentionally empty.
+    }
+
+    @Override
+    public void registerInstanceOf(DexType type) {
+      // Intentionally empty.
+    }
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
index 627d84f..1b1d37c 100644
--- a/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
+++ b/src/main/java/com/android/tools/r8/shaking/AppInfoWithLiveness.java
@@ -31,7 +31,7 @@
 import com.android.tools.r8.graph.FieldAccessInfoCollectionImpl;
 import com.android.tools.r8.graph.FieldResolutionResult;
 import com.android.tools.r8.graph.GraphLens;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.graph.InstantiatedSubTypeInfo;
 import com.android.tools.r8.graph.LookupResult.LookupResultSuccess;
 import com.android.tools.r8.graph.LookupTarget;
@@ -993,17 +993,9 @@
   }
 
   public AppInfoWithLiveness rewrittenWithLens(
-      DirectMappedDexApplication application, NestedGraphLens lens) {
+      DirectMappedDexApplication application, NonIdentityGraphLens lens) {
     assert checkIfObsolete();
-    // The application has already been rewritten with all of lens' parent lenses. Therefore, we
-    // temporarily replace lens' parent lens with an identity lens to avoid the overhead of
-    // traversing the entire lens chain upon each lookup during the rewriting.
-    return lens.withAlternativeParentLens(
-        GraphLens.getIdentityLens(), () -> createRewrittenAppInfoWithLiveness(application, lens));
-  }
 
-  private AppInfoWithLiveness createRewrittenAppInfoWithLiveness(
-      DirectMappedDexApplication application, NestedGraphLens lens) {
     // Switchmap classes should never be affected by renaming.
     assert lens.assertDefinitionsNotModified(
         switchMaps.keySet().stream()
diff --git a/src/main/java/com/android/tools/r8/shaking/KeepInfoCollection.java b/src/main/java/com/android/tools/r8/shaking/KeepInfoCollection.java
index efc610a..d2cbf66 100644
--- a/src/main/java/com/android/tools/r8/shaking/KeepInfoCollection.java
+++ b/src/main/java/com/android/tools/r8/shaking/KeepInfoCollection.java
@@ -16,7 +16,7 @@
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexReference;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.graph.ProgramField;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.shaking.KeepFieldInfo.Joiner;
@@ -162,7 +162,7 @@
   @Deprecated
   public abstract void forEachPinnedField(Consumer<DexField> consumer);
 
-  public abstract KeepInfoCollection rewrite(NestedGraphLens lens, InternalOptions options);
+  public abstract KeepInfoCollection rewrite(NonIdentityGraphLens lens, InternalOptions options);
 
   public abstract KeepInfoCollection mutate(Consumer<MutableKeepInfoCollection> mutator);
 
@@ -198,7 +198,7 @@
     }
 
     @Override
-    public KeepInfoCollection rewrite(NestedGraphLens lens, InternalOptions options) {
+    public KeepInfoCollection rewrite(NonIdentityGraphLens lens, InternalOptions options) {
       Map<DexType, KeepClassInfo> newClassInfo = new IdentityHashMap<>(keepClassInfo.size());
       keepClassInfo.forEach(
           (type, info) -> {
diff --git a/src/main/java/com/android/tools/r8/synthesis/SynthesizingContext.java b/src/main/java/com/android/tools/r8/synthesis/SynthesizingContext.java
index c03022e..380c2dc 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SynthesizingContext.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SynthesizingContext.java
@@ -4,7 +4,7 @@
 package com.android.tools.r8.synthesis;
 
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.origin.Origin;
 
 /**
@@ -29,7 +29,7 @@
     this.origin = origin;
   }
 
-  SynthesizingContext rewrite(NestedGraphLens lens) {
+  SynthesizingContext rewrite(NonIdentityGraphLens lens) {
     DexType rewritten = lens.lookupType(type);
     return rewritten == type ? this : new SynthesizingContext(type, origin);
   }
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
index d104c4d..2996794 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticFinalization.java
@@ -12,6 +12,7 @@
 import com.android.tools.r8.graph.DexType;
 import com.android.tools.r8.graph.GraphLens;
 import com.android.tools.r8.graph.GraphLens.Builder;
+import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
 import com.android.tools.r8.shaking.MainDexClasses;
 import com.android.tools.r8.utils.DescriptorUtils;
 import com.android.tools.r8.utils.InternalOptions;
@@ -99,7 +100,7 @@
     Map<DexType, EquivalenceGroup<SyntheticMethodDefinition>> equivalences =
         computeActualEquivalences(potentialEquivalences, options.itemFactory);
 
-    Builder lensBuilder = GraphLens.builder();
+    Builder lensBuilder = NestedGraphLens.builder();
     List<DexProgramClass> newProgramClasses = new ArrayList<>();
     List<DexProgramClass> finalSyntheticClasses = new ArrayList<>();
     buildLensAndProgram(
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
index b3c179f..ef45975 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticItems.java
@@ -10,7 +10,7 @@
 import com.android.tools.r8.graph.DexItemFactory;
 import com.android.tools.r8.graph.DexProgramClass;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
 import com.android.tools.r8.synthesis.SyntheticFinalization.Result;
 import com.google.common.collect.ImmutableList;
@@ -233,7 +233,8 @@
         nextSyntheticId);
   }
 
-  public CommittedItems commitRewrittenWithLens(DexApplication application, NestedGraphLens lens) {
+  public CommittedItems commitRewrittenWithLens(
+      DexApplication application, NonIdentityGraphLens lens) {
     // Rewrite the previously committed synthetic types.
     ImmutableSet<DexType> rewrittenLegacyTypes = lens.rewriteTypes(this.legacySyntheticTypes);
     ImmutableMap.Builder<DexType, SyntheticReference> rewrittenItems = ImmutableMap.builder();
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticMethodReference.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticMethodReference.java
index 326eee0..6070e49 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticMethodReference.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticMethodReference.java
@@ -6,7 +6,7 @@
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexMethod;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.graph.ProgramMethod;
 import java.util.function.Function;
 
@@ -40,7 +40,7 @@
   }
 
   @Override
-  SyntheticReference rewrite(NestedGraphLens lens) {
+  SyntheticReference rewrite(NonIdentityGraphLens lens) {
     SynthesizingContext context = getContext().rewrite(lens);
     DexMethod rewritten = lens.lookupMethod(method);
     return context == getContext() && rewritten == method
diff --git a/src/main/java/com/android/tools/r8/synthesis/SyntheticReference.java b/src/main/java/com/android/tools/r8/synthesis/SyntheticReference.java
index 9981624..9b95337 100644
--- a/src/main/java/com/android/tools/r8/synthesis/SyntheticReference.java
+++ b/src/main/java/com/android/tools/r8/synthesis/SyntheticReference.java
@@ -5,7 +5,7 @@
 
 import com.android.tools.r8.graph.DexClass;
 import com.android.tools.r8.graph.DexType;
-import com.android.tools.r8.graph.GraphLens.NestedGraphLens;
+import com.android.tools.r8.graph.GraphLens.NonIdentityGraphLens;
 import com.android.tools.r8.origin.Origin;
 import java.util.function.Function;
 
@@ -37,5 +37,5 @@
 
   abstract DexType getHolder();
 
-  abstract SyntheticReference rewrite(NestedGraphLens lens);
+  abstract SyntheticReference rewrite(NonIdentityGraphLens lens);
 }
diff --git a/src/main/java/com/android/tools/r8/utils/InternalOptions.java b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
index ae7a9d7..e3a85c8 100644
--- a/src/main/java/com/android/tools/r8/utils/InternalOptions.java
+++ b/src/main/java/com/android/tools/r8/utils/InternalOptions.java
@@ -1218,6 +1218,7 @@
     public boolean allowClassInlinerGracefulExit =
         System.getProperty("com.android.tools.r8.disallowClassInlinerGracefulExit") == null;
     public boolean reportUnusedProguardConfigurationRules = false;
+    public boolean alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding = true;
     public boolean alwaysUsePessimisticRegisterAllocation = false;
     public boolean enableCheckCastAndInstanceOfRemoval = true;
     public boolean enableDeadSwitchCaseElimination = true;
diff --git a/src/test/java/com/android/tools/r8/repackage/RepackagingWithNonReboundFieldReferenceTest.java b/src/test/java/com/android/tools/r8/repackage/RepackagingWithNonReboundFieldReferenceTest.java
index c8a9d82..65e3e84 100644
--- a/src/test/java/com/android/tools/r8/repackage/RepackagingWithNonReboundFieldReferenceTest.java
+++ b/src/test/java/com/android/tools/r8/repackage/RepackagingWithNonReboundFieldReferenceTest.java
@@ -7,14 +7,13 @@
 import static com.android.tools.r8.shaking.ProguardConfigurationParser.FLATTEN_PACKAGE_HIERARCHY;
 import static com.android.tools.r8.shaking.ProguardConfigurationParser.REPACKAGE_CLASSES;
 import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.assertTrue;
 
-import com.android.tools.r8.CompilationFailedException;
 import com.android.tools.r8.TestBase;
 import com.android.tools.r8.TestParameters;
-import com.android.tools.r8.repackage.RepackageWithMainDexListTest.TestClass;
 import com.android.tools.r8.repackage.testclasses.RepackagingWithNonReboundFieldReferenceTestClasses;
 import com.android.tools.r8.repackage.testclasses.RepackagingWithNonReboundFieldReferenceTestClasses.B;
+import com.android.tools.r8.utils.BooleanUtils;
 import com.google.common.collect.ImmutableList;
 import java.util.List;
 import org.junit.Test;
@@ -27,45 +26,50 @@
 
   private static final String REPACKAGE_DIR = "foo";
 
+  private final boolean alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding;
   private final String flattenPackageHierarchyOrRepackageClasses;
   private final TestParameters parameters;
 
-  @Parameters(name = "{1}, kind: {0}")
+  @Parameters(name = "{2}, reuse FieldAccessInfoCollection: {0}, kind: {1}")
   public static List<Object[]> data() {
     return buildParameters(
+        BooleanUtils.values(),
         ImmutableList.of(FLATTEN_PACKAGE_HIERARCHY, REPACKAGE_CLASSES),
         getTestParameters().withAllRuntimesAndApiLevels().build());
   }
 
   public RepackagingWithNonReboundFieldReferenceTest(
-      String flattenPackageHierarchyOrRepackageClasses, TestParameters parameters) {
+      boolean alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding,
+      String flattenPackageHierarchyOrRepackageClasses,
+      TestParameters parameters) {
+    this.alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding =
+        alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding;
     this.flattenPackageHierarchyOrRepackageClasses = flattenPackageHierarchyOrRepackageClasses;
     this.parameters = parameters;
   }
 
   @Test
   public void test() throws Exception {
-    try {
-      testForR8(parameters.getBackend())
-          .addInnerClasses(getClass(), RepackagingWithNonReboundFieldReferenceTestClasses.class)
-          .addKeepMainRule(TestClass.class)
-          .addKeepRules(
-              "-" + flattenPackageHierarchyOrRepackageClasses + " \"" + REPACKAGE_DIR + "\"")
-          .addOptionsModification(
-              options -> {
-                assertFalse(options.testing.enableExperimentalRepackaging);
-                options.testing.enableExperimentalRepackaging = true;
-              })
-          .enableMemberValuePropagationAnnotations()
-          .enableNoVerticalClassMergingAnnotations()
-          .setMinApi(parameters.getApiLevel())
-          .compile();
-
-      // TODO(b/168282032): Support lens rewriting of non-rebound references in the writer.
-      fail();
-    } catch (CompilationFailedException exception) {
-      // Ignore.
-    }
+    testForR8(parameters.getBackend())
+        .addInnerClasses(getClass(), RepackagingWithNonReboundFieldReferenceTestClasses.class)
+        .addKeepMainRule(TestClass.class)
+        .addKeepRules(
+            "-" + flattenPackageHierarchyOrRepackageClasses + " \"" + REPACKAGE_DIR + "\"")
+        .addOptionsModification(
+            options -> {
+              assertTrue(
+                  options.testing.alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding);
+              options.testing.alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding =
+                  alwaysUseExistingFieldAccessInfoCollectionInMemberRebinding;
+              assertFalse(options.testing.enableExperimentalRepackaging);
+              options.testing.enableExperimentalRepackaging = true;
+            })
+        .enableMemberValuePropagationAnnotations()
+        .enableNoVerticalClassMergingAnnotations()
+        .setMinApi(parameters.getApiLevel())
+        .compile()
+        .run(parameters.getRuntime(), TestClass.class)
+        .assertSuccessWithOutputLines("Hello world!");
   }
 
   static class TestClass {