diff --git a/src/vm/wren_core.c b/src/vm/wren_core.c index 33029188..51c1e9ab 100644 --- a/src/vm/wren_core.c +++ b/src/vm/wren_core.c @@ -869,7 +869,7 @@ DEF_PRIMITIVE(string_contains) ObjString* string = AS_STRING(args[0]); ObjString* search = AS_STRING(args[1]); - RETURN_BOOL(wrenStringFind(string, search) != UINT32_MAX); + RETURN_BOOL(wrenStringFind(string, search, 0) != UINT32_MAX); } DEF_PRIMITIVE(string_endsWith) @@ -893,7 +893,19 @@ DEF_PRIMITIVE(string_indexOf) ObjString* string = AS_STRING(args[0]); ObjString* search = AS_STRING(args[1]); - uint32_t index = wrenStringFind(string, search); + uint32_t index = wrenStringFind(string, search, 0); + RETURN_NUM(index == UINT32_MAX ? -1 : (int)index); +} + +DEF_PRIMITIVE(string_indexOf_with_startIndex) +{ + if (!validateString(vm, args[1], "Argument")) return false; + + ObjString* string = AS_STRING(args[0]); + ObjString* search = AS_STRING(args[1]); + uint32_t startIndex = AS_NUM(args[2]); + + uint32_t index = wrenStringFind(string, search, startIndex); RETURN_NUM(index == UINT32_MAX ? -1 : (int)index); } @@ -1252,6 +1264,7 @@ void wrenInitializeCore(WrenVM* vm) PRIMITIVE(vm->stringClass, "contains(_)", string_contains); PRIMITIVE(vm->stringClass, "endsWith(_)", string_endsWith); PRIMITIVE(vm->stringClass, "indexOf(_)", string_indexOf); + PRIMITIVE(vm->stringClass, "indexOf(_,_)", string_indexOf_with_startIndex); PRIMITIVE(vm->stringClass, "iterate(_)", string_iterate); PRIMITIVE(vm->stringClass, "iterateByte_(_)", string_iterateByte); PRIMITIVE(vm->stringClass, "iteratorValue(_)", string_iteratorValue); diff --git a/src/vm/wren_value.c b/src/vm/wren_value.c index 36526559..06bfc85b 100644 --- a/src/vm/wren_value.c +++ b/src/vm/wren_value.c @@ -854,13 +854,16 @@ Value wrenStringCodePointAt(WrenVM* vm, ObjString* string, uint32_t index) } // Uses the Boyer-Moore-Horspool string matching algorithm. -uint32_t wrenStringFind(ObjString* haystack, ObjString* needle) +uint32_t wrenStringFind(ObjString* haystack, ObjString* needle, uint32_t startIndex) { // Corner case, an empty needle is always found. if (needle->length == 0) return 0; // If the needle is longer than the haystack it won't be found. - if (needle->length > haystack->length) return UINT32_MAX; + if (needle->length > (haystack->length - startIndex)) return UINT32_MAX; + + // If the startIndex is too far it also won't be found. + if (startIndex >= haystack->length) return UINT32_MAX; // Pre-calculate the shift table. For each character (8-bit value), we // determine how far the search window can be advanced if that character is @@ -890,18 +893,18 @@ uint32_t wrenStringFind(ObjString* haystack, ObjString* needle) // Slide the needle across the haystack, looking for the first match or // stopping if the needle goes off the end. char lastChar = needle->value[needleEnd]; - uint32_t range = haystack->length - needle->length; + uint32_t range = (haystack->length - startIndex) - needle->length; for (uint32_t index = 0; index <= range; ) { // Compare the last character in the haystack's window to the last character // in the needle. If it matches, see if the whole needle matches. - char c = haystack->value[index + needleEnd]; + char c = haystack->value[startIndex + (index + needleEnd)]; if (lastChar == c && - memcmp(haystack->value + index, needle->value, needleEnd) == 0) + memcmp(haystack->value + startIndex + index, needle->value, needleEnd) == 0) { // Found a match. - return index; + return index + startIndex; } // Otherwise, slide the needle forward. diff --git a/src/vm/wren_value.h b/src/vm/wren_value.h index 9d4cd3ee..d85297a9 100644 --- a/src/vm/wren_value.h +++ b/src/vm/wren_value.h @@ -728,7 +728,7 @@ Value wrenStringCodePointAt(WrenVM* vm, ObjString* string, uint32_t index); // Search for the first occurence of [needle] within [haystack] and returns its // zero-based offset. Returns `UINT32_MAX` if [haystack] does not contain // [needle]. -uint32_t wrenStringFind(ObjString* haystack, ObjString* needle); +uint32_t wrenStringFind(ObjString* haystack, ObjString* needle, uint32_t startIndex); // Creates a new open upvalue pointing to [value] on the stack. ObjUpvalue* wrenNewUpvalue(WrenVM* vm, Value* value); diff --git a/test/core/string/index_of.wren b/test/core/string/index_of.wren index 22fb8f79..c6b3e02d 100644 --- a/test/core/string/index_of.wren +++ b/test/core/string/index_of.wren @@ -5,6 +5,12 @@ System.print("abcd".indexOf("abcd")) // expect: 0 System.print("abcd".indexOf("abcde")) // expect: -1 System.print("abab".indexOf("ab")) // expect: 0 +System.print("abcd".indexOf("cd", 0)) // expect: 2 +System.print("abcd".indexOf("cd", 1)) // expect: 2 +System.print("abcd".indexOf("cd", 2)) // expect: 2 +System.print("abcd".indexOf("cd", 3)) // expect: -1 +System.print("abcd".indexOf("cd", 10)) // expect: -1 + // More complex cases. System.print("abcdefabcdefg".indexOf("defg")) // expect: 9 System.print("abcdabcdabcd".indexOf("dab")) // expect: 3