erts, kernel: Fix erts_debug:size/1 for hashmaps

This commit introduces two BIFs:
* erts_internal:map_type/1
* erts_internal:map_hashmap_children/1

erts_internal:map_hashmap_children/1 is only intended for use within
erts_debug:size/1 since the internal hashmap node is not allowed to
leak anywhere.

3 files changed, 108 insertions, 10 deletions

diff --git a/erts/emulator/beam/bif.tab b/erts/emulator/beam/bif.tab
index b4e821a986..c56a108b34 100644
--- a/erts/emulator/beam/bif.tab
+++ b/erts/emulator/beam/bif.tab
@@ -157,6 +157,8 @@ bif erts_internal:request_system_task/3
 bif erts_internal:check_process_code/2
 
 bif erts_internal:map_to_tuple_keys/1
+bif erts_internal:map_type/1
+bif erts_internal:map_hashmap_children/1
 
 # inet_db support
 bif erlang:port_set_data/2
diff --git a/erts/emulator/beam/erl_map.c b/erts/emulator/beam/erl_map.c
index 16293668ad..0e24f2e1b1 100644
--- a/erts/emulator/beam/erl_map.c
+++ b/erts/emulator/beam/erl_map.c
@@ -2540,6 +2540,81 @@ BIF_RETTYPE erts_internal_map_to_tuple_keys_1(BIF_ALIST_1) {
     BIF_ERROR(BIF_P, BADARG);
 }
 
