Support non-divergent list building

1 files changed, 20 insertions, 5 deletions

diff --git a/runtime/src/shed/namespaces/Std/Lists.c b/runtime/src/shed/namespaces/Std/Lists.c
index 41760d83..f4a01a6d 100644
--- a/runtime/src/shed/namespaces/Std/Lists.c
+++ b/runtime/src/shed/namespaces/Std/Lists.c
@@ -58,10 +58,8 @@ ShedObject shed__runtime__namespace__Std__Lists__method__ArrayBuilder__size(Shed
 }
 
 ShedObject shed__runtime__namespace__Std__Lists__method__ArrayBuilder__unsafeAdd(ShedClosurePointer closurePointer, ShedObject arg0, ShedObject arg1) {
-    // TODO: track whether array has been built or added to
     struct ShedArrayBuilder* array_builder = (struct ShedArrayBuilder*) arg0;
-    array_builder->array->elements[array_builder->array->size] = arg1;
-    array_builder->array->size += 1;
+    array_builder->array->elements[array_builder->size] = arg1;
 
     struct ShedArrayBuilder* new_array_builder = shed__runtime__gc_alloc(sizeof(struct ShedArrayBuilder), 0);
     new_array_builder->object_header.mark = 0b110;
@@ -78,7 +76,24 @@ ShedObject shed__runtime__namespace__Std__Lists__method__ArrayBuilder__unsafeEnl
 }
 
 ShedObject shed__runtime__namespace__Std__Lists__method__ArrayBuilder__build(ShedClosurePointer closurePointer, ShedObject arg0) {
-    // TODO: track whether array has been built or added to
     struct ShedArrayBuilder* array_builder = (struct ShedArrayBuilder*) arg0;
-    return (ShedObject) array_builder->array;
+    if (array_builder->array->size == 0) {
+        array_builder->array->size = array_builder->size;
+        return (ShedObject) array_builder->array;
+    } else if (array_builder->size == array_builder->array->size) {
+        return (ShedObject) array_builder->array;
+    } else {
+        struct ShedArray* array = shed__runtime__gc_alloc(
+            sizeof(struct ShedArray) + sizeof(ShedObject) * array_builder->size,
+            0
+        );
+        array->object_header.mark = 0b110;
+        array->size = array_builder->size;
+
+        for (PatioInt element_index = 0; element_index < array_builder->size; element_index++) {
+            array->elements[element_index] = array_builder->array->elements[element_index];
+        }
+
+        return (ShedObject) array;
+    }
 }