Extend StringBuilder optimization to non-string type arguments.

Bug: 114002137
Change-Id: I1c9d355e98d4e63fce184b839ccdb9b1258a3edf
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 3ba841d..0b445a6 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
@@ -16,6 +16,7 @@
 import com.android.tools.r8.ir.analysis.type.TypeLatticeElement;
 import com.android.tools.r8.ir.code.BasicBlock;
 import com.android.tools.r8.ir.code.BasicBlock.ThrowingInfo;
+import com.android.tools.r8.ir.code.ConstNumber;
 import com.android.tools.r8.ir.code.ConstString;
 import com.android.tools.r8.ir.code.DominatorTree;
 import com.android.tools.r8.ir.code.DominatorTree.Assumption;
@@ -26,7 +27,9 @@
 import com.android.tools.r8.ir.code.InvokeMethod;
 import com.android.tools.r8.ir.code.InvokeMethodWithReceiver;
 import com.android.tools.r8.ir.code.InvokeVirtual;
+import com.android.tools.r8.ir.code.NumberConversion;
 import com.android.tools.r8.ir.code.Value;
+import com.android.tools.r8.ir.code.ValueType;
 import com.android.tools.r8.logging.Log;
 import com.android.tools.r8.utils.StringUtils;
 import com.google.common.annotations.VisibleForTesting;
@@ -85,7 +88,7 @@
   private int numberOfBuildersThatEscape = 0;
   private int numberOfBuildersWhoseResultIsInterned = 0;
   private int numberOfBuildersWithNonTrivialStateChange = 0;
-  private int numberOfBuildersWithNonStringArg = 0;
+  private int numberOfBuildersWithUnsupportedArg = 0;
   private int numberOfBuildersWithMergingPoints = 0;
   private int numberOfBuildersWithNonDeterministicArg = 0;
   private int numberOfDeadBuilders = 0;
@@ -126,7 +129,7 @@
     Log.info(getClass(),
         "# builders w/ non-trivial state change: %s", numberOfBuildersWithNonTrivialStateChange);
     Log.info(getClass(),
-        "# builders w/ non-string arg: %s", numberOfBuildersWithNonStringArg);
+        "# builders w/ unsupported arg: %s", numberOfBuildersWithUnsupportedArg);
     Log.info(getClass(),
         "# builders w/ merging points: %s", numberOfBuildersWithMergingPoints);
     Log.info(getClass(),
@@ -354,8 +357,9 @@
               continue;
             }
             assert invoke.inValues().size() == 2;
