From a15c9223bd12b25f72337d78efc20aed8294d5bf Mon Sep 17 00:00:00 2001
From: Fred Dushin <fred@dushin.net>
Date: Sat, 15 Jul 2023 10:19:28 -0400
Subject: [PATCH] Add support for `interop_iolist_fold`

Signed-off-by: Fred Dushin <fred@dushin.net>
---
 src/libAtomVM/interop.c | 93 +++++++++++++++++------------------------
 src/libAtomVM/interop.h |  3 ++
 2 files changed, 41 insertions(+), 55 deletions(-)

diff --git a/src/libAtomVM/interop.c b/src/libAtomVM/interop.c
index e8013a11c..b04b00281 100644
--- a/src/libAtomVM/interop.c
+++ b/src/libAtomVM/interop.c
@@ -176,19 +176,16 @@ term interop_proplist_get_value_default(term list, term key, term default_value)
     return default_value;
 }
 
-InteropFunctionResult interop_iolist_size(term t, size_t *size)
+inline InteropFunctionResult interop_iolist_fold(term t, interop_iolist_fold_fun fold_fun, void *accum)
 {
     if (term_is_binary(t)) {
-        *size = term_binary_size(t);
-        return InteropOk;
+        return fold_fun(t, accum);
     }
 
     if (UNLIKELY(!term_is_list(t))) {
         return InteropBadArg;
     }
 
-    unsigned long acc = 0;
-
     struct TempStack temp_stack;
     if (UNLIKELY(temp_stack_init(&temp_stack) != TempStackOk)) {
         return InteropMemoryAllocFail;
@@ -200,9 +197,14 @@ InteropFunctionResult interop_iolist_size(term t, size_t *size)
     }
 
     while (!temp_stack_is_empty(&temp_stack)) {
-        if (term_is_integer(t)) {
-            acc++;
-            t = temp_stack_pop(&temp_stack);
+        if (term_is_integer(t) || term_is_binary(t)) {
+            InteropFunctionResult result = fold_fun(t, accum);
+            if (UNLIKELY(result != InteropOk)) {
+                temp_stack_destroy(&temp_stack);
+                return result;
+            } else {
+                t = temp_stack_pop(&temp_stack);
+            }
 
         } else if (term_is_nil(t)) {
             t = temp_stack_pop(&temp_stack);
@@ -214,10 +216,6 @@ InteropFunctionResult interop_iolist_size(term t, size_t *size)
             }
             t = term_get_list_head(t);
 
-        } else if (term_is_binary(t)) {
-            acc += term_binary_size(t);
-            t = temp_stack_pop(&temp_stack);
-
         } else {
             temp_stack_destroy(&temp_stack);
             return InteropBadArg;
@@ -226,60 +224,45 @@ InteropFunctionResult interop_iolist_size(term t, size_t *size)
 
     temp_stack_destroy(&temp_stack);
 
-    *size = acc;
     return InteropOk;
 }
 
-InteropFunctionResult interop_write_iolist(term t, char *p)
+static inline InteropFunctionResult size_fold_fun(term t, void *accum)
 {
-    if (term_is_binary(t)) {
-        int len = term_binary_size(t);
-        memcpy(p, term_binary_data(t), len);
-        return InteropOk;
-    }
-
-    struct TempStack temp_stack;
-    if (UNLIKELY(temp_stack_init(&temp_stack) != TempStackOk)) {
-        return InteropMemoryAllocFail;
-    }
-
-    if (UNLIKELY(temp_stack_push(&temp_stack, t) != TempStackOk)) {
-        temp_stack_destroy(&temp_stack);
-        return InteropMemoryAllocFail;
+    size_t *size = (size_t *) accum;
+    if (term_is_integer(t)) {
+        *size += 1;
+    } else if (term_is_binary(t)) {
+        *size += term_binary_size(t);
     }
+    return InteropOk;
+}
 
-    while (!temp_stack_is_empty(&temp_stack)) {
-        if (term_is_integer(t)) {
-            *p = term_to_int(t);
-            p++;
-            t = temp_stack_pop(&temp_stack);
-
-        } else if (term_is_nil(t)) {
-            t = temp_stack_pop(&temp_stack);
-
-        } else if (term_is_nonempty_list(t)) {
-            if (UNLIKELY(temp_stack_push(&temp_stack, term_get_list_tail(t)) != TempStackOk)) {
-                temp_stack_destroy(&temp_stack);
-                return InteropMemoryAllocFail;
-            }
-            t = term_get_list_head(t);
-
-        } else if (term_is_binary(t)) {
-            int len = term_binary_size(t);
-            memcpy(p, term_binary_data(t), len);
-            p += len;
-            t = temp_stack_pop(&temp_stack);
+InteropFunctionResult interop_iolist_size(term t, size_t *size)
+{
+    *size = 0;
+    return interop_iolist_fold(t, size_fold_fun, size);
+}
 
-        } else {
-            temp_stack_destroy(&temp_stack);
-            return InteropBadArg;
-        }
+static inline InteropFunctionResult write_string_fold_fun(term t, void *accum)
+{
+    char **p = (char **) accum;
+    if (term_is_integer(t)) {
+        **p = term_to_int(t);
+        (*p)++;
+    } else if (term_is_binary(t)) {
+        int len = term_binary_size(t);
+        memcpy(*p, term_binary_data(t), len);
+        *p += len;
     }
-
-    temp_stack_destroy(&temp_stack);
     return InteropOk;
 }
 
+InteropFunctionResult interop_write_iolist(term t, char *p)
+{
+    return interop_iolist_fold(t, write_string_fold_fun, (void *) &p);
+}
+
 term interop_map_get_value(GlobalContext *glb, term map, term key)
 {
     return interop_map_get_value_default(glb, map, key, term_nil());
diff --git a/src/libAtomVM/interop.h b/src/libAtomVM/interop.h
index 987612196..c1d92b836 100644
--- a/src/libAtomVM/interop.h
+++ b/src/libAtomVM/interop.h
@@ -53,6 +53,8 @@ typedef struct
     int i_val;
 } AtomStringIntPair;
 
+typedef InteropFunctionResult (*interop_iolist_fold_fun)(term t, void *accum);
+
 char *interop_term_to_string(term t, int *ok);
 char *interop_binary_to_string(term binary);
 char *interop_list_to_string(term list, int *ok);
@@ -65,6 +67,7 @@ term interop_map_get_value_default(GlobalContext *glb, term map, term key, term
 
 NO_DISCARD InteropFunctionResult interop_iolist_size(term t, size_t *size);
 NO_DISCARD InteropFunctionResult interop_write_iolist(term t, char *p);
+NO_DISCARD InteropFunctionResult interop_iolist_fold(term t, interop_iolist_fold_fun fold_fun, void *accum);
 
 /**
  * @brief Finds on a table the first matching atom string.