Implement data flow analysis over an abstraction of the IR

Change-Id: I4a81e523188054d3493bf1f44ba0cfdb4d6210eb
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/AbstractTransferFunction.java b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/AbstractTransferFunction.java
index af7a836..a3013e3 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/AbstractTransferFunction.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/AbstractTransferFunction.java
@@ -4,19 +4,17 @@
 
 package com.android.tools.r8.ir.analysis.framework.intraprocedural;
 
-import com.android.tools.r8.ir.code.BasicBlock;
-import com.android.tools.r8.ir.code.Instruction;
-
 /**
  * A transfer function that defines the abstract semantics of the instructions in the program
  * according to some abstract state {@link StateType}.
  */
-public interface AbstractTransferFunction<StateType extends AbstractState<StateType>> {
+public interface AbstractTransferFunction<
+    Block, Instruction, StateType extends AbstractState<StateType>> {
 
   TransferFunctionResult<StateType> apply(Instruction instruction, StateType state);
 
   default StateType computeBlockEntryState(
-      BasicBlock block, BasicBlock predecessor, StateType predecessorExitState) {
+      Block block, Block predecessor, StateType predecessorExitState) {
     return predecessorExitState;
   }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/ControlFlowGraph.java b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/ControlFlowGraph.java
new file mode 100644
index 0000000..41fcd44
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/ControlFlowGraph.java
@@ -0,0 +1,42 @@
+// Copyright (c) 2022, 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.ir.analysis.framework.intraprocedural;
+
+import com.android.tools.r8.utils.TraversalContinuation;
+import com.google.common.collect.Iterables;
+import java.util.Collection;
+import java.util.function.BiFunction;
+
+public interface ControlFlowGraph<Block, Instruction> {
+
+  Collection<Block> getPredecessors(Block block);
+
+  Collection<Block> getSuccessors(Block block);
+
+  default boolean hasUniquePredecessor(Block block) {
+    return getPredecessors(block).size() == 1;
+  }
+
+  default Block getUniquePredecessor(Block block) {
+    assert hasUniquePredecessor(block);
+    return Iterables.getOnlyElement(getPredecessors(block));
+  }
+
+  default boolean hasUniqueSuccessor(Block block) {
+    return getSuccessors(block).size() == 1;
+  }
+
+  default boolean hasUniqueSuccessorWithUniquePredecessor(Block block) {
+    return hasUniqueSuccessor(block) && getPredecessors(getUniqueSuccessor(block)).size() == 1;
+  }
+
+  default Block getUniqueSuccessor(Block block) {
+    assert hasUniqueSuccessor(block);
+    return Iterables.getOnlyElement(getSuccessors(block));
+  }
+
+  <T> TraversalContinuation<T> traverseInstructions(
+      Block block, BiFunction<Instruction, T, TraversalContinuation<T>> fn, T initialValue);
+}
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/DataflowAnalysisResult.java b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/DataflowAnalysisResult.java
index a8a1832..8974bbf 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/DataflowAnalysisResult.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/DataflowAnalysisResult.java
@@ -24,8 +24,8 @@
     return false;
   }
 
-  public <StateType extends AbstractState<StateType>>
-      SuccessfulDataflowAnalysisResult<StateType> asSuccessfulAnalysisResult() {
+  public <Block, StateType extends AbstractState<StateType>>
+      SuccessfulDataflowAnalysisResult<Block, StateType> asSuccessfulAnalysisResult() {
     return null;
   }
 
@@ -33,12 +33,13 @@
     return false;
   }
 
-  public static class SuccessfulDataflowAnalysisResult<StateType extends AbstractState<StateType>>
+  public static class SuccessfulDataflowAnalysisResult<
+          Block, StateType extends AbstractState<StateType>>
       extends DataflowAnalysisResult {
 
-    private final Map<BasicBlock, StateType> blockExitStates;
+    private final Map<Block, StateType> blockExitStates;
 
-    public SuccessfulDataflowAnalysisResult(Map<BasicBlock, StateType> blockExitStates) {
+    public SuccessfulDataflowAnalysisResult(Map<Block, StateType> blockExitStates) {
       this.blockExitStates = blockExitStates;
     }
 
@@ -57,7 +58,7 @@
 
     @SuppressWarnings("unchecked")
     @Override
-    public SuccessfulDataflowAnalysisResult<StateType> asSuccessfulAnalysisResult() {
+    public SuccessfulDataflowAnalysisResult<Block, StateType> asSuccessfulAnalysisResult() {
       return this;
     }
   }
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraProceduralDataflowAnalysisBase.java b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraProceduralDataflowAnalysisBase.java
new file mode 100644
index 0000000..9b015d2
--- /dev/null
+++ b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraProceduralDataflowAnalysisBase.java
@@ -0,0 +1,147 @@
+// Copyright (c) 2022, 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.ir.analysis.framework.intraprocedural;
+
+import com.android.tools.r8.ir.analysis.framework.intraprocedural.DataflowAnalysisResult.FailedDataflowAnalysisResult;
+import com.android.tools.r8.ir.analysis.framework.intraprocedural.DataflowAnalysisResult.SuccessfulDataflowAnalysisResult;
+import com.android.tools.r8.utils.Timing;
+import com.android.tools.r8.utils.TraversalContinuation;
+import com.android.tools.r8.utils.WorkList;
+import java.util.IdentityHashMap;
+import java.util.Map;
+
+/**
+ * This defines a simple fixpoint solver for running an intraprocedural dataflow analysis.
+ *
+ * <p>The solver computes an {@link AbstractState} for each {@link Block} using the {@link
+ * AbstractTransferFunction} which defines the abstract semantics for each instruction.
+ *
+ * <p>Once the fixpoint is reached the analysis returns a {@link SuccessfulDataflowAnalysisResult}.
+ * If the supplied {@link AbstractTransferFunction} returns a {@link FailedTransferFunctionResult}
+ * for a given instruction and abstract state, then the analysis return a {@link
+ * FailedDataflowAnalysisResult}.
+ */
+public class IntraProceduralDataflowAnalysisBase<
+    Block, Instruction, StateType extends AbstractState<StateType>> {
+
+  final StateType bottom;
+
+  final ControlFlowGraph<Block, Instruction> cfg;
+
+  // The transfer function that defines the abstract semantics for each instruction.
+  final AbstractTransferFunction<Block, Instruction, StateType> transfer;
+
+  // The state of the analysis.
+  final Map<Block, StateType> blockExitStates = new IdentityHashMap<>();
+
+  // The entry states for each block that satisfies the predicate
+  // shouldCacheBlockEntryStateFor(block). These entry states can be computed from the exit states
+  // of the predecessors, but doing so can be expensive when a block has many predecessors.
+  final Map<Block, StateType> blockEntryStatesCache = new IdentityHashMap<>();
+
+  public IntraProceduralDataflowAnalysisBase(
+      StateType bottom,
+      ControlFlowGraph<Block, Instruction> cfg,
+      AbstractTransferFunction<Block, Instruction, StateType> transfer) {
+    this.bottom = bottom;
+    this.cfg = cfg;
+    this.transfer = transfer;
+  }
+
+  public DataflowAnalysisResult run(Block root) {
+    return run(root, Timing.empty());
+  }
+
+  public DataflowAnalysisResult run(Block root, Timing timing) {
+    return run(WorkList.newIdentityWorkList(root), timing);
+  }
+
+  private DataflowAnalysisResult run(WorkList<Block> worklist, Timing timing) {
+    while (worklist.hasNext()) {
+      Block initialBlock = worklist.next();
+      Block block = initialBlock;
+      Block end = null;
+      // Compute the abstract state upon entry to the basic block, by joining all the predecessor
+      // exit states.
+      StateType state =
+          timing.time("Compute block entry state", () -> computeBlockEntryState(initialBlock));
+
+      timing.begin("Compute transfers");
+      do {
+        TraversalContinuation<StateType> traversalContinuation =
+            cfg.traverseInstructions(
+                block,
+                (instruction, previousState) -> {
+                  TransferFunctionResult<StateType> transferResult =
+                      transfer.apply(instruction, previousState);
+                  if (transferResult.isFailedTransferResult()) {
+                    timing.end();
+                    return TraversalContinuation.doBreak();
+                  }
+                  assert transferResult.isAbstractState();
+                  return TraversalContinuation.doContinue(transferResult.asAbstractState());
+                },
+                state);
+        if (traversalContinuation.isBreak()) {
+          return new FailedDataflowAnalysisResult();
+        }
+        state = traversalContinuation.getValue();
+        if (cfg.hasUniqueSuccessorWithUniquePredecessor(block)) {
+          block = cfg.getUniqueSuccessor(block);
+        } else {
+          end = block;
+          block = null;
+        }
+      } while (block != null);
+      timing.end();
+
+      // Update the block exit state, and re-enqueue all successor blocks if the abstract state
+      // changed.
+      if (setBlockExitState(end, state)) {
+        worklist.addAllIgnoringSeenSet(cfg.getSuccessors(end));
+      }
+
+      // Add the computed exit state to the entry state of each successor that satisfies the
+      // predicate shouldCacheBlockEntryStateFor(successor).
+      updateBlockEntryStateCacheForSuccessors(end, state);
+    }
+    return new SuccessfulDataflowAnalysisResult<>(blockExitStates);
+  }
+
+  StateType computeBlockEntryState(Block block) {
+    if (shouldCacheBlockEntryStateFor(block)) {
+      return blockEntryStatesCache.getOrDefault(block, bottom);
+    }
+    StateType result = bottom;
+    for (Block predecessor : cfg.getPredecessors(block)) {
+      StateType edgeState =
+          transfer.computeBlockEntryState(
+              block, predecessor, blockExitStates.getOrDefault(predecessor, bottom));
+      result = result.join(edgeState);
+    }
+    return result;
+  }
+
+  boolean setBlockExitState(Block block, StateType state) {
+    assert !cfg.hasUniqueSuccessorWithUniquePredecessor(block);
+    StateType previous = blockExitStates.put(block, state);
+    assert previous == null || state.isGreaterThanOrEquals(previous);
+    return !state.equals(previous);
+  }
+
+  void updateBlockEntryStateCacheForSuccessors(Block block, StateType state) {
+    for (Block successor : cfg.getSuccessors(block)) {
+      if (shouldCacheBlockEntryStateFor(successor)) {
+        StateType edgeState = transfer.computeBlockEntryState(successor, block, state);
+        StateType previous = blockEntryStatesCache.getOrDefault(successor, bottom);
+        blockEntryStatesCache.put(successor, previous.join(edgeState));
+      }
+    }
+  }
+
+  boolean shouldCacheBlockEntryStateFor(Block block) {
+    return cfg.getPredecessors(block).size() > 2;
+  }
+}
diff --git a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraproceduralDataflowAnalysis.java b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraproceduralDataflowAnalysis.java
index 4f00ef0..13a9c42 100644
--- a/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraproceduralDataflowAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/analysis/framework/intraprocedural/IntraproceduralDataflowAnalysis.java
@@ -4,130 +4,17 @@
 
 package com.android.tools.r8.ir.analysis.framework.intraprocedural;
 
-import com.android.tools.r8.ir.analysis.framework.intraprocedural.DataflowAnalysisResult.FailedDataflowAnalysisResult;
-import com.android.tools.r8.ir.analysis.framework.intraprocedural.DataflowAnalysisResult.SuccessfulDataflowAnalysisResult;
 import com.android.tools.r8.ir.code.BasicBlock;
+import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
-import com.android.tools.r8.utils.Timing;
-import com.android.tools.r8.utils.WorkList;
-import java.util.IdentityHashMap;
-import java.util.Map;
 
-/**
- * This defines a simple fixpoint solver for running an intraprocedural dataflow analysis.
- *
- * <p>The solver computes an {@link AbstractState} for each {@link BasicBlock} using the {@link
- * AbstractTransferFunction} which defines the abstract semantics for each instruction.
- *
- * <p>Once the fixpoint is reached the analysis returns a {@link SuccessfulDataflowAnalysisResult}.
- * If the supplied {@link AbstractTransferFunction} returns a {@link FailedTransferFunctionResult}
- * for a given instruction and abstract state, then the analysis return a {@link
- * FailedDataflowAnalysisResult}.
- */
-public class IntraproceduralDataflowAnalysis<StateType extends AbstractState<StateType>> {
-
-  private final StateType bottom;
-
-  // The transfer function that defines the abstract semantics for each instruction.
-  private final AbstractTransferFunction<StateType> transfer;
-
-  // The state of the analysis.
-  private final Map<BasicBlock, StateType> blockExitStates = new IdentityHashMap<>();
-
-  // The entry states for each block that satisfies the predicate
-  // shouldCacheBlockEntryStateFor(block). These entry states can be computed from the exit states
-  // of the predecessors, but doing so can be expensive when a block has many predecessors.
-  private final Map<BasicBlock, StateType> blockEntryStatesCache = new IdentityHashMap<>();
+public class IntraproceduralDataflowAnalysis<StateType extends AbstractState<StateType>>
+    extends IntraProceduralDataflowAnalysisBase<BasicBlock, Instruction, StateType> {
 
   public IntraproceduralDataflowAnalysis(
-      StateType bottom, AbstractTransferFunction<StateType> transfer) {
-    this.bottom = bottom;
-    this.transfer = transfer;
-  }
-
-  public DataflowAnalysisResult run(BasicBlock root) {
-    return run(root, Timing.empty());
-  }
-
-  public DataflowAnalysisResult run(BasicBlock root, Timing timing) {
-    return run(WorkList.newIdentityWorkList(root), timing);
-  }
-
-  private DataflowAnalysisResult run(WorkList<BasicBlock> worklist, Timing timing) {
-    while (worklist.hasNext()) {
-      BasicBlock initialBlock = worklist.next();
-      BasicBlock block = initialBlock;
-      BasicBlock end = null;
-      // Compute the abstract state upon entry to the basic block, by joining all the predecessor
-      // exit states.
-      StateType state =
-          timing.time("Compute block entry state", () -> computeBlockEntryState(initialBlock));
-
-      timing.begin("Compute transfers");
-      do {
-        for (Instruction instruction : block.getInstructions()) {
-          TransferFunctionResult<StateType> transferResult = transfer.apply(instruction, state);
-          if (transferResult.isFailedTransferResult()) {
-            timing.end();
-            return new FailedDataflowAnalysisResult();
-          }
-          assert transferResult.isAbstractState();
-          state = transferResult.asAbstractState();
-        }
-        if (block.hasUniqueSuccessorWithUniquePredecessor()) {
-          block = block.getUniqueSuccessor();
-        } else {
-          end = block;
-          block = null;
-        }
-      } while (block != null);
-      timing.end();
-
-      // Update the block exit state, and re-enqueue all successor blocks if the abstract state
-      // changed.
-      if (setBlockExitState(end, state)) {
-        worklist.addAllIgnoringSeenSet(end.getSuccessors());
-      }
-
-      // Add the computed exit state to the entry state of each successor that satisfies the
-      // predicate shouldCacheBlockEntryStateFor(successor).
-      updateBlockEntryStateCacheForSuccessors(end, state);
-    }
-    return new SuccessfulDataflowAnalysisResult<>(blockExitStates);
-  }
-
-  private StateType computeBlockEntryState(BasicBlock block) {
-    if (shouldCacheBlockEntryStateFor(block)) {
-      return blockEntryStatesCache.getOrDefault(block, bottom);
-    }
-    StateType result = bottom;
-    for (BasicBlock predecessor : block.getPredecessors()) {
-      StateType edgeState =
-          transfer.computeBlockEntryState(
-              block, predecessor, blockExitStates.getOrDefault(predecessor, bottom));
-      result = result.join(edgeState);
-    }
-    return result;
-  }
-
-  private boolean setBlockExitState(BasicBlock block, StateType state) {
-    assert !block.hasUniqueSuccessorWithUniquePredecessor();
-    StateType previous = blockExitStates.put(block, state);
-    assert previous == null || state.isGreaterThanOrEquals(previous);
-    return !state.equals(previous);
-  }
-
-  private void updateBlockEntryStateCacheForSuccessors(BasicBlock block, StateType state) {
-    for (BasicBlock successor : block.getSuccessors()) {
-      if (shouldCacheBlockEntryStateFor(successor)) {
-        StateType edgeState = transfer.computeBlockEntryState(successor, block, state);
-        StateType previous = blockEntryStatesCache.getOrDefault(successor, bottom);
-        blockEntryStatesCache.put(successor, previous.join(edgeState));
-      }
-    }
-  }
-
-  private boolean shouldCacheBlockEntryStateFor(BasicBlock block) {
-    return block.getPredecessors().size() > 2;
+      StateType bottom,
+      IRCode code,
+      AbstractTransferFunction<BasicBlock, Instruction, StateType> transfer) {
+    super(bottom, code, transfer);
   }
 }
diff --git a/src/main/java/com/android/tools/r8/ir/code/IRCode.java b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
index 01b08a8..a67618b 100644
--- a/src/main/java/com/android/tools/r8/ir/code/IRCode.java
+++ b/src/main/java/com/android/tools/r8/ir/code/IRCode.java
@@ -16,6 +16,7 @@
 import com.android.tools.r8.graph.classmerging.MergedClassesCollection;
 import com.android.tools.r8.ir.analysis.TypeChecker;
 import com.android.tools.r8.ir.analysis.VerifyTypesHelper;
+import com.android.tools.r8.ir.analysis.framework.intraprocedural.ControlFlowGraph;
 import com.android.tools.r8.ir.analysis.type.ClassTypeElement;
 import com.android.tools.r8.ir.analysis.type.Nullability;
 import com.android.tools.r8.ir.analysis.type.TypeElement;
@@ -33,6 +34,7 @@
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.SetUtils;
 import com.android.tools.r8.utils.StringUtils;
+import com.android.tools.r8.utils.TraversalContinuation;
 import com.android.tools.r8.utils.WorkList;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -55,13 +57,14 @@
 import java.util.NoSuchElementException;
 import java.util.Queue;
 import java.util.Set;
+import java.util.function.BiFunction;
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
-public class IRCode implements ValueFactory {
+public class IRCode implements ControlFlowGraph<BasicBlock, Instruction>, ValueFactory {
 
   private static final int MAX_MARKING_COLOR = 0x40000000;
 
@@ -1363,6 +1366,29 @@
     return blocks;
   }
 
+  @Override
+  public Collection<BasicBlock> getPredecessors(BasicBlock block) {
+    return block.getPredecessors();
+  }
+
+  @Override
+  public Collection<BasicBlock> getSuccessors(BasicBlock block) {
+    return block.getSuccessors();
+  }
+
+  @Override
+  public <T> TraversalContinuation<T> traverseInstructions(
+      BasicBlock block, BiFunction<Instruction, T, TraversalContinuation<T>> fn, T initialValue) {
+    TraversalContinuation<T> traversalContinuation = TraversalContinuation.doContinue(initialValue);
+    for (Instruction instruction : block.getInstructions()) {
+      traversalContinuation = fn.apply(instruction, traversalContinuation.getValue());
+      if (traversalContinuation.shouldBreak()) {
+        return traversalContinuation;
+      }
+    }
+    return traversalContinuation;
+  }
+
   /**
    * Returns the set of blocks that are reachable from the given source. The source itself is only
    * included if there is a path from the given block to itself.
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/ClassInlinerMethodConstraintAnalysis.java b/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/ClassInlinerMethodConstraintAnalysis.java
index 854188e..f45344f 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/ClassInlinerMethodConstraintAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/ClassInlinerMethodConstraintAnalysis.java
@@ -26,8 +26,8 @@
     // Analyze code.
     IntraproceduralDataflowAnalysis<ParameterUsages> analysis =
         new IntraproceduralDataflowAnalysis<>(
-            ParameterUsages.bottom(), new TransferFunction(appView, method, code));
-    SuccessfulDataflowAnalysisResult<ParameterUsages> result =
+            ParameterUsages.bottom(), code, new TransferFunction(appView, method, code));
+    SuccessfulDataflowAnalysisResult<?, ParameterUsages> result =
         timing.time(
             "Data flow analysis",
             () -> analysis.run(code.entryBlock(), timing).asSuccessfulAnalysisResult());
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/TransferFunction.java b/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/TransferFunction.java
index 65c05e6..0e4a8dd 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/TransferFunction.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/classinliner/analysis/TransferFunction.java
@@ -52,7 +52,8 @@
 import com.google.common.collect.Sets;
 import java.util.Set;
 
-class TransferFunction implements AbstractTransferFunction<ParameterUsages> {
+class TransferFunction
+    implements AbstractTransferFunction<BasicBlock, Instruction, ParameterUsages> {
 
   private static final AliasedValueConfiguration aliasedValueConfiguration =
       AssumeAndCheckCastAliasedValueConfiguration.getInstance();
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderAppendFlowAnalysis.java b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderAppendFlowAnalysis.java
index 80066f3..1179c6d 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderAppendFlowAnalysis.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderAppendFlowAnalysis.java
@@ -12,6 +12,7 @@
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.IntraproceduralDataflowAnalysis;
 import com.android.tools.r8.ir.analysis.framework.intraprocedural.TransferFunctionResult;
 import com.android.tools.r8.ir.code.BasicBlock;
+import com.android.tools.r8.ir.code.IRCode;
 import com.android.tools.r8.ir.code.Instruction;
 import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.ir.code.InvokeVirtual;
@@ -39,10 +40,10 @@
    * loop.
    */
   static boolean hasAppendInstructionInLoop(
-      Value builder, StringBuilderOptimizationConfiguration configuration) {
+      IRCode code, Value builder, StringBuilderOptimizationConfiguration configuration) {
     IntraproceduralDataflowAnalysis<AbstractStateImpl> analysis =
         new IntraproceduralDataflowAnalysis<>(
-            AbstractStateImpl.bottom(), new TransferFunction(builder, configuration));
+            AbstractStateImpl.bottom(), code, new TransferFunction(builder, configuration));
     DataflowAnalysisResult result = analysis.run(builder.definition.getBlock());
     return result.isFailedAnalysisResult();
   }
@@ -121,7 +122,8 @@
    * <p>If a call to {@code toString()} on the builder i seen, then the abstract state is reset to
    * bottom.
    */
-  private static class TransferFunction implements AbstractTransferFunction<AbstractStateImpl> {
+  private static class TransferFunction
+      implements AbstractTransferFunction<BasicBlock, Instruction, AbstractStateImpl> {
 
     private final Value builder;
     private final StringBuilderOptimizationConfiguration configuration;
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
index b851f1e..7e0b884 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizer.java
@@ -722,7 +722,7 @@
         return null;
       }
       if (StringBuilderAppendFlowAnalysis.hasAppendInstructionInLoop(
-          builder, optimizationConfiguration)) {
+          code, builder, optimizationConfiguration)) {
         return null;
       }
       return StringUtils.join("", contents);