-            Value arg = invoke.inValues().get(1);
-            String addition = extractConstantArgument(arg);
+            Value arg = invoke.inValues().get(1).getAliasedValue();
+            DexType argType = invoke.getInvokedMethod().proto.parameters.values[0];
+            String addition = extractConstantArgument(arg, argType);
             Map<Instruction, BuilderState> perInstrState = getBuilderState(builder);
             BuilderState dominantState = findDominantState(dominatorTree, perInstrState, instr);
             if (dominantState != null) {
@@ -380,8 +384,9 @@
             if (!candidateBuilders.contains(builder)) {
               continue;
             }
-            Value arg = invoke.inValues().get(1);
-            String addition = extractConstantArgument(arg);
+            Value arg = invoke.inValues().get(1).getAliasedValue();
+            DexType argType = invoke.getInvokedMethod().proto.parameters.values[0];
+            String addition = extractConstantArgument(arg, argType);
             Map<Instruction, BuilderState> perInstrState = getBuilderState(builder);
             BuilderState dominantState = findDominantState(dominatorTree, perInstrState, instr);
             if (dominantState != null) {
@@ -414,17 +419,86 @@
       return this;
     }
 
-    private String extractConstantArgument(Value arg) {
+    private String extractConstantArgument(Value arg, DexType argType) {
       String addition = ANY_STRING;
       if (!arg.isPhi()) {
-        // TODO(b/114002137): Improve arg extraction and type conversion.
         if (arg.definition.isConstString()) {
           addition = arg.definition.asConstString().getValue().toString();
+        } else if (arg.definition.isConstNumber() || arg.definition.isNumberConversion()) {
+          Number number = extractConstantNumber(arg);
+          if (number == null) {
+            return addition;
+          }
+          if (arg.getTypeLattice().isPrimitive()) {
+            if (argType == factory.booleanType) {
+              addition = String.valueOf(number.intValue() != 0);
+            } else if (argType == factory.byteType) {
+              addition = String.valueOf(number.byteValue());
+            } else if (argType == factory.shortType) {
+              addition = String.valueOf(number.shortValue());
+            } else if (argType == factory.charType) {
+              addition = String.valueOf((char) number.intValue());
+            } else if (argType == factory.intType) {
+              addition = String.valueOf(number.intValue());
+            } else if (argType == factory.longType) {
+              addition = String.valueOf(number.longValue());
+            } else if (argType == factory.floatType) {
+              addition = String.valueOf(number.floatValue());
+            } else if (argType == factory.doubleType) {
+              addition = String.valueOf(number.doubleValue());
+            }
+          } else if (arg.getTypeLattice().isNullType()) {
+            assert number.intValue() == 0;
+            addition = "null";
+          }
         }
       }
       return addition;
     }
 
+    private Number extractConstantNumber(Value arg) {
+      assert !arg.isPhi();
+      if (arg.definition.isConstNumber()) {
+        ConstNumber cst = arg.definition.asConstNumber();
+        if (cst.outType() == ValueType.LONG) {
+          return cst.getLongValue();
+        } else if (cst.outType() == ValueType.FLOAT) {
+          return cst.getFloatValue();
+        } else if (cst.outType() == ValueType.DOUBLE) {
+          return cst.getDoubleValue();
+        } else {
+          assert cst.outType() == ValueType.INT || cst.outType() == ValueType.OBJECT;
+          return cst.getIntValue();
+        }
+      } else if (arg.definition.isNumberConversion()) {
+        NumberConversion conversion = arg.definition.asNumberConversion();
+        assert conversion.inValues().size() == 1;
+        Number temp = extractConstantNumber(conversion.inValues().get(0));
+        if (temp == null) {
+          return null;
+        }
+        DexType conversionType = conversion.to.dexTypeFor(factory);
+        if (conversionType == factory.booleanType) {
+          return temp.intValue() != 0 ? 1 : 0;
+        } else if (conversionType == factory.byteType) {
+          return temp.byteValue();
+        } else if (conversionType == factory.shortType) {
+          return temp.shortValue();
+        } else if (conversionType == factory.charType) {
+          return temp.intValue();
+        } else if (conversionType == factory.intType) {
+          return temp.intValue();
+        } else if (conversionType == factory.longType) {
+          return temp.longValue();
+        } else if (conversionType == factory.floatType) {
+          return temp.floatValue();
+        } else if (conversionType == factory.doubleType) {
+          return temp.doubleValue();
+        }
+      }
+      return null;
+    }
+
     private BuilderState findDominantState(
         DominatorTree dominatorTree,
         Map<Instruction, BuilderState> perInstrState,
@@ -690,11 +764,15 @@
         numberOfBuildersWithNonTrivialStateChange++;
         return false;
       }
-      for (DexType argType : invokedMethod.proto.parameters.values) {
-        if (!canHandleArgumentType(argType)) {
-          numberOfBuildersWithNonStringArg++;
-          return false;
-        }
+      assert invoke.inValues().size() == 2;
+      TypeLatticeElement argType = invoke.inValues().get(1).getTypeLattice();
+      if (!argType.isPrimitive() && !argType.isClassType() && !argType.isNullType()) {
+        numberOfBuildersWithUnsupportedArg++;
+        return false;
+      }
+      if (argType.isClassType()) {
+        DexType argClassType = argType.asClassTypeLatticeElement().getClassType();
+        return canHandleArgumentType(argClassType);
       }
       return true;
     }
@@ -707,8 +785,6 @@
 
     private boolean canHandleArgumentType(DexType argType) {
       // TODO(b/113859361): passed to another builder should be an eligible case.
-      // TODO(b/114002137): Improve arg extraction and type conversion.
-      //   For now, skip any append(arg) that receives non-string types.
       return argType == factory.stringType || argType == factory.charSequenceType;
     }
   }
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 b21941a..c434aac 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
@@ -16,6 +16,7 @@
 import com.android.tools.r8.ir.optimize.string.StringBuilderOptimizer.BuilderState;
 import java.util.Map;
 import java.util.function.Consumer;
+import org.junit.Ignore;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.Parameterized;
@@ -85,9 +86,13 @@
     buildAndCheckIR(
         "nonStringArgs",
         checkOptimizerStates(appView, optimizer -> {
-          // TODO(b/114002137): Improve arg extraction and type conversion.
-          assertEquals(0, optimizer.analysis.builderStates.size());
-          assertEquals(0, optimizer.analysis.simplifiedBuilders.size());
+          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());
         }));
   }
 
