Implement `ConstrainSingle` in `MessageChainBuilder`

mirai-core/src/commonMain/kotlin/net.mamoe.mirai/message/data/MessageChain.kt

@@ -466,20 +466,35 @@ internal fun Sequence<SingleMessage>.constrainSingleMessages(): List<SingleMessa
 internal fun Iterable<SingleMessage>.constrainSingleMessages(): List<SingleMessage> {
     val list = ArrayList<SingleMessage>()
 
+    var firstConstrainIndex = -1
+
     for (singleMessage in this) {
         if (singleMessage is ConstrainSingle<*>) {
-            val key = singleMessage.key
-            val index = list.indexOfFirst { it is ConstrainSingle<*> && it.key == key }
-            if (index != -1) {
-                list[index] = singleMessage
-                continue
+            if (firstConstrainIndex == -1) {
+                firstConstrainIndex = list.size // we are going to add one
+            } else {
+                val key = singleMessage.key
+                val index = list.indexOfFirst(firstConstrainIndex) { it is ConstrainSingle<*> && it.key == key }
+                if (index != -1) {
+                    list[index] = singleMessage
+                    continue
+                }
             }
         }
+
         list.add(singleMessage)
     }
     return list
 }
 
+internal inline fun <T> List<T>.indexOfFirst(offset: Int, predicate: (T) -> Boolean): Int {
+    for (index in offset..this.lastIndex) {
+        if (predicate(this[index]))
+            return index
+    }
+    return -1
+}
+
 /**
  * 使用 [Collection] 作为委托的 [MessageChain]
  */

mirai-core/src/commonMain/kotlin/net.mamoe.mirai/message/data/builder.kt

@@ -13,9 +13,9 @@
 
 package net.mamoe.mirai.message.data
 
+import net.mamoe.mirai.utils.MiraiExperimentalAPI
 import kotlin.jvm.JvmMultifileClass
 import kotlin.jvm.JvmName
-import kotlin.jvm.JvmOverloads
 import kotlin.jvm.JvmSynthetic
 
 /**
@@ -38,42 +38,50 @@ inline fun buildMessageChain(initialSize: Int, block: MessageChainBuilder.() ->
     return MessageChainBuilder(initialSize).apply(block).asMessageChain()
 }
 
-/**
- * 使用特定的容器构建一个 [MessageChain]
- *
- * @see MessageChainBuilder
- */
-@JvmSynthetic
-inline fun buildMessageChain(
-    container: MutableCollection<SingleMessage>,
-    block: MessageChainBuilder.() -> Unit
-): MessageChain {
-    return MessageChainBuilder(container).apply(block).asMessageChain()
-}
-
 /**
  * [MessageChain] 构建器.
  * 多个连续的 [String] 会被连接为单个 [PlainText] 以优化性能.
+ * **注意:** 无并发安全性.
  *
  * @see buildMessageChain 推荐使用
  * @see asMessageChain 完成构建
  */
-class MessageChainBuilder
-@JvmOverloads constructor(
-    private val container: MutableCollection<SingleMessage> = mutableListOf()
-) : MutableCollection<SingleMessage> by container, Appendable {
+@OptIn(MiraiExperimentalAPI::class)
+class MessageChainBuilder private constructor(
+    private val container: MutableList<SingleMessage>
+) : MutableList<SingleMessage> by container, Appendable {
+    constructor() : this(mutableListOf())
     constructor(initialSize: Int) : this(ArrayList<SingleMessage>(initialSize))
 
+    private var firstConstrainSingleIndex = -1
+
+    private fun addAndCheckConstrainSingle(element: SingleMessage): Boolean {
+        if (element is ConstrainSingle<*>) {
+            if (firstConstrainSingleIndex == -1) {
+                firstConstrainSingleIndex = container.size
+                return container.add(element)
+            }
+            val key = element.key
+
+            container[container.indexOfFirst(firstConstrainSingleIndex) { it is ConstrainSingle<*> && it.key == key }] =
+                element
+            return true
+        } else {
+            return container.add(element)
+        }
+    }
+
     override fun add(element: SingleMessage): Boolean {
         flushCache()
-        return container.add(element)
+        return addAndCheckConstrainSingle(element)
     }
 
     fun add(element: Message): Boolean {
         flushCache()
         @Suppress("UNCHECKED_CAST")
         return when (element) {
-            is SingleMessage -> container.add(element)
+            is ConstrainSingle<*> -> addAndCheckConstrainSingle(element)
+            is SingleMessage -> container.add(element) // no need to constrain
             is Iterable<*> -> this.addAll(element.flatten())
             else -> error("stub")
         }
@@ -81,13 +89,13 @@ class MessageChainBuilder
 
     override fun addAll(elements: Collection<SingleMessage>): Boolean {
         flushCache()
-        return container.addAll(elements.flatten())
+        return addAll(elements.flatten())
     }
 
     @JvmName("addAllFlatten") // erased generic type cause declaration clash
     fun addAll(elements: Collection<Message>): Boolean {
         flushCache()
-        return container.addAll(elements.flatten())
+        return addAll(elements.flatten())
     }
 
     operator fun Message.unaryPlus() {
@@ -146,6 +154,6 @@ class MessageChainBuilder
 
     fun asMessageChain(): MessageChain {
         this.flushCache()
-        return (this as MutableCollection<SingleMessage>).asMessageChain()
+        return MessageChainImplByCollection(this) // fast-path, no need to constrain
     }
 }

mirai-core/src/commonTest/kotlin/net/mamoe/mirai/message.data/ConstrainSingleTest.kt

@@ -15,44 +15,44 @@ import kotlin.test.assertEquals
 import kotlin.test.assertSame
 
 
-internal class ConstrainSingleTest {
+@OptIn(MiraiExperimentalAPI::class)
+internal class TestConstrainSingleMessage : ConstrainSingle<TestConstrainSingleMessage>, Any() {
+    companion object Key : Message.Key<TestConstrainSingleMessage> {
+        override val typeName: String
+            get() = "TestMessage"
+    }
 
-    @OptIn(MiraiExperimentalAPI::class)
-    internal class TestMessage : ConstrainSingle<TestMessage>, Any() {
-        companion object Key : Message.Key<TestMessage> {
-            override val typeName: String
-                get() = "TestMessage"
-        }
-
-        override fun toString(): String = super.toString()
-
-        override fun contentToString(): String {
-            TODO("Not yet implemented")
-        }
-
-        override val key: Message.Key<TestMessage>
-            get() = Key
-        override val length: Int
-            get() = TODO("Not yet implemented")
-
-        override fun get(index: Int): Char {
-            TODO("Not yet implemented")
-        }
-
-        override fun subSequence(startIndex: Int, endIndex: Int): CharSequence {
-            TODO("Not yet implemented")
-        }
-
-        override fun compareTo(other: String): Int {
-            TODO("Not yet implemented")
-        }
+    override fun toString(): String = "<TestConstrainSingleMessage#${super.hashCode()}>"
+
+    override fun contentToString(): String {
+        TODO("Not yet implemented")
+    }
+
+    override val key: Message.Key<TestConstrainSingleMessage>
+        get() = Key
+    override val length: Int
+        get() = TODO("Not yet implemented")
+
+    override fun get(index: Int): Char {
+        TODO("Not yet implemented")
     }
 
+    override fun subSequence(startIndex: Int, endIndex: Int): CharSequence {
+        TODO("Not yet implemented")
+    }
+
+    override fun compareTo(other: String): Int {
+        TODO("Not yet implemented")
+    }
+}
+
+internal class ConstrainSingleTest {
+
     @OptIn(MiraiExperimentalAPI::class)
     @Test
     fun testConstrainSingleInPlus() {
-        val new = TestMessage()
-        val combined = TestMessage() + new
+        val new = TestConstrainSingleMessage()
+        val combined = TestConstrainSingleMessage() + new
 
         assertEquals(combined.left, EmptyMessageChain)
         assertSame(combined.tail, new)
@@ -60,10 +60,10 @@ internal class ConstrainSingleTest {
 
     @Test // net.mamoe.mirai/message/data/MessageChain.kt:441
     fun testConstrainSingleInSequence() {
-        val last = TestMessage()
+        val last = TestConstrainSingleMessage()
         val sequence: Sequence<SingleMessage> = sequenceOf(
-            TestMessage(),
-            TestMessage(),
+            TestConstrainSingleMessage(),
+            TestConstrainSingleMessage(),
             last
         )
 
@@ -74,11 +74,11 @@ internal class ConstrainSingleTest {
 
     @Test // net.mamoe.mirai/message/data/MessageChain.kt:441
     fun testConstrainSingleOrderInSequence() {
-        val last = TestMessage()
+        val last = TestConstrainSingleMessage()
         val sequence: Sequence<SingleMessage> = sequenceOf(
-            TestMessage(), // last should replace here
+            TestConstrainSingleMessage(), // last should replace here
             PlainText("test"),
-            TestMessage(),
+            TestConstrainSingleMessage(),
             last
         )
 

mirai-core/src/commonTest/kotlin/net/mamoe/mirai/message.data/MessageChainBuilderTest.kt

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2020 Mamoe Technologies and contributors.
+ *
+ * 此源代码的使用受 GNU AFFERO GENERAL PUBLIC LICENSE version 3 许可证的约束, 可以在以下链接找到该许可证.
+ * Use of this source code is governed by the GNU AGPLv3 license that can be found through the following link.
+ *
+ * https://github.com/mamoe/mirai/blob/master/LICENSE
+ */
+
+package net.mamoe.mirai.message.data
+
+import kotlin.test.Test
+import kotlin.test.assertEquals
+
+internal class MessageChainBuilderTest {
+    @Test
+    fun testConcat() {
+        val chain = buildMessageChain {
+            +"test"
+            +" "
+            +PlainText("foo")
+            +" "
+            +(PlainText("bar") + " goo ")
+            +buildMessageChain {
+                +"1"
+                +"2"
+                +"3"
+            }
+        }
+
+        assertEquals("test foo bar goo 123", chain.toString())
+    }
+
+    @Test
+    fun testConstrain() {
+        val lastSingle = TestConstrainSingleMessage()
+
+        val chain = buildMessageChain {
+            +"test"
+            +TestConstrainSingleMessage()
+            +TestConstrainSingleMessage()
+            +PlainText("foo")
+            +TestConstrainSingleMessage()
+            +lastSingle
+        }
+
+        assertEquals(chain.size, 3)
+        assertEquals("test${lastSingle}foo", chain.toString())
+    }
+}