Improve support for string templates

The initial implementation passed through the entire string unmodified, this allows formatting the Java expressions inside the `\{...}`. See #1010 Co-authored-by: butterunderflow <azhong.934@gmail.com> PiperOrigin-RevId: 592940163
google · Dec 21, 2023 · 38de9c4 · 38de9c4
1 parent 8afdfca
commit 38de9c4
Show file tree

Hide file tree

Showing 6 changed files with 103 additions and 29 deletions.
diff --git a/core/src/main/java/com/google/googlejavaformat/java/JavacTokens.java b/core/src/main/java/com/google/googlejavaformat/java/JavacTokens.java
@@ -28,6 +28,11 @@
 import com.sun.tools.javac.parser.Tokens.TokenKind;
 import com.sun.tools.javac.parser.UnicodeReader;
 import com.sun.tools.javac.util.Context;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Objects;
 import java.util.Set;
 
@@ -83,22 +88,53 @@ static boolean isStringFragment(TokenKind kind) {
     return STRINGFRAGMENT != null && Objects.equals(kind, STRINGFRAGMENT);
   }
 
-  /** Lex the input and return a list of {@link RawTok}s. */
-  public static ImmutableList<RawTok> getTokens(
-      String source, Context context, Set<TokenKind> stopTokens) {
+  private static ImmutableList<Token> readAllTokens(
+      String source, Context context, Set<Integer> nonTerminalStringFragments) {
     if (source == null) {
       return ImmutableList.of();
     }
     ScannerFactory fac = ScannerFactory.instance(context);
     char[] buffer = (source + EOF_COMMENT).toCharArray();
     Scanner scanner =
         new AccessibleScanner(fac, new CommentSavingTokenizer(fac, buffer, buffer.length));
+    List<Token> tokens = new ArrayList<>();
+    do {
+      scanner.nextToken();
+      tokens.add(scanner.token());
+    } while (scanner.token().kind != TokenKind.EOF);
+    for (int i = 0; i < tokens.size(); i++) {
+      if (isStringFragment(tokens.get(i).kind)) {
+        int start = i;
+        while (isStringFragment(tokens.get(i).kind)) {
+          i++;
+        }
+        for (int j = start; j < i - 1; j++) {
+          nonTerminalStringFragments.add(tokens.get(j).pos);
+        }
+      }
+    }
+    // A string template is tokenized as a series of STRINGFRAGMENT tokens containing the string
+    // literal values, followed by the tokens for the template arguments. For the formatter, we
+    // want the stream of tokens to appear in order by their start position.
+    if (Runtime.version().feature() >= 21) {
+      Collections.sort(tokens, Comparator.comparingInt(t -> t.pos));
+    }
+    return ImmutableList.copyOf(tokens);
+  }
+
+  /** Lex the input and return a list of {@link RawTok}s. */
+  public static ImmutableList<RawTok> getTokens(
+      String source, Context context, Set<TokenKind> stopTokens) {
+    if (source == null) {
+      return ImmutableList.of();
+    }
+    Set<Integer> nonTerminalStringFragments = new HashSet<>();
+    ImmutableList<Token> javacTokens = readAllTokens(source, context, nonTerminalStringFragments);
+
     ImmutableList.Builder<RawTok> tokens = ImmutableList.builder();
     int end = source.length();
     int last = 0;
-    do {
-      scanner.nextToken();
-      Token t = scanner.token();
+    for (Token t : javacTokens) {
       if (t.comments != null) {
         for (Comment c : Lists.reverse(t.comments)) {
           if (last < c.getSourcePos(0)) {
@@ -118,27 +154,12 @@ public static ImmutableList<RawTok> getTokens(
       if (last < t.pos) {
         tokens.add(new RawTok(null, null, last, t.pos));
       }
-      int pos = t.pos;
-      int endPos = t.endPos;
       if (isStringFragment(t.kind)) {
-        // A string template is tokenized as a series of STRINGFRAGMENT tokens containing the string
-        // literal values, followed by the tokens for the template arguments. For the formatter, we
-        // want the stream of tokens to appear in order by their start position, and also to have
-        // all the content from the original source text (including leading and trailing ", and the
-        // \ escapes from template arguments). This logic processes the token stream from javac to
-        // meet those requirements.
-        while (isStringFragment(t.kind)) {
-          endPos = t.endPos;
-          scanner.nextToken();
-          t = scanner.token();
-        }
-        // Read tokens for the string template arguments, until we read the end of the string
-        // template. The last token in a string template is always a trailing string fragment. Use
-        // lookahead to defer reading the token after the template until the next iteration of the
-        // outer loop.
-        while (scanner.token(/* lookahead= */ 1).endPos < endPos) {
-          scanner.nextToken();
-          t = scanner.token();
+        int endPos = t.endPos;
+        int pos = t.pos;
+        if (nonTerminalStringFragments.contains(t.pos)) {
+          // Include the \ escape from \{...} in the preceding string fragment
+          endPos++;
         }
         tokens.add(new RawTok(source.substring(pos, endPos), t.kind, pos, endPos));
         last = endPos;
@@ -151,7 +172,7 @@ public static ImmutableList<RawTok> getTokens(
                 t.endPos));
         last = t.endPos;
       }
-    } while (scanner.token().kind != TokenKind.EOF);
+    }
     if (last < end) {
       tokens.add(new RawTok(null, null, last, end));
     }

diff --git a/core/src/main/java/com/google/googlejavaformat/java/java21/Java21InputAstVisitor.java b/core/src/main/java/com/google/googlejavaformat/java/java21/Java21InputAstVisitor.java
@@ -82,11 +82,20 @@ public Void visitDeconstructionPattern(DeconstructionPatternTree node, Void unus
 
   @SuppressWarnings("preview")
   @Override
-  public Void visitStringTemplate(StringTemplateTree node, Void aVoid) {
+  public Void visitStringTemplate(StringTemplateTree node, Void unused) {
     sync(node);
+    builder.open(plusFour);
     scan(node.getProcessor(), null);
     token(".");
     token(builder.peekToken().get());
+    for (int i = 0; i < node.getFragments().size() - 1; i++) {
+      token("{");
+      builder.breakOp();
+      scan(node.getExpressions().get(i), null);
+      token("}");
+      token(builder.peekToken().get());
+    }
+    builder.close();
     return null;
   }
 

diff --git a/core/src/test/java/com/google/googlejavaformat/java/FormatterIntegrationTest.java b/core/src/test/java/com/google/googlejavaformat/java/FormatterIntegrationTest.java
@@ -60,7 +60,8 @@ public class FormatterIntegrationTest {
               "SwitchUnderscore",
               "I880",
               "Unnamed",
-              "I981")
+              "I981",
+              "StringTemplate")
           .build();
 
   @Parameters(name = "{index}: {0}")

diff --git a/core/src/test/java/com/google/googlejavaformat/java/FormatterTest.java b/core/src/test/java/com/google/googlejavaformat/java/FormatterTest.java
@@ -18,6 +18,7 @@
 import static com.google.common.truth.Truth.assertWithMessage;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.junit.Assert.fail;
+import static org.junit.Assume.assumeTrue;
 
 import com.google.common.base.Joiner;
 import com.google.common.io.CharStreams;
@@ -492,4 +493,27 @@ public void removeTrailingTabsInComments() throws Exception {
                 + "  }\n"
                 + "}\n");
   }
+
+  @Test
+  public void stringTemplateTests() throws Exception {
+    assumeTrue(Runtime.version().feature() >= 21);
+    assertThat(
+            new Formatter()
+                .formatSource(
+                    "public class Foo {\n"
+                        + "    String test(){\n"
+                        + "        var simple = STR.\"mytemplate1XXXX \\{exampleXXXX.foo()}yyy\";\n"
+                        + "        var nested = STR.\"template \\{example.  foo()+"
+                        + " STR.\"templateInner\\{  example}\"}xxx }\";\n"
+                        + "    }\n"
+                        + "}\n"))
+        .isEqualTo(
+            "public class Foo {\n"
+                + "  String test() {\n"
+                + "    var simple = STR.\"mytemplate1XXXX \\{exampleXXXX.foo()}yyy\";\n"
+                + "    var nested = STR.\"template \\{example.foo() +"
+                + " STR.\"templateInner\\{example}\"}xxx }\";\n"
+                + "  }\n"
+                + "}\n");
+  }
 }
diff --git a/core/src/test/resources/com/google/googlejavaformat/java/testdata/StringTemplate.input b/core/src/test/resources/com/google/googlejavaformat/java/testdata/StringTemplate.input
@@ -0,0 +1,10 @@
+public class StringTemplates {
+    void test(){
+        var m = STR."template \{example}xxx";
+        var nested = STR."template \{example.foo()+ STR."templateInner\{example}"}xxx }";
+        var nestNested = STR."template \{example0.
+                    foo() +
+                    STR."templateInner\{example1.test(STR."\{example2
+                        }")}"}xxx }";
+    }
+}
diff --git a/core/src/test/resources/com/google/googlejavaformat/java/testdata/StringTemplate.output b/core/src/test/resources/com/google/googlejavaformat/java/testdata/StringTemplate.output
@@ -0,0 +1,9 @@
+public class StringTemplates {
+  void test() {
+    var m = STR."template \{example}xxx";
+    var nested = STR."template \{example.foo() + STR."templateInner\{example}"}xxx }";
+    var nestNested =
+        STR."template \{
+            example0.foo() + STR."templateInner\{example1.test(STR."\{example2}")}"}xxx }";
+  }
+}