@@ -96,12 +101,17 @@
     buildAndCheckIR(
         "typeConversion",
         checkOptimizerStates(appView, optimizer -> {
-          // TODO(b/114002137): Improve arg extraction and type conversion.
-          assertEquals(0, optimizer.analysis.builderStates.size());
-          assertEquals(0, optimizer.analysis.simplifiedBuilders.size());
+          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, "0.14 0 false null", true);
+          }
+          assertEquals(1, optimizer.analysis.simplifiedBuilders.size());
         }));
   }
 
+  @Ignore("TODO(b/113859361): passed to another builder should be an eligible case.")
   @Test
   public void testNestedBuilders_appendBuilderItself() throws Exception {
     buildAndCheckIR(
@@ -139,7 +149,7 @@
     buildAndCheckIR(
         "simplePhi",
         checkOptimizerStates(appView, optimizer -> {
-          // TODO(b/114002137): Improve arg extraction and type conversion.
+          assertEquals(0, optimizer.analysis.builderStates.size());
           assertEquals(0, optimizer.analysis.simplifiedBuilders.size());
         }));
   }
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 3f726f9..86cd835 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
@@ -69,6 +69,8 @@
       int expectedStringCountInTrivialSequence,
       int expectedStringCountInBuilderWithInitialValue,
       int expectedStringCountInBuilderWithCapacity,
+      int expectedStringCountInNonStringArgs,
+      int expectedStringCountInTypeConversion,
       int expectedStringCountInNestedBuilderAppendItself,
       int expectedStringCountInNestedBuilderAppendResult)
       throws Exception {
@@ -97,13 +99,13 @@
     assertThat(method, isPresent());
     count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(0, count);
+    assertEquals(expectedStringCountInNonStringArgs, count);
 
     method = mainClass.uniqueMethodWithName("typeConversion");
     assertThat(method, isPresent());
     count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(0, count);
+    assertEquals(expectedStringCountInTypeConversion, count);
 
     method = mainClass.uniqueMethodWithName("nestedBuilders_appendBuilderItself");
     assertThat(method, isPresent());
@@ -121,7 +123,7 @@
     assertThat(method, isPresent());
     count = Streams.stream(method.iterateInstructions(
         i -> i.isConstString(JumboStringMode.ALLOW))).count();
-    assertEquals(5, count);
+    assertEquals(4, count);
 
     method = mainClass.uniqueMethodWithName("phiAtInit");
     assertThat(method, isPresent());
@@ -160,7 +162,7 @@
             .addOptionsModification(this::configure)
             .run(parameters.getRuntime(), MAIN)
             .assertSuccessWithOutput(JAVA_OUTPUT);
-    test(result, 3, 3, 2, 4, 4);
+    test(result, 3, 3, 0, 0, 0, 3, 3);
 
     result =
         testForD8()
@@ -172,7 +174,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, 3, 2, 4, 3);
+    test(result, 3, 3, 0, 0, 0, 3, 3);
   }
 
   @Test
@@ -189,7 +191,7 @@
             .addOptionsModification(this::configure)
             .run(parameters.getRuntime(), MAIN)
             .assertSuccessWithOutput(JAVA_OUTPUT);
-    test(result, 1, 1, 1, 4, 3);
+    test(result, 1, 1, 1, 1, 1, 3, 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 da409fc..19c7129 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
@@ -25,9 +25,8 @@
   @NeverInline
   public static void builderWithCapacity() {
     StringBuilder builder = new StringBuilder(3);
-    // TODO(b/114002137): switch to use the integer.
-    builder.append("4");
-    builder.append("2");
+    builder.append(4);
+    builder.append(2);
     System.out.println(builder.toString());
   }
 
@@ -64,8 +63,7 @@
     b1.append(",");
     StringBuilder b2 = new StringBuilder();
     b2.append("R");
-    // TODO(b/114002137): switch to use the integer.
-    b2.append("8");
+    b2.append(8);
     b1.append(b2);
     System.out.println(b1.toString());
   }
@@ -77,8 +75,7 @@
     b1.append(",");
     StringBuilder b2 = new StringBuilder();
     b2.append("R");
-    // TODO(b/114002137): switch to use the integer.
-    b2.append("8");
+    b2.append(8);
     b1.append(b2.toString());
     System.out.println(b1.toString());
   }
@@ -95,8 +92,7 @@
     } else {
       builder.append("R");
     }
-    // TODO(b/114002137): switch to use the integer.
-    builder.append("8");
+    builder.append(8);
     System.out.println(builder.toString());
   }