change unsafe immutable hash table iteration

racket/src/racket/src/hash.c

@@ -2759,6 +2759,116 @@ Scheme_Object *scheme_hash_tree_next_pos(Scheme_Hash_Tree *tree, mzlonglong pos)
 # define HAMT_TRAVERSE_NEXT(i) ((i)+1)
 #endif
 
+XFORM_NONGCING static void hamt_subtree_at_index(Scheme_Hash_Tree *ht, mzlonglong pos,
+                                                 Scheme_Hash_Tree **_subtree, int *_i, int *_popcount)
+{
+  int popcount, i;
+  Scheme_Hash_Tree *sub;
+
+  while (1) {
+    popcount = hamt_popcount(ht->bitmap);
+    i = HAMT_TRAVERSE_INIT(popcount);
+    while (1) {
+      if (HASHTR_SUBTREEP(ht->els[i])
+          || HASHTR_COLLISIONP(ht->els[i])) {
+        sub = (Scheme_Hash_Tree *)ht->els[i];
+        if (pos < sub->count) {
+          ht = sub;
+          break; /* to outer loop */
+        } else
+          pos -= sub->count;
+      } else {
+        if (!pos) {
+          *_subtree = ht;
+          *_i = i;
+          if (_popcount) *_popcount = popcount;
+          return;
+        }
+        --pos;
+      }
+      i = HAMT_TRAVERSE_NEXT(i);
+    }
+  }
+}
+
+/* We have two different implementations of Scheme_Hash_Tree traversal
+   as exposed by `unsafe-immutable-hash-next`, etc.:
+
+   * Consecutive integers from 0 as the index values (the same as
+     `scheme_hash_tree_next`), or
+
+   * Subtree lists (for the spine of a subtree up to the root) plus
+     offset into a subtree.
+
+   The second one is more direct, but requires some allocation. To
+   avoid allocation for small trees, it uses an integer encoding of a
+   path.
+
+   The first implementation is better for small trees, and the second
+   is better for very large trees, but there's not a big difference.
+   Small trees dominate for macro expansion, which is a big use of
+   hash trees, so that gives the first implementation the edge.
+
+   Microbenchmark:
+
+    (let loop ([size 2])
+      (define ht (for/hasheq ([i (in-range size)])
+                   (values i i)))
+      (define c (quotient 10000000 size))
+      (printf "~s: " size)
+      (time
+       (for ([j (in-range c)])
+         (for/fold ([v #f]) ([k (in-immutable-hash-keys ht)])
+           k)))
+      (unless (= size (expt 2 20))
+        (loop (* size 2))))
+
+ */
+
+#if 1
+/* ------ Implementation 1 of hash-tree traversal ------ */
+
+Scheme_Object *scheme_unsafe_hash_tree_start(Scheme_Hash_Tree *ht)
+{
+  ht = resolve_placeholder(ht);
+
+  if (!ht->count)
+    return scheme_false;
+  else
+    return scheme_make_integer(0);
+}
+
+XFORM_NONGCING void scheme_unsafe_hash_tree_subtree(Scheme_Object *obj, Scheme_Object *args,
+                                                    Scheme_Hash_Tree **_subtree, int *_i)
+{
+  Scheme_Hash_Tree *ht;
+
+  if (SCHEME_NP_CHAPERONEP(obj))
+    obj = SCHEME_CHAPERONE_VAL(obj);
+  ht = (Scheme_Hash_Tree *)obj;
+  ht = resolve_placeholder(ht);
+
+  hamt_subtree_at_index(ht, SCHEME_INT_VAL(args), _subtree, _i, NULL);
+}
+
+XFORM_NONGCING Scheme_Object *scheme_unsafe_hash_tree_access(Scheme_Hash_Tree *subtree, int i)
+{
+  return _mzHAMT_VAL(subtree, i, hamt_popcount(subtree->bitmap));
+}
+
+Scheme_Object *scheme_unsafe_hash_tree_next(Scheme_Hash_Tree *ht, Scheme_Object *args)
+{
+  intptr_t i = SCHEME_INT_VAL(args) + 1;
+  ht = resolve_placeholder(ht);
+  if (i < ht->count)
+    return scheme_make_integer(i);
+  else
+    return scheme_false;
+}
+
+#else
+/* ------ Implementation 2 of hash-tree traversal ------ */
+
 #define mzHAMT_MAX_INDEX_LEVEL 4 /* For the compressed form of the index */
 
 Scheme_Object *make_index_frame(Scheme_Hash_Tree *ht, intptr_t i, Scheme_Object *rest)
