xref: /aosp_15_r20/external/leakcanary2/shark-graph/src/main/java/shark/internal/SortedBytesMap.kt (revision d9e8da70d8c9df9a41d7848ae506fb3115cae6e6)

<lambda>null1 package shark.internal
2 
3 import shark.internal.hppc.LongObjectPair
4 import shark.internal.hppc.to
5 
6 /**
7  * A read only map of `id` => `byte array` sorted by id, where `id` is a long if [longIdentifiers]
8  * is true and an int otherwise. Each entry has a value byte array of size [bytesPerValue].
9  *
10  * Instances are created by [UnsortedByteEntries]
11  *
12  * [get] and [contains] perform a binary search to locate a specific entry by key.
13  */
14 internal class SortedBytesMap(
15   private val longIdentifiers: Boolean,
16   private val bytesPerValue: Int,
17   private val sortedEntries: ByteArray
18 ) {
19   private val bytesPerKey = if (longIdentifiers) 8 else 4
20   private val bytesPerEntry = bytesPerKey + bytesPerValue
21 
22   val size = sortedEntries.size / bytesPerEntry
23 
24   operator fun get(key: Long): ByteSubArray? {
25     val keyIndex = binarySearch(key)
26     if (keyIndex < 0) {
27       return null
28     }
29     return getAtIndex(keyIndex)
30   }
31 
32   fun indexOf(key: Long): Int {
33     return binarySearch(key)
34   }
35 
36   fun getAtIndex(keyIndex: Int): ByteSubArray {
37     val valueIndex = keyIndex * bytesPerEntry + bytesPerKey
38     return ByteSubArray(sortedEntries, valueIndex, bytesPerValue, longIdentifiers)
39   }
40 
41   operator fun contains(key: Long): Boolean {
42     val keyIndex = binarySearch(key)
43     return keyIndex >= 0
44   }
45 
46   fun entrySequence(): Sequence<LongObjectPair<ByteSubArray>> {
47     return (0 until size).asSequence()
48       .map { keyIndex ->
49         val valueIndex = keyIndex * bytesPerEntry + bytesPerKey
50         keyAt(keyIndex) to ByteSubArray(sortedEntries, valueIndex, bytesPerValue, longIdentifiers)
51       }
52   }
53 
54   private fun binarySearch(
55     key: Long
56   ): Int {
57     val startIndex = 0
58     val endIndex = size
59     var lo = startIndex
60     var hi = endIndex - 1
61     while (lo <= hi) {
62       val mid = (lo + hi).ushr(1)
63       val midVal = keyAt(mid)
64       when {
65         midVal < key -> lo = mid + 1
66         midVal > key -> hi = mid - 1
67         else -> return mid
68       }
69     }
70     return lo.inv()
71   }
72 
73   fun keyAt(index: Int): Long {
74     val keyIndex = index * bytesPerEntry
75     return if (longIdentifiers) {
76       sortedEntries.readLong(keyIndex)
77     } else {
78       sortedEntries.readInt(keyIndex).toLong()
79     }
80   }
81 }