Add array:counts()

4 files changed, 21 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);
diff --git a/compile.c b/compile.c
index 0cd34235..30868d52 100644
--- a/compile.c
+++ b/compile.c
@@ -2028,6 +2028,10 @@ CORD compile(env_t *env, ast_t *ast)
                 CORD 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, ", $SetInfo(", compile_type_info(env, item_t), "))");
+            } else if (streq(call->name, "counts")) {
+                CORD self = compile_to_pointer_depth(env, call->self, 0, false);
+                (void)compile_arguments(env, ast, NULL, call->args);
+                return CORD_all("Array$counts(", self, ", ", compile_type_info(env, self_value_t), ")");
             } else code_err(ast, "There is no '%s' method for arrays", call->name);
         }
         case SetType: {
diff --git a/typecheck.c b/typecheck.c
index 00839898..d8691544 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -707,6 +707,7 @@ type_t *get_type(env_t *env, ast_t *ast)
             else if (streq(call->name, "by")) return self_value_t;
             else if (streq(call->name, "reversed")) return self_value_t;
             else if (streq(call->name, "unique")) return Type(SetType, .item_type=Match(self_value_t, ArrayType)->item_type);
+            else if (streq(call->name, "counts")) return Type(TableType, .key_type=Match(self_value_t, ArrayType)->item_type, .value_type=INT_TYPE);
             else if (streq(call->name, "heapify")) return Type(VoidType);
             else if (streq(call->name, "heap_push")) return Type(VoidType);
             else if (streq(call->name, "heap_pop")) return Match(self_value_t, ArrayType)->item_type;