Improved diagnostic rendering

cli/ballerina-cli/build.gradle

@@ -34,6 +34,7 @@ configurations {
 
 dependencies {
 
+    implementation project(':ballerina-parser')
     implementation project(':ballerina-lang')
     implementation project(':ballerina-runtime')
     implementation project(':ballerina-tools-api')

cli/ballerina-cli/src/main/java/io/ballerina/cli/task/CompileTask.java

@@ -18,11 +18,15 @@
 
 package io.ballerina.cli.task;
 
+import io.ballerina.cli.utils.AnnotateDiagnostics;
 import io.ballerina.cli.utils.BuildTime;
 import io.ballerina.projects.CodeGeneratorResult;
 import io.ballerina.projects.CodeModifierResult;
+import io.ballerina.projects.Document;
 import io.ballerina.projects.JBallerinaBackend;
 import io.ballerina.projects.JvmTarget;
+import io.ballerina.projects.ModuleName;
+import io.ballerina.projects.Package;
 import io.ballerina.projects.PackageCompilation;
 import io.ballerina.projects.PackageResolution;
 import io.ballerina.projects.Project;
@@ -42,8 +46,12 @@
 import org.wso2.ballerinalang.util.RepoUtils;
 
 import java.io.PrintStream;
+import java.nio.file.Paths;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 import static io.ballerina.cli.launcher.LauncherUtils.createLauncherException;
@@ -56,6 +64,7 @@
  * @since 2.0.0
  */
 public class CompileTask implements Task {
+
     private final transient PrintStream out;
     private final transient PrintStream err;
     private final boolean compileForBalPack;
@@ -208,14 +217,39 @@ public void execute(Project project) {
             if (project.buildOptions().dumpBuildTime()) {
                 BuildTime.getInstance().codeGenDuration = System.currentTimeMillis() - start;
             }
+            // HashSet to keep track of the diagnostics to avoid duplicate diagnostics
+            Set<String> diagnosticSet = new HashSet<>();
+            // HashMap for documents based on filename
+            Map<String, Document> documentMap = new HashMap<>();
+            int terminalWidth = AnnotateDiagnostics.getTerminalWidth();
+            boolean colorEnabled = terminalWidth != 0;
 
+            Package currentPackage = project.currentPackage();
+            currentPackage.moduleIds().forEach(moduleId -> {
+                currentPackage.module(moduleId).documentIds().forEach(documentId -> {
+                    Document document = currentPackage.module(moduleId).document(documentId);
+                    documentMap.put(getDocumentPath(document.module().moduleName(), document.name()), document);
+                });
+                currentPackage.module(moduleId).testDocumentIds().forEach(documentId -> {
+                    Document document = currentPackage.module(moduleId).document(documentId);
+                    documentMap.put(getDocumentPath(document.module().moduleName(), document.name()), document);
+                });
+            });
             // Report package compilation and backend diagnostics
             diagnostics.addAll(jBallerinaBackend.diagnosticResult().diagnostics(false));
             diagnostics.forEach(d -> {
                 if (d.diagnosticInfo().code() == null || (!d.diagnosticInfo().code().equals(
                         ProjectDiagnosticErrorCode.BUILT_WITH_OLDER_SL_UPDATE_DISTRIBUTION.diagnosticId()) &&
                         !d.diagnosticInfo().code().startsWith(TOOL_DIAGNOSTIC_CODE_PREFIX))) {
-                    err.println(d);
+                    if (diagnosticSet.add(d.toString())) {
+                        Document document = documentMap.get(d.location().lineRange().fileName());
+                        if (document != null) {
+                            err.println(AnnotateDiagnostics.renderDiagnostic(d, document,
+                                    terminalWidth == 0 ? 999 : terminalWidth, colorEnabled));
+                        } else {
+                            err.println(AnnotateDiagnostics.renderDiagnostic(d, colorEnabled));
+                        }
+                    }
                 }
             });
             // Report build tool execution diagnostics
@@ -239,6 +273,13 @@ public void execute(Project project) {
         }
     }
 
+    private String getDocumentPath(ModuleName moduleName, String documentName) {
+        if (moduleName.isDefaultModuleName()) {
+            return documentName;
+        }
+        return Paths.get("modules", moduleName.moduleNamePart(), documentName).toString();
+    }
+
     private boolean isPackCmdForATemplatePkg(Project project) {
         return compileForBalPack && project.currentPackage().manifest().template();
     }

cli/ballerina-cli/src/main/java/io/ballerina/cli/utils/AnnotateDiagnostics.java

@@ -0,0 +1,209 @@
+/*
+ * Copyright (c) 2024, WSO2 LLC. (https://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package io.ballerina.cli.utils;
+
+import io.ballerina.compiler.internal.diagnostics.StringDiagnosticProperty;
+import io.ballerina.projects.Document;
+import io.ballerina.projects.internal.PackageDiagnostic;
+import io.ballerina.tools.diagnostics.Diagnostic;
+import io.ballerina.tools.diagnostics.DiagnosticSeverity;
+import io.ballerina.tools.diagnostics.Location;
+import io.ballerina.tools.text.TextDocument;
+import org.jline.jansi.Ansi;
+import org.jline.terminal.TerminalBuilder;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+
+import static io.ballerina.cli.utils.DiagnosticAnnotation.NEW_LINE;
+import static io.ballerina.cli.utils.DiagnosticAnnotation.SEVERITY_COLORS;
+import static io.ballerina.cli.utils.DiagnosticAnnotation.getColoredString;
+
+/**
+ * This class is used to generate diagnostic annotations from diagnostics.
+ *
+ * @since 2201.9.0
+ */
+public class AnnotateDiagnostics {
+
+    private static final String COMPILER_ERROR_PREFIX = "BCE";
+    private static final int SYNTAX_ERROR_CODE_THRESHOLD = 1000;
+    private static final int MISSING_TOKEN_KEYWORD_CODE_THRESHOLD = 400;
+    private static final int INVALID_TOKEN_CODE = 600;
+
+    public static Ansi renderDiagnostic(Diagnostic diagnostic, Document document, int terminalWidth,
+                                        boolean colorEnabled) {
+
+        String diagnosticCode = diagnostic.diagnosticInfo().code();
+        if (diagnostic instanceof PackageDiagnostic && diagnosticCode != null &&
+                diagnosticCode.startsWith(COMPILER_ERROR_PREFIX)) {
+            int diagnosticCodeNumber = Integer.parseInt(diagnosticCode.substring(3));
+            if (diagnosticCodeNumber < SYNTAX_ERROR_CODE_THRESHOLD) {
+                PackageDiagnostic packageDiagnostic = (PackageDiagnostic) diagnostic;
+                return Ansi.ansi()
+                        .render(diagnosticToString(diagnostic, colorEnabled) + NEW_LINE + getSyntaxDiagnosticAnnotation(
+                                document, packageDiagnostic, diagnosticCodeNumber, terminalWidth, colorEnabled));
+            }
+        }
+        DiagnosticAnnotation diagnosticAnnotation = getDiagnosticLineFromSyntaxAPI(
+                document, diagnostic.location(), diagnostic.diagnosticInfo().severity(), terminalWidth, colorEnabled);
+        return Ansi.ansi().render(diagnosticToString(diagnostic, colorEnabled) + NEW_LINE + diagnosticAnnotation);
+
+    }
+
+    public static int getTerminalWidth() {
+        try {
+            return TerminalBuilder.builder().dumb(true).build().getWidth();
+        } catch (IOException e) {
+            return 999;
+        }
+    }
+
+    public static Ansi renderDiagnostic(Diagnostic diagnostic, boolean colorEnabled) {
+        return Ansi.ansi().render(diagnosticToString(diagnostic, colorEnabled));
+    }
+
+    private static String diagnosticToString(Diagnostic diagnostic, boolean colorEnabled) {
+        DiagnosticSeverity severity = diagnostic.diagnosticInfo().severity();
+        String severityString = severity.toString();
+        String color = SEVERITY_COLORS.get(severity);
+        String message = diagnostic.toString().substring(severityString.length());
+        String code = diagnostic.diagnosticInfo().code();
+        boolean isMultiline = diagnostic.message().contains(NEW_LINE);
+        String formatString = getColoredString("%s", color, colorEnabled) + "%s" +
+                (code != null ? (isMultiline ? NEW_LINE + "(%s)" : " (%s)") : "");
+
+        return String.format(formatString, severityString, message, code != null ? code : "");
+    }
+
+    private static DiagnosticAnnotation getDiagnosticLineFromSyntaxAPI(Document document, Location location,
+                                                                       DiagnosticSeverity severity, int terminalWidth,
+                                                                       boolean colorEnabled) {
+        TextDocument textDocument = document.textDocument();
+        int startOffset = location.lineRange().startLine().offset();
+        int endOffset = location.lineRange().endLine().offset();
+        int startLine = location.lineRange().startLine().line();
+        int endLine = location.lineRange().endLine().line();
+        boolean isMultiline = startLine != endLine;
+        int length = isMultiline ? textDocument.line(startLine).length() - startOffset : endOffset - startOffset;
+
+        return new DiagnosticAnnotation(
+                getLines(textDocument, startLine, endLine),
+                startOffset,
+                length == 0 ? 1 : length,
+                isMultiline,
+                endOffset,
+                startLine + 1,
+                severity,
+                DiagnosticAnnotation.DiagnosticAnnotationType.REGULAR,
+                terminalWidth, colorEnabled);
+    }
+
+    private static DiagnosticAnnotation getSyntaxDiagnosticAnnotation(Document document,
+                                                                      PackageDiagnostic packageDiagnostic,
+                                                                      int diagnosticCode, int terminalWidth,
+                                                                      boolean colorEnabled) {
+        TextDocument textDocument = document.textDocument();
+        Location location = packageDiagnostic.location();
+        int startLine = location.lineRange().startLine().line();
+        int startOffset = location.lineRange().startLine().offset();
+        int padding = 0;
+        int endLine = location.lineRange().endLine().line();
+        int endOffset = location.lineRange().endLine().offset();
+        String color = SEVERITY_COLORS.get(DiagnosticSeverity.ERROR);
+
+        if (diagnosticCode < MISSING_TOKEN_KEYWORD_CODE_THRESHOLD) {
+            StringDiagnosticProperty strProperty = (StringDiagnosticProperty) packageDiagnostic.properties().get(0);
+            String lineString = textDocument.line(startLine).text();
+            String missingTokenString = getColoredString(strProperty.value(), color, colorEnabled);
+            if (startOffset < lineString.length() && lineString.charAt(startOffset) != ' ') {
+                missingTokenString = missingTokenString + " ";
+            }
+            if (startOffset > 0 && lineString.charAt(startOffset - 1) != ' ') {
+                missingTokenString = " " + missingTokenString;
+                padding++;
+            }
+
+            String lineWithMissingToken = lineString.substring(0, startOffset) + missingTokenString +
+                    lineString.substring(startOffset);
+            List<String> lines = new ArrayList<>();
+            lines.add(lineWithMissingToken);
+            return new DiagnosticAnnotation(
+                    lines,
+                    padding + startOffset,
+                    strProperty.value().length(),
+                    false,
+                    0,
+                    startLine + 1,
+                    DiagnosticSeverity.ERROR,
+                    DiagnosticAnnotation.DiagnosticAnnotationType.MISSING,
+                    terminalWidth, colorEnabled);
+        }
+
+        if (diagnosticCode == INVALID_TOKEN_CODE) {
+            List<String> lines = getLines(textDocument, startLine, endLine);
+            if (lines.size() > 1) {
+                String annotatedLine1 = lines.get(0).substring(0, startOffset) +
+                        getColoredString(lines.get(0).substring(startOffset), color, colorEnabled);
+                String annotatedLine2 =
+                        getColoredString(lines.get(lines.size() - 1).substring(0, endOffset), color, colorEnabled) +
+                                lines.get(lines.size() - 1).substring(endOffset);
+                lines.set(0, annotatedLine1);
+                lines.set(lines.size() - 1, annotatedLine2);
+                return new DiagnosticAnnotation(
+                        lines,
+                        startOffset,
+                        textDocument.line(startLine).length() - location.lineRange().startLine().offset(),
+                        true,
+                        endOffset,
+                        startLine + 1,
+                        DiagnosticSeverity.ERROR,
+                        DiagnosticAnnotation.DiagnosticAnnotationType.INVALID,
+                        terminalWidth, colorEnabled);
+            }
+            String line = lines.get(0);
+            String annotatedLine = line.substring(0, startOffset) +
+                    getColoredString(line.substring(startOffset, endOffset), color, colorEnabled) +
+                    line.substring(endOffset);
+            lines.set(0, annotatedLine);
+            return new DiagnosticAnnotation(
+                    lines,
+                    startOffset,
+                    endOffset - startOffset,
+                    false,
+                    0,
+                    startLine + 1,
+                    DiagnosticSeverity.ERROR,
+                    DiagnosticAnnotation.DiagnosticAnnotationType.INVALID,
+                    terminalWidth, colorEnabled);
+        }
+        return getDiagnosticLineFromSyntaxAPI(document, location, DiagnosticSeverity.ERROR, terminalWidth,
+                colorEnabled);
+    }
+
+    private static List<String> getLines(TextDocument textDocument, int start, int end) {
+        List<String> lines = new ArrayList<>();
+        for (int i = start; i <= end; i++) {
+            lines.add(textDocument.line(i).text());
+        }
+        return lines;
+    }
+
+}