feat: adjust selection to a meaningful code snippet (LLM-1202)

GitOrigin-RevId: 29e9d00c789ca7d64327a149323f248c1f01e050

platform/util/src/com/intellij/util/KtUtils.kt

@@ -56,4 +56,61 @@ inline fun <T1 : AutoCloseable?, T2 : AutoCloseable?, T3 : AutoCloseable, R>
     callsInPlace(block, InvocationKind.AT_MOST_ONCE)
   }
   return first.invoke().use { o1 -> second.invoke().use { o2 -> third.invoke().use { o3 -> block(o1, o2, o3) } } }
-}
+}
+
+/**
+ * @return a new sequence that contains the elements from the original sequence
+ * up to and including the element for which the [predicate] returns false.
+ */
+fun <T> Sequence<T>.takeWhileInclusive(predicate: (T) -> Boolean): Sequence<T> {
+  return TakeWhileInclusiveSequence(this, predicate)
+}
+
+internal class TakeWhileInclusiveSequence<T>(
+  private val sequence: Sequence<T>,
+  private val predicate: (T) -> Boolean
+) : Sequence<T> {
+
+  enum class NextState {
+    UNKNOWN,
+    DONE,
+    CONTINUE
+  }
+
+  override fun iterator(): Iterator<T> = object : Iterator<T> {
+    val iterator = sequence.iterator()
+    var currentState = NextState.UNKNOWN
+    var nextState: NextState = NextState.UNKNOWN
+    var nextItem: T? = null
+
+    private fun calcNext() {
+      if (iterator.hasNext()) {
+        val item = iterator.next()
+        nextItem = item
+        currentState = NextState.CONTINUE
+        nextState = if (predicate(item)) NextState.UNKNOWN else NextState.DONE
+      }
+      else {
+        currentState = NextState.DONE
+        nextState = NextState.DONE
+      }
+    }
+
+    override fun next(): T {
+      if (!hasNext()) throw NoSuchElementException()
+      @Suppress("UNCHECKED_CAST")
+      val result = nextItem as T
+
+      // Clean next to avoid keeping reference on yielded instance
+      nextItem = null
+      currentState = nextState
+      return result
+    }
+
+    override fun hasNext(): Boolean {
+      if (currentState == NextState.UNKNOWN)
+        calcNext() // will change nextState
+      return currentState == NextState.CONTINUE
+    }
+  }
+}