Move small number-key maps to the direct index from the second key (#2204)

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/6310-small-number-map-direct-index.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/6310-small-number-map-direct-index.ts"
+ ],
+ "diags": []
+}

packages/runtime/src/scr_map.c

@@ -1205,9 +1205,11 @@ static __attribute__((noinline)) void scr_map_set_num(ScrMap *m, uint64_t hash,
 }
 
 /* A number-key map with at least this many entries moves to the direct
- * index when its keys allow it, so small integer keys are indexed before
- * the linear scan ends. */
-#define SCR_MAP_DENSE_AFTER 4
+ * index when its keys allow it: from its second key, so a map of a few
+ * small integer keys (a histogram over 0..3) reads through the inline
+ * direct index rather than the out-of-line linear scan. A map with one
+ * key keeps its linear storage. */
+#define SCR_MAP_DENSE_AFTER 1
 
 /* Whether appending one more entry would grow the lookup index rather than
  * reuse tombstones: scr_map_reserve_append compacts a linear map with a

packages/runtime/src/scr_map.test.c

@@ -184,9 +184,11 @@ static void iterate_while_mutating(void) {
 
 static void modes(void) {
   ScrMap *map = scr_map_new(SCR_MAP_KEY_F64, SCR_MAP_VAL_F64, NULL);
-  for (int i = 1; i <= 4; i++) scr_map_set_f64_f64(map, i, i);
-  assert(!map->ndense && !map->nbuckets); /* still a linear scan */
-  scr_map_set_f64_f64(map, 5, 5);
+  scr_map_set_f64_f64(map, 1, 1);
+  assert(!map->ndense && !map->nbuckets); /* one key: still a linear scan */
+  scr_map_set_f64_f64(map, 2, 2);
+  assert(map->ndense >= 16 && !map->nbuckets); /* from the second small integer key */
+  for (int i = 3; i <= 5; i++) scr_map_set_f64_f64(map, i, i);
   assert(map->ndense >= 16 && !map->nbuckets);
   for (int i = 6; i < 1000; i++) scr_map_set_f64_f64(map, i, i);
   assert(map->ndense >= 1000 && map->ndense <= 1000 + 1000 / 4 + 64); /* near the largest key */

packages/runtime/test/test_map.c

