From ba26b452bebaf841810ac6a2aa0441822ce71dc9 Mon Sep 17 00:00:00 2001
From: Tony Robalik <trobalik@squareup.com>
Date: Mon, 16 Dec 2024 13:06:18 -0800
Subject: [PATCH] fix: add plugins block below imports, if any.

---
 .../cash/recipes/plugins/PluginMutator.kt     | 28 +++++++++++--
 .../cash/recipes/plugins/PluginMutatorTest.kt | 41 +++++++++++++++++--
 2 files changed, 62 insertions(+), 7 deletions(-)

diff --git a/recipes/plugins/src/main/kotlin/cash/recipes/plugins/PluginMutator.kt b/recipes/plugins/src/main/kotlin/cash/recipes/plugins/PluginMutator.kt
index 2720757..297a591 100644
--- a/recipes/plugins/src/main/kotlin/cash/recipes/plugins/PluginMutator.kt
+++ b/recipes/plugins/src/main/kotlin/cash/recipes/plugins/PluginMutator.kt
@@ -13,6 +13,7 @@ import cash.grammar.kotlindsl.utils.Whitespace
 import cash.grammar.kotlindsl.utils.Whitespace.trimGently
 import cash.grammar.utils.ifNotEmpty
 import cash.recipes.plugins.exception.NonNormalizedScriptException
+import com.squareup.cash.grammar.KotlinParser.ImportListContext
 import com.squareup.cash.grammar.KotlinParser.NamedBlockContext
 import com.squareup.cash.grammar.KotlinParser.PostfixUnaryExpressionContext
 import com.squareup.cash.grammar.KotlinParser.ScriptContext
@@ -46,11 +47,14 @@ public class PluginMutator private constructor(
   private val pluginsToRemove: Set<String>,
 ) : KotlinParserBaseListener() {
 
+  // TODO(tsr): use System.lineSeparator() universally
+  private val newline = System.lineSeparator()
   private val rewriter = Rewriter(tokens)
   private val smartIndent = SmartIndent(tokens)
   private val terminalNewlines = Whitespace.countTerminalNewlines(tokens)
 
   private val blockStack = ArrayDeque<NamedBlockContext>()
+  private var importsList: ImportListContext? = null
   private var pluginsBlock: NamedBlockContext? = null
 
   private val pluginIds = mutableSetOf<String>()
@@ -63,7 +67,14 @@ public class PluginMutator private constructor(
       throw KotlinParseException.withErrors(it)
     }
 
-    return rewriter.text.trimGently(terminalNewlines)
+    return rewriter.text
+      // Simplifies conditional logic elsewhere. We never want two blank lines in a row.
+      .replace("$newline$newline$newline", "$newline$newline")
+      .trimGently(terminalNewlines)
+  }
+
+  override fun enterImportList(ctx: ImportListContext) {
+    importsList = ctx
   }
 
   override fun enterNamedBlock(ctx: NamedBlockContext) {
@@ -143,21 +154,30 @@ public class PluginMutator private constructor(
 
     rewriter.insertBefore(
       tokenToInsertBefore,
-      contentToAdd.joinToString(separator = "\n") + "\n"
+      contentToAdd.joinToString(separator = newline) + newline
     )
   }
 
   private fun addNewPluginBlock(script: ScriptContext) {
     if (pluginsToAdd.isEmpty()) return
 
+    val hasImports = importsList != null
+
     val pluginsContentToAdd: List<String> = pluginContentToAdd(emptyList())
     val pluginBlockContent = buildString {
+      if (hasImports) {
+        appendLine()
+      }
+
       appendLine("plugins {")
       pluginsContentToAdd.forEach { appendLine(it) }
-      append("}")
+      appendLine("}")
+      appendLine()
     }
 
-    rewriter.insertBefore(script.start, "$pluginBlockContent\n\n")
+    val insertLocation = importsList?.stop ?: script.start
+
+    rewriter.insertBefore(insertLocation, pluginBlockContent)
   }
 
   private fun removeBlockIdPlugins(ctx: List<ParserRuleContext>) {
diff --git a/recipes/plugins/src/test/kotlin/cash/recipes/plugins/PluginMutatorTest.kt b/recipes/plugins/src/test/kotlin/cash/recipes/plugins/PluginMutatorTest.kt
index c862fac..c8dcda5 100644
--- a/recipes/plugins/src/test/kotlin/cash/recipes/plugins/PluginMutatorTest.kt
+++ b/recipes/plugins/src/test/kotlin/cash/recipes/plugins/PluginMutatorTest.kt
@@ -92,6 +92,41 @@ internal class PluginMutatorTest {
     assertThat(rewrittenContent).isEqualTo(expectedContent)
   }
 
+  @Test
+  fun `can add plugins when no existing plugin block is present, below imports`() {
+    val pluginToAdd = setOf("foo", "bar")
+    val pluginsToRemove = emptySet<String>()
+
+    val buildScript =
+      """
+        import magic
+        
+        extension {
+          foo = bar
+        }
+      """.trimIndent()
+
+    // When...
+    val pluginMutator = PluginMutator.of(buildScript, pluginToAdd, pluginsToRemove)
+    val rewrittenContent = pluginMutator.rewritten()
+    val expectedContent =
+      """
+        import magic
+        
+        plugins {
+          id("foo")
+          id("bar")
+        }
+
+        extension {
+          foo = bar
+        }
+      """.trimIndent()
+
+    // Then...
+    assertThat(rewrittenContent).isEqualTo(expectedContent)
+  }
+
   @Test
   fun `skip adding plugins that are already present`() {
     val pluginToAdd = setOf("cash.server")
@@ -184,7 +219,7 @@ internal class PluginMutatorTest {
 
   @Test
   fun `throw error when same plugins are in add and remove sets`() {
-    val pluginToAdd =  setOf("bar")
+    val pluginToAdd = setOf("bar")
     val pluginToRemove = setOf("foo", "bar")
 
     val buildScript =
@@ -202,7 +237,7 @@ internal class PluginMutatorTest {
 
   @Test
   fun `throws error for non-normalized script, with non-BlockId plugins`() {
-    val pluginToAdd =  setOf("foo")
+    val pluginToAdd = setOf("foo")
     val pluginToRemove = setOf("bar")
 
     val buildScript =
@@ -220,7 +255,7 @@ internal class PluginMutatorTest {
 
   @Test
   fun `throws error for non-normalized script, with apply plugins`() {
-    val pluginToAdd =  setOf("foo")
+    val pluginToAdd = setOf("foo")
     val pluginToRemove = setOf("bar")
 
     val buildScript =