Implement sum extension function for UArrays (KT-28779)

libraries/tools/kotlin-stdlib-gen/src/templates/Arrays.kt

@@ -968,7 +968,7 @@ object ArrayOps : TemplateGroupBase() {
                     PrimitiveType.Char ->
                         body { "return withType(\"CharArray\", fillFrom(this, ${primitive}Array(newSize)))" }
                     PrimitiveType.Long ->
-                        body { "return withType(\"LongArray\", arrayCopyResize(this, newSize, ZERO))" }
+                        body { "return withType(\"LongArray\", arrayCopyResize(this, newSize, ${primitive!!.zero()}))" }
                     else ->
                         body { "return fillFrom(this, ${primitive}Array(newSize))" }
                 }

libraries/tools/kotlin-stdlib-gen/src/templates/Numeric.kt

@@ -5,31 +5,65 @@
 
 package templates
 
+import templates.Family.*
+
 object Numeric : TemplateGroupBase() {
 
     init {
         defaultBuilder {
             sequenceClassification(SequenceClass.terminal)
+            specialFor(ArraysOfUnsigned) {
+                since("1.3")
+                annotation("@ExperimentalUnsignedTypes")
+            }
         }
     }
 
     private val numericPrimitivesDefaultOrder = PrimitiveType.defaultPrimitives intersect PrimitiveType.numericPrimitives
+    private val summablePrimitives = numericPrimitivesDefaultOrder + PrimitiveType.unsignedPrimitives
 
     val f_sum = fn("sum()") {
-        Family.defaultFamilies.forEach { family -> include(family, numericPrimitivesDefaultOrder) }
+        listOf(Iterables, Sequences, ArraysOfObjects).forEach { include(it, summablePrimitives) }
+        include(ArraysOfPrimitives, numericPrimitivesDefaultOrder)
+        include(ArraysOfUnsigned)
     } builder {
+        val p = primitive!!
 
         doc { "Returns the sum of all elements in the ${f.collection}." }
-        returns("SUM")
-        platformName("sumOf<T>")
-        body {
-            """
-            var sum: SUM = ZERO
-            for (element in this) {
-                sum += element
+        returns(p.sumType().name)
+
+        specialFor(ArraysOfUnsigned) {
+            inlineOnly()
+
+            body {
+                if (p == p.sumType())
+                    "return storage.sum().to${p.sumType().name}()"
+                else
+                    "return sumBy { it.to${p.sumType().name}() }"
+            }
+        }
+        specialFor(Iterables, Sequences, ArraysOfObjects, ArraysOfPrimitives) {
+            platformName("sumOf<T>")
+
+            if (p.isUnsigned()) {
+                require(f != ArraysOfPrimitives) { "Arrays of unsigneds are separate from arrays of primitives." }
+                specialFor(Iterables) { sourceFile(SourceFile.UCollections) }
+                specialFor(Sequences) { sourceFile(SourceFile.USequences) }
+                specialFor(ArraysOfObjects) { sourceFile(SourceFile.UArrays) }
+
+                since("1.3")
+                annotation("@ExperimentalUnsignedTypes")
+            }
+
+            body {
+                """
+                var sum: ${p.sumType().name} = ${p.sumType().zero()}
+                for (element in this) {
+                    sum += element
+                }
+                return sum
+                """
             }
-            return sum
-            """
         }
     }
 

libraries/tools/kotlin-stdlib-gen/src/templates/dsl/CommonTypes.kt

@@ -75,6 +75,22 @@ enum class PrimitiveType {
 fun PrimitiveType.isIntegral(): Boolean = this in PrimitiveType.integralPrimitives
 fun PrimitiveType.isNumeric(): Boolean = this in PrimitiveType.numericPrimitives
 fun PrimitiveType.isFloatingPoint(): Boolean = this in PrimitiveType.floatingPointPrimitives
+fun PrimitiveType.isUnsigned(): Boolean = this in PrimitiveType.unsignedPrimitives
+
+fun PrimitiveType.sumType() = when (this) {
+    PrimitiveType.Byte, PrimitiveType.Short, PrimitiveType.Char -> PrimitiveType.Int
+    PrimitiveType.UByte, PrimitiveType.UShort -> PrimitiveType.UInt
+    else -> this
+}
+
+fun PrimitiveType.zero() = when (this) {
+    PrimitiveType.Double -> "0.0"
+    PrimitiveType.Float -> "0.0f"
+    PrimitiveType.Long -> "0L"
+    PrimitiveType.ULong -> "0uL"
+    in PrimitiveType.unsignedPrimitives -> "0u"
+    else -> "0"
+}
 
 enum class Inline {
     No,

libraries/tools/kotlin-stdlib-gen/src/templates/dsl/MemberBuilder.kt

@@ -194,28 +194,19 @@ class MemberBuilder(
                     "RECEIVER" -> receiver
                     "SELF" -> self
                     "PRIMITIVE" -> primitive?.name ?: token
-                    "SUM" -> {
-                        when (primitive) {
-                            PrimitiveType.Byte, PrimitiveType.Short, PrimitiveType.Char -> "Int"
-                            else -> primitive
-                        }
-                    }
-                    "ZERO" -> when (primitive) {
-                        PrimitiveType.Double -> "0.0"
-                        PrimitiveType.Float -> "0.0f"
-                        PrimitiveType.Long -> "0L"
-                        else -> "0"
-                    }
                     "ONE" -> when (primitive) {
                         PrimitiveType.Double -> "1.0"
                         PrimitiveType.Float -> "1.0f"
                         PrimitiveType.Long -> "1L"
+                        PrimitiveType.ULong -> "1uL"
+                        in PrimitiveType.unsignedPrimitives -> "1u"
                         else -> "1"
                     }
                     "-ONE" -> when (primitive) {
                         PrimitiveType.Double -> "-1.0"
                         PrimitiveType.Float -> "-1.0f"
                         PrimitiveType.Long -> "-1L"
+                        in PrimitiveType.unsignedPrimitives -> error("-ONE is not in the domain of unsigned primitives")
                         else -> "-1"
                     }
                     "TCollection" -> {

libraries/tools/kotlin-stdlib-gen/src/templates/dsl/SourceFile.kt

@@ -14,6 +14,7 @@ enum class SourceFile(jvmClassName: String? = null, val multifile: Boolean = tru
     Sets(packageName = "kotlin.collections"),
     Maps(packageName = "kotlin.collections"),
     Sequences(packageName = "kotlin.sequences"),
+    USequences(packageName = "kotlin.sequences"),
     Ranges(packageName = "kotlin.ranges"),
     URanges(packageName = "kotlin.ranges"),
     Comparisons(packageName = "kotlin.comparisons"),