erts: Implement map iterator using a stack

This version does not work great as the subtrees
created are not proper hash maps. Also it is not
all that performant as the extra allocations to
keep the stack there is expensive.

4 files changed, 84 insertions, 5 deletions

diff --git a/erts/emulator/beam/bif.tab b/erts/emulator/beam/bif.tab
index f7b4451890..4ee5af14bc 100644
--- a/erts/emulator/beam/bif.tab
+++ b/erts/emulator/beam/bif.tab
@@ -691,3 +691,4 @@ bif erts_internal:maps_to_list/2
 #
 
 bif erlang:iolist_to_iovec/1
+bif erts_internal:map_get_index/2
\ No newline at end of file
diff --git a/erts/emulator/beam/erl_map.c b/erts/emulator/beam/erl_map.c
index f0c54e05f7..37fcedbc6d 100644
--- a/erts/emulator/beam/erl_map.c
+++ b/erts/emulator/beam/erl_map.c
@@ -90,7 +90,6 @@ static BIF_RETTYPE hashmap_merge(Process *p, Eterm nodeA, Eterm nodeB, int swap_
                                  struct HashmapMergeContext_*);
 static Export hashmap_merge_trap_export;
 static BIF_RETTYPE maps_merge_trap_1(BIF_ALIST_1);
-static Uint hashmap_subtree_size(Eterm node);
 static Eterm hashmap_to_list(Process *p, Eterm map, Sint n);
 static Eterm hashmap_keys(Process *p, Eterm map);
 static Eterm hashmap_values(Process *p, Eterm map);
@@ -1519,7 +1518,7 @@ trap:  /* Yield */
     return trap_ret;
 }
 
-static Uint hashmap_subtree_size(Eterm node) {
+Uint hashmap_subtree_size(Eterm node) {
     DECLARE_WSTACK(stack);
     Uint size = 0;
 
@@ -3060,6 +3059,75 @@ static Eterm hashmap_bld_tuple_uint(Uint **hpp, Uint *szp, Uint n, Uint nums[])
     return res;
 }
 
+static Uint32
+index_to_hash(Sint idx, Uint32 hval)
+{
+    Uint32 new_hval = 1;
+    while (1) {
+        if (hval & 0x1) {
+            idx--;
+            if (!idx) break;
+        }
+        new_hval = new_hval << 1;
+        hval = hval >> 1;
+    }
+    return new_hval;
+}
+
+BIF_RETTYPE erts_internal_map_get_index_2(BIF_ALIST_2) {
+    Sint idx;
+    Uint32 hval, sz;
+    Eterm *ptr, hdr;
+
+    if (!is_hashmap(BIF_ARG_2) || is_not_small(BIF_ARG_1))
+        BIF_ERROR(BIF_P, BADARG);
+
+    idx = signed_val(BIF_ARG_1);
+
+    if (idx < 1 || idx > 17)
+        BIF_ERROR(BIF_P, BADARG);
+
+    ASSERT(is_boxed(BIF_ARG_2));
+
+    ptr = boxed_val(BIF_ARG_2);
+again:
+    hdr = *ptr;
+
+    switch(hdr & _HEADER_MAP_SUBTAG_MASK) {
+    case HAMT_SUBTAG_HEAD_ARRAY:
+        ptr++;
+        sz = 16;
+        hval = 0xffff;
+        break;
+    case HAMT_SUBTAG_HEAD_BITMAP:
+        ptr++;
+        if (ptr[0] == -1) {
+            ptr = boxed_val(ptr[1]);
+            goto again;
+        }
+    case HAMT_SUBTAG_NODE_BITMAP:
+        hval = MAP_HEADER_VAL(hdr);
+        sz = hashmap_bitcount(hval);
+        ASSERT(sz < 17);
+        break;
+    default:
+        erts_exit(ERTS_ABORT_EXIT, "bad header");
+    }
+
+    if (idx <= sz) {
+        if (is_list(ptr[idx]))
+            return ptr[idx];
+        else {
+            Eterm *hp = HAlloc(BIF_P, HAMT_HEAD_BITMAP_SZ(1));
+            hp[0] = MAP_HEADER_HAMT_HEAD_BITMAP(index_to_hash(idx, hval));
+            hp[1] = -1;
+            hp[2] = ptr[idx];
+            return make_hashmap(hp);
+        }
+    } else {
+        return am_none;
+    }
+}
 
 /* implementation of builtin emulations */
 
diff --git a/erts/emulator/beam/erl_map.h b/erts/emulator/beam/erl_map.h
index c3ccf80b85..d910e98398 100644
--- a/erts/emulator/beam/erl_map.h
+++ b/erts/emulator/beam/erl_map.h
@@ -56,7 +56,7 @@ typedef struct flatmap_s {
 /* the head-node is a bitmap or array with an untagged size */
 
 
-#define hashmap_size(x)               (((hashmap_head_t*) hashmap_val(x))->size)
+#define hashmap_size(x)               (((hashmap_head_t*) hashmap_val(x))->size == -1 ? hashmap_subtree_size(x) : ((hashmap_head_t*) hashmap_val(x))->size)
 #define hashmap_make_hash(Key)        make_internal_hash(Key, 0)
 
 #define hashmap_restore_hash(Heap,Lvl,Key) \
@@ -106,6 +106,7 @@ Eterm  erts_hashmap_from_ks_and_vs_extra(ErtsHeapFactory *factory,
                                          Eterm *ks, Eterm *vs, Uint n,
 					 Eterm k, Eterm v);
 
+Uint hashmap_subtree_size(Eterm node);
 const Eterm *erts_maps_get(Eterm key, Eterm map);
 
 const Eterm *erts_hashmap_get(Uint32 hx, Eterm key, Eterm map);
diff --git a/lib/stdlib/src/maps.erl b/lib/stdlib/src/maps.erl
index 7ed32a3988..0da12cf8e2 100644
--- a/lib/stdlib/src/maps.erl
+++ b/lib/stdlib/src/maps.erl
@@ -32,7 +32,7 @@
          new/0, put/3, remove/2, take/2,
          to_list/1, update/3, values/1]).
 
--opaque iterator() :: [{term(), term()}].
+-opaque iterator() :: [{term(), term()}] | [[pos_integer() | map()]].
 
 -export_type([iterator/0]).
 
@@ -279,7 +279,10 @@ size(Val) ->
       Iterator :: iterator().
 
 iterator(Map) when is_map(Map) ->
-    maps:to_list(Map);
+    case maps:size(Map) < 42 of
+        true -> maps:to_list(Map);
+        false -> [[1 | Map]]
+    end;
 iterator(M) ->
     erlang:error(error_type(M),[M]).
 
@@ -289,6 +292,12 @@ iterator(M) ->
       V :: term(),
       NextIterator :: iterator().
 next([]) -> none;
+next([[Index | Map] | St]) ->
+    case erts_internal:map_get_index(Index, Map) of
+        SubMap when is_map(SubMap) -> next([[1 | SubMap], [Index+1|Map] | St]);
+        [K|V] -> {K,V,[[Index+1 | Map] | St]};
+        none -> next(St)
+    end;
 next([{K, V} | I]) ->
     {K, V, I}.