Shamrock: 快速序列化/反序列化 Protobuf

Signed-off-by: 白池 <[email protected]>
whitechi73 · Feb 23, 2024 · b4c40e2 · b4c40e2
1 parent 9bbcc2f
commit b4c40e2
Show file tree

Hide file tree

Showing 54 changed files with 380 additions and 175 deletions.
diff --git a/annotations/build.gradle.kts b/annotations/build.gradle.kts
@@ -6,4 +6,9 @@ plugins {
 java {
     sourceCompatibility = JavaVersion.VERSION_17
     targetCompatibility = JavaVersion.VERSION_17
+}
+
+dependencies {
+    implementation(DEPENDENCY_PROTOBUF)
+    implementation(kotlinx("serialization-protobuf", "1.6.2"))
 }
diff --git a/annotations/src/main/java/moe/fuqiuluo/symbols/Protobuf.kt b/annotations/src/main/java/moe/fuqiuluo/symbols/Protobuf.kt
@@ -0,0 +1,16 @@
+package moe.fuqiuluo.symbols
+
+import kotlinx.serialization.decodeFromByteArray
+import kotlinx.serialization.protobuf.ProtoBuf
+
+import kotlin.reflect.KClass
+
+interface Protobuf<T: Protobuf<T>>
+
+inline fun <reified T: Protobuf<T>> KClass<T>.decode(data: ByteArray): T {
+    return ProtoBuf.decodeFromByteArray(data)
+}
+
+inline fun <reified T: Protobuf<T>> ByteArray.decodeProtobuf(to: KClass<T>? = null): T {
+    return ProtoBuf.decodeFromByteArray(this)
+}
diff --git a/processor/build.gradle.kts b/processor/build.gradle.kts
@@ -1,6 +1,7 @@
 plugins {
     kotlin("jvm")
     id("com.google.devtools.ksp") version "1.9.21-1.0.15"
+    kotlin("plugin.serialization") version "1.9.21"
 }
 
 ksp {
@@ -14,5 +15,8 @@ dependencies {
     implementation("com.google.devtools.ksp:symbol-processing-api:1.9.21-1.0.15")
     implementation("com.squareup:kotlinpoet:1.14.2")
 
+    implementation(DEPENDENCY_PROTOBUF)
+    implementation(kotlinx("serialization-protobuf", "1.6.2"))
+
     ksp("dev.zacsweers.autoservice:auto-service-ksp:1.1.0")
 }
diff --git a/processor/src/main/java/moe/fuqiuluo/ksp/impl/OneBotHandlerProcessor.kt b/processor/src/main/java/moe/fuqiuluo/ksp/impl/OneBotHandlerProcessor.kt
@@ -6,6 +6,7 @@ package moe.fuqiuluo.ksp.impl
 import com.google.devtools.ksp.KspExperimental
 import com.google.devtools.ksp.getAnnotationsByType
 import com.google.devtools.ksp.getClassDeclarationByName
+import com.google.devtools.ksp.getKotlinClassByName
 import com.google.devtools.ksp.processing.CodeGenerator
 import com.google.devtools.ksp.processing.Dependencies
 import com.google.devtools.ksp.processing.KSPLogger
@@ -26,18 +27,18 @@ class OneBotHandlerProcessor(
 ): SymbolProcessor {
     override fun process(resolver: Resolver): List<KSAnnotated> {
         val ActionManagerNode = resolver.getClassDeclarationByName("moe.fuqiuluo.shamrock.remote.action.ActionManager")
-        if (ActionManagerNode == null) {
-            logger.error("OneBotHandlerProcessor: ActionManager not found")
-            return emptyList()
-        }
+            ?: resolver.getKotlinClassByName("moe.fuqiuluo.shamrock.remote.action.ActionManager")
+            ?: resolver.getClassDeclarationByName("ActionManager")
         val symbols = resolver.getSymbolsWithAnnotation(OneBotHandler::class.qualifiedName!!)
         val unableToProcess = symbols.filterNot { it.validate() }
-        val oneBotHandlers = (symbols.filter {
-            it is KSClassDeclaration && it.validate()  && it.classKind == ClassKind.OBJECT
-        } as Sequence<KSClassDeclaration>).toList()
+        if (ActionManagerNode != null) {
+            val oneBotHandlers = (symbols.filter {
+                it is KSClassDeclaration && it.validate() && it.classKind == ClassKind.OBJECT
+            } as Sequence<KSClassDeclaration>).toList()
 
-        if (oneBotHandlers.isNotEmpty()) {
-            ActionManagerNode.accept(ActionManagerVisitor(oneBotHandlers), Unit)
+            if (oneBotHandlers.isNotEmpty()) {
+                ActionManagerNode.accept(ActionManagerVisitor(oneBotHandlers), Unit)
+            }
         }
 
         return unableToProcess.toList()

diff --git a/processor/src/main/java/moe/fuqiuluo/ksp/impl/ProtobufProcessor.kt b/processor/src/main/java/moe/fuqiuluo/ksp/impl/ProtobufProcessor.kt
@@ -0,0 +1,77 @@
+@file:Suppress("UNCHECKED_CAST")
+package moe.fuqiuluo.ksp.impl
+
+import com.google.devtools.ksp.isInternal
+import com.google.devtools.ksp.isPrivate
+import com.google.devtools.ksp.processing.CodeGenerator
+import com.google.devtools.ksp.processing.Dependencies
+import com.google.devtools.ksp.processing.KSPLogger
+import com.google.devtools.ksp.processing.Resolver
+import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.symbol.ClassKind
+import com.google.devtools.ksp.symbol.KSAnnotated
+import com.google.devtools.ksp.symbol.KSClassDeclaration
+import com.google.devtools.ksp.symbol.KSDeclaration
+import com.google.devtools.ksp.validate
+import com.squareup.kotlinpoet.FileSpec
+import kotlinx.serialization.Serializable
+
+class ProtobufProcessor(
+    private val codeGenerator: CodeGenerator,
+    private val logger: KSPLogger
+): SymbolProcessor {
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        val symbols = resolver.getSymbolsWithAnnotation(Serializable::class.qualifiedName!!)
+        val unableToProcess = symbols.filterNot { it.validate() }
+        val actions = (symbols.filter {
+            it is KSClassDeclaration && it.validate()  && it.classKind == ClassKind.CLASS
+        } as Sequence<KSClassDeclaration>).filter {
+            it.superTypes.any { superType ->
+                superType.resolve().declaration.qualifiedName?.asString() == "moe.fuqiuluo.symbols.Protobuf"
+            }
+        }.toList()
+
+        if (actions.isNotEmpty()) {
+            actions.forEach { clz ->
+                if (clz.isInternal()) return@forEach
+                if (clz.isPrivate()) return@forEach
+
+                val packageName = "protobuf.auto"
+                val fileSpecBuilder = FileSpec.scriptBuilder("FastProtobuf", packageName)
+
+                fileSpecBuilder.addImport("kotlinx.serialization.protobuf", "ProtoBuf")
+                fileSpecBuilder.addImport("kotlinx.serialization", "decodeFromByteArray")
+                fileSpecBuilder.addImport("kotlinx.serialization", "encodeToByteArray")
+
+                if (clz.parentDeclaration != null) {
+                    fileSpecBuilder.addImport(clz.importPackage, clz.simpleName.asString())
+                } else {
+                    fileSpecBuilder.addImport(clz.packageName.asString(), clz.simpleName.asString())
+                }
+                if (clz.typeParameters.isNotEmpty()) {
+                    val genericType = clz.typeParameters.joinToString(", ") { it.name.asString() }
+                    fileSpecBuilder.addStatement("""inline fun <$genericType> ${clz.simpleName.asString()}<$genericType>.toByteArray() = ProtoBuf.encodeToByteArray(this)""")
+                } else {
+                    fileSpecBuilder.addStatement("inline fun ${clz.simpleName.asString()}.toByteArray() = ProtoBuf.encodeToByteArray(this)")
+                }
+
+                codeGenerator.createNewFile(
+                    dependencies = Dependencies.ALL_FILES,
+                    packageName = packageName,
+                    fileName = clz.simpleName.asString() + "\$FP"
+                ).use { outputStream ->
+                    outputStream.writer().use {
+                        fileSpecBuilder.build().writeTo(it)
+                    }
+                }
+            }
+        }
+
+        return unableToProcess.toList()
+    }
+
+    private val KSDeclaration.importPackage: String
+        get() = if (parentDeclaration != null) {
+            parentDeclaration!!.importPackage + "." + parentDeclaration!!.simpleName.asString()
+        } else packageName.asString()
+}
diff --git a/processor/src/main/java/moe/fuqiuluo/ksp/providers/ProtobufProvider.kt b/processor/src/main/java/moe/fuqiuluo/ksp/providers/ProtobufProvider.kt
@@ -0,0 +1,17 @@
+package moe.fuqiuluo.ksp.providers
+
+import com.google.auto.service.AutoService
+import com.google.devtools.ksp.processing.SymbolProcessor
+import com.google.devtools.ksp.processing.SymbolProcessorEnvironment
+import com.google.devtools.ksp.processing.SymbolProcessorProvider
+import moe.fuqiuluo.ksp.impl.ProtobufProcessor
+
+@AutoService(SymbolProcessorProvider::class)
+class ProtobufProvider: SymbolProcessorProvider {
+    override fun create(environment: SymbolProcessorEnvironment): SymbolProcessor {
+        return ProtobufProcessor(
+            environment.codeGenerator,
+            environment.logger
+        )
+    }
+}
diff --git a/protobuf/build.gradle.kts b/protobuf/build.gradle.kts
@@ -2,6 +2,7 @@ plugins {
     id("com.android.library")
     id("org.jetbrains.kotlin.android")
     kotlin("plugin.serialization") version "1.9.21"
+    id("com.google.devtools.ksp") version "1.9.21-1.0.15"
 }
 
 android {
@@ -38,4 +39,7 @@ dependencies {
     implementation(kotlinx("serialization-protobuf", "1.6.2"))
     implementation(kotlinx("serialization-json", "1.6.2"))
 
+    implementation(project(":annotations"))
+
+    ksp(project(":processor"))
 }
diff --git a/protobuf/src/main/java/protobuf/fav/WeiyunComm.kt b/protobuf/src/main/java/protobuf/fav/WeiyunComm.kt
@@ -2,9 +2,10 @@ package protobuf.fav
 
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 @Serializable
 data class WeiyunComm(
     @ProtoNumber(1) val req: WeiyunCommonReq? = null,
     @ProtoNumber(2) val resp: WeiyunCommonResp? = null
-)
+): Protobuf<WeiyunComm>
diff --git a/protobuf/src/main/java/protobuf/fav/WeiyunMsgHead.kt b/protobuf/src/main/java/protobuf/fav/WeiyunMsgHead.kt
@@ -6,6 +6,7 @@ package protobuf.fav
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 @Serializable
 data class WeiyunMsgHead(
@@ -27,4 +28,5 @@ data class WeiyunMsgHead(
     @ProtoNumber(103) val promptMsg: String? = null,
     @ProtoNumber(111) val totalSpace: ULong = ULong.MIN_VALUE,
     @ProtoNumber(112) val usedSpace: ULong = ULong.MIN_VALUE,
-)
+): Protobuf<WeiyunMsgHead>
+
diff --git a/protobuf/src/main/java/protobuf/guild/GetGuildFeeds.kt b/protobuf/src/main/java/protobuf/guild/GetGuildFeeds.kt
@@ -6,6 +6,7 @@ import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 import protobuf.qweb.QWebExtInfo
 
 @Serializable
@@ -19,15 +20,15 @@ data class GetGuildFeedsReq(
     @ProtoNumber(7) var u7: Int? = null,
     @ProtoNumber(8) var u8: Int? = null,
     @ProtoNumber(9) var u9: ByteArray? = null,
-)
+): Protobuf<GetGuildFeedsReq>
 
 @Serializable
 data class GetGuildFeedsRsp(
     @ProtoNumber(1) var vecFeed: List<StFeed>? = null,
     @ProtoNumber(2) var isFinish: Int = 0,
     //@ProtoNumber(3) var feedAttchInfo: ByteArray? = null,
     //@ProtoNumber(4) var traceId: String? = null,
-)
+): Protobuf<GetGuildFeedsRsp>
 
 @Serializable
 data class StFeed(

diff --git a/protobuf/src/main/java/protobuf/lightapp/AdaptShareInfo.kt b/protobuf/src/main/java/protobuf/lightapp/AdaptShareInfo.kt
@@ -0,0 +1,37 @@
+@file:OptIn(ExperimentalSerializationApi::class)
+package protobuf.lightapp
+
+import kotlinx.serialization.ExperimentalSerializationApi
+import kotlinx.serialization.Serializable
+import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
+
+@Serializable
+data class AdaptShareInfoReq(
+    //@ProtoNumber(1) var extInfo: Any? = null,
+    @ProtoNumber(2) var appid: String? = null,
+    @ProtoNumber(3) var title: String? = null,
+    @ProtoNumber(4) var desc: String? = null,
+    @ProtoNumber(5) var time: ULong? = null,
+    @ProtoNumber(6) var scene: UInt? = null,
+    @ProtoNumber(7) var templetType: UInt? = null,
+    @ProtoNumber(8) var businessType: UInt? = null,
+    @ProtoNumber(9) var picUrl: String? = null,
+    @ProtoNumber(10) var vidUrl: String? = null,
+    @ProtoNumber(11) var jumpUrl: String? = null,
+    @ProtoNumber(12) var iconUrl: String? = null,
+    @ProtoNumber(13) var verType: UInt? = null,
+    @ProtoNumber(14) var shareType: UInt? = null,
+    @ProtoNumber(15) var versionId: String? = null,
+    @ProtoNumber(16) var withShareTicket: UInt? = null,
+    @ProtoNumber(17) var webURL: String? = null,
+    //@ProtoNumber(18) var appidRich: Any? = null,
+    @ProtoNumber(19) var template: Template? = null,
+    //@ProtoNumber(20) var rcvOpenId: Any? = null,
+): Protobuf<AdaptShareInfoReq>
+
+@Serializable
+data class Template(
+    @ProtoNumber(1) var templateId: UInt? = null,
+    @ProtoNumber(2) var templateData: ByteArray? = null,
+)
diff --git a/protobuf/src/main/java/protobuf/message/longmsg/LongMsg.kt b/protobuf/src/main/java/protobuf/message/longmsg/LongMsg.kt
@@ -4,6 +4,7 @@ package protobuf.message.longmsg
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 
 @Serializable
@@ -39,14 +40,14 @@ data class LongMsgReq(
     @ProtoNumber(1) val recvInfo: RecvLongMsgInfo? = null,
     @ProtoNumber(2) val sendInfo: SendLongMsgInfo? = null,
     @ProtoNumber(15) val setting: LongMsgSettings? = null,
-)
+): Protobuf<LongMsgReq>
 
 @Serializable
 data class LongMsgRsp(
     @ProtoNumber(1) val recvResult: RecvLongMsgResult? = null,
     @ProtoNumber(2) val sendResult: SendLongMsgResult? = null,
     @ProtoNumber(15) val setting: LongMsgSettings? = null
-) {
+): Protobuf<LongMsgRsp> {
     companion object {
         @Serializable
         data class SendLongMsgResult(

diff --git a/protobuf/src/main/java/protobuf/message/longmsg/LongMsgPayload.kt b/protobuf/src/main/java/protobuf/message/longmsg/LongMsgPayload.kt
@@ -5,6 +5,7 @@ package protobuf.message.longmsg
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 import protobuf.message.MessageBody
 import protobuf.message.MessageContent
 import protobuf.message.MessageHead
@@ -29,4 +30,4 @@ data class LongMsgAction(
 @Serializable
 data class LongMsgPayload(
     @ProtoNumber(2) val action: List<LongMsgAction>? = null
-)
+): Protobuf<LongMsgPayload>
diff --git a/protobuf/src/main/java/protobuf/message/multimedia/RichMediaForPicData.kt b/protobuf/src/main/java/protobuf/message/multimedia/RichMediaForPicData.kt
@@ -4,12 +4,13 @@ package protobuf.message.multimedia
 import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 @Serializable
 data class RichMediaForPicData(
     @ProtoNumber(1) val info: MediaInfo?,
     @ProtoNumber(2) val display: DisplayMediaInfo?,
-) {
+): Protobuf<RichMediaForPicData> {
     companion object {
         @Serializable
         data class MediaInfo(

diff --git a/protobuf/src/main/java/protobuf/msg/SendMessage.kt b/protobuf/src/main/java/protobuf/msg/SendMessage.kt
@@ -2,6 +2,7 @@ package protobuf.msg
 
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 @Serializable
 class PbSendMsgReq(
@@ -19,7 +20,7 @@ class PbSendMsgReq(
     //@ProtoNumber(12)  var msgCtrl: MsgCtrl? = null,
     //@ProtoNumber(13)  var receipt_req: ReceiptReq? = null,
     //@ProtoNumber(14)  var multi_send_seq: UInt = 0u,
-)
+): Protobuf<PbSendMsgReq>
 
 @Serializable
 data class ContentHead(

diff --git a/protobuf/src/main/java/protobuf/oidb/TrpcOidb.kt b/protobuf/src/main/java/protobuf/oidb/TrpcOidb.kt
@@ -2,11 +2,12 @@ package protobuf.oidb
 
 import kotlinx.serialization.Serializable
 import kotlinx.serialization.protobuf.ProtoNumber
+import moe.fuqiuluo.symbols.Protobuf
 
 @Serializable
 data class TrpcOidb(
     @ProtoNumber(1) val cmd: Int = Int.MIN_VALUE,
     @ProtoNumber(2) val service: Int = Int.MIN_VALUE,
     @ProtoNumber(4) val buffer: ByteArray,
     @ProtoNumber(12) val flag: Int = Int.MIN_VALUE,
-)
+): Protobuf<TrpcOidb>