+/*
+ * erts_internal:map_type/1
+ *
+ * Used in erts_debug:size/1
+ */
+
+BIF_RETTYPE erts_internal_map_type_1(BIF_ALIST_1) {
+    DECL_AM(hashmap);
+    DECL_AM(hashmap_node);
+    DECL_AM(flatmap);
+    if (is_flatmap(BIF_ARG_1)) {
+	BIF_RET(AM_flatmap);
+    } else if (is_hashmap(BIF_ARG_1)) {
+        Eterm hdr = *(boxed_val(BIF_ARG_1));
+        ASSERT(is_header(hdr));
+        switch (hdr & _HEADER_MAP_SUBTAG_MASK) {
+            case HAMT_SUBTAG_HEAD_ARRAY:
+            case HAMT_SUBTAG_HEAD_BITMAP:
+                BIF_RET(AM_hashmap);
+            case HAMT_SUBTAG_NODE_ARRAY:
+            case HAMT_SUBTAG_NODE_BITMAP:
+                BIF_RET(AM_hashmap_node);
+            default:
+                erl_exit(1, "bad header");
+        }
+    }
+    BIF_ERROR(BIF_P, BADARG);
+}
+
+/*
+ * erts_internal:map_hashmap_children/1
+ *
+ * Used in erts_debug:size/1
+ */
+
+BIF_RETTYPE erts_internal_map_hashmap_children_1(BIF_ALIST_1) {
+    if (is_hashmap(BIF_ARG_1)) {
+        Eterm node = BIF_ARG_1;
+        Eterm *ptr, hdr, *hp, res = NIL;
+        Uint  sz = 0;
+        ptr = boxed_val(node);
+        hdr = *ptr;
+
+        ASSERT(is_header(hdr));
+
+        switch(hdr & _HEADER_MAP_SUBTAG_MASK) {
+            case HAMT_SUBTAG_NODE_ARRAY:
+                sz   = 16;
+                ptr += 1;
+                break;
+            case HAMT_SUBTAG_NODE_BITMAP:
+                sz   = hashmap_bitcount(MAP_HEADER_VAL(hdr));
+                ptr += 1;
+                break;
+            case HAMT_SUBTAG_HEAD_BITMAP:
+                sz   = hashmap_bitcount(MAP_HEADER_VAL(hdr));
+                ptr += 2;
+                break;
+            case HAMT_SUBTAG_HEAD_ARRAY:
+                sz   = 16;
+                ptr += 2;
+                break;
+            default:
+                erl_exit(1, "bad header\r\n");
+                break;
+        }
+        ASSERT(sz < 17);
+        hp = HAlloc(BIF_P, 2*sz);
+        while(sz--) { res = CONS(hp, *ptr++, res); hp += 2; }
+        BIF_RET(res);
+    }
+    BIF_ERROR(BIF_P, BADARG);
+}
+
+
 static Eterm hashmap_info(Process *p, Eterm node) {
     Eterm *hp;
     Eterm res = NIL, info = NIL;
diff --git a/lib/kernel/src/erts_debug.erl b/lib/kernel/src/erts_debug.erl
index ef605d0bfe..a82d9e750f 100644
--- a/lib/kernel/src/erts_debug.erl
+++ b/lib/kernel/src/erts_debug.erl
@@ -164,8 +164,10 @@ set_internal_state(_, _) ->
 
 -spec size(term()) -> non_neg_integer().
 
+-record(s, {seen, maps}).
+
 size(Term) ->
-    {Sum,_} = size(Term, gb_trees:empty(), 0),
+    {Sum,_} = size(Term, #s{seen=gb_trees:empty(),maps=[]}, 0),
     Sum.
 
 size([H|T]=Term, Seen0, Sum0) ->
@@ -209,10 +211,24 @@ tuple_size(I, Sz, Tuple, Seen0, Sum0) ->
     tuple_size(I+1, Sz, Tuple, Seen, Sum).
 
 map_size(Map,Seen0,Sum0) ->
-    Kt = erts_internal:map_to_tuple_keys(Map),
-    Vs = maps:values(Map),
-    {Sum1,Seen1} = size(Kt,Seen0,Sum0),
-    fold_size(Vs,Seen1,Sum1+length(Vs)+3).
+    %% Danger:
+    %% The internal nodes from erts_internal:map_hashmap_children/1
+    %% is not allowed to leak anywhere. They are only allowed in
+    %% containers (cons cells and tuples, not maps), in gc and
+    %% in erts_debug:same/2
+    case erts_internal:map_type(Map) of
+        flatmap ->
+            Kt = erts_internal:map_to_tuple_keys(Map),
+            Vs = maps:values(Map),
+            {Sum1,Seen1} = size(Kt,Seen0,Sum0),
+            fold_size(Vs,Seen1,Sum1+length(Vs)+3);
+        hashmap ->
+            Cs = erts_internal:map_hashmap_children(Map),
+            fold_size(Cs,Seen0,Sum0+length(Cs)+2);
+        hashmap_node ->
+            Cs = erts_internal:map_hashmap_children(Map),
+            fold_size(Cs,Seen0,Sum0+length(Cs)+1)
+    end.
 
 fun_size(Fun, Seen, Sum) ->
     case erlang:fun_info(Fun, type) of
@@ -229,13 +245,18 @@ fold_size([H|T], Seen0, Sum0) ->
     fold_size(T, Seen, Sum);
 fold_size([], Seen, Sum) -> {Sum,Seen}.
 
-remember_term(Term, Seen) ->
-    case gb_trees:lookup(Term, Seen) of
-	none -> gb_trees:insert(Term, [Term], Seen);
+remember_term(Term, #s{maps=Ms}=S) when is_map(Term) ->
+    case is_term_seen(Term, Ms) of
+        false -> S#s{maps=[Term|Ms]};
+        true  -> seen
+    end;
+remember_term(Term, #s{seen=T}=S) ->
+    case gb_trees:lookup(Term,T) of
+	none -> S#s{seen=gb_trees:insert(Term,[Term],T)};
 	{value,Terms} ->
 	    case is_term_seen(Term, Terms) of
-		false -> gb_trees:update(Term, [Term|Terms], Seen);
-		true -> seen
+		false -> S#s{seen=gb_trees:update(Term,[Term|Terms],T)};
+		true  -> seen
 	    end
     end.