Fix formatting of comments above import declarations

misc/formatter/modules/formatter-core/src/main/java/org/ballerinalang/formatter/core/FormatterUtils.java

@@ -16,7 +16,13 @@
 package org.ballerinalang.formatter.core;
 
 import io.ballerina.compiler.syntax.tree.ImportDeclarationNode;
+import io.ballerina.compiler.syntax.tree.ImportOrgNameNode;
+import io.ballerina.compiler.syntax.tree.Minutiae;
+import io.ballerina.compiler.syntax.tree.MinutiaeList;
 import io.ballerina.compiler.syntax.tree.Node;
+import io.ballerina.compiler.syntax.tree.NodeFactory;
+import io.ballerina.compiler.syntax.tree.SyntaxKind;
+import io.ballerina.compiler.syntax.tree.Token;
 import io.ballerina.tools.text.LineRange;
 import org.apache.commons.lang3.builder.CompareToBuilder;
 
@@ -79,4 +85,81 @@ static void sortImportDeclarations(List<ImportDeclarationNode> importDeclaration
                         node2.moduleName().stream().map(node -> node.toString().trim()).collect(Collectors.joining()))
                 .toComparison());
     }
+
+    /**
+     * Swap leading minutiae of the first import in original code and the first import when sorted.
+     *
+     * @param firstImportNode First ImportDeclarationNode in original code
+     * @param importNodes     Sorted formatted ImportDeclarationNode nodes
+     */
+    static void swapLeadingMinutiae(ImportDeclarationNode firstImportNode, List<ImportDeclarationNode> importNodes) {
+        int prevFirstImportIndex = -1;
+        String firstImportOrgName = firstImportNode.orgName().map(ImportOrgNameNode::toSourceCode).orElse("");
+        String firstImportModuleName = getImportModuleName(firstImportNode);
+        for (int i = 0; i < importNodes.size(); i++) {
+            if (doesImportMatch(firstImportOrgName, firstImportModuleName, importNodes.get(i))) {
+                prevFirstImportIndex = i;
+                break;
+            }
+        }
+        if (prevFirstImportIndex == 0) {
+            return;
+        }
+
+        // remove comments from the previous first import
+        ImportDeclarationNode prevFirstImportNode = importNodes.get(prevFirstImportIndex);
+        MinutiaeList prevFirstLeadingMinutiae = prevFirstImportNode.leadingMinutiae();
+
+        MinutiaeList prevFirstNewLeadingMinutiae = NodeFactory.createEmptyMinutiaeList();
+        Minutiae prevFirstMinutiae = prevFirstLeadingMinutiae.get(0);
+        if (prevFirstMinutiae.kind() == SyntaxKind.END_OF_LINE_MINUTIAE) {
+            // if the prevFirstImport now is the first of a group of imports, handle the added leading newline
+            prevFirstNewLeadingMinutiae = prevFirstNewLeadingMinutiae.add(prevFirstMinutiae);
+            prevFirstLeadingMinutiae = prevFirstLeadingMinutiae.remove(0);
+        }
+
+        importNodes.set(prevFirstImportIndex,
+                modifyImportDeclLeadingMinutiae(prevFirstImportNode, prevFirstNewLeadingMinutiae));
+
+        if (!hasEmptyLineAtEnd(prevFirstLeadingMinutiae)) {
+            // adds a newline after prevFirstImport's leading minutiae if not present
+            prevFirstLeadingMinutiae =
+                    prevFirstLeadingMinutiae.add(NodeFactory.createEndOfLineMinutiae(System.lineSeparator()));
+        }
+
+        ImportDeclarationNode newFirstImportNode = importNodes.get(0);
+        MinutiaeList newFirstLeadingMinutiae = newFirstImportNode.importKeyword().leadingMinutiae();
+        for (int i = 0; i < newFirstLeadingMinutiae.size(); i++) {
+            Minutiae minutiae = newFirstLeadingMinutiae.get(i);
+            if (i == 0 && minutiae.kind() == SyntaxKind.END_OF_LINE_MINUTIAE) {
+                // since we added a new line after prevFirstImport's leading minutiae we can skip th newline here
+                continue;
+            }
+            prevFirstLeadingMinutiae = prevFirstLeadingMinutiae.add(minutiae);
+        }
+
+        importNodes.set(0, modifyImportDeclLeadingMinutiae(newFirstImportNode, prevFirstLeadingMinutiae));
+    }
+
+    private static ImportDeclarationNode modifyImportDeclLeadingMinutiae(ImportDeclarationNode importDecl,
+                                                                         MinutiaeList leadingMinutiae) {
+        Token importToken = importDecl.importKeyword();
+        Token modifiedImportToken = importToken.modify(leadingMinutiae, importToken.trailingMinutiae());
+        return importDecl.modify().withImportKeyword(modifiedImportToken).apply();
+    }
+
+    private static boolean doesImportMatch(String orgName, String moduleName, ImportDeclarationNode importDeclNode) {
+        String importDeclOrgName = importDeclNode.orgName().map(ImportOrgNameNode::toSourceCode).orElse("");
+        return orgName.equals(importDeclOrgName) && moduleName.equals(getImportModuleName(importDeclNode));
+    }
+
+    private static String getImportModuleName(ImportDeclarationNode importDeclarationNode) {
+        return importDeclarationNode.moduleName().stream().map(Node::toSourceCode).collect(Collectors.joining("."));
+    }
+
+    private static boolean hasEmptyLineAtEnd(MinutiaeList minutiaeList) {
+        int size = minutiaeList.size();
+        return minutiaeList.get(size - 1).kind() == SyntaxKind.END_OF_LINE_MINUTIAE &&
+                minutiaeList.get(size - 2).kind() == SyntaxKind.END_OF_LINE_MINUTIAE;
+    }
 }

