From 78dc73710e5dca40da47ffcd281bbdc7a489758d Mon Sep 17 00:00:00 2001
From: Nacho Lopez <nacho@nlopez.io>
Date: Mon, 27 Nov 2023 23:00:50 +0100
Subject: [PATCH] Take for loops into account for MultipleContentEmitters rule

---
 .../compose/rules/MultipleContentEmitters.kt  | 24 ++++++++++++--
 .../MultipleContentEmittersCheckTest.kt       | 29 +++++++++++++++++
 .../MultipleContentEmittersCheckTest.kt       | 32 +++++++++++++++++++
 3 files changed, 82 insertions(+), 3 deletions(-)

diff --git a/rules/common/src/main/kotlin/io/nlopez/compose/rules/MultipleContentEmitters.kt b/rules/common/src/main/kotlin/io/nlopez/compose/rules/MultipleContentEmitters.kt
index 1120a6b9..5a4e2378 100644
--- a/rules/common/src/main/kotlin/io/nlopez/compose/rules/MultipleContentEmitters.kt
+++ b/rules/common/src/main/kotlin/io/nlopez/compose/rules/MultipleContentEmitters.kt
@@ -9,8 +9,10 @@ import io.nlopez.rules.core.util.emitsContent
 import io.nlopez.rules.core.util.findChildrenByClass
 import io.nlopez.rules.core.util.hasReceiverType
 import io.nlopez.rules.core.util.isComposable
+import org.jetbrains.kotlin.psi.KtBlockExpression
 import org.jetbrains.kotlin.psi.KtCallExpression
 import org.jetbrains.kotlin.psi.KtFile
+import org.jetbrains.kotlin.psi.KtForExpression
 import org.jetbrains.kotlin.psi.KtFunction
 
 class MultipleContentEmitters : ComposeKtVisitor {
@@ -69,11 +71,27 @@ class MultipleContentEmitters : ComposeKtVisitor {
     companion object {
         internal val KtFunction.directUiEmitterCount: Int
             get() = bodyBlockExpression?.let { block ->
-                block.statements
-                    .filterIsInstance<KtCallExpression>()
-                    .count { it.emitsContent }
+                // If there's content emitted in a for loop, we assume there's at
+                // least two iterations and thus count any emitters in them as multiple
+                val forLoopCount = when {
+                    block.forLoopHasUiEmitters -> 2
+                    else -> 0
+                }
+                block.directUiEmitterCount + forLoopCount
             } ?: 0
 
+        internal val KtBlockExpression.forLoopHasUiEmitters: Boolean
+            get() = statements.filterIsInstance<KtForExpression>().any {
+                when (val body = it.body) {
+                    is KtBlockExpression -> body.directUiEmitterCount > 0
+                    is KtCallExpression -> body.emitsContent
+                    else -> false
+                }
+            }
+
+        internal val KtBlockExpression.directUiEmitterCount: Int
+            get() = statements.filterIsInstance<KtCallExpression>().count { it.emitsContent }
+
         internal fun KtFunction.indirectUiEmitterCount(mapping: Map<KtFunction, Int>): Int {
             val bodyBlock = bodyBlockExpression ?: return 0
             return bodyBlock.statements
diff --git a/rules/detekt/src/test/kotlin/io/nlopez/compose/rules/detekt/MultipleContentEmittersCheckTest.kt b/rules/detekt/src/test/kotlin/io/nlopez/compose/rules/detekt/MultipleContentEmittersCheckTest.kt
index 0469f282..2ef973c8 100644
--- a/rules/detekt/src/test/kotlin/io/nlopez/compose/rules/detekt/MultipleContentEmittersCheckTest.kt
+++ b/rules/detekt/src/test/kotlin/io/nlopez/compose/rules/detekt/MultipleContentEmittersCheckTest.kt
@@ -166,4 +166,33 @@ class MultipleContentEmittersCheckTest {
             .hasStartSourceLocation(2, 5)
         assertThat(errors.first()).hasMessage(MultipleContentEmitters.MultipleContentEmittersDetected)
     }
+
+    @Test
+    fun `for loops are captured`() {
+        @Language("kotlin")
+        val code = """
+            @Composable
+            fun MultipleContent(texts: List<String>, modifier: Modifier = Modifier) {
+                for (text in texts) {
+                    Text(text)
+                }
+            }
+            @Composable
+            fun MultipleContent(otherTexts: List<String>, modifier: Modifier = Modifier) {
+                Text("text 1")
+                for (otherText in otherTexts) {
+                    Text(otherText)
+                }
+            }
+        """.trimIndent()
+        val errors = rule.lint(code)
+        assertThat(errors)
+            .hasStartSourceLocations(
+                SourceLocation(2, 5),
+                SourceLocation(8, 5),
+            )
+        for (error in errors) {
+            assertThat(error).hasMessage(MultipleContentEmitters.MultipleContentEmittersDetected)
+        }
+    }
 }
diff --git a/rules/ktlint/src/test/kotlin/io/nlopez/compose/rules/ktlint/MultipleContentEmittersCheckTest.kt b/rules/ktlint/src/test/kotlin/io/nlopez/compose/rules/ktlint/MultipleContentEmittersCheckTest.kt
index cb452df0..b8a5af26 100644
--- a/rules/ktlint/src/test/kotlin/io/nlopez/compose/rules/ktlint/MultipleContentEmittersCheckTest.kt
+++ b/rules/ktlint/src/test/kotlin/io/nlopez/compose/rules/ktlint/MultipleContentEmittersCheckTest.kt
@@ -171,4 +171,36 @@ class MultipleContentEmittersCheckTest {
             ),
         )
     }
+
+    @Test
+    fun `for loops are captured`() {
+        @Language("kotlin")
+        val code = """
+            @Composable
+            fun MultipleContent(texts: List<String>, modifier: Modifier = Modifier) {
+                for (text in texts) {
+                    Text(text)
+                }
+            }
+            @Composable
+            fun MultipleContent(otherTexts: List<String>, modifier: Modifier = Modifier) {
+                Text("text 1")
+                for (otherText in otherTexts) {
+                    Text(otherText)
+                }
+            }
+        """.trimIndent()
+        emittersRuleAssertThat(code).hasLintViolationsWithoutAutoCorrect(
+            LintViolation(
+                line = 2,
+                col = 5,
+                detail = MultipleContentEmitters.MultipleContentEmittersDetected,
+            ),
+            LintViolation(
+                line = 8,
+                col = 5,
+                detail = MultipleContentEmitters.MultipleContentEmittersDetected,
+            ),
+        )
+    }
 }