support CRLF

openrewrite · Oct 30, 2023 · 66bf010 · 66bf010
1 parent 1bab1ff
commit 66bf010
Show file tree

Hide file tree

Showing 3 changed files with 103 additions and 5 deletions.
diff --git a/src/main/java/org/openrewrite/kotlin/KotlinParser.java b/src/main/java/org/openrewrite/kotlin/KotlinParser.java
@@ -425,15 +425,18 @@ public CompiledSource parse(List<Parser.Input> sources, Disposable disposable, E
                 fileName = source.getPath().toString();
             }
 
-            VirtualFile vFile = new LightVirtualFile(fileName, KotlinFileType.INSTANCE, StringUtilRt.convertLineSeparators(source.getSource(ctx).readFully()));
+            String sourceText = source.getSource(ctx).readFully();
+            List<Integer> cRLFLocations = getCRLFLocations(sourceText);
+
+            VirtualFile vFile = new LightVirtualFile(fileName, KotlinFileType.INSTANCE, StringUtilRt.convertLineSeparators(sourceText));
             final FileViewProvider fileViewProvider = new SingleRootFileViewProvider(
                     PsiManager.getInstance(environment.getProject()),
                     vFile
             );
             KtFile file = (KtFile) fileViewProvider.getPsi(KotlinLanguage.INSTANCE);
             assert file != null;
             ktFiles.add(file);
-            kotlinSources.add(new KotlinSource(source, file));
+            kotlinSources.add(new KotlinSource(source, file, cRLFLocations));
         }
 
         BaseDiagnosticsCollector diagnosticsReporter = DiagnosticReporterFactory.INSTANCE.createReporter(false);
@@ -611,4 +614,25 @@ private ApiVersion getApiVersion(KotlinLanguageLevel languageLevel) {
                 throw new IllegalArgumentException("Unknown language level: " + languageLevel);
         }
     }
+
+    private List<Integer> getCRLFLocations(String source) {
+        if (source.isEmpty()) {
+            return emptyList();
+        }
+        List<Integer> cRLFIndices = new ArrayList<>();
+        int pos = 0;
+        for (int i = 0; i < source.length(); i++) {
+            char currentChar = source.charAt(i);
+            if (currentChar == '\r') {
+                // Check if the next character is '\n' (CRLF)
+                if (i + 1 < source.length() && source.charAt(i + 1) == '\n') {
+                    cRLFIndices.add(pos);
+                    i++; // Skip the next character ('\n')
+                }
+            }
+            pos++;
+        }
+
+        return cRLFIndices;
+    }
 }
diff --git a/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java b/src/main/java/org/openrewrite/kotlin/internal/KotlinTreeParserVisitor.java
@@ -18,6 +18,7 @@
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.kotlin.KtNodeTypes;
 import org.jetbrains.kotlin.com.intellij.lang.ASTNode;
+import org.jetbrains.kotlin.com.intellij.openapi.util.TextRange;
 import org.jetbrains.kotlin.com.intellij.psi.*;
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.LeafPsiElement;
 import org.jetbrains.kotlin.com.intellij.psi.impl.source.tree.PsiErrorElementImpl;
@@ -87,6 +88,7 @@ public class KotlinTreeParserVisitor extends KtVisitor<J, ExecutionContext> {
     private final Boolean charsetBomMarked;
     private final Stack<KtElement> ownerStack = new Stack<>();
     private final ExecutionContext executionContext;
+    private final List<Integer> cRLFLocations;
 
     public KotlinTreeParserVisitor(KotlinSource kotlinSource,
                                    PsiElementAssociations psiElementAssociations,
@@ -103,6 +105,7 @@ public KotlinTreeParserVisitor(KotlinSource kotlinSource,
         charsetBomMarked = stream.isCharsetBomMarked();
         ownerStack.push(kotlinSource.getKtFile());
         executionContext = ctx;
+        cRLFLocations = kotlinSource.getCRLFLocations();
     }
 
     public K.CompilationUnit parse() {
@@ -3765,12 +3768,26 @@ private Space space(PsiElement node) {
         for (; node != null; node = next(node)) {
             PsiElement finalNode = node;
             if (isWhiteSpace(node)) {
+                String whiteSpace = node.getText();
+
+                // replace `\n` to CRLF back if it's CRLF in the source
+                TextRange range = node.getTextRange();
+                int left = findFirstGreaterOrEqual(cRLFLocations, range.getStartOffset());
+                int right = left != -1 ? findFirstLessOrEqual(cRLFLocations, range.getEndOffset(), left) : -1;
+                boolean hasCRLF = left != -1 && left <= right;
+
+                if (hasCRLF) {
+                    for (int i = right; i >= left; i--) {
+                        whiteSpace = replaceNewLineWithCRLF(whiteSpace, cRLFLocations.get(i) - range.getStartOffset());
+                    }
+                }
+
                 if (space == null) {
-                    space = Space.build(node.getText(), emptyList());
+                    space = Space.build(whiteSpace, emptyList());
                 } else {
                     if (isWhiteSpace(preNode)) {
                         // merge space
-                        space = space.withWhitespace(space.getWhitespace() + node.getText());
+                        space = space.withWhitespace(space.getWhitespace() + whiteSpace);
                     } else {
                         space = space.withComments(ListUtils.mapLast(space.getComments(), c -> c.withSuffix(finalNode.getText())));
                     }
@@ -3856,4 +3873,60 @@ private PsiElement findTrailingComma(PsiElement element) {
         return null;
     }
 
+    private int findFirstGreaterOrEqual(List<Integer> sequence, int target) {
+        int left = 0;
+        int right = sequence.size() - 1;
+        int resultIndex = -1;
+
+        while (left <= right) {
+            int mid = left + (right - left) / 2;
+            int midValue = sequence.get(mid);
+
+            if (midValue >= target) {
+                resultIndex = mid;
+                right = mid - 1;
+            } else {
+                left = mid + 1;
+            }
+        }
+
+        return resultIndex;
+    }
+
+    private int findFirstLessOrEqual(List<Integer> sequence, int target, int left) {
+        int right = sequence.size() - 1;
+        int resultIndex = -1;
+
+        while (left <= right) {
+            int mid = left + (right - left) / 2;
+            int midValue = sequence.get(mid);
+
+            if (midValue <= target) {
+                resultIndex = mid;
+                left = mid + 1;
+            } else {
+                right = mid - 1;
+            }
+        }
+
+        return resultIndex;
+    }
+
+    private String replaceNewLineWithCRLF(String source, int index) {
+        if (index < 0 || index >= source.length()) {
+            return source;
+        }
+
+        StringBuilder sb = new StringBuilder(source);
+        char charAtIndex = source.charAt(index);
+        if (charAtIndex != '\n') {
+            return source;
+        }
+
+        sb.setCharAt(index, '\r');
+        sb.insert(index + 1, '\n');
+
+        return sb.toString();
+    }
+
 }
diff --git a/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinSource.kt b/src/main/kotlin/org/openrewrite/kotlin/internal/KotlinSource.kt
@@ -30,7 +30,8 @@ import kotlin.collections.ArrayDeque
 @Getter
 class KotlinSource(
     var input: Parser.Input,
-    val ktFile: KtFile
+    val ktFile: KtFile,
+    val cRLFLocations : List<Int>
 ) {
     val nodes: Map<Int, ASTNode>