@@ -422,12 +422,14 @@ static void test_live_iteration_indices(void) {
 /* A small (linear) map that must grow its lookup index while an iteration
  * is active may not compact: entry indices, tombstones and the append
  * position stay put. Integer keys in a compact range move to the direct
- * index (dense) from the fifth entry, any other number keys to the bucket
+ * index (dense) from the second entry, any other number keys to the bucket
  * table from the ninth; both transitions are covered by running the same
- * script with `offset` 0 and 0.5. */
+ * script with `offset` 0 and 0.5. The walk deletes `victim`, an existing
+ * key, then appends the key that makes the index grow. */
 static void run_small_storage_transitions(double offset) {
   bool dense = offset == 0;
-  int n = dense ? 4 : 8; /* entries a linear map holds before its index grows */
+  int n = dense ? 1 : 8; /* entries a linear map holds before its index grows */
+  int victim = dense ? 0 : 1;
   ScrMap *m = scr_map_new(SCR_MAP_KEY_F64, SCR_MAP_VAL_F64, NULL);
   for (int i = 0; i < n; i++) scr_map_set_f64_f64(m, i + offset, i + 10);
   check(m->nbuckets == 0 && m->buckets == NULL && m->ndense == 0,
@@ -441,22 +443,23 @@ static void run_small_storage_transitions(double offset) {
   ScrMap *copy = scr_map_clone(m, false, NULL);
   check(copy->nbuckets == 0 && copy->ndense == 0, "small clone preserves compact storage");
   scr_map_iter_enter(m);
-  scr_map_delete_f64(m, 1 + offset);
+  scr_map_delete_f64(m, victim + offset);
   scr_map_set_f64_f64(m, n + offset, 14);
   check(dense ? m->ndense > 0 && m->nbuckets == 0 : m->nbuckets > 0 && m->ndense == 0,
         "live iteration promotes without moving entry indices");
-  check(m->nentries == (uint32_t)n + 1 && !scr_map_iter_live(m, 1) &&
+  check(m->nentries == (uint32_t)n + 1 && !scr_map_iter_live(m, victim) &&
             scr_map_iter_key_f64(m, n) == n + offset,
         "promotion preserves tombstones and append position");
   for (int i = n + 1; i < 64; i++) scr_map_set_f64_f64(m, i + offset, i + 10);
-  check(m->nentries == 64 && !scr_map_iter_live(m, 1) && scr_map_iter_key_f64(m, 63) == 63 + offset,
+  check(m->nentries == 64 && !scr_map_iter_live(m, victim) && scr_map_iter_key_f64(m, 63) == 63 + offset,
         "growth during iteration keeps entry indices");
   double out = 0;
-  check(scr_map_get_f64_f64(m, offset, &out) && out == 10,
-        "original key survives promotion and growth");
-  check(!scr_map_has_f64(m, 1 + offset), "deleted small key stays absent after rehash");
+  check(dense ? scr_map_get_f64_f64(m, n + offset, &out) && out == 14
+              : scr_map_get_f64_f64(m, offset, &out) && out == 10,
+        "keys written before the promotion survive it and the growth");
+  check(!scr_map_has_f64(m, victim + offset), "deleted small key stays absent after rehash");
   scr_map_iter_exit(m);
-  check(scr_map_has_f64(copy, 1 + offset) && !scr_map_has_f64(copy, n + offset),
+  check(scr_map_has_f64(copy, victim + offset) && !scr_map_has_f64(copy, n + offset),
         "small clone remains independent");
   scr_map_clear(copy);
   scr_map_iter_enter(copy);

tests/corpus/6310-small-number-map-direct-index.ts

@@ -0,0 +1,77 @@
+// Number-keyed Maps and Sets with two to four small integer keys use the
+// direct index. Results must match Node through overwrites, deletes and
+// re-adds, a fractional, negative or large key that moves the map to
+// hashing, clear, and mutation during a for-of or forEach walk.
+
+const hist = new Map<number, number>();
+const codes = [0, 3, 1, 3, 2, 0, 3, 3, 1];
+for (const c of codes) hist.set(c, (hist.get(c) ?? 0) + 1);
+console.log([...hist].map(([k, v]) => `${k}:${v}`).join(" "), hist.size, hist.has(4), hist.get(-0));
+
+// Delete and re-add: the re-added key moves to the end.
+hist.delete(3);
+hist.set(3, 9);
+hist.delete(0);
+console.log([...hist.keys()].join(","), hist.get(3), hist.has(0));
+
+// A key the direct index cannot hold moves the map to hashing.
+for (const odd of [0.5, -1, 2 ** 40]) {
+  const m = new Map<number, string>([
+    [1, "a"],
+    [2, "b"],
+  ]);
+  m.set(odd, "c");
+  m.set(1, "A");
+  console.log([...m].map(([k, v]) => `${k}=${v}`).join(" "), m.get(2), m.get(odd), m.size);
+}
+
+// Two keys, then clear, then reuse.
+const pair = new Map<number, number>();
+pair.set(5, 50);
+pair.set(7, 70);
+pair.clear();
+pair.set(7, 1);
+pair.set(1, 2);
+console.log([...pair].map(([k, v]) => `${k}=${v}`).join(" "), pair.size);
+
+// A walk that deletes ahead and adds keys.
+const walk = new Map<number, number>([
+  [0, 0],
+  [1, 10],
+  [2, 20],
+]);
+const seen: number[] = [];
+for (const [k, v] of walk) {
+  seen.push(v);
+  if (k === 0) {
+    walk.delete(2);
+    walk.set(3, 30);
+  }
+  if (k === 3) walk.set(0, -1);
+}
+console.log(seen.join(","), [...walk].map(([k, v]) => `${k}=${v}`).join(" "));
+
+// Sets of small integers, forEach with deletion and addition.
+const set = new Set<number>([2, 0]);
+set.add(1);
+set.add(2);
+const order: number[] = [];
+set.forEach((v) => {
+  order.push(v);
+  if (v === 2) {
+    set.delete(1);
+    set.add(3);
+  }
+});
+console.log(order.join(","), [...set].join(","), set.has(1), set.size);
+
+// Many small maps, each with a few keys.
+let total = 0;
+for (let i = 0; i < 1000; i++) {
+  const m = new Map<number, number>();
+  m.set(i % 3, i);
+  m.set((i + 1) % 3, 1);
+  if (i % 5 === 0) m.delete(i % 3);
+  total += (m.get(i % 3) ?? 100) + (m.get((i + 1) % 3) ?? 0) + m.size;
+}
+console.log(total);