@@ -2948,38 +3058,21 @@ Scheme_Object *scheme_unsafe_hash_tree_next(Scheme_Hash_Tree *ht, Scheme_Object
   }
 }
 
+#endif
+
 XFORM_NONGCING static void hamt_at_index(Scheme_Hash_Tree *ht, mzlonglong pos,
                                          Scheme_Object **_key, Scheme_Object **_val, uintptr_t *_code)
 {
   int popcount, i;
   Scheme_Hash_Tree *sub;
 
-  while (1) {
+  hamt_subtree_at_index(ht, pos, &sub, &i, &popcount);
-    popcount = hamt_popcount(ht->bitmap);
+
-    i = HAMT_TRAVERSE_INIT(popcount);
+  *_key = sub->els[i];
-    while (1) {
-      if (HASHTR_SUBTREEP(ht->els[i])
-          || HASHTR_COLLISIONP(ht->els[i])) {
-        sub = (Scheme_Hash_Tree *)ht->els[i];
-        if (pos < sub->count) {
-          ht = sub;
-          break; /* to outer loop */
-        } else
-          pos -= sub->count;
-      } else {
-        if (!pos) {
-          *_key = ht->els[i];
   if (_val)
-            *_val = _mzHAMT_VAL(ht, i, popcount);
+    *_val = _mzHAMT_VAL(sub, i, popcount);
   if (_code)
-            *_code = _mzHAMT_CODE(ht, i, popcount);
+    *_code = _mzHAMT_CODE(sub, i, popcount);
-          return;
-        }
-        --pos;
-      }
-      i = HAMT_TRAVERSE_NEXT(i);
-    }
-  }
 }
 
 int scheme_hash_tree_index(Scheme_Hash_Tree *ht, mzlonglong pos, Scheme_Object **_key, Scheme_Object **_val)

racket/src/racket/src/optimize.c