misc/formatter/modules/formatter-core/src/main/java/org/ballerinalang/formatter/core/FormattingTreeModifier.java

@@ -260,6 +260,7 @@
 
 import static org.ballerinalang.formatter.core.FormatterUtils.isInlineRange;
 import static org.ballerinalang.formatter.core.FormatterUtils.sortImportDeclarations;
+import static org.ballerinalang.formatter.core.FormatterUtils.swapLeadingMinutiae;
 
 /**
  * A formatter implementation that updates the minutiae of a given tree according to the ballerina formatting
@@ -290,7 +291,7 @@ public FormattingTreeModifier(FormattingOptions options, LineRange lineRange) {
 
     @Override
     public ModulePartNode transform(ModulePartNode modulePartNode) {
-        NodeList<ImportDeclarationNode> imports = sortAndGroupImportDeclarationNodes(modulePartNode.imports());
+        NodeList<ImportDeclarationNode> imports = arrangeAndFormatImportDeclarations(modulePartNode.imports());
         NodeList<ModuleMemberDeclarationNode> members =
                 formatMemberDeclarations(modulePartNode.members(), n -> isMultilineModuleMember(n));
         Token eofToken = formatToken(modulePartNode.eofToken(), 0, 0);
@@ -651,7 +652,7 @@ public RecordFieldWithDefaultValueNode transform(RecordFieldWithDefaultValueNode
     @Override
     public ImportDeclarationNode transform(ImportDeclarationNode importDeclarationNode) {
         boolean prevPreservedNewLine = env.hasPreservedNewline;
-        setPreserveNewline(false);
+        setPreserveNewline(hasLeadingComments(importDeclarationNode));
         Token importKeyword = formatToken(importDeclarationNode.importKeyword(), 1, 0);
         setPreserveNewline(prevPreservedNewLine);
         boolean hasPrefix = importDeclarationNode.prefix().isPresent();
@@ -4781,8 +4782,13 @@ private boolean matchesMinutiaeKind(Minutiae minutiae, SyntaxKind kind) {
         return minutiae != null && minutiae.kind() == kind;
     }
 
-    private NodeList<ImportDeclarationNode> sortAndGroupImportDeclarationNodes(
+    private NodeList<ImportDeclarationNode> arrangeAndFormatImportDeclarations(
             NodeList<ImportDeclarationNode> importDeclarationNodes) {
+        if (importDeclarationNodes.isEmpty()) {
+            return importDeclarationNodes;
+        }
+
+        ImportDeclarationNode firstImport = importDeclarationNodes.get(0);
         // moduleImports would collect only module level imports if grouping is enabled,
         // and would collect all imports otherwise
         List<ImportDeclarationNode> moduleImports = new ArrayList<>();
@@ -4818,6 +4824,12 @@ private NodeList<ImportDeclarationNode> sortAndGroupImportDeclarationNodes(
         imports.addAll(moduleImportNodes.stream().collect(Collectors.toList()));
         imports.addAll(stdLibImportNodes.stream().collect(Collectors.toList()));
         imports.addAll(thirdPartyImportNodes.stream().collect(Collectors.toList()));
+
+        if (hasLeadingComments(firstImport)) {
+            // This is to ensure license header remains at top of the file
+            swapLeadingMinutiae(firstImport, imports);
+        }
+
         return NodeFactory.createNodeList(imports);
     }
 
@@ -4831,4 +4843,14 @@ private boolean isScopedFunctionArgument(FunctionArgumentNode functionArgumentNo
         }
         return false;
     }
+
+    private boolean hasLeadingComments(Node node) {
+        MinutiaeList minutiaeList = node.leadingMinutiae();
+        for (Minutiae minutiae: minutiaeList) {
+            if (minutiae.kind() == SyntaxKind.COMMENT_MINUTIAE) {
+                return true;
+            }
+        }
+        return false;
+    }
 }

misc/formatter/modules/formatter-core/src/test/java/org/ballerinalang/formatter/core/declarations/ImportDeclarationsTest.java

@@ -52,11 +52,18 @@ public Object[][] dataProvider() {
     @DataProvider(name = "test-file-provider-custom")
     @Override
     public Object[][] dataProviderWithCustomTests(Method testName) {
+        String testResourceDir = this.getTestResourceDir();
         switch (testName.getName()) {
             case "test":
                 return new Object[][] {
-                        {"import_declaration_1.bal", this.getTestResourceDir()},
-                        {"import_declaration_2.bal", this.getTestResourceDir()}
+                        {"import_declaration_1.bal", testResourceDir},
+                        {"import_declaration_2.bal", testResourceDir},
+                        {"import_declaration_6.bal", testResourceDir},
+                        {"import_declaration_7.bal", testResourceDir},
+                        {"import_declaration_8.bal", testResourceDir},
+                        {"import_declaration_9.bal", testResourceDir},
+                        {"import_declaration_10.bal", testResourceDir},
+                        {"import_declaration_11.bal", testResourceDir}
                 };
             case "testWithCustomOptions":
                 FormattingOptions optionWithNoGrouping = FormattingOptions.builder()
@@ -70,9 +77,9 @@ public Object[][] dataProviderWithCustomTests(Method testName) {
                                 ImportFormattingOptions.builder().setGroupImports(false).setSortImports(false).build())
                         .build();
                 return new Object[][] {
-                        {"import_declaration_3.bal", this.getTestResourceDir(), optionWithNoGrouping},
-                        {"import_declaration_4.bal", this.getTestResourceDir(), optionWithNoSorting},
-                        {"import_declaration_5.bal", this.getTestResourceDir(), optionWithNoGroupingAndSorting}
+                        {"import_declaration_3.bal", testResourceDir, optionWithNoGrouping},
+                        {"import_declaration_4.bal", testResourceDir, optionWithNoSorting},
+                        {"import_declaration_5.bal", testResourceDir, optionWithNoGroupingAndSorting}
                 };
         }
         return null;

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_1.bal

@@ -1,7 +1,9 @@
 import ballerina/io as console;
 import ballerina/lang.'int;
 import ballerina/log as logger;
+
 //Test comment
+
 //Another comment
 import ballerina/math;
 

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_10.bal

@@ -0,0 +1,24 @@
+// Copyright (c) 2024 WSO2 LLC. (http://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.
+
+// check this
+
+import ballerina/jballerina.java as _;
+import ballerinax/twilio as _;
+
+function bar() {
+
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_11.bal

@@ -0,0 +1,20 @@
+// comment 1
+// comment 2
+
+// module imports
+import bas;
+import fos;
+
+import ballerina/io;
+import ballerina/time;
+import ballerinax/http;
+import ballerinax/oracledb;
+
+// bar
+import abc/bar;
+
+// foo
+import abc/foo;
+
+function baz() {
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_5.bal

@@ -1,4 +1,5 @@
 import greeter.email;
+
 // First line of comments
 // Second line of comments
 import ballerina/lang.'int;

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_6.bal

@@ -0,0 +1,17 @@
+// Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+//
+// This software is the property of WSO2 LLC. and its suppliers, if any.
+// Dissemination of any information or reproduction of any material contained
+// herein in any form is strictly forbidden, unless permitted by WSO2 expressly.
+// You may not alter or remove any copyright or other notice from copies of this content.
+
+import cp_configuration_service.config;
+import cp_configuration_service.utils;
+
+import ballerina/http;
+import ballerina/time;
+import ballerina/uuid;
+
+function bar() {
+    string a = "a";
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_7.bal

@@ -0,0 +1,23 @@
+// Copyright (c) 2024, WSO2 LLC. (http://www.wso2.com).
+//
+// This software is the property of WSO2 LLC. and its suppliers, if any.
+// Dissemination of any information or reproduction of any material contained
+// herein in any form is strictly forbidden, unless permitted by WSO2 expressly.
+// You may not alter or remove any copyright or other notice from copies of this content.
+
+// module imports
+import bas;
+import fos;
+
+import ballerina/log;
+import ballerina/time;
+import ballerinax/jdbc;
+import ballerinax/oracledb;
+
+import abc/bar;
+import abc/foo;
+import cde/baz;
+
+function bar() {
+    string a = "a";
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_8.bal

@@ -0,0 +1,40 @@
+// Copyright (c) 2024 WSO2 LLC. (http://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.
+
+// just a random comment
+// for two lines part of this import
+
+// comment 3
+// comment 4
+
+// module imports
+import bas;
+import fos;
+
+import ballerina/io;
+import ballerina/time;
+import ballerinax/http;
+import ballerinax/oracledb;
+
+// bar
+import abc/bar;
+
+// foo
+import abc/foo;
+
+function baz() {
+
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/assert/import_declaration_9.bal

@@ -0,0 +1,8 @@
+// check this
+
+import ballerina/jballerina.java as _;
+import ballerinax/twilio as _;
+
+function foz() {
+
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/source/import_declaration_10.bal

@@ -0,0 +1,23 @@
+// Copyright (c) 2024 WSO2 LLC. (http://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.
+
+// check this
+import ballerinax/twilio as _;
+import ballerina/jballerina.java as _;
+
+function bar() {
+
+}

misc/formatter/modules/formatter-core/src/test/resources/declarations/import/source/import_declaration_11.bal

@@ -0,0 +1,20 @@
+
+// comment 1
+// comment 2
+import ballerina/io;
+import ballerinax/oracledb;
+import ballerinax/http;
+import ballerina/time;
+
+// module imports
+import bas;
+import fos;
+
+// foo
+import abc/foo;
+
+// bar
+import abc/bar;
+
+function baz() {
+}