aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--builtins/array.c18
-rw-r--r--builtins/array.h3
-rw-r--r--compile.c9
-rw-r--r--typecheck.c1
4 files changed, 31 insertions, 0 deletions
diff --git a/builtins/array.c b/builtins/array.c
index a6e0dca2..abc9f923 100644
--- a/builtins/array.c
+++ b/builtins/array.c
@@ -634,4 +634,22 @@ public void Array$heapify(array_t *heap, closure_t comparison, int64_t padded_it
ARRAY_DECREF(*heap);
}
+public Int_t Array$binary_search(array_t array, void *target, closure_t comparison)
+{
+ typedef int32_t (*cmp_fn_t)(void*, void*, void*);
+ int64_t lo = 0, hi = array.length-1;
+ while (lo <= hi) {
+ int64_t mid = (lo + hi) / 2;
+ int32_t cmp = ((cmp_fn_t)comparison.fn)(
+ array.data + array.stride*mid, target, comparison.userdata);
+ if (cmp == 0)
+ return I(mid+1);
+ else if (cmp < 0)
+ lo = mid + 1;
+ else if (cmp > 0)
+ hi = mid - 1;
+ }
+ return I(lo+1); // Return the index where the target would be inserted
+}
+
// vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/builtins/array.h b/builtins/array.h
index a3305ca3..88184f8a 100644
--- a/builtins/array.h
+++ b/builtins/array.h
@@ -85,5 +85,8 @@ void Array$heap_pop(array_t *heap, closure_t comparison, int64_t padded_item_siz
#define Array$heap_pop_value(heap, comparison, padded_item_size, type) \
({ array_t *_heap = heap; if (_heap->length == 0) fail("Attempt to pop from an empty array"); \
type value = *(type*)_heap->data; Array$heap_pop(_heap, comparison, padded_item_size); value; })
+Int_t Array$binary_search(array_t array, void *target, closure_t comparison);
+#define Array$binary_search_value(array, target, comparison) \
+ ({ __typeof(target) _target = target; Array$binary_search(array, &_target, comparison); })
// vim: ts=4 sw=0 et cino=L2,l1,(0,W4,m1,\:0
diff --git a/compile.c b/compile.c
index 516ca36d..f7589ea7 100644
--- a/compile.c
+++ b/compile.c
@@ -2166,6 +2166,15 @@ CORD compile(env_t *env, ast_t *ast)
arg_t *arg_spec = new(arg_t, .name="by", .type=Type(ClosureType, .fn=fn_t), .default_val=default_cmp);
CORD arg_code = compile_arguments(env, ast, arg_spec, call->args);
return CORD_all("Array$heap_pop_value(", self, ", ", arg_code, ", ", padded_item_size, ", ", compile_type(item_t), ")");
+ } else if (streq(call->name, "binary_search")) {
+ CORD self = compile_to_pointer_depth(env, call->self, 0, false);
+ type_t *item_ptr = Type(PointerType, .pointed=item_t, .is_stack=true);
+ type_t *fn_t = Type(FunctionType, .args=new(arg_t, .name="x", .type=item_ptr, .next=new(arg_t, .name="y", .type=item_ptr)),
+ .ret=Type(IntType, .bits=32));
+ ast_t *default_cmp = FakeAST(InlineCCode, .code=CORD_all("((closure_t){.fn=generic_compare, .userdata=(void*)", compile_type_info(env, item_t), "})"), .type=NewTypeAST(NULL, NULL, NULL, FunctionTypeAST));
+ arg_t *arg_spec = new(arg_t, .name="target", .type=item_t, .next=new(arg_t, .name="by", .type=Type(ClosureType, .fn=fn_t), .default_val=default_cmp));
+ CORD arg_code = compile_arguments(env, ast, arg_spec, call->args);
+ return CORD_all("Array$binary_search_value(", self, ", ", arg_code, ")");
} else if (streq(call->name, "clear")) {
CORD self = compile_to_pointer_depth(env, call->self, 1, false);
(void)compile_arguments(env, ast, NULL, call->args);
diff --git a/typecheck.c b/typecheck.c
index 1100aedc..2814bd03 100644
--- a/typecheck.c
+++ b/typecheck.c
@@ -728,6 +728,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, "heap_push")) return Type(VoidType);
else if (streq(call->name, "heap_pop")) return Match(self_value_t, ArrayType)->item_type;
+ else if (streq(call->name, "binary_search")) return INT_TYPE;
else code_err(ast, "There is no '%s' method for arrays", call->name);
}
case SetType: {