From 9d365b1db321196b3c02f063942cbe387da1bc63 Mon Sep 17 00:00:00 2001
From: Andrei Shikov <ashikov@google.com>
Date: Wed, 8 Nov 2023 12:22:54 +0000
Subject: [PATCH] Use ObjectIntMap for lazy key-index map

Using primitive collection to avoid boxing integers on each key-index map update.
(cherry picked from https://android-review.googlesource.com/q/commit:88c07301a8825a671643c470f8a690893c7483c1)
Merged-In: I138202435b9876ccef47912ffa1267d1b46b2acc
Change-Id: I138202435b9876ccef47912ffa1267d1b46b2acc
---
 compose/foundation/foundation/build.gradle           |  1 +
 .../foundation/lazy/layout/LazyLayoutKeyIndexMap.kt  | 12 ++++++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/compose/foundation/foundation/build.gradle b/compose/foundation/foundation/build.gradle
index ae0efec408d19..b9f07fc323d3f 100644
--- a/compose/foundation/foundation/build.gradle
+++ b/compose/foundation/foundation/build.gradle
@@ -35,6 +35,7 @@ androidXMultiplatform {
         commonMain {
             dependencies {
                 implementation(libs.kotlinStdlibCommon)
+                api(project(':collection:collection'))
                 api(project(':compose:animation:animation'))
                 api(project(':compose:runtime:runtime'))
                 api(project(':compose:ui:ui'))
diff --git a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutKeyIndexMap.kt b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutKeyIndexMap.kt
index 386b38990f9d4..437e6a6ad7c72 100644
--- a/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutKeyIndexMap.kt
+++ b/compose/foundation/foundation/src/commonMain/kotlin/androidx/compose/foundation/lazy/layout/LazyLayoutKeyIndexMap.kt
@@ -16,6 +16,9 @@
 
 package androidx.compose.foundation.lazy.layout
 
+import androidx.collection.MutableObjectIntMap
+import androidx.collection.ObjectIntMap
+import androidx.collection.emptyObjectIntMap
 import androidx.compose.foundation.ExperimentalFoundationApi
 
 /**
@@ -54,7 +57,7 @@ internal class NearestRangeKeyIndexMap(
     nearestRange: IntRange,
     intervalContent: LazyLayoutIntervalContent<*>
 ) : LazyLayoutKeyIndexMap {
-    private val map: Map<Any, Int>
+    private val map: ObjectIntMap<Any>
     private val keys: Array<Any?>
     private val keysStartIndex: Int
 
@@ -66,13 +69,14 @@ internal class NearestRangeKeyIndexMap(
         check(first >= 0) { "negative nearestRange.first" }
         val last = minOf(nearestRange.last, list.size - 1)
         if (last < first) {
-            map = emptyMap()
+            map = emptyObjectIntMap()
             keys = emptyArray()
             keysStartIndex = 0
         } else {
-            keys = arrayOfNulls<Any?>(last - first + 1)
+            val size = last - first + 1
+            keys = arrayOfNulls<Any?>(size)
             keysStartIndex = first
-            map = hashMapOf<Any, Int>().also { map ->
+            map = MutableObjectIntMap<Any>(size).also { map ->
                 list.forEach(
                     fromIndex = first,
                     toIndex = last,