iolist_size/1: Fix truncation of result

iolist_size/1 would silently return a truncated result for iolists
whose size exceeded the word size. For example:

  iolist(lists:duplicate(256, <<0:(1 bsl 24)/unit:8>>)).

would return 0 (instead of 4294967296 or 1 bsl 32).

Rewrite iolist_size/1 to accumulate the size in a bignum if
necessary and result the correct result.

2 files changed, 191 insertions, 18 deletions

diff --git a/erts/emulator/beam/bif.c b/erts/emulator/beam/bif.c
index 6b1ce823cb..e886792ba7 100644
--- a/erts/emulator/beam/bif.c
+++ b/erts/emulator/beam/bif.c
@@ -2168,20 +2168,146 @@ BIF_RETTYPE tl_1(BIF_ALIST_1)
 /**********************************************************************/
 /* return the size of an I/O list */
 
-BIF_RETTYPE iolist_size_1(BIF_ALIST_1)
+static Eterm
+accumulate(Eterm acc, Uint size)
 {
-    Sint size = io_list_len(BIF_ARG_1);
+    if (is_non_value(acc)) {
+	/*
+	 * There is no pre-existing accumulator. Allocate a
+	 * bignum buffer with one extra word to be used if
+	 * the bignum grows in the future.
+	 */
+	Eterm* hp = (Eterm *) erts_alloc(ERTS_ALC_T_TEMP_TERM,
+					 (BIG_UINT_HEAP_SIZE+1) *
+					 sizeof(Eterm));
+	return uint_to_big(size, hp);
+    } else {
+	Eterm* big;
+	int need_heap;
 
-    if (size == -1) {
-	BIF_ERROR(BIF_P, BADARG);
-    } else if (IS_USMALL(0, (Uint) size)) {
-	BIF_RET(make_small(size));
+	/*
+	 * Add 'size' to 'acc' in place. There is always one
+	 * extra word allocated in case the bignum grows by one word.
+	 */
+	big = big_val(acc);
+	need_heap = BIG_NEED_SIZE(BIG_SIZE(big));
+	acc = big_plus_small(acc, size, big);
+	if (BIG_NEED_SIZE(big_size(acc)) > need_heap) {
+	    /*
+	     * The extra word has been consumed. Grow the
+	     * allocation by one word.
+	     */
+	    big = (Eterm *) erts_realloc(ERTS_ALC_T_TEMP_TERM,
+					 big_val(acc),
+					 (need_heap+1) * sizeof(Eterm));
+	    acc = make_big(big);
+	}
+	return acc;
+    }
+}
+
+static Eterm
+consolidate(Process* p, Eterm acc, Uint size)
+{
+    Eterm* hp;
+
+    if (is_non_value(acc)) {
+	return erts_make_integer(size, p);
     } else {
-	Eterm* hp = HAlloc(BIF_P, BIG_UINT_HEAP_SIZE);
-	BIF_RET(uint_to_big(size, hp));
+	Eterm* big;
+	Uint sz;
+	Eterm res;
+	
+	acc = accumulate(acc, size);
+	big = big_val(acc);
+	sz = BIG_NEED_SIZE(BIG_SIZE(big));
+	hp = HAlloc(p, sz);
+	res = make_big(hp);
+	while (sz--) {
+	    *hp++ = *big++;
+	}
+	erts_free(ERTS_ALC_T_TEMP_TERM, (void *) big_val(acc));
+	return res;
     }
 }
 
+BIF_RETTYPE iolist_size_1(BIF_ALIST_1)
+{
+    Eterm obj, hd;
+    Eterm* objp;
+    Uint size = 0;
+    Uint cur_size;
+    Uint new_size;
+    Eterm acc = THE_NON_VALUE;
+    DECLARE_ESTACK(s);
+
+    obj = BIF_ARG_1;
+    goto L_again;
+
+    while (!ESTACK_ISEMPTY(s)) {
+	obj = ESTACK_POP(s);
+    L_again:
+	if (is_list(obj)) {
+	L_iter_list:
+	    objp = list_val(obj);
+	    hd = CAR(objp);
+	    obj = CDR(objp);
+	    /* Head */
+	    if (is_byte(hd)) {
+		size++;
+		if (size == 0) {
+		    acc = accumulate(acc, (Uint) -1);
+		    size = 1;
+		}
+	    } else if (is_binary(hd) && binary_bitsize(hd) == 0) {
+		cur_size = binary_size(hd);
+		if ((new_size = size + cur_size) >= size) {
+		    size = new_size;
+		} else {
+		    acc = accumulate(acc, size);
+		    size = cur_size;
+		}
+	    } else if (is_list(hd)) {
+		ESTACK_PUSH(s, obj);
+		obj = hd;
+		goto L_iter_list;
+	    } else if (is_not_nil(hd)) {
+		goto L_type_error;
+	    }
+	    /* Tail */
+	    if (is_list(obj)) {
+		goto L_iter_list;
+	    } else if (is_binary(obj) && binary_bitsize(obj) == 0) {
+		cur_size = binary_size(obj);
+		if ((new_size = size + cur_size) >= size) {
+		    size = new_size;
+		} else {
+		    acc = accumulate(acc, size);
+		    size = cur_size;
+		}
+	    } else if (is_not_nil(obj)) {
+		goto L_type_error;
+	    }
+	} else if (is_binary(obj) && binary_bitsize(obj) == 0) {
+	    cur_size = binary_size(obj);
+	    if ((new_size = size + cur_size) >= size) {
+		size = new_size;
+	    } else {
+		acc = accumulate(acc, size);
+		size = cur_size;
+	    }
+	} else if (is_not_nil(obj)) {
+	    goto L_type_error;
+	}
+    }
+
+    DESTROY_ESTACK(s);
+    BIF_RET(consolidate(BIF_P, acc, size));
+
+ L_type_error:
+    DESTROY_ESTACK(s);
+    BIF_ERROR(BIF_P, BADARG);
+}
 
 /**********************************************************************/
 
diff --git a/erts/emulator/test/binary_SUITE.erl b/erts/emulator/test/binary_SUITE.erl
index b7b2aaf89e..401b3c23b6 100644
--- a/erts/emulator/test/binary_SUITE.erl
+++ b/erts/emulator/test/binary_SUITE.erl
@@ -517,18 +517,65 @@ external_size_1(Term, Size0, Limit) when Size0 < Limit ->
 external_size_1(_, _, _) -> ok.
 
 t_iolist_size(Config) when is_list(Config) ->
-    %% Build a term whose external size only fits in a big num (on 32-bit CPU).
-    Bin = iolist_to_binary(lists:seq(0, 254)),
-    ?line ok = t_iolist_size_1(Bin, 0, 16#7FFFFFFF),
-    ?line ok = t_iolist_size_1(make_unaligned_sub_binary(Bin), 0, 16#7FFFFFFF).
+    ?line Seed = now(),
+    ?line io:format("Seed: ~p", [Seed]),
+    ?line random:seed(Seed),
+    ?line Base = <<0:(1 bsl 20)/unit:8>>,
+    ?line Powers = [1 bsl N || N <- lists:seq(2, 37)],
+    ?line Sizes0 = [[N - random:uniform(N div 2),
+		     lists:seq(N-2, N+2),
+		     N+N div 2,
+		     N + random:uniform(N div 2)] ||
+		       N <- Powers],
+    %% Test sizes around 1^32 more thoroughly.
+    FourGigs = 1 bsl 32,
+    ?line Sizes1 = [FourGigs+N || N <- lists:seq(-8, 40)] ++ Sizes0,
+    ?line Sizes2 = lists:flatten(Sizes1),
+    ?line Sizes = lists:usort(Sizes2),
+    io:format("~p sizes:", [length(Sizes)]),
+    io:format("~p\n", [Sizes]),
+    ?line [Sz = iolist_size(build_iolist(Sz, Base)) || Sz <- Sizes],
+    ok.
 
-t_iolist_size_1(IOList, Size0, Limit) when Size0 < Limit ->
-    case iolist_size(IOList) of
-	Size when is_integer(Size), Size0 < Size ->
-	    io:format("~p", [Size]),
-	    t_iolist_size_1([IOList|IOList], Size, Limit)
+build_iolist(N, Base) when N < 16 ->
+    case random:uniform(3) of
+	1 ->
+	    <<Bin:N/binary,_/binary>> = Base,
+	    Bin;
+	_ ->
+	    lists:seq(1, N)
     end;
-t_iolist_size_1(_, _, _) -> ok.
+build_iolist(N, Base) when N =< byte_size(Base) ->
+    case random:uniform(3) of
+	1 ->
+	    <<Bin:N/binary,_/binary>> = Base,
+	    Bin;
+	2 ->
+	    <<Bin:N/binary,_/binary>> = Base,
+	    [Bin];
+	3 ->
+	    case N rem 2 of
+		0 ->
+		    L = build_iolist(N div 2, Base),
+		    [L,L];
+		1 ->
+		    L = build_iolist(N div 2, Base),
+		    [L,L,45]
+	    end
+    end;
+build_iolist(N0, Base) ->
+    Small = random:uniform(15),
+    Seq = lists:seq(1, Small),
+    N = N0 - Small,
+    case N rem 2 of
+	0 ->
+	    L = build_iolist(N div 2, Base),
+	    [L,L|Seq];
+	1 ->
+	    L = build_iolist(N div 2, Base),
+	    [47,L,L|Seq]
+    end.
+
 
 bad_binary_to_term_2(doc) -> "OTP-4053.";
 bad_binary_to_term_2(suite) -> [];