@@ -10385,7 +10385,7 @@ static Scheme_Object *get_linklet_or_instance_for_import_key(Optimize_Info *info
   if (!cross || !cross->get_import)
     return NULL;
 
-  v = scheme_hash_tree_get(cross->linklets, key);
+  v = scheme_eq_hash_tree_get(cross->linklets, key);
   if (!v) {
     a[0] = key;
     v = scheme_apply_multi(cross->get_import, 1, a);
@@ -10455,7 +10455,7 @@ static Scheme_Object *get_import_inline_or_shape(Optimize_Info *info, Scheme_IR_
   if (!info->cross || (var->instance_pos < 0))
     return NULL;
 
-  key = scheme_hash_tree_get(info->cross->import_keys, scheme_make_integer(var->instance_pos));
+  key = scheme_eq_hash_tree_get(info->cross->import_keys, scheme_make_integer(var->instance_pos));
   if (!key)
     return NULL;
 
@@ -10606,10 +10606,10 @@ Scheme_Object *scheme_optimize_add_import_variable(Optimize_Info *info, Scheme_O
   if (SCHEME_FALSEP(linklet_key))
     return NULL;
 
-  pos = scheme_hash_tree_get(info->cross->rev_import_keys, linklet_key);
+  pos = scheme_eq_hash_tree_get(info->cross->rev_import_keys, linklet_key);
   MZ_ASSERT(pos);
 
-  syms = (Scheme_Hash_Tree *)scheme_hash_tree_get(info->cross->import_syms, pos);
+  syms = (Scheme_Hash_Tree *)scheme_eq_hash_tree_get(info->cross->import_syms, pos);
   if (!syms) {
     syms = empty_eq_hash_tree;
     if (SCHEME_INT_VAL(pos) < SCHEME_VEC_SIZE(info->linklet->importss)) {
@@ -10626,7 +10626,7 @@ Scheme_Object *scheme_optimize_add_import_variable(Optimize_Info *info, Scheme_O
     info->cross->import_syms = ht;
   }
 
-  var_pos = scheme_hash_tree_get(syms, symbol);
+  var_pos = scheme_eq_hash_tree_get(syms, symbol);
   if (!var_pos) {
     var_pos = scheme_make_integer(syms->count >> 1);
     syms = scheme_hash_tree_set(syms, symbol, var_pos);
@@ -10646,7 +10646,7 @@ Scheme_Object *scheme_optimize_get_import_key(Optimize_Info *info, Scheme_Object
   Scheme_Object *next_keys, *key, *pos;
   Scheme_Hash_Tree *ht;
 
-  next_keys = scheme_hash_tree_get(info->cross->import_next_keys, linklet_key);
+  next_keys = scheme_eq_hash_tree_get(info->cross->import_next_keys, linklet_key);
   if (!next_keys) {
     /* chaining is not supported by the compilation client */
     return NULL;
@@ -10655,7 +10655,7 @@ Scheme_Object *scheme_optimize_get_import_key(Optimize_Info *info, Scheme_Object
   MZ_ASSERT(instance_pos < SCHEME_VEC_SIZE(next_keys));
 
   key = SCHEME_VEC_ELS(next_keys)[instance_pos];
-  pos = scheme_hash_tree_get(info->cross->rev_import_keys, key);
+  pos = scheme_eq_hash_tree_get(info->cross->rev_import_keys, key);
   if (!pos) {
     /* Add this linklet as an import */
     pos = scheme_make_integer(info->cross->import_keys->count);
@@ -10720,7 +10720,7 @@ static void record_optimize_shapes(Optimize_Info *info, Scheme_Linklet *linklet,
 
     /* Add imported variables for each instance */
     for (i = 0; i < SCHEME_VEC_SIZE(linklet->importss); i++) {
-      ht = (Scheme_Hash_Tree *)scheme_hash_tree_get(info->cross->import_syms, scheme_make_integer(i));
+      ht = (Scheme_Hash_Tree *)scheme_eq_hash_tree_get(info->cross->import_syms, scheme_make_integer(i));
       if (ht && ((ht->count >> 1) > SCHEME_VEC_SIZE(SCHEME_VEC_ELS(linklet->importss)[i]))) {
         Scheme_Object *sym;
         v = scheme_make_vector((ht->count >> 1), NULL);
@@ -10776,7 +10776,7 @@ static void record_optimize_shapes(Optimize_Info *info, Scheme_Linklet *linklet,
     used = 0;
     for (i = 0; i < SCHEME_VEC_SIZE(linklet->importss); i++) {
       if (!SCHEME_INTP(SCHEME_VEC_ELS(linklet->importss)[i])) {
-        v = scheme_hash_tree_get(info->cross->import_keys, scheme_make_integer(i));
+        v = scheme_eq_hash_tree_get(info->cross->import_keys, scheme_make_integer(i));
         MZ_ASSERT(v);
         SCHEME_VEC_ELS((*_import_keys))[used++] = v;
       }
@@ -10791,9 +10791,9 @@ static void record_optimize_shapes(Optimize_Info *info, Scheme_Linklet *linklet,
     k = 0;
     for (i = 0; i < SCHEME_VEC_SIZE(linklet->importss); i++) {
       if (!SCHEME_INTP(SCHEME_VEC_ELS(linklet->importss)[i])) {
-        v = scheme_hash_tree_get(info->cross->import_keys, scheme_make_integer(i));
+        v = scheme_eq_hash_tree_get(info->cross->import_keys, scheme_make_integer(i));
         if (v)
-          v = scheme_hash_tree_get(info->cross->linklets, v);
+          v = scheme_eq_hash_tree_get(info->cross->linklets, v);
         in_linklet = ((v && SAME_TYPE(SCHEME_TYPE(v), scheme_linklet_type)) ? (Scheme_Linklet *)v : NULL);
         in_instance = ((v && SAME_TYPE(SCHEME_TYPE(v), scheme_instance_type)) ? (Scheme_Instance *)v : NULL);
         MZ_ASSERT(!in_linklet || SAME_TYPE(in_linklet->so.type, scheme_linklet_type));