bytes index_of can now take another byte as argument

blade-lang · Sep 26, 2024 · 1eca5fb · 1eca5fb
1 parent bef62d9
commit 1eca5fb
Showing 1 changed file with 18 additions and 6 deletions.
diff --git a/src/bytes.c b/src/bytes.c
@@ -111,20 +111,32 @@ DECLARE_BYTES_METHOD(extend) {
 
 DECLARE_BYTES_METHOD(index_of) {
   ENFORCE_ARG_RANGE(index_of, 1, 2);
-  ENFORCE_ARG_TYPE(index_of, 0, IS_NUMBER);
+  ENFORCE_ARG_TYPES(index_of, 0, IS_NUMBER, IS_BYTES);
   b_obj_bytes *bytes = AS_BYTES(METHOD_OBJECT);
-  uint8_t needle = AS_NUMBER(args[0]);
 
   int start_index = 0;
   if(arg_count == 2) {
     ENFORCE_ARG_TYPE(index_of, 1, IS_NUMBER);
     start_index = AS_NUMBER(args[1]);
   }
 
-  if(bytes->bytes.count > 0 && start_index >= 0 && start_index < bytes->bytes.count - 1) {
-    for (int i = start_index; i < bytes->bytes.count; i++) {
-      if (bytes->bytes.bytes[i] == needle) {
-        RETURN_NUMBER(i);
+  if(IS_NUMBER(args[0])) {
+    uint8_t needle = AS_NUMBER(args[0]);
+    if(bytes->bytes.count > 0 && start_index >= 0 && start_index < bytes->bytes.count - 1) {
+      for (int i = start_index; i < bytes->bytes.count; i++) {
+        if (bytes->bytes.bytes[i] == needle) {
+          RETURN_NUMBER(i);
+        }
+      }
+    }
+  } else {
+    b_obj_bytes *needle = AS_BYTES(args[0]);
+
+    if(bytes->bytes.count > 0 && start_index >= 0 && start_index < bytes->bytes.count - 1) {
+      for (int i = start_index; i < bytes->bytes.count; i++) {
+        if (memcmp(bytes->bytes.bytes + i, needle->bytes.bytes, needle->bytes.count) == 0) {
+          RETURN_NUMBER(i);
+        }
       }
     }
   }