Distinguish StringBuilder init with initial values v.s. capacity.

Bug: 114002137
Change-Id: If6a119189db9c95278b1467901ecf47e2b59423d
diff --git a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizationConfiguration.java b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizationConfiguration.java
index 994b495..d1612c3 100644
--- a/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizationConfiguration.java
+++ b/src/main/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizationConfiguration.java
@@ -14,6 +14,8 @@
 
   boolean isBuilderInit(DexMethod method);
 
+  boolean isBuilderInitWithInitialValue(InvokeMethod invoke);
+
   boolean isAppendMethod(DexMethod method);
 
   boolean isSupportedAppendMethod(InvokeMethod invoke);
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 7f2faef..e4446c6 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
@@ -367,27 +367,23 @@
           }
 
           if (instr.isInvokeDirect()
-              && optimizationConfiguration.isBuilderInit(
-                  instr.asInvokeDirect().getInvokedMethod())) {
+              && optimizationConfiguration.isBuilderInitWithInitialValue(instr.asInvokeDirect())) {
             InvokeDirect invoke = instr.asInvokeDirect();
             Value builder = invoke.getReceiver();
             if (!candidateBuilders.contains(builder)) {
               continue;
             }
-            // builder initialization with the initial content.
-            if (invoke.inValues().size() > 1) {
-              assert invoke.inValues().size() == 2;
-              Value arg = invoke.inValues().get(1);
-              String addition = extractConstantArgument(arg);
-              Map<Instruction, BuilderState> perInstrState = getBuilderState(builder);
-              BuilderState dominantState = findDominantState(dominatorTree, perInstrState, instr);
-              if (dominantState != null) {
-                BuilderState currentState = dominantState.createChild(addition);
-                perInstrState.put(instr, currentState);
-              } else {
-                // TODO(b/114002137): if we want to utilize partial results, don't remove it here.
-                candidateBuilders.remove(builder);
-              }
+            assert invoke.inValues().size() == 2;
+            Value arg = invoke.inValues().get(1);
+            String addition = extractConstantArgument(arg);
+            Map<Instruction, BuilderState> perInstrState = getBuilderState(builder);
+            BuilderState dominantState = findDominantState(dominatorTree, perInstrState, instr);
+            if (dominantState != null) {
+              BuilderState currentState = dominantState.createChild(addition);
+              perInstrState.put(instr, currentState);
+            } else {
+              // TODO(b/114002137): if we want to utilize partial results, don't remove it here.
+              candidateBuilders.remove(builder);
             }
             continue;
           }
@@ -689,6 +685,13 @@
     }
 
     @Override
+    public boolean isBuilderInitWithInitialValue(InvokeMethod invoke) {
+      return isBuilderInit(invoke.getInvokedMethod())
+          && invoke.inValues().size() == 2
+          && !invoke.inValues().get(1).getTypeLattice().isPrimitive();
+    }
+
+    @Override
     public boolean isAppendMethod(DexMethod method) {
       return factory.stringBuilderMethods.isAppendMethod(method)
           || factory.stringBufferMethods.isAppendMethod(method);
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizerAnalysisTest.java b/src/test/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizerAnalysisTest.java
index 371aeb0..a77978c 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizerAnalysisTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/string/StringBuilderOptimizerAnalysisTest.java
@@ -54,6 +54,36 @@
   }
 
   @Test
