fix(level-hash/persistent): restructure keymap file for minimal disk usage

A keymap file of ~3.75MB can store 491520 value addresses (slots). The previous keymap layout resulted in internal fragmentation where a block (4KB in this case) was allocated for a single value address, resulting in wasted disk space. The new restructured layout also makes it possible to resize the level hash easily.
2026-10-16 07:49:44 +08:00 · 2024-07-31 10:42:45 +05:30
parent 3270f0e014
commit c300fc87e6
12 changed files with 448 additions and 625 deletions
@@ -21,8 +21,28 @@ package com.itsaky.androidide.levelhash
 * Level hash is a write-optimized and high-performance hashing index scheme with cost-efficient
 * resizing and low-overhead consistency guarantee for persistent memory.
 *
+ * **Level size and bucket size**
+ *
+ * The level size is used to determine the total capacity of the level hash,
+ * while the bucket size is the number of slots in each bucket. These two
+ * properties define the overall performance (load factor) of the level hash
+ * (in terms of hash collisions and insertion failures). As a result, the values
+ * for these properties must be chosen carefully.
+ *
+ * ```
+ * Load factor = occupied slots / total slots
+ * ```
+ *
+ * When the load factor of the level hash increases, the hash collisions increase
+ * as well. This results in decreased access performance and insertion failures.
+ * It is **recommended** to enable [AbstractLevelHashBuilder.autoExpand] which
+ * will automatically expand the level hash when the load factor reaches a certain
+ * threshold (specifically, when the top level is full).
+ *
 * @param K The key type.
 * @param V The value type.
