Import get inserted at correct index with newline space (#4242)

Closes: #4241

Co-authored-by: Jente Sondervorst <[email protected]>

Author: Jenson3210

rewrite-java/src/main/java/org/openrewrite/java/style/ImportLayoutStyle.java

@@ -151,6 +151,7 @@ public List<JRightPadded<J.Import>> addImport(List<JRightPadded<J.Import>> origi
                     // Find the import in the original list to establish insertion position.
                     for (int j = 0; j < originalImports.size(); j++) {
                         if (after != null && after.getElement().equals(originalImports.get(j).getElement()) && addToBlock.accept(after)) {
+                            insertPosition = j;
                             break;
                         } else if (before.getElement().equals(originalImports.get(j).getElement())) {
                             insertPosition = j + 1;

rewrite-java/src/test/java/org/openrewrite/java/style/ImportLayoutStyleTest.java

@@ -121,4 +121,50 @@ void addImportInPresenceOfDuplicateOtherImport() {
                 .containsExactlyInAnyOrder(
                         import1, import1, new JRightPadded<>(importToAdd, Space.EMPTY, Markers.EMPTY));
     }
+
+    @Test
+    @Issue("https://github.com/openrewrite/rewrite/issues/4241")
+    void addImportWithNewLineInUnsortedImportList() {
+        ImportLayoutStyle style = new ImportLayoutStyle(
+          Integer.MAX_VALUE, Integer.MAX_VALUE, Collections.emptyList(), Collections.emptyList());
+        JRightPadded<J.Import> import0 = new JRightPadded<>(
+          new J.Import(
+            randomId(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            new JLeftPadded<>(Space.SINGLE_SPACE, true, Markers.EMPTY),
+            TypeTree.build("pkg.AAA.MEMBER_0").withPrefix(Space.SINGLE_SPACE),
+            null),
+          Space.EMPTY,
+          Markers.EMPTY);
+        JRightPadded<J.Import> import1 = new JRightPadded<>(
+          new J.Import(
+            randomId(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            new JLeftPadded<>(Space.SINGLE_SPACE, true, Markers.EMPTY),
+            TypeTree.build("pkg.Clazz.MEMBER_3").withPrefix(Space.SINGLE_SPACE),
+            null),
+          Space.EMPTY,
+          Markers.EMPTY);
+        JRightPadded<J.Import> import3 = new JRightPadded<>(
+          new J.Import(
+            randomId(),
+            Space.EMPTY,
+            Markers.EMPTY,
+            new JLeftPadded<>(Space.SINGLE_SPACE, true, Markers.EMPTY),
+            TypeTree.build("pkg.Clazz.MEMBER_1").withPrefix(Space.SINGLE_SPACE),
+            null),
+          Space.EMPTY,
+          Markers.EMPTY);
+        J.Import importToAdd = new J.Import(
+          randomId(),
+          Space.EMPTY,
+          Markers.EMPTY,
+          new JLeftPadded<>(Space.EMPTY, true, Markers.EMPTY),
+          TypeTree.build("pkg.Clazz.MEMBER_2").withPrefix(Space.SINGLE_SPACE),
+          null);
+
+        assertThat(style.addImport(List.of(import0, import1, import3), importToAdd, null, Collections.emptyList()).get(1).getElement().getPrefix()).isEqualTo(Space.format("\n"));
+    }
 }