+  public void testBuilderWithInitialValue() throws Exception {
+    buildAndCheckIR(
+        "builderWithInitialValue",
+        checkOptimizerStates(appView, optimizer -> {
+          assertEquals(1, optimizer.analysis.builderStates.size());
+          for (Value builder : optimizer.analysis.builderStates.keySet()) {
+            Map<Instruction, BuilderState> perBuilderState =
+                optimizer.analysis.builderStates.get(builder);
+            checkBuilderState(optimizer, perBuilderState, "Hello,R8", true);
+          }
+          assertEquals(1, optimizer.analysis.simplifiedBuilders.size());
+        }));
+  }
+
+  @Test
+  public void testBuilderWithCapacity() throws Exception {
+    buildAndCheckIR(
+        "builderWithCapacity",
+        checkOptimizerStates(appView, optimizer -> {
+          assertEquals(1, optimizer.analysis.builderStates.size());
+          for (Value builder : optimizer.analysis.builderStates.keySet()) {
+            Map<Instruction, BuilderState> perBuilderState =
+                optimizer.analysis.builderStates.get(builder);
+            checkBuilderState(optimizer, perBuilderState, "42", true);
+          }
+          assertEquals(1, optimizer.analysis.simplifiedBuilders.size());
+        }));
+  }
+
+  @Test
   public void testNonStringArgs() throws Exception {
     buildAndCheckIR(
         "nonStringArgs",
@@ -242,6 +272,13 @@
     }
 
     @Override
+    public boolean isBuilderInitWithInitialValue(InvokeMethod invoke) {
+      return isBuilderInit(invoke.getInvokedMethod())
+          && invoke.inValues().size() == 2
+          && !invoke.inValues().get(1).getTypeLattice().isPrimitive();
+    }
+
+    @Override
     public boolean isAppendMethod(DexMethod method) {
       return isBuilderType(method.holder) && method.name.toString().equals("append");
     }
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTest.java b/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTest.java
index 1e4e3d7..3f726f9 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTest.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTest.java
@@ -30,6 +30,8 @@
   private static final Class<?> MAIN = StringConcatenationTestClass.class;
   private static final String JAVA_OUTPUT = StringUtils.lines(
       "xyz",
+      "Hello,R8",
+      "42",
       "42",
       "0.14 0 false null",
       "Hello,R8",
@@ -64,9 +66,11 @@
 
   private void test(
       TestRunResult result,
-      int expectedStringCount1,
-      int expectedStringCount2,
-      int expectedStringCount3)
+      int expectedStringCountInTrivialSequence,
+      int expectedStringCountInBuilderWithInitialValue,
+      int expectedStringCountInBuilderWithCapacity,
+      int expectedStringCountInNestedBuilderAppendItself,
+      int expectedStringCountInNestedBuilderAppendResult)
       throws Exception {
     CodeInspector codeInspector = result.inspector();
     ClassSubject mainClass = codeInspector.clazz(MAIN);
@@ -75,7 +79,19 @@
     assertThat(method, isPresent());
     long count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(expectedStringCount1, count);
+    assertEquals(expectedStringCountInTrivialSequence, count);
+
+    method = mainClass.uniqueMethodWithName("builderWithInitialValue");
+    assertThat(method, isPresent());
+    count = Streams.stream(method.iterateInstructions(
+        i -> i.isConstString(JumboStringMode.ALLOW))).count();
+    assertEquals(expectedStringCountInBuilderWithInitialValue, count);
+
+    method = mainClass.uniqueMethodWithName("builderWithCapacity");
+    assertThat(method, isPresent());
+    count = Streams.stream(method.iterateInstructions(
+        i -> i.isConstString(JumboStringMode.ALLOW))).count();
+    assertEquals(expectedStringCountInBuilderWithCapacity, count);
 
     method = mainClass.uniqueMethodWithName("nonStringArgs");
     assertThat(method, isPresent());
@@ -93,13 +109,13 @@
     assertThat(method, isPresent());
     count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(expectedStringCount2, count);
+    assertEquals(expectedStringCountInNestedBuilderAppendItself, count);
 
     method = mainClass.uniqueMethodWithName("nestedBuilders_appendBuilderResult");
     assertThat(method, isPresent());
     count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(expectedStringCount3, count);
+    assertEquals(expectedStringCountInNestedBuilderAppendResult, count);
 
     method = mainClass.uniqueMethodWithName("simplePhi");
     assertThat(method, isPresent());
@@ -144,7 +160,7 @@
             .addOptionsModification(this::configure)
             .run(parameters.getRuntime(), MAIN)
             .assertSuccessWithOutput(JAVA_OUTPUT);
-    test(result, 3, 4, 4);
+    test(result, 3, 3, 2, 4, 4);
 
     result =
         testForD8()
@@ -156,7 +172,7 @@
             .assertSuccessWithOutput(JAVA_OUTPUT);
     // TODO(b/114002137): The lack of subtyping made the escape analysis to regard
     //    StringBuilder#toString as an alias-introducing instruction.
-    test(result, 3, 4, 3);
+    test(result, 3, 3, 2, 4, 3);
   }
 
   @Test
@@ -173,7 +189,7 @@
             .addOptionsModification(this::configure)
             .run(parameters.getRuntime(), MAIN)
             .assertSuccessWithOutput(JAVA_OUTPUT);
-    test(result, 1, 4, 3);
+    test(result, 1, 1, 1, 4, 3);
   }
 
   // TODO(b/114002137): Once enabled, remove this test-specific setting.
diff --git a/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTestClass.java b/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTestClass.java
index 3caff25..da409fc 100644
--- a/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTestClass.java
+++ b/src/test/java/com/android/tools/r8/ir/optimize/string/StringConcatenationTestClass.java
@@ -15,6 +15,23 @@
   }
 
   @NeverInline
+  public static void builderWithInitialValue() {
+    StringBuilder builder = new StringBuilder("Hello");
+    builder.append(",");
+    builder.append("R8");
+    System.out.println(builder.toString());
+  }
+
+  @NeverInline
+  public static void builderWithCapacity() {
+    StringBuilder builder = new StringBuilder(3);
+    // TODO(b/114002137): switch to use the integer.
+    builder.append("4");
+    builder.append("2");
+    System.out.println(builder.toString());
+  }
+
+  @NeverInline
   public static void nonStringArgs() {
     StringBuilder builder = new StringBuilder();
     builder.append(4);
@@ -121,6 +138,8 @@
 
   public static void main(String[] args) {
     trivialSequence();
+    builderWithInitialValue();
+    builderWithCapacity();
     nonStringArgs();
     typeConversion();
     nestedBuilders_appendBuilderItself();