+ *
+ * @author Akash Yadav
 */
 interface LevelHash<K : Any, V : Any?> : AutoCloseable {

@@ -127,6 +147,11 @@ interface LevelHash<K : Any, V : Any?> : AutoCloseable {

  companion object {

+    /**
+     * The default level size.
+     */
+    const val LEVEL_SIZE_DEFAULT = 3
+
    /**
     * The default level size.
     */
@@ -163,12 +188,12 @@ interface LevelHash<K : Any, V : Any?> : AutoCloseable {
   * An [InsertionError] which is thrown when the level hash is full and unable
   * to insert new values.
   */
-  class OverflowException : InsertionError("Level hash has reached is capacity")
+  class OverflowException(capacity: Int, occupied: Int) : InsertionError("Level hash has reached is capacity. cap=$capacity occ=$occupied")

  /**
   * Thrown when the level hash cannot be resized.
   */
-  class ResizeFailure(msg: String) : InsertionError(msg)
+  open class ResizeFailure(msg: String) : InsertionError(msg)

  /**
   * Thrown when *update* operation fails to find the slot in a [LevelHash].
@@ -42,7 +42,7 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
  protected val seeds: Pair<HashT, HashT>,
 ) : LevelHash<K, V> {

-  final override var levelSize: Int = levelSize
+  override var levelSize: Int = levelSize
    set(value) {
      require(value <= 30) { "Level size must be <= 30" }
      field = value
@@ -119,13 +119,13 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo

  override fun insert(key: K, value: V): Boolean {

-    if (this.levelItemCounts.sum() == this.totalSlotCount) {
+    if (this.levelItemCounts[0] == this.topLevelBucketCount * this.bucketSize) {
      if (autoExpand) {
        check(expand(1)) {
          "Failed to expand the level hash"
        }
      } else {
-        throw OverflowException()
+        throw OverflowException(this.totalSlotCount, this.levelItemCounts.sum())
      }
    }

@@ -266,33 +266,39 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
    for (oldBuckIdx in 0..<(this.topLevelBucketCount shr 1)) {
      for (oldSlotIdx in 0..<this.bucketSize) {
        val oldSlot = getSlot(Level.BOTTOM, oldBuckIdx, oldSlotIdx)
-        if (oldSlot.isOccupied()) {
-          val fhash = fhash(oldSlot.key)
-          val shash = shash(oldSlot.key)
-          val fidx = buckIdxCap(fhash, newTopLevelCapacity)
-          val sidx = buckIdxCap(shash, newTopLevelCapacity)
-
-          var insertSuccess = false
-          for (newSlotIdx in 0..<this.bucketSize) {
-            if (tryMoveToInterim(oldSlot, fidx, newSlotIdx) || tryMoveToInterim(
-                oldSlot, sidx, newSlotIdx)
-            ) {
-              insertSuccess = true
-              newLevelItemCount++
-              break
-            }
-          }
-
-          if (!insertSuccess) {
-            throw ResizeFailure("Failed to move slot to the interim level")
-          }
-
-          oldSlot.reset(null, null)
+        if (!oldSlot.isOccupied()) {
+          continue
        }
+
+        val fhash = fhash(oldSlot.key)
+        val shash = shash(oldSlot.key)
+        val fidx = buckIdxCap(fhash, newTopLevelCapacity)
+        val sidx = buckIdxCap(shash, newTopLevelCapacity)
+
+        var insertSuccess = false
+        for (newSlotIdx in 0..<this.bucketSize) {
+          if (moveForExpansion(oldSlot, fidx, newSlotIdx) || moveForExpansion(
+              oldSlot, sidx, newSlotIdx)
+          ) {
+            insertSuccess = true
+            newLevelItemCount++
+            break
+          }
+        }
+
+        if (!insertSuccess) {
+          throw ResizeFailure("Failed to move slot to the interim level")
+        }
+
+        oldSlot.reset(null, null)
      }
    }

+    this.levelSize = levelSize
+    this.topLevelBucketCount = 2.0.pow(levelSize).toInt()
+
    onExpand(levelSize, newLevelItemCount)
+
    this.levelItemCounts[1] = this.levelItemCounts[0]
    this.levelItemCounts[0] = newLevelItemCount
    this.expandCount++
@@ -306,7 +312,7 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
   *
   * @param bucketCount The bucket count of the interim level.
   */
-  protected abstract fun prepareExpansion(bucketCount: Int)
+  protected open fun prepareExpansion(bucketCount: Int) {}

  /**
   * Called to notify that the expand operation was successful.
@@ -315,8 +321,7 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
   * @param interimItemCount The number of slot that are occupied in the
   * interim level.
   */
-  protected abstract fun onExpand(newLevelSize: Int, interimItemCount: Int
-  )
+  protected open fun onExpand(newLevelSize: Int, interimItemCount: Int) {}

  /**
   * Move the given slot to the interim level at the given slot in the given bucket.
@@ -330,8 +335,8 @@ internal abstract class AbstractLevelHash<K : Any, V : Any?> internal constructo
   * slot should be moved.
   * @return Whether the slot was moved to the interim level.
   */
-  protected abstract fun tryMoveToInterim(slot: LevelSlot<K, V>,
-                                          bucketIdx: Int, slotIdx: Int
+  protected abstract fun moveForExpansion(slot: LevelSlot<K, V>,
+                                  bucketIdx: Int, slotIdx: Int
  ): Boolean

  /**
@@ -95,8 +95,8 @@ internal class InMemoryLevelHash<K : Any, V : Any?>(
    this.interimLevel = List(bucketCount) { InMemoryLevelBucket(bucketSize) }
  }

-  override fun tryMoveToInterim(slot: LevelSlot<K, V>, bucketIdx: Int,
-                                slotIdx: Int
+  override fun moveForExpansion(slot: LevelSlot<K, V>, bucketIdx: Int,
+                        slotIdx: Int
  ): Boolean {
    val interimLevel = checkNotNull(interimLevel)
    slot as InMemoryLevelSlot<K, V>
@@ -113,9 +113,6 @@ internal class InMemoryLevelHash<K : Any, V : Any?>(
  }

  override fun onExpand(newLevelSize: Int, interimItemCount: Int) {
-    this.levelSize = newLevelSize
-    this.topLevelBucketCount = 2.0.pow(newLevelSize).toInt()
-
    // after expand, interim level becomes the new top level
    // and the current top level becomes the new bottom level
    // the current bottom level is released
@@ -1,121 +0,0 @@
-/*
- *  This file is part of AndroidIDE.
- *
- *  AndroidIDE is free software: you can redistribute it and/or modify
- *  it under the terms of the GNU General Public License as published by
- *  the Free Software Foundation, either version 3 of the License, or
- *  (at your option) any later version.
- *
- *  AndroidIDE is distributed in the hope that it will be useful,
- *  but WITHOUT ANY WARRANTY; without even the implied warranty of
- *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *  GNU General Public License for more details.
- *
- *  You should have received a copy of the GNU General Public License
- *   along with AndroidIDE.  If not, see <https://www.gnu.org/licenses/>.
- */
-
-package com.itsaky.androidide.levelhash.internal
-
-import com.itsaky.androidide.levelhash.RandomAccessIO
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO.Companion.KEYMAP_SEGMENT_SIZE_BYTES
-import com.itsaky.androidide.levelhash.util.MappedRandomAccessIO
-import org.slf4j.LoggerFactory
-import java.io.DataInput
-import java.io.DataOutput
-import java.io.RandomAccessFile
-import java.nio.MappedByteBuffer
-import java.nio.channels.FileChannel
-
-/**
- * A segment of the key map file.
- */
-internal class KeymapSegment(
-  val file: RandomAccessFile,
-  val path: String,
-  val segmentOffset: Long,
-  val segmentLength: Long,
-  val io: MappedRandomAccessIO = MappedRandomAccessIO(),
-) : DataInput by io, DataOutput by io, RandomAccessIO by io, AutoCloseable {
-
-  lateinit var mappedBuffer: MappedByteBuffer
-
-  init {
-    remap(segmentOffset, segmentLength, true)
-  }
-
-  /**
-   * Map a region of the backing file into memory, unmapping previous mappings,
-   * if any.
-   */
-  @Suppress("SameParameterValue")
-  private fun remap(position: Long, size: Long, forceRemap: Boolean = false
-  ) {
-    if (!forceRemap && (segmentOffset == position && segmentLength == size)) {
-      // the given region is already mapped into memory
-      return
-    }
-
-    unmap()
-
-    this.mappedBuffer =
-      file.channel.map(FileChannel.MapMode.READ_WRITE, segmentOffset,
-        segmentLength)
-    this.io.reset(mappedBuffer, segmentOffset, segmentLength)
-  }
-
-  /**
-   * Unmap the memory mapped region, if mapped.
-   */
-
-  override fun close() {
-    unmap()
-  }
-
-  /**
-   * Unmap the memory mapped buffer. This [KeymapSegment] becomes invalid after
-   * this operation and must not be used further unless remapped.
-   *
-   * @return `true` if the unmapping request was successful, `false` otherwise.
-   */
-  fun unmap(): Boolean {
-    if (!::mappedBuffer.isInitialized) {
-      return false
-    }
-
-    try {
-      val cleanerF = mappedBuffer.javaClass.getDeclaredField("cleaner")
-      cleanerF.isAccessible = true
-      val cleaner = cleanerF.get(mappedBuffer)
-
-      val cleanM = cleaner.javaClass.getDeclaredMethod("clean")
-      cleanM.isAccessible = true
-
-      cleanM.invoke(cleaner)
-      return true
-    } catch (err: Throwable) {
-      log.error("Failed to unmap keymap segment. file={}, offset={}, size={}",
-        path, segmentOffset, segmentLength, err)
-      err.printStackTrace()
-      return false
-    }
-  }
-
-  companion object {
-
-    private val log = LoggerFactory.getLogger(KeymapSegment::class.java)
-
-    /**
-     * Create a new segment.
-     *
-     * @param offset The offset of the segment in the keymap file.
-     * @param file The file to create the segment in.
-     * @param path The path of the file (used for debugging).
-     */
-    fun create(offset: Long, file: RandomAccessFile, path: String
-    ): KeymapSegment {
-      return KeymapSegment(file = file, path = path, segmentOffset = offset,
-        segmentLength = KEYMAP_SEGMENT_SIZE_BYTES)
-    }
-  }
-}
@@ -1,208 +0,0 @@
-/*
- *  This file is part of AndroidIDE.
- *
- *  AndroidIDE is free software: you can redistribute it and/or modify
- *  it under the terms of the GNU General Public License as published by
- *  the Free Software Foundation, either version 3 of the License, or
- *  (at your option) any later version.
- *
- *  AndroidIDE is distributed in the hope that it will be useful,
- *  but WITHOUT ANY WARRANTY; without even the implied warranty of
- *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- *  GNU General Public License for more details.
- *
- *  You should have received a copy of the GNU General Public License
- *   along with AndroidIDE.  If not, see <https://www.gnu.org/licenses/>.
- */
-
-package com.itsaky.androidide.levelhash.internal
-
-import com.itsaky.androidide.levelhash.HashT
-import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.MAGIC_NUMBER_SIZE_BYTES
-import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_INT
-import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_LONG
-import com.itsaky.androidide.levelhash.util.FileUtils
-import org.slf4j.LoggerFactory
-import java.io.File
-import java.io.RandomAccessFile
-
-/**
- * @author Akash Yadav
- */
-internal class PersistentKeymapIO(indexFile: File
-) : AutoCloseable {
-  // TODO(itsaky): use LFU cache to limit the number of segments mapped in memory at once.
-  /**
-   * The segments for the keymap file.
-   */
-  private val keymapSegments = arrayOfNulls<KeymapSegment>(KEYMAP_MAX_SEGMENTS)
-  private val keymapFile = File(indexFile.parentFile, "${indexFile.name}._i")
-  private val raKeymapFile by lazy { RandomAccessFile(keymapFile, "rw") }
-
-  private val segmentCount by lazy {
-    if (!keymapFile.exists() || !keymapFile.isFile || keymapFile.length() <= KEYMAP_HEADER_SIZE_BYTES) {
-      // there should be 1 segment initially
-      return@lazy 1
-    }
-
-    raKeymapFile.seek(KEYMAP_HEADER_SIZE_BYTES)
-    val count = raKeymapFile.readLong()
-    raKeymapFile.seek(0)
-    return@lazy count
-  }
-
-  init {
-    BinaryFileUtils.initSparseFile(
-      file = keymapFile,
-      magicNumber = KEYMAP_MAGIC_NUMBER
-    ) { raf ->
-      var position = MAGIC_NUMBER_SIZE_BYTES
-      raf.seek(position)
-      raf.writeInt(0)
-      position += SIZE_INT
-
-      raf.seek(position)
-      raf.writeLong(KEYMAP_SEGMENT_SIZE_BYTES)
-    }
-  }
-
-  fun getSegmentOffset(index: Int): Long {
-    return KEYMAP_HEADER_SIZE_BYTES + index * KEYMAP_SEGMENT_SIZE_BYTES
-  }
-
-  /**
-   * Get or create a segment for the keymap file.
-   */
-  fun getOrCreateSegment(index: Int): KeymapSegment {
-    var segment = keymapSegments[index]
-    if (segment == null) {
-      segment = KeymapSegment.create(getSegmentOffset(index), raKeymapFile,
-        keymapFile.canonicalPath)
-      keymapSegments[index] = segment
-    }
-    return segment
-  }
-
-  fun segmentAndPosForPos(
-    levelIdx: Int,
-    bucketIdx: Int,
-    slotIdx: Int,
-  ): Pair<KeymapSegment, Long> {
-    val hash = hash(levelIdx, bucketIdx, slotIdx)
-    val segmentIdx = segmentIndex(hash)
-    val positionInSegment = positionInSegment(hash)
-    val segment = getOrCreateSegment(segmentIdx)
-    return segment to positionInSegment
-  }
-
-  fun deallocate(offset: Long, len: Long
-  ) {
-    FileUtils.deallocate(raKeymapFile, offset, len)
-  }
-
-  fun deallocate(segment: KeymapSegment, positionInSegment: Long, len: Long
-  ) {
-    FileUtils.deallocate(raKeymapFile,
-      segment.segmentOffset + positionInSegment, len)
-  }
-
-  override fun close() {
-    keymapSegments.forEach { segment ->
-      try {
-        segment?.unmap()
-      } catch (err: Throwable) {
-        log.error("Failed to unmap segment", err)
-      }
-    }
-
-    try {
-      raKeymapFile.close()
-    } catch (e: Throwable) {
-      log.error("Failed to close keymap file", e)
-    }
-  }
-
-  companion object {
-
-    private val log = LoggerFactory.getLogger(PersistentKeymapIO::class.java)
-
-    /**
-     * The number of bytes in the header of the keymap file that is used to store
-     * the segment count.
-     */
-    internal const val KEYMAP_HEADER_SEGCOUNT_BYTES = SIZE_INT
-
-    /**
-     * The number of bytes in the header of the keymap file that is used to store
-     * the size of the segment structures.
-     */
-    internal const val KEYMAP_HEADER_SEGSIZE_BYTES = SIZE_LONG
-
-    /**
-     * The number of bytes used to store the header of the keymap file.
-     */
-    internal const val KEYMAP_HEADER_SIZE_BYTES: Long =
-      MAGIC_NUMBER_SIZE_BYTES + KEYMAP_HEADER_SEGCOUNT_BYTES + KEYMAP_HEADER_SEGSIZE_BYTES
-
-    /**
-     * The size of one segment in the key map file.
-     */
-    internal const val KEYMAP_SEGMENT_SIZE_BYTES: Long = 1024 * 1024 // 1MB
-
-    /**
-     * The number of bytes that are used to store an entry in a key map file.
-     */
-    internal const val KEYMAP_ENTRY_BYTES: Long = SIZE_LONG
-
-    /**
-     * The number of entries that can be stored in a segment in the key map file.
-     */
-    internal const val KEYMAP_ENTRIES_PER_SEGMENT =
-      KEYMAP_SEGMENT_SIZE_BYTES / KEYMAP_ENTRY_BYTES
-
-    /**
-     * The maximum number of segments that can be stored for an index.
-     */
-    internal const val KEYMAP_MAX_SEGMENTS = 100
-
-    /**
-     * The maximum number of entries that can be stored in all of the keymap segments combined.
-     */
-    const val KEYMAP_MAX_ENTRIES =
-      KEYMAP_MAX_SEGMENTS * KEYMAP_ENTRIES_PER_SEGMENT
-
-    /**
-     * Magic number that is used as the file signature to identify the keymap file.
-     */
-    internal const val KEYMAP_MAGIC_NUMBER = 0x414944584B
-
-    /**
-     * Get the ID of the segment that should be used to store the store the
-     * given position hash.
-     */
-    private fun segmentIndex(positionHash: HashT): Int {
-      return (positionHash % KEYMAP_MAX_SEGMENTS.toULong()).toInt()
-    }
-
-    /**
-     * Get the position of the entry in the segment where the entry should be
-     * written. The returned position is in bytes from the beginning of the
-     * segment.
-     */
-    private fun positionInSegment(positionHash: HashT): Long {
-      return ((positionHash % KEYMAP_ENTRIES_PER_SEGMENT.toULong()).toLong() * KEYMAP_ENTRY_BYTES)
-    }
-
-    /**
-     * Get the hash of the position coordinates.
-     */
-    private fun hash(levelNum: Int, bucketIdx: Int, slotIdx: Int
-    ): HashT {
-      var hash = 17L
-      hash = 31L * hash + levelNum
-      hash = 31L * hash + bucketIdx
-      hash = 31L * hash + slotIdx
-      return hash.toULong()
-    }
-  }
-}
@@ -22,7 +22,6 @@ import com.itsaky.androidide.levelhash.HashT
 import com.itsaky.androidide.levelhash.LevelHashFn
 import com.itsaky.androidide.levelhash.LevelSlot
 import java.io.File
-import kotlin.math.pow

 /**
 * An file-based persistent implementation of [AbstractLevelHash].
@@ -39,35 +38,22 @@ internal class PersistentLevelHash<K : Any, V : Any?>(
  keyExternalizer: DataExternalizer<K>,
  valueExternalizer: DataExternalizer<V>,
  indexFile: File,
-) : AbstractLevelHash<K, V>(
-  levelSize = levelSize,
-  bucketSize = bucketSize,
-  uniqueKeys = uniqueKeys,
-  autoExpand = autoExpand,
-  levelHashFn = levelHashFn,
-  seeds = seeds
-) {
+) : AbstractLevelHash<K, V>(levelSize = levelSize, bucketSize = bucketSize,
+  uniqueKeys = uniqueKeys, autoExpand = autoExpand, levelHashFn = levelHashFn,
+  seeds = seeds) {

-  private val io =
-    PersistentLevelHashIO(indexFile, keyExternalizer, valueExternalizer)
+  private val io = PersistentLevelHashIO(indexFile, levelSize, bucketSize, keyExternalizer,
+    valueExternalizer)
+
+  override var levelSize: Int
+    get() = io.metaIo.levelSize
+    set(value) {
+      io.metaIo.levelSize = value
+    }

  override fun getSlot(levelIdx: Int, bucketIdx: Int, slotIdx: Int
  ): LevelSlot<K, V> {
-    return object : LevelSlot<K, V> {
-      override val key: K
-        get() = checkNotNull(io.readKey(levelIdx, bucketIdx, slotIdx))
-
-      override val value: V?
-        get() = io.readValue(levelIdx, bucketIdx, slotIdx)
-
-      override fun isOccupied(): Boolean {
-        return io.isOccupied(levelIdx, bucketIdx, slotIdx)
-      }
-
-      override fun reset(key: K?, value: V?) {
-        io.writeEntry(levelIdx, bucketIdx, slotIdx, key, value)
-      }
-    }
+    return PersistentLevelSlot(levelIdx, bucketIdx, slotIdx, io)
  }

  override fun clear() {
@@ -78,22 +64,31 @@ internal class PersistentLevelHash<K : Any, V : Any?>(
    io.close()
  }

-  private var interimBucketCount = -1
-  override fun prepareExpansion(bucketCount: Int) {
-    check(this.interimBucketCount == -1)
-    this.interimBucketCount = bucketCount
-  }
-
-  override fun tryMoveToInterim(
-    slot: LevelSlot<K, V>,
-    bucketIdx: Int,
-    slotIdx: Int
+  override fun moveForExpansion(slot: LevelSlot<K, V>, bucketIdx: Int,
+                                slotIdx: Int
  ): Boolean {
-    TODO()
+    slot as PersistentLevelSlot<K, V>
+    return false
  }

-  override fun onExpand(newLevelSize: Int, interimItemCount: Int) {
-    this.levelSize = newLevelSize
-    this.topLevelBucketCount = 2.0.pow(newLevelSize).toInt()
+  private class PersistentLevelSlot<K : Any, V : Any?>(
+    val levelIdx: Int,
+    val bucketIdx: Int,
+    val slotIdx: Int,
+    private val io: PersistentLevelHashIO<K, V>,
+  ) : LevelSlot<K, V> {
+
+    override val key: K
+      get() = checkNotNull(io.readKey(levelIdx, bucketIdx, slotIdx))
+
+    override val value: V?
+      get() = io.readValue(levelIdx, bucketIdx, slotIdx)
+
+    override fun isOccupied(): Boolean =
+      io.isOccupied(levelIdx, bucketIdx, slotIdx)
+
+    override fun reset(key: K?, value: V?) {
+      io.writeEntry(levelIdx, bucketIdx, slotIdx, key, value)
+    }
  }
 }
@@ -18,13 +18,9 @@
 package com.itsaky.androidide.levelhash.internal

 import com.itsaky.androidide.levelhash.DataExternalizer
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO.Companion.KEYMAP_HEADER_SIZE_BYTES
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO.Companion.KEYMAP_MAX_SEGMENTS
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO.Companion.KEYMAP_SEGMENT_SIZE_BYTES
-import com.itsaky.androidide.levelhash.seekBoolean
+import com.itsaky.androidide.levelhash.LevelHash.ResizeFailure
 import com.itsaky.androidide.levelhash.seekInt
 import com.itsaky.androidide.levelhash.seekShort
-import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_BOOLEAN
 import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_INT
 import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_LONG
 import com.itsaky.androidide.levelhash.util.DataExternalizers.SIZE_SHORT
@@ -34,6 +30,7 @@ import org.slf4j.LoggerFactory
 import java.io.File
 import java.io.RandomAccessFile
 import java.nio.channels.FileChannel
+import kotlin.math.pow

 private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)

@@ -51,41 +48,45 @@ private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)
 *
 * ## Keymap
 *
- * - The keymap file header occupies [PersistentKeymapIO.KEYMAP_HEADER_SIZE_BYTES]
- *    bytes.
- * - Each segment is occupies [PersistentKeymapIO.KEYMAP_SEGMENT_SIZE_BYTES]
- *    bytes.
- * - There can be at most [PersistentKeymapIO.KEYMAP_MAX_SEGMENTS] segments in
- *    the keymap file.
- * - Each segment can store at most [PersistentKeymapIO.KEYMAP_ENTRIES_PER_SEGMENT]
- *    key-value pairs.
- *
- * Considering the default values, a keymap file can:
- * - be `8 + 4 + 1024 * 1024 * 100 = ~100 MB` in size.
- * - hold `(1024 * 1024 * 100) / 16 = 6553600 entries`
- *
 * Structure of the keymap file:
 *
 * ```
 * keymap {
 *   u64 magic_number;
- *   segment segments[segment_count];
+ *   level levels[2];
+ *   level interim_level?;
+ * }
 *
- *   segment {
- *       u64 value_address[];
- *   }
+ * level {
+ *    bucket buckets[2^(level_size-level_idx)];
+ * }
+ *
+ * bucket {
+ *    u64 slots[bucket_size];
 * }
 * ```
 *
- * - `segment_count` - The number of segments in the file.
- * - `segment_size` - The size of each segment.
- * - `segments` - The segments of the file. There are `segment_count` entries,
- *    each of exactly `segment_size` bytes.
- * - `value_address` - A 1-based 64-bit offset pointing the value entry in the values
- *    file. Note that the offset is 1-based because 0 is considered as an invalid
- *    value for the address. This helps in recognizing if an entry in the keymap
- *    points to a valid address or not. We can't really fill up the keymap entries
- *    with a negative offset as that would just take up unnecessary space.
+ * The `keymap` struct contains fields :
+ * - `magic_number` - Magic number for uniquely identifying the keymap file.
+ * - `levels` - The levels of the keymap. There are 2 levels, the top level (index 0)
+ *    and the bottom level (index 1).
+ * - `interim_level` - The temporary level that is used to move the slots from
+ *    the bottom level to the yet-to-be top level during expansion process.
+ *
+ * The `level` struct contains fields :
+ * - `buckets` - The buckets of the level. The number of buckets depends on the
+ *    index of the level and can be calculated using `2^(level_size-level_idx)`
+ *    where `level_size` is the size of the level hash and `level_idx` is the
+ *    index of the current level.
+ *
+ * The `bucket` struct contains fields :
+ * - `slots` - An array of size `bucket_size` that contains the 1-based, 64-bit
+ *    address of the the value entry in the values file. Note that the offset is
+ *    1-based because 0 is considered as an invalid value for the address (hence,
+ *    if the value is 0, then that slot is considered as an empty slot). This
+ *    helps in recognizing if an entry in the keymap points to a valid address
+ *    or not. We can't really fill up the keymap entries with a negative offset
+ *    as that would just take up unnecessary space.
 *
 * ## Values
 *
@@ -97,7 +98,6 @@ private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)
 *    value values[];
 *
 *    value {
- *      u8 is_occupied;
 *      u16 entry_size;
 *      u32 key_size;
 *      u8 key[key_size];
@@ -112,16 +112,13 @@ private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)
 * - `values` - The value entries.
 *
 * Each `value` entry contains fields :
- * - `is_occupied` - A 1-byte boolean value (0=false or 1=true) that indicates
- *    if the entry is occupied or not.
- * - `entry_size` - The size of the entry in bytes (excluding the size of
- *    `is_occupied` and `entry_size` itself).
+ * - `entry_size` - The size of the entry in bytes (excluding the size of this
+ *    `entry_size` field).
 * - `key_size` - The size of the key in bytes.
 * - `key` - The key of `key_size` bytes.
 * - `value_size` - The size of the value in bytes.
 * - `value` - The value of `value_size` bytes.
 *
- *
 * ## Metadata
 *
 * Structure of the metadata file :
@@ -133,8 +130,10 @@ private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)
 *    u64 values_first_entry;
 *    u64 values_next_entry;
 *    u64 values_file_size_bytes;
- *    u32 km_segment_count;
- *    u64 km_segment_size;
+ *    u32 km_level_size;
+ *    u32 km_bucket_size;
+ *    u64 km_l0_addr;
+ *    u64 km_l1_addr;
 * }
 * ```
 *
@@ -144,13 +143,17 @@ private val logger = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)
 * - `values_first_entry` - The address of the first entry in the values file.
 * - `values_next_entry` - The address of the next entry in the values file.
 * - `values_file_size_bytes` - The (occupied) size of the values file in bytes.
- * - `km_segment_count` - The number of segments in the keymap file.
- * - `km_segment_size` - The size of each segment in the keymap file.
+ * - `km_level_size` - The level size of the level hash.
+ * - `km_bucket_size` - The bucket size of the level hash.
+ * - `km_l0_addr` - Address of the level 0 (top level) in the keymap.
+ * - `km_l1_addr` - Address of the level 1 (bottom level) in the keymap.
 *
 * @author Akash Yadav
 */
 internal class PersistentLevelHashIO<K : Any, V : Any?>(
  indexFile: File,
+  levelSize: Int,
+  bucketSize: Int,
  private val keyExternalizer: DataExternalizer<K>,
  private val valueExternalizer: DataExternalizer<V>,
 ) : AutoCloseable {
@@ -161,32 +164,67 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  // as the values file is randomly accessed, and by multiple functions, or
  // even multiple threads, the current position cannot be determined

-  private val metaIo =
-    PersistentMetaIO(File(indexFile.parentFile, "${indexFile.name}._meta"))
-  private val keymapIo = PersistentKeymapIO(indexFile)
+  internal val metaIo =
+    PersistentMetaIO(File(indexFile.parentFile, "${indexFile.name}._meta"), levelSize, bucketSize)
+  private val keymapFile = File(indexFile.parentFile, "${indexFile.name}._i")
+  private val raKeymapFile by lazy { RandomAccessFile(keymapFile, "rw") }
  private val raIndexFile by lazy { RandomAccessFile(indexFile, "rw") }

+  private val keymapIo = MappedRandomAccessIO()
  private var valIo = MappedRandomAccessIO()

+  private val topLevelBucketCount: Long
+    get() = 2.0.pow(metaIo.levelSize).toLong()
+
+  private val topLevelSizeBytes: Long
+    get() = topLevelBucketCount * metaIo.bucketSize * KEYMAP_ENTRY_SIZE_BYTES
+
+  private val totalBucketCount: Long
+    get() = topLevelBucketCount + (topLevelBucketCount shr 1)
+
+  private val totalSlotCount: Long
+    get() = totalBucketCount * metaIo.bucketSize
+
  init {
-    BinaryFileUtils.initSparseFile(
-      file = indexFile,
-      magicNumber = VALUES_MAGIC_NUMBER
-    )
+    BinaryFileUtils.initSparseFile(file = keymapFile,
+      magicNumber = KEYMAP_MAGIC_NUMBER)
+    BinaryFileUtils.initSparseFile(file = indexFile,
+      magicNumber = VALUES_MAGIC_NUMBER)

    // the header region is not memory mapped
-    log.debug("mmap values file: offset={}, size={}", VALUES_HEADER_SIZE_BYTES,
-      metaIo.valuesFileSize)
-    val buffer = raIndexFile.channel.map(FileChannel.MapMode.READ_WRITE,
+    var buffer = raIndexFile.channel.map(FileChannel.MapMode.READ_WRITE,
      VALUES_HEADER_SIZE_BYTES, metaIo.valuesFileSize)
    valIo.reset(buffer, 0L, metaIo.valuesFileSize)
+
+    val keymapSize = metaIo.keymapSize
+    buffer = raKeymapFile.channel.map(FileChannel.MapMode.READ_WRITE,
+      KEYMAP_HEADER_SIZE_BYTES, keymapSize)
+    keymapIo.reset(buffer, 0L, keymapSize)
  }

-  private fun getValuesRealOffset(offset: Long) =
-    VALUES_HEADER_SIZE_BYTES + offset
+  private fun realValueOffset(offset: Long) = VALUES_HEADER_SIZE_BYTES + offset

-  private fun valuesDeallocate(offset: Long, len: Long) {
-    FileUtils.deallocate(raIndexFile, getValuesRealOffset(offset), len)
+  private fun realKeymapOffset(offset: Long) = KEYMAP_HEADER_SIZE_BYTES + offset
+
+  private fun valuesDeallocate(offset: Long, len: Long) =
+    FileUtils.deallocate(raIndexFile, realValueOffset(offset), len)
+
+  private fun keymapDeallocate(offset: Long, len: Long) =
+    FileUtils.deallocate(raKeymapFile, realKeymapOffset(offset), len)
+
+  private fun slotAddress(
+    levelIdx: Int,
+    bucketIdx: Int,
+    slotIdx: Int,
+  ): Long {
+    val levelPos = when (levelIdx) {
+      0 -> metaIo.l0Addr
+      1 -> metaIo.l1Addr
+      else -> throw IllegalArgumentException("Invalid level index: $levelIdx")
+    }
+
+    return levelPos + (KEYMAP_ENTRY_SIZE_BYTES * metaIo.bucketSize * bucketIdx) +
+      (KEYMAP_ENTRY_SIZE_BYTES * slotIdx)
  }

  /**
@@ -202,12 +240,11 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    bucketIdx: Int,
    slotIdx: Int,
  ): Long {
-    val (segment, positionInSegment) = keymapIo.segmentAndPosForPos(levelIdx,
-      bucketIdx, slotIdx)
-    segment.position(positionInSegment)
-    return segment.readLong().also { address ->
+    val slotAddr = slotAddress(levelIdx, bucketIdx, slotIdx)
+    keymapIo.position(slotAddr)
+    return keymapIo.readLong().also { address ->
      check(
-        address >= 0) { "Invalid value address in keymap at $positionInSegment: $address" }
+        address >= 0) { "Invalid value address in keymap at $slotAddr: $address" }
    }
  }

@@ -231,12 +268,8 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    bucketIdx: Int,
    slotIdx: Int,
  ): Boolean {
-    if (!moveToSlot(levelNum, bucketIdx, slotIdx)) {
-      return false
-    }
-
-    // token must be set and key size must be > 0
-    return valIo.readBoolean() && valIo.readInt() > 0
+    // 0 == no value entry
+    return valueAddressForPos(levelNum, bucketIdx, slotIdx) > 0L
  }

  /**
@@ -253,12 +286,6 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
      return null
    }

-    val isOccupied = valIo.readBoolean()
-
-    if (!isOccupied) {
-      return null
-    }
-
    if (valIo.readUnsignedShort() <= 0L) {
      // entry size is 0, so the slot is empty
      return null
@@ -341,21 +368,21 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    key: K?,
    value: V?,
  ) {
-    val (segment, positionInSegment) = keymapIo.segmentAndPosForPos(levelNum,
-      bucketIdx, slotIdx)
+    // TODO: expand values file in case of buffer overflow
+    val slotAddr = slotAddress(levelNum, bucketIdx, slotIdx)

    if (key == null) {
      // delete entry
-      deleteEntry(segment, positionInSegment)
+      deleteEntry(slotAddr)
      return
    }

-    segment.position(positionInSegment)
-    val existingValueAddr = segment.readLong()
+    keymapIo.position(slotAddr)
+    val existingValueAddr = keymapIo.readLong()
    val (isUpdate, existingEntrySize) = run {
      if (existingValueAddr > 0L) {
        valIo.position(existingValueAddr - 1)
-        valIo.readBoolean() to valIo.readUnsignedShort().toLong()
+        true to valIo.readUnsignedShort().toLong()
      } else {
        false to 0L
      }
@@ -364,8 +391,7 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
    val tokenAddr = metaIo.valuesNextEntry
    valIo.position(tokenAddr)

-    // leave space for isOccupied token and entry size
-    valIo.seekBoolean()
+    // leave space for entry size
    valIo.seekShort()

    val keyValPos = valIo.position()
@@ -382,20 +408,19 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(

    // then set the token
    valIo.position(tokenAddr)
-    valIo.writeBoolean(true)
    valIo.writeShort((entrySize and 0xFFFF).toInt())

    // reset to the final position
    valIo.position(finalAddr)

    // then update the address in the keymap
-    segment.position(positionInSegment)
-    segment.writeLong(tokenAddr + 1)
+    keymapIo.position(slotAddr)
+    keymapIo.writeLong(tokenAddr + 1)

    metaIo.valuesNextEntry = finalAddr

    if (isUpdate) {
-      var size: Long = SIZE_BOOLEAN + SIZE_SHORT
+      var size: Long = SIZE_SHORT
      if (existingEntrySize > 0L) {
        size += existingEntrySize
      }
@@ -406,12 +431,11 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  }

  private fun deleteEntry(
-    segment: KeymapSegment,
-    positionInSegment: Long,
+    slotAddr: Long,
    readValue: Boolean = false,
  ): V? {
-    segment.position(positionInSegment)
-    val valueAddr = segment.readLong()
+    keymapIo.position(slotAddr)
+    val valueAddr = keymapIo.readLong()

    if (valueAddr == 0L) {
      return null
@@ -419,23 +443,12 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(

    // reading this region again will return 0, which is considered
    // a null pointer
-    keymapIo.deallocate(segment, positionInSegment, SIZE_LONG)
+    keymapDeallocate(slotAddr, SIZE_LONG)

    valIo.position(valueAddr - 1)
+    valIo.seekShort() // seek over entrySize

-    var entrySize = 0L
-
-    if (!valIo.readBoolean()) {
-      // slot is not occupied
-      return null
-    }
-
-    entrySize += SIZE_BOOLEAN // isOccupied
-
-    // seek over entrySize
-    valIo.seekShort()
-
-    entrySize += SIZE_SHORT // entrySize
+    var entrySize = SIZE_SHORT

    val keySize = valIo.readInt()
    entrySize += SIZE_INT
@@ -468,20 +481,11 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  }

  fun clear() {
-    keymapIo.deallocate(KEYMAP_HEADER_SIZE_BYTES,
-      KEYMAP_SEGMENT_SIZE_BYTES * KEYMAP_MAX_SEGMENTS)
-    valuesDeallocate(VALUES_HEADER_SIZE_BYTES, metaIo.valuesFileSize)
-  }
+    var keymapSize = topLevelSizeBytes
+    keymapSize += topLevelSizeBytes shr 1

-  @Suppress("UNUSED", "UNUSED_PARAMETER")
-  fun tryMoveToInterim(
-    levelNum: Int,
-    bucketIdx: Int,
-    slotIdx: Int,
-    key: K?,
-    value: V?,
-  ): V? {
-    TODO()
+    keymapDeallocate(KEYMAP_HEADER_SIZE_BYTES, keymapSize)
+    valuesDeallocate(VALUES_HEADER_SIZE_BYTES, metaIo.valuesFileSize)
  }

  private fun <T> writeSlotKeyOrVal(obj: T?, externalizer: DataExternalizer<T>
@@ -510,29 +514,21 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
  }

  override fun close() {
-
-    try {
-      keymapIo.close()
-    } catch (e: Throwable) {
-      log.error("Failed to close keymap file", e)
-    }
-
-    try {
-      metaIo.close()
-    } catch (e: Throwable) {
-      log.error("Failed to close meta file", e)
-    }
-
-    try {
-      raIndexFile.close()
-    } catch (e: Throwable) {
-      logger.error("Failed to close index file", e)
-    }
+    closeAndLogErrs(metaIo, keymapIo, valIo, raKeymapFile, raIndexFile)
  }

-
  companion object {

+    private fun closeAndLogErrs(vararg closeable: AutoCloseable) {
+      closeable.forEachIndexed { index, autoCloseable ->
+        try {
+          autoCloseable.close()
+        } catch (e: Throwable) {
+          logger.error("Failed to close [{}]", index, e)
+        }
+      }
+    }
+
    private val log = LoggerFactory.getLogger(PersistentLevelHashIO::class.java)

    /**
@@ -557,8 +553,32 @@ internal class PersistentLevelHashIO<K : Any, V : Any?>(
     */
    internal const val VALUES_INITIAL_SIZE_BYTES = 512L * 1024L

+    /**
+     * The number of bytes used to store the header of the keymap file.
+     */
+    internal const val KEYMAP_HEADER_SIZE_BYTES: Long = MAGIC_NUMBER_SIZE_BYTES
+
+    /**
+     * The number of bytes that are used to store an entry in a key map file.
+     */
+    internal const val KEYMAP_ENTRY_SIZE_BYTES: Long = SIZE_LONG
+
+    /**
+     * Magic number that is used as the file signature to identify the keymap file.
+     */
+    internal const val KEYMAP_MAGIC_NUMBER = 0x414944584B
+
    internal const val LEVEL_VALUES_VERSION = 1
    internal const val LEVEL_KEYMAP_VERSION = 1
  }

+  /**
+   * Thrown during a move operation when the destination slot is already occupied.
+   */
+  class DestinationOccupiedException(
+    levelIdx: Int,
+    bucketIdx: Int,
+    slotIdx: Int,
+  ) : ResizeFailure(
+    "Destination slot is already occupied: levelIdx=$levelIdx, bucketIdx=$bucketIdx, slotIdx=$slotIdx")
 }
@@ -18,7 +18,8 @@
 package com.itsaky.androidide.levelhash.internal

 import androidx.collection.MutableLongLongMap
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO.Companion.KEYMAP_SEGMENT_SIZE_BYTES
+import com.itsaky.androidide.levelhash.LevelHash
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.KEYMAP_ENTRY_SIZE_BYTES
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.LEVEL_KEYMAP_VERSION
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.LEVEL_VALUES_VERSION
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.VALUES_HEADER_SIZE_BYTES
@@ -29,70 +30,67 @@ import org.slf4j.LoggerFactory
 import java.io.EOFException
 import java.io.File
 import java.io.RandomAccessFile
+import kotlin.math.min
+import kotlin.math.pow

 /**
 * @author Akash Yadav
 */
-internal class PersistentMetaIO(private val metaFile: File
+internal class PersistentMetaIO(private val metaFile: File,
+                                levelSize: Int,
+                                bucketSize: Int,
 ) : AutoCloseable {

  private val raFile by lazy { RandomAccessFile(metaFile, "rw") }
  private val cache = MutableLongLongMap(META__FIELD_COUNT)

  var valuesVersion: Int
-    get() {
-      return readInt(VAL__VERSION__OFFSET)
-    }
-    set(value) {
-      writeInt(VAL__VERSION__OFFSET, value)
-    }
+    get() = readInt(VAL__VERSION__OFFSET)
+    set(value) = writeInt(VAL__VERSION__OFFSET, value)

  var keymapVersion: Int
-    get() {
-      return readInt(KM__VERSION__OFFSET)
-    }
-    set(value) {
-      writeInt(KM__VERSION__OFFSET, value)
-    }
+    get() = readInt(KM__VERSION__OFFSET)
+    set(value) = writeInt(KM__VERSION__OFFSET, value)

  var valuesFirstEntry: Long
-    get() {
-      return readLong(VAL__FIRST_ENTRY__OFFSET)
-    }
-    set(value) {
-      writeLong(VAL__FIRST_ENTRY__OFFSET, value)
-    }
+    get() = readLong(VAL__FIRST_ENTRY__OFFSET)
+    set(value) = writeLong(VAL__FIRST_ENTRY__OFFSET, value)

  var valuesNextEntry: Long
-    get() {
-      return readLong(VAL__NEXT_ENTRY__OFFSET)
-    }
-    set(value) {
-      writeLong(VAL__NEXT_ENTRY__OFFSET, value)
-    }
+    get() = readLong(VAL__NEXT_ENTRY__OFFSET)
+    set(value) = writeLong(VAL__NEXT_ENTRY__OFFSET, value)

  var valuesFileSize: Long
-    get() {
-      return readLong(VAL__VAL_SIZE__OFFSET)
-    }
-    set(value) {
-      writeLong(VAL__VAL_SIZE__OFFSET, value)
-    }
+    get() = readLong(VAL__VAL_SIZE__OFFSET)
+    set(value) = writeLong(VAL__VAL_SIZE__OFFSET, value)

-  var keymapSegmentCount: Int
-    get() {
-      return readInt(KM__SEGMENT_COUNT__OFFSET)
-    }
-    set(value) {
-      writeInt(KM__SEGMENT_COUNT__OFFSET, value)
-    }
+  var levelSize: Int
+    get() = readInt(KM__LEVEL_SIZE__OFFSET)
+    set(value) = writeInt(KM__LEVEL_SIZE__OFFSET, value)

-  var keymapSegmentSize: Long
+  var bucketSize: Int
+    get() = readInt(KM__BUCKET_SIZE__OFFSET)
+    set(value) = writeInt(KM__BUCKET_SIZE__OFFSET, value)
+
+  var l0Addr: Long
+    get() = readLong(KM__L0_ADDR__OFFSET)
+    set(value) { writeLong(KM__L0_ADDR__OFFSET, value) }
+
+  var l1Addr: Long
+    get() = readLong(KM__L1_ADDR__OFFSET)
+    set(value) { writeLong(KM__L1_ADDR__OFFSET, value) }
+
+  val keymapSize: Long
    get() {
-      return readLong(KM__SEGMENT_SIZE__OFFSET)
-    }
-    set(value) {
-      writeLong(KM__SEGMENT_SIZE__OFFSET, value)
+      // for this to work, l0 and l1 must be in a contiguosly allocated space
+      // in the keymap file
+      val l0Addr = this.l0Addr
+      val l1Addr = this.l1Addr
+      val l0SizeBytes = 2.0.pow(levelSize).toLong() * bucketSize * KEYMAP_ENTRY_SIZE_BYTES
+      var size = min(l0Addr, l1Addr)
+      size += l0SizeBytes
+      size += l0SizeBytes shr 1
+      return size
    }

  init {
@@ -110,17 +108,21 @@ internal class PersistentMetaIO(private val metaFile: File
      valuesFileSize = VALUES_INITIAL_SIZE_BYTES
    }
    if (valuesFirstEntry == 0L) {
-      valuesFirstEntry = VALUES_HEADER_SIZE_BYTES
+      valuesFirstEntry = 0
    }
    if (valuesNextEntry == 0L) {
-      valuesNextEntry = valuesFirstEntry
+      valuesNextEntry = 0
    }
-    if (keymapSegmentCount == 0) {
-      // initially, there should be at least one segment
-      keymapSegmentCount = 1
+    if (this.levelSize != levelSize) {
+      this.levelSize = levelSize
    }
-    if (keymapSegmentSize == 0L) {
-      keymapSegmentSize = KEYMAP_SEGMENT_SIZE_BYTES
+    if (this.bucketSize != bucketSize) {
+      this.bucketSize = bucketSize
+    }
+    // default value of l0Addr is 0
+    // only the value of l1Addr should be updated
+    if (this.l1Addr == 0L) {
+      this.l1Addr = 2.0.pow(levelSize).toLong() * bucketSize * KEYMAP_ENTRY_SIZE_BYTES
    }
  }

@@ -223,26 +225,34 @@ internal class PersistentMetaIO(private val metaFile: File
      VAL__NEXT_ENTRY__OFFSET + VAL__NEXT_ENTRY__SIZE_BYTES
    private const val VAL__VAL_SIZE__SIZE_BYTES = SIZE_LONG

-    private const val KM__SEGMENT_COUNT__OFFSET =
+    private const val KM__LEVEL_SIZE__OFFSET =
      VAL__VAL_SIZE__OFFSET + VAL__VAL_SIZE__SIZE_BYTES
-    private const val KM__SEGMENT_COUNT__SIZE_BYTES = SIZE_INT
+    private const val KM__LEVEL_SIZE__SIZE_BYTES = SIZE_INT

-    private const val KM__SEGMENT_SIZE__OFFSET =
-      KM__SEGMENT_COUNT__OFFSET + KM__SEGMENT_COUNT__SIZE_BYTES
-    private const val KM__SEGMENT_SIZE__SIZE_BYTES = SIZE_LONG
+    private const val KM__BUCKET_SIZE__OFFSET =
+      KM__LEVEL_SIZE__OFFSET + KM__LEVEL_SIZE__SIZE_BYTES
+    private const val KM__BUCKET_SIZE__SIZE_BYTES = SIZE_INT
+
+    private const val KM__L0_ADDR__OFFSET =
+      KM__BUCKET_SIZE__OFFSET + KM__BUCKET_SIZE__SIZE_BYTES
+    private const val KM__L0_ADDR__SIZE_BYTES = SIZE_LONG
+
+    private const val KM__L1_ADDR__OFFSET =
+      KM__L0_ADDR__OFFSET + KM__L0_ADDR__SIZE_BYTES
+    private const val KM__L1_ADDR__SIZE_BYTES = SIZE_LONG

    // TODO: The properties below must be updated when a new field is added to the meta file

    /**
     * The number of fields in the meta file.
     */
-    private const val META__FIELD_COUNT = 7
+    private const val META__FIELD_COUNT = 9

    /**
     * The size of the meta file in bytes.
     */
    // Offset of the last field + its size
    private const val META__SIZE_BYTES =
-      KM__SEGMENT_SIZE__OFFSET + KM__SEGMENT_SIZE__SIZE_BYTES
+      KM__L1_ADDR__OFFSET + KM__L1_ADDR__SIZE_BYTES
  }
 }
@@ -25,6 +25,7 @@ import com.itsaky.androidide.levelhash.seekFloat
 import com.itsaky.androidide.levelhash.seekInt
 import com.itsaky.androidide.levelhash.seekLong
 import com.itsaky.androidide.levelhash.seekShort
+import org.slf4j.LoggerFactory
 import java.io.DataInput
 import java.io.DataInputStream
 import java.io.DataOutput
@@ -38,32 +39,44 @@ import kotlin.math.min
 *
 * @author Akash Yadav
 */
-internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO {
+internal class MappedRandomAccessIO : DataInput, DataOutput, RandomAccessIO, AutoCloseable {

+  private var positionOffset = 0L
  private var position = 0L
  private var size = 0L

  private lateinit var buffer: MappedByteBuffer

+  companion object {
+    private val log = LoggerFactory.getLogger(MappedRandomAccessIO::class.java)
+  }
+
  internal fun buf(): MappedByteBuffer = buf(this.position)

  internal fun buf(position: Long): MappedByteBuffer {
    return this.buffer.apply { position(position.toInt()) }
  }

-  fun reset(buffer: MappedByteBuffer, position: Long, length: Long) {
+  fun reset(buffer: MappedByteBuffer, position: Long, length: Long, positionOffset: Long = 0L) {
    this.buffer = buffer
    this.position = position
    this.size = length
+    this.positionOffset = positionOffset
+  }
+
+  override fun close() {
+    try {
+      buffer.unmap()
+    } catch (err: Throwable) {
+      log.error("Failed to close memory-mapped buffer", err)
+    }
  }

  override fun position(position: Long) {
    val currentPos = position()
    if (currentPos == position) return
    require(position >= 0L) { "position must be >= 0, position=$position" }
-    require(
-      position < size) { "position must be < size:$size, position=$position" }
-    this.position = position
+    this.position = this.positionOffset + position
  }

  override fun position(): Long {
@@ -0,0 +1,47 @@
+/*
+ *  This file is part of AndroidIDE.
+ *
+ *  AndroidIDE is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  AndroidIDE is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *   along with AndroidIDE.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.itsaky.androidide.levelhash.util
+
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+import java.nio.MappedByteBuffer
+
+/**
+ * Unmap the memory mapped buffer. This [KeymapSegment] becomes invalid after
+ * this operation and must not be used further unless remapped.
+ *
+ * @return `true` if the unmapping request was successful, `false` otherwise.
+ */
+fun MappedByteBuffer.unmap(): Boolean {
+  try {
+    val cleanerF = javaClass.getDeclaredField("cleaner")
+    cleanerF.isAccessible = true
+    val cleaner = cleanerF.get(this)
+
+    val cleanM = cleaner.javaClass.getDeclaredMethod("clean")
+    cleanM.isAccessible = true
+
+    cleanM.invoke(cleaner)
+    return true
+  } catch (err: Throwable) {
+    LoggerFactory.getLogger(Logger.ROOT_LOGGER_NAME)
+      .error("Failed to unmap keymap segment")
+    err.printStackTrace()
+    return false
+  }
+}
@@ -56,9 +56,9 @@ class PersistentLevelHashTest {
      .keyExternalizer(DataExternalizers.STRING)
      .valueExternalizer(valueExternalizer)
      .indexFile(indexFile)
-      .apply(conf)
      .levelSize(2)
      .hashFn(::stringHash)
+      .apply { conf(this as PersistentHashBuilder<String, V>) }
      .build()

    assertThat(hash).isNotNull()
@@ -100,6 +100,17 @@ class PersistentLevelHashTest {
    }
  }

+  @Test
+  fun `test values file expands and remaps in case of buffer overflow`() {
+    persistentStringHash("huge") { autoExpand(false).levelSize(15).bucketSize(10) }.use { hash ->
+      for (i in 0..<hash.totalSlotCount) {
+        val key = "key${i}"
+        val value = "value${i}"
+        assertThat(hash.insert(key, value)).isTrue()
+      }
+    }
+  }
+
  @Test
  fun testRemoval() {
    persistentStringHash("remove").use { hash ->
@@ -183,4 +194,28 @@ class PersistentLevelHashTest {
      assertThat(hash["long"]).isEqualTo("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
    }
  }
+
+  @Test
+  fun `test level hash expand`() {
+    persistentStringHash(
+      fileName = "expand",
+      valueExternalizer = DataExternalizers.STRING.nullable()
+    ).use { hash ->
+      for (i in 0..<hash.totalSlotCount) {
+        val key = "key${i}"
+        val value = "value${i}"
+        assertThat(hash.insert(key, value)).isTrue()
+      }
+
+      // expand the level hash
+      hash.expand(1)
+
+      // then check that all existing entries are still present
+      for (i in 0..<hash.totalSlotCount) {
+        val key = "key${i}"
+        val value = "value${i}"
+        assertThat(hash[key]).isEqualTo(value)
+      }
+    }
+  }
 }
@@ -18,13 +18,20 @@
 package com.itsaky.androidide.levelhash

 import com.google.common.truth.Truth.assertThat
-import com.itsaky.androidide.levelhash.internal.PersistentKeymapIO
+import com.itsaky.androidide.levelhash.LevelHash.Companion.BUCKET_SIZE_DEFAULT
+import com.itsaky.androidide.levelhash.LevelHash.Companion.LEVEL_SIZE_DEFAULT
 import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.KEYMAP_ENTRY_SIZE_BYTES
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.LEVEL_KEYMAP_VERSION
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.LEVEL_VALUES_VERSION
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.VALUES_HEADER_SIZE_BYTES
+import com.itsaky.androidide.levelhash.internal.PersistentLevelHashIO.Companion.VALUES_INITIAL_SIZE_BYTES
 import com.itsaky.androidide.levelhash.internal.PersistentMetaIO
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
 import java.io.File
+import kotlin.math.pow

 /**
 * @author Akash Yadav
@@ -43,40 +50,36 @@ class PersistentMetaIOTest {
      metaFile.parentFile!!.mkdirs()
    }

-    return PersistentMetaIO(metaFile)
+    return PersistentMetaIO(metaFile = metaFile,
+      levelSize = LEVEL_SIZE_DEFAULT,
+      bucketSize = BUCKET_SIZE_DEFAULT)
  }

  @Test
  fun `test meta init with default values`() {
    val io = createMetaIo("init-with-default")
-    assertThat(io.valuesVersion).isEqualTo(PersistentLevelHashIO.LEVEL_VALUES_VERSION)
-    assertThat(io.keymapVersion).isEqualTo(PersistentLevelHashIO.LEVEL_KEYMAP_VERSION)
-    assertThat(io.valuesFirstEntry).isEqualTo(PersistentLevelHashIO.VALUES_HEADER_SIZE_BYTES)
-    assertThat(io.valuesNextEntry).isEqualTo(PersistentLevelHashIO.VALUES_HEADER_SIZE_BYTES)
-    assertThat(io.valuesFileSize).isEqualTo(PersistentLevelHashIO.VALUES_INITIAL_SIZE_BYTES)
-    assertThat(io.keymapSegmentCount).isEqualTo(1)
-    assertThat(io.keymapSegmentSize).isEqualTo(PersistentKeymapIO.KEYMAP_SEGMENT_SIZE_BYTES)
+    assertThat(io.valuesVersion).isEqualTo(LEVEL_VALUES_VERSION)
+    assertThat(io.keymapVersion).isEqualTo(LEVEL_KEYMAP_VERSION)
+    assertThat(io.valuesFirstEntry).isEqualTo(VALUES_HEADER_SIZE_BYTES)
+    assertThat(io.valuesNextEntry).isEqualTo(VALUES_HEADER_SIZE_BYTES)
+    assertThat(io.valuesFileSize).isEqualTo(VALUES_INITIAL_SIZE_BYTES)
+    assertThat(io.levelSize).isEqualTo(LEVEL_SIZE_DEFAULT)
+    assertThat(io.bucketSize).isEqualTo(BUCKET_SIZE_DEFAULT)
+    assertThat(io.l0Addr).isEqualTo(0L)
+    assertThat(io.l1Addr).isEqualTo(2.0.pow(LEVEL_SIZE_DEFAULT).toLong() * BUCKET_SIZE_DEFAULT * KEYMAP_ENTRY_SIZE_BYTES)
  }

  @Test
  fun `test meta init with existing file`() {
    createMetaIo(name = "init-with-existing").use { io ->
-      assertThat(io.valuesVersion).isEqualTo(PersistentLevelHashIO.LEVEL_VALUES_VERSION)
-      assertThat(io.keymapVersion).isEqualTo(PersistentLevelHashIO.LEVEL_KEYMAP_VERSION)
-      assertThat(io.valuesFirstEntry).isEqualTo(PersistentLevelHashIO.VALUES_HEADER_SIZE_BYTES)
-      assertThat(io.valuesNextEntry).isEqualTo(PersistentLevelHashIO.VALUES_HEADER_SIZE_BYTES)
-      assertThat(io.valuesFileSize).isEqualTo(PersistentLevelHashIO.VALUES_INITIAL_SIZE_BYTES)
-      assertThat(io.keymapSegmentCount).isEqualTo(1)
-      assertThat(io.keymapSegmentSize).isEqualTo(PersistentKeymapIO.KEYMAP_SEGMENT_SIZE_BYTES)
-
      io.valuesVersion = 2
      io.keymapVersion = 3
      io.valuesFirstEntry = 100
      io.valuesNextEntry = 200
      io.valuesFileSize = 1024

-      io.keymapSegmentCount = 10
-      io.keymapSegmentSize = 1024
+      io.levelSize = 10
+      io.bucketSize = 20
    }

    createMetaIo(name = "init-with-existing", createNew = false).use { io ->
@@ -85,8 +88,10 @@ class PersistentMetaIOTest {
      assertThat(io.valuesFirstEntry).isEqualTo(100)
      assertThat(io.valuesNextEntry).isEqualTo(200)
      assertThat(io.valuesFileSize).isEqualTo(1024)
-      assertThat(io.keymapSegmentCount).isEqualTo(10)
-      assertThat(io.keymapSegmentSize).isEqualTo(1024)
+      assertThat(io.levelSize).isEqualTo(10)
+      assertThat(io.bucketSize).isEqualTo(20)
+      assertThat(io.l0Addr).isEqualTo(0L)
+      assertThat(io.l1Addr).isEqualTo(2.0.pow(io.levelSize).toLong() * io.bucketSize * KEYMAP_ENTRY_SIZE_BYTES)
    }
  }
 }