From 6563215b667b96c7d43f16da34a87004eabbc1cf Mon Sep 17 00:00:00 2001
From: Liam Miller-Cushon <cushon@google.com>
Date: Tue, 5 Apr 2022 12:26:16 -0700
Subject: [PATCH] Make `AppliedFix` more efficient

PiperOrigin-RevId: 439646636
---
 .../JavacErrorDescriptionListener.java        |  3 +-
 .../google/errorprone/fixes/AppliedFix.java   | 34 ++++++++++++++-----
 2 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/check_api/src/main/java/com/google/errorprone/JavacErrorDescriptionListener.java b/check_api/src/main/java/com/google/errorprone/JavacErrorDescriptionListener.java
index d7fa1c7403d..8e11a0d0a00 100644
--- a/check_api/src/main/java/com/google/errorprone/JavacErrorDescriptionListener.java
+++ b/check_api/src/main/java/com/google/errorprone/JavacErrorDescriptionListener.java
@@ -83,7 +83,8 @@ private JavacErrorDescriptionListener(
     checkNotNull(endPositions);
     try {
       CharSequence sourceFileContent = sourceFile.getCharContent(true);
-      fixToAppliedFix = fix -> AppliedFix.fromSource(sourceFileContent, endPositions).apply(fix);
+      AppliedFix.Applier applier = AppliedFix.fromSource(sourceFileContent, endPositions);
+      fixToAppliedFix = applier::apply;
     } catch (IOException e) {
       throw new UncheckedIOException(e);
     }
diff --git a/check_api/src/main/java/com/google/errorprone/fixes/AppliedFix.java b/check_api/src/main/java/com/google/errorprone/fixes/AppliedFix.java
index 72a413e347d..928e45d9f3c 100644
--- a/check_api/src/main/java/com/google/errorprone/fixes/AppliedFix.java
+++ b/check_api/src/main/java/com/google/errorprone/fixes/AppliedFix.java
@@ -18,12 +18,18 @@
 
 import static com.google.common.base.Preconditions.checkArgument;
 
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
 import com.sun.tools.javac.tree.EndPosTable;
 import java.io.IOException;
 import java.io.LineNumberReader;
 import java.io.StringReader;
 import java.util.HashSet;
+import java.util.NavigableMap;
 import java.util.Set;
+import java.util.TreeMap;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 /**
@@ -52,10 +58,12 @@ public boolean isRemoveLine() {
   public static class Applier {
     private final CharSequence source;
     private final EndPosTable endPositions;
+    private final Supplier<NavigableMap<Integer, Integer>> lineOffsets;
 
     public Applier(CharSequence source, EndPosTable endPositions) {
       this.source = source;
       this.endPositions = endPositions;
+      this.lineOffsets = Suppliers.memoize(() -> lineOffsets(source.toString()));
     }
 
     /**
@@ -80,15 +88,7 @@ public AppliedFix apply(Fix suggestedFix) {
         replaced.replace(repl.startPosition(), repl.endPosition(), repl.replaceWith());
 
         // Find the line number(s) being modified
-        // TODO: this could be more efficient
-        try {
-          LineNumberReader lineNumberReader =
-              new LineNumberReader(new StringReader(source.toString()));
-          lineNumberReader.skip(repl.startPosition());
-          modifiedLines.add(lineNumberReader.getLineNumber());
-        } catch (IOException e) {
-          // impossible since source is in-memory
-        }
+        modifiedLines.add(lineOffsets.get().floorEntry(repl.startPosition()).getValue());
       }
 
       // Not sure this is really the right behavior, but otherwise we can end up with an infinite
@@ -138,4 +138,20 @@ private static Set<Replacement> descending(Set<Replacement> set) {
   public static Applier fromSource(CharSequence source, EndPosTable endPositions) {
     return new Applier(source, endPositions);
   }
+
+  private static final Pattern NEWLINE = Pattern.compile("\\R");
+
+  /** Returns the start offsets of the lines in the input. */
+  private static NavigableMap<Integer, Integer> lineOffsets(String input) {
+    NavigableMap<Integer, Integer> lines = new TreeMap<>();
+    int line = 0;
+    int idx = 0;
+    lines.put(idx, line++);
+    Matcher matcher = NEWLINE.matcher(input);
+    while (matcher.find(idx)) {
+      idx = matcher.end();
+      lines.put(idx, line++);
+    }
+    return lines;
+  }
 }