core: generate fast collection for primitive types

core/kt-fast-collections-generator/src/main/kotlin/CollectionGenerator.kt

@@ -71,6 +71,10 @@ class StorageType(
             "UInt" -> return "0u"
             "Long" -> return "0L"
             "ULong" -> return "0uL"
+            "Float" -> return "0.0f"
+            "Double" -> return "0.0"
+            "Boolean" -> return "false"
+            "Char" -> return "' '"
         }
         throw RuntimeException("unknown primitive type")
     }

core/kt-fast-collections-generator/src/main/kotlin/PrimitiveSymbolProcessor.kt

@@ -0,0 +1,46 @@
+package fr.sncf.osrd.fast_collections.generator
+
+import com.google.devtools.ksp.processing.*
+import com.google.devtools.ksp.symbol.*
+import com.google.devtools.ksp.validate
+
+private const val ANNOTATION_PACKAGE = "fr.sncf.osrd.fast_collections"
+private const val ANNOTATION_SIMPLE_NAME = "PrimitiveCollections"
+private const val ANNOTATION_QUALIFIED_NAME = "${ANNOTATION_PACKAGE}.${ANNOTATION_SIMPLE_NAME}"
+
+private class PrimitiveSymbolProcessor(val context: GeneratorContext) : SymbolProcessor {
+    override fun process(resolver: Resolver): List<KSAnnotated> {
+        val symbols = resolver.getSymbolsWithAnnotation(ANNOTATION_QUALIFIED_NAME)
+        val invalidSymbols = arrayListOf<KSAnnotated>()
+        for (symbol in symbols) {
+            if (!symbol.validate()) {
+                invalidSymbols.add(symbol)
+                continue
+            }
+            val file = symbol as KSFile
+            for (annotation in file.annotations) {
+                if (annotation.shortName.getShortName() != ANNOTATION_SIMPLE_NAME)
+                    continue
+
+                val primitiveType = (annotation.arguments[0].value as KSType).declaration
+
+                val fromPrimitive = "(%s)"
+                val toPrimitive = "(%s)"
+                val toPrimitiveFun: (String) -> String = { toPrimitive.format(it) }
+                val fromPrimitiveFun: (String) -> String = { fromPrimitive.format(it) }
+                val collections = (annotation.arguments[1].value as List<*>).map { it as String }.toList()
+                generateCollections(context, file, primitiveType, primitiveType, toPrimitiveFun, fromPrimitiveFun, collections)
+            }
+        }
+        return invalidSymbols
+    }
+}
+
+class PrimitiveCollectionsProcessorProvider : SymbolProcessorProvider {
+    override fun create(
+        environment: SymbolProcessorEnvironment
+    ): SymbolProcessor {
+        val context = GeneratorContext(environment.codeGenerator, environment.logger)
+        return PrimitiveSymbolProcessor(context)
+    }
+}

core/kt-fast-collections-generator/src/main/kotlin/WrapperSymbolProcessor.kt

@@ -4,9 +4,9 @@ import com.google.devtools.ksp.processing.*
 import com.google.devtools.ksp.symbol.*
 import com.google.devtools.ksp.validate
 
-const val ANNOTATION_PACKAGE = "fr.sncf.osrd.fast_collections"
-const val ANNOTATION_SIMPLE_NAME = "PrimitiveWrapperCollections"
-const val ANNOTATION_QUALIFIED_NAME = "${ANNOTATION_PACKAGE}.${ANNOTATION_SIMPLE_NAME}"
+private const val ANNOTATION_PACKAGE = "fr.sncf.osrd.fast_collections"
+private const val ANNOTATION_SIMPLE_NAME = "PrimitiveWrapperCollections"
+private const val ANNOTATION_QUALIFIED_NAME = "${ANNOTATION_PACKAGE}.${ANNOTATION_SIMPLE_NAME}"
 
 fun addGenerator(generators: MutableSet<CollectionGenerator>, generatorId: String) {
     val generator = GENERATORS[generatorId] ?: throw RuntimeException("unknown generator: $generatorId")
@@ -19,10 +19,14 @@ fun addGenerator(generators: MutableSet<CollectionGenerator>, generatorId: Strin
 fun getArrayType(primitive: KSDeclaration): String {
     val primitiveQName = primitive.qualifiedName!!.asString()
     when (primitiveQName) {
-        "kotlin.Int" -> return "IntArray"
-        "kotlin.UInt" -> return "UIntArray"
-        "kotlin.Long" -> return "LongArray"
-        "kotlin.ULong" -> return "ULongArray"
+        "kotlin.Int" -> return "kotlin.IntArray"
+        "kotlin.UInt" -> return "kotlin.UIntArray"
+        "kotlin.Long" -> return "kotlin.LongArray"
+        "kotlin.ULong" -> return "kotlin.ULongArray"
+        "kotlin.Float" -> return "kotlin.FloatArray"
+        "kotlin.Double" -> return "kotlin.DoubleArray"
+        "kotlin.Boolean" -> return "kotlin.BooleanArray"
+        "kotlin.Char" -> return "kotlin.CharArray"
     }
     throw NotImplementedError("unsupported primitive type: $primitiveQName")
 }

core/kt-fast-collections-generator/src/main/resources/META-INF/services/com.google.devtools.ksp.processing.SymbolProcessorProvider

@@ -1 +1,2 @@
 fr.sncf.osrd.fast_collections.generator.PrimitiveWrapperCollectionsProcessorProvider
+fr.sncf.osrd.fast_collections.generator.PrimitiveCollectionsProcessorProvider

core/kt-fast-collections/src/main/kotlin/PrimitiveCollections.kt

@@ -0,0 +1,41 @@
+@file:PrimitiveCollections(
+    primitive = Int::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = UInt::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = Long::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = ULong::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = Float::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = Double::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = Boolean::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+@file:PrimitiveCollections(
+    primitive = Char::class,
+    collections = ["Array", "ArrayList", "ArraySortedSet"],
+)
+
+package fr.sncf.osrd.fast_collections