From 63c5cc1f525a66d1c14976dbd3fce3b7b0928a7f Mon Sep 17 00:00:00 2001
From: Stephan Schroevers <stephan.schroevers@teampicnic.com>
Date: Sat, 21 Sep 2024 16:04:05 -0700
Subject: [PATCH] Optimize VisitorState#getConstantExpression

By avoiding a second pass over the string.

Fixes #4586

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/google/error-prone/pull/4586 from PicnicSupermarket:sschroevers/optimize-getConstantExpression ff02145af6127f50940b5017ed1efd6183fc367e
PiperOrigin-RevId: 677291745
---
 .../com/google/errorprone/VisitorState.java   | 31 +++++++++----------
 .../RobolectricShadowDirectlyOn.java          |  2 +-
 .../flogger/FloggerArgumentToString.java      |  2 +-
 .../flogger/FloggerStringConcatenation.java   |  2 +-
 .../google/errorprone/VisitorStateTest.java   |  2 ++
 5 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/check_api/src/main/java/com/google/errorprone/VisitorState.java b/check_api/src/main/java/com/google/errorprone/VisitorState.java
index 1139bdc9c7a..21ce3e7307f 100644
--- a/check_api/src/main/java/com/google/errorprone/VisitorState.java
+++ b/check_api/src/main/java/com/google/errorprone/VisitorState.java
@@ -51,6 +51,7 @@
 import com.sun.tools.javac.tree.JCTree.JCCompilationUnit;
 import com.sun.tools.javac.tree.TreeMaker;
 import com.sun.tools.javac.util.Context;
+import com.sun.tools.javac.util.Convert;
 import com.sun.tools.javac.util.Name;
 import com.sun.tools.javac.util.Names;
 import com.sun.tools.javac.util.Options;
@@ -770,24 +771,20 @@ private static final class SharedState {
    * Like {@link Elements#getConstantExpression}, but doesn't over-escape single quotes in strings.
    */
   public String getConstantExpression(Object value) {
-    String escaped = getElements().getConstantExpression(value);
-    if (value instanceof String) {
-      // Don't escape single-quotes in string literals
-      StringBuilder sb = new StringBuilder();
-      for (int i = 0; i < escaped.length(); i++) {
-        char c = escaped.charAt(i);
-        if (c == '\\' && i + 1 < escaped.length()) {
-          char next = escaped.charAt(++i);
-          if (next != '\'') {
-            sb.append(c);
-          }
-          sb.append(next);
-        } else {
-          sb.append(c);
-        }
+    if (!(value instanceof CharSequence str)) {
+      return getElements().getConstantExpression(value);
+    }
+
+    // Don't escape single-quotes in string literals.
+    StringBuilder sb = new StringBuilder("\"");
+    for (int i = 0; i < str.length(); i++) {
+      char c = str.charAt(i);
+      if (c == '\'') {
+        sb.append('\'');
+      } else {
+        sb.append(Convert.quote(c));
       }
-      return sb.toString();
     }
-    return escaped;
+    return sb.append('"').toString();
   }
 }
diff --git a/core/src/main/java/com/google/errorprone/bugpatterns/RobolectricShadowDirectlyOn.java b/core/src/main/java/com/google/errorprone/bugpatterns/RobolectricShadowDirectlyOn.java
index 5e9bada9f6c..7b3f264b42d 100644
--- a/core/src/main/java/com/google/errorprone/bugpatterns/RobolectricShadowDirectlyOn.java
+++ b/core/src/main/java/com/google/errorprone/bugpatterns/RobolectricShadowDirectlyOn.java
@@ -73,7 +73,7 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
     MethodSymbol symbol = getSymbol(parent);
     String argReplacement =
         Streams.concat(
-                Stream.of(state.getConstantExpression(symbol.getSimpleName().toString())),
+                Stream.of(state.getConstantExpression(symbol.getSimpleName())),
                 Streams.zip(
                     symbol.getParameters().stream(),
                     parent.getArguments().stream(),
diff --git a/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerArgumentToString.java b/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerArgumentToString.java
index 62c00eb1b46..dc0d70184dd 100644
--- a/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerArgumentToString.java
+++ b/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerArgumentToString.java
@@ -351,7 +351,7 @@ Description unwrapArguments(
     if (!fixed) {
       return NO_MATCH;
     }
-    fix.replace(tree.getArguments().get(0), state.getConstantExpression(sb.toString()));
+    fix.replace(tree.getArguments().get(0), state.getConstantExpression(sb));
     return describeMatch(tree, fix.build());
   }
 
diff --git a/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerStringConcatenation.java b/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerStringConcatenation.java
index ddf535ec835..0b5ea9a0f74 100644
--- a/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerStringConcatenation.java
+++ b/core/src/main/java/com/google/errorprone/bugpatterns/flogger/FloggerStringConcatenation.java
@@ -112,7 +112,7 @@ protected Void defaultAction(Tree tree, Void unused) {
         tree,
         SuggestedFix.replace(
             argument,
-            state.getConstantExpression(formatString.toString())
+            state.getConstantExpression(formatString)
                 + ", "
                 + formatArguments.stream().map(state::getSourceForNode).collect(joining(", "))));
   }
diff --git a/core/src/test/java/com/google/errorprone/VisitorStateTest.java b/core/src/test/java/com/google/errorprone/VisitorStateTest.java
index b0cde5a0ca8..579f92c04de 100644
--- a/core/src/test/java/com/google/errorprone/VisitorStateTest.java
+++ b/core/src/test/java/com/google/errorprone/VisitorStateTest.java
@@ -106,6 +106,8 @@ public void getConstantExpression() {
     assertThat(visitorState.getConstantExpression("hello \n world"))
         .isEqualTo("\"hello \\n world\"");
     assertThat(visitorState.getConstantExpression('\'')).isEqualTo("'\\''");
+    assertThat(visitorState.getConstantExpression(new StringBuilder("hello ' world")))
+        .isEqualTo("\"hello ' world\"");
   }
 
   // The following is taken from ErrorProneJavacPluginTest. There may be an easier way.