Add array:counts()

2 files changed, 16 insertions, 0 deletions

diff --git a/builtins/array.c b/builtins/array.c
index aa6929d5..c5d52380 100644
--- a/builtins/array.c
+++ b/builtins/array.c
@@ -14,6 +14,7 @@
 #include "functions.h"
 #include "halfsiphash.h"
 #include "integers.h"
+#include "table.h"
 #include "types.h"
 #include "util.h"
 
@@ -219,6 +220,20 @@ public void *Array$random(array_t arr)
     return arr.data + arr.stride*index;
 }
 
+public table_t Array$counts(array_t arr, const TypeInfo *type)
+{
+    table_t counts = {};
+    const TypeInfo count_type = {.size=sizeof(table_t), .align=__alignof__(table_t),
+        .tag=TableInfo, .TableInfo.key=type->ArrayInfo.item, .TableInfo.value=&$Int};
+    for (int64_t i = 0; i < arr.length; i++) {
+        void *key = arr.data + i*arr.stride;
+        int64_t *count = Table$get(counts, key, &count_type);
+        int64_t val = count ? *count + 1 : 1;
+        Table$set(&counts, key, &val, &count_type);
+    }
+    return counts;
+}
+
 public array_t Array$sample(array_t arr, int64_t n, array_t weights, int64_t padded_item_size)
 {
     if (arr.length == 0 || n <= 0)
diff --git a/builtins/array.h b/builtins/array.h
index cefea807..4acf2664 100644
--- a/builtins/array.h
+++ b/builtins/array.h
@@ -63,6 +63,7 @@ void Array$shuffle(array_t *arr, int64_t padded_item_size);
 void *Array$random(array_t arr);
 #define Array$random_value(arr, t) ({ array_t _arr = arr; if (_arr.length == 0) fail("Cannot get a random value from an empty array!"); *(t*)Array$random(_arr); })
 array_t Array$sample(array_t arr, int64_t n, array_t weights, int64_t padded_item_size);
+table_t Array$counts(array_t arr, const TypeInfo *type);
 void Array$clear(array_t *array);
 void Array$compact(array_t *arr, int64_t padded_item_size);
 bool Array$contains(array_t array, void *item, const TypeInfo *type);