Add array:pop()

6 files changed, 77 insertions, 4 deletions

diff --git a/compile.c b/compile.c
index 94ef0513..2cc3284e 100644
--- a/compile.c
+++ b/compile.c
@@ -3013,6 +3013,12 @@ CORD compile(env_t *env, ast_t *ast)
                 self = compile_to_pointer_depth(env, call->self, 0, false);
                 (void)compile_arguments(env, ast, NULL, call->args);
                 return CORD_all("Table$from_entries(", self, ", Set$info(", compile_type_info(env, item_t), "))");
+            } else if (streq(call->name, "pop")) {
+                EXPECT_POINTER("an", "array");
+                arg_t *arg_spec = new(arg_t, .name="index", .type=INT_TYPE, .default_val=FakeAST(Int, "-1"));
+                CORD index = compile_arguments(env, ast, arg_spec, call->args);
+                return CORD_all("Array$pop(", self, ", ", index, ", ", compile_type(item_t), ", _, ",
+                                promote_to_optional(item_t, "_"), ", ", compile_none(item_t), ", ", padded_item_size, ")");
             } else if (streq(call->name, "counts")) {
                 self = compile_to_pointer_depth(env, call->self, 0, false);
                 (void)compile_arguments(env, ast, NULL, call->args);
diff --git a/docs/arrays.md b/docs/arrays.md
index 76372cb9..3bb15bdf 100644
--- a/docs/arrays.md
+++ b/docs/arrays.md
@@ -608,6 +608,44 @@ arr:insert_all([99, 100], at=2)
 
 ---
 
+### `pop`
+
+**Description:**  
+Removes and returns an item from the array. If the given index is present in
+the array, the item at that index will be removed and the array will become one
+element shorter.
+
+**Signature:**  
+```tomo
+func pop(arr: &[T], index: Int = -1 -> T?)
+```
+
+**Parameters:**
+
+- `arr`: The array to remove an item from.
+- `index`: The index from which to remove the item (default: the last item).
+
+**Returns:**  
+`none` if the array is empty or the given index does not exist in the array,
+otherwise the item at the given index.
+
+**Example:**  
+```tomo
+>> arr := [10, 20, 30, 40]
+
+>> arr:pop()
+= 40
+>> arr
+= &[10, 20, 30]
+
+>> arr:pop(index=2)
+= 20
+>> arr
+= &[10, 30]
+```
+
+---
+
 ### `random`
 
 **Description:**  
diff --git a/stdlib/arrays.c b/stdlib/arrays.c
index e5ffbfd4..1f4b8252 100644
--- a/stdlib/arrays.c
+++ b/stdlib/arrays.c
@@ -217,7 +217,7 @@ public void Array$remove_item(Array_t *arr, void *item, Int_t max_removals, cons
     }
 }
 
-public Int_t Array$find(Array_t arr, void *item, const TypeInfo_t *type)
+public OptionalInt_t Array$find(Array_t arr, void *item, const TypeInfo_t *type)
 {
     const TypeInfo_t *item_type = type->ArrayInfo.item;
     for (int64_t i = 0; i < arr.length; i++) {
@@ -227,7 +227,7 @@ public Int_t Array$find(Array_t arr, void *item, const TypeInfo_t *type)
     return NONE_INT;
 }
 
-public Int_t Array$first(Array_t arr, Closure_t predicate)
+public OptionalInt_t Array$first(Array_t arr, Closure_t predicate)
 {
     bool (*is_good)(void*, void*) = (void*)predicate.fn;
     for (int64_t i = 0; i < arr.length; i++) {
diff --git a/stdlib/arrays.h b/stdlib/arrays.h
index 5e0ca7e4..f8921e0d 100644
--- a/stdlib/arrays.h
+++ b/stdlib/arrays.h
@@ -65,9 +65,21 @@ void Array$insert_all(Array_t *arr, Array_t to_insert, Int_t index, int64_t padd
 void Array$remove_at(Array_t *arr, Int_t index, Int_t count, int64_t padded_item_size);
 void Array$remove_item(Array_t *arr, void *item, Int_t max_removals, const TypeInfo_t *type);
 #define Array$remove_item_value(arr, item_expr, max, type) ({ __typeof(item_expr) item = item_expr; Array$remove_item(arr, &item, max, type); })
-Int_t Array$find(Array_t arr, void *item, const TypeInfo_t *type);
+
+#define Array$pop(arr_expr, index_expr, item_type, nonnone_var, nonnone_expr, none_expr, padded_item_size) ({ \
+    Array_t *arr = arr_expr; \
+    Int_t index = index_expr; \
+    int64_t index64 = Int_to_Int64(index, false); \
+    int64_t off = index64 + (index64 < 0) * (arr->length + 1) - 1; \
+    (off >= 0 && off < arr->length) ? ({ \
+        item_type nonnone_var = *(item_type*)(arr->data + off*arr->stride); \
+        Array$remove_at(arr, index, I_small(1), padded_item_size); \
+        nonnone_expr; \
+    }) : none_expr; })
+
+OptionalInt_t Array$find(Array_t arr, void *item, const TypeInfo_t *type);
 #define Array$find_value(arr, item_expr, type) ({ __typeof(item_expr) item = item_expr; Array$find(arr, &item, type); })
-Int_t Array$first(Array_t arr, Closure_t predicate);
+OptionalInt_t Array$first(Array_t arr, Closure_t predicate);
 void Array$sort(Array_t *arr, Closure_t comparison, int64_t padded_item_size);
 Array_t Array$sorted(Array_t arr, Closure_t comparison, int64_t padded_item_size);
 void Array$shuffle(Array_t *arr, RNG_t rng, int64_t padded_item_size);
diff --git a/test/arrays.tm b/test/arrays.tm
index 8bc5a7f6..e32ee598 100644
--- a/test/arrays.tm
+++ b/test/arrays.tm
@@ -168,3 +168,19 @@ func main():
 	= none : Int?
 	>> [4, 5, 6]:first(func(i:&Int): i:is_prime())
 	= 2 : Int?
+
+	do:
+		>> nums := &[10, 20, 30, 40, 50]
+		>> nums:pop()
+		= 50
+		>> nums
+		= &[10, 20, 30, 40]
+		>> nums:pop(2)
+		= 20
+		>> nums
+		= &[10, 30, 40]
+		>> nums:clear()
+		>> nums
+		= &[]
+		>> nums:pop()
+		= none
diff --git a/typecheck.c b/typecheck.c
index 0bca1e3a..afbc8df4 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -829,6 +829,7 @@ type_t *get_type(env_t *env, ast_t *ast)
             else if (streq(call->name, "heapify")) return Type(VoidType);
             else if (streq(call->name, "insert")) return Type(VoidType);
             else if (streq(call->name, "insert_all")) return Type(VoidType);
+            else if (streq(call->name, "pop")) return Type(OptionalType, .type=item_type);
             else if (streq(call->name, "random")) return item_type;
             else if (streq(call->name, "remove_at")) return Type(VoidType);
             else if (streq(call->name, "remove_item")) return Type(VoidType);