diff --git a/stdlib/src/builtin/string_literal.mojo b/stdlib/src/builtin/string_literal.mojo index 377173470f8..30d671ba9c6 100644 --- a/stdlib/src/builtin/string_literal.mojo +++ b/stdlib/src/builtin/string_literal.mojo @@ -591,10 +591,13 @@ struct StringLiteral( writer.write(self.as_string_slice()) - fn find(self, substr: StringLiteral, start: Int = 0) -> Int: + fn find[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: """Finds the offset of the first occurrence of `substr` starting at `start`. If not found, returns -1. + Parameters: + T: The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -602,12 +605,15 @@ struct StringLiteral( Returns: The offset of `substr` relative to the beginning of the string. """ - return StringRef(self).find(substr, start=start) + return self.as_string_slice().find(substr, start=start) - fn rfind(self, substr: StringLiteral, start: Int = 0) -> Int: + fn rfind[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: """Finds the offset of the last occurrence of `substr` starting at `start`. If not found, returns -1. + Parameters: + T: The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -615,7 +621,7 @@ struct StringLiteral( Returns: The offset of `substr` relative to the beginning of the string. """ - return StringRef(self).rfind(substr, start=start) + return self.as_string_slice().rfind(substr, start=start) fn replace(self, old: StringLiteral, new: StringLiteral) -> StringLiteral: """Return a copy of the string with all occurrences of substring `old` diff --git a/stdlib/src/collections/string.mojo b/stdlib/src/collections/string.mojo index f1b34ffff92..36948d7949e 100644 --- a/stdlib/src/collections/string.mojo +++ b/stdlib/src/collections/string.mojo @@ -1707,10 +1707,13 @@ struct String( """ return substr.as_string_slice() in self.as_string_slice() - fn find(self, substr: String, start: Int = 0) -> Int: + fn find[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: """Finds the offset of the first occurrence of `substr` starting at `start`. If not found, returns -1. + Parameters: + T: The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -1718,13 +1721,15 @@ struct String( Returns: The offset of `substr` relative to the beginning of the string. """ + return self.as_string_slice().find(substr, start) - return self.as_string_slice().find(substr.as_string_slice(), start) - - fn rfind(self, substr: String, start: Int = 0) -> Int: + fn rfind[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: """Finds the offset of the last occurrence of `substr` starting at `start`. If not found, returns -1. + Parameters: + The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -1732,10 +1737,7 @@ struct String( Returns: The offset of `substr` relative to the beginning of the string. """ - - return self.as_string_slice().rfind( - substr.as_string_slice(), start=start - ) + return self.as_string_slice().rfind(substr, start=start) fn isspace(self) -> Bool: """Determines whether every character in the given String is a diff --git a/stdlib/src/utils/string_slice.mojo b/stdlib/src/utils/string_slice.mojo index 03ed9be4909..d91d542bfb9 100644 --- a/stdlib/src/utils/string_slice.mojo +++ b/stdlib/src/utils/string_slice.mojo @@ -885,10 +885,14 @@ struct StringSlice[is_mutable: Bool, //, origin: Origin[is_mutable]]( """ return _FormatCurlyEntry.format(self, args) - fn find(ref self, substr: StringSlice, start: Int = 0) -> Int: + # FIXME(#3526): this should return unicode codepoint offsets + fn find[T: Stringlike, //](ref self, substr: T, start: Int = 0) -> Int: """Finds the offset of the first occurrence of `substr` starting at `start`. If not found, returns `-1`. + Parameters: + T: The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -918,10 +922,14 @@ struct StringSlice[is_mutable: Bool, //, origin: Origin[is_mutable]]( return int(loc) - int(self.unsafe_ptr()) - fn rfind(self, substr: StringSlice, start: Int = 0) -> Int: + # FIXME(#3526): this should return unicode codepoint offsets + fn rfind[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: """Finds the offset of the last occurrence of `substr` starting at `start`. If not found, returns `-1`. + Parameters: + T: The Stringlike type. + Args: substr: The substring to find. start: The offset from which to find. @@ -930,9 +938,9 @@ struct StringSlice[is_mutable: Bool, //, origin: Origin[is_mutable]]( The offset of `substr` relative to the beginning of the string. """ if not substr: - return len(self) + return self.byte_length() - if len(self) < len(substr) + start: + if self.byte_length() < substr.byte_length() + start: return -1 # The substring to search within, offset from the beginning if `start` @@ -1091,6 +1099,43 @@ struct StringSlice[is_mutable: Bool, //, origin: Origin[is_mutable]]( # ===-----------------------------------------------------------------------===# +trait Stringlike(CollectionElement, CollectionElementNew): + """Trait intended to be used as a generic entrypoint for all String-like + types.""" + + fn byte_length(self) -> Int: + """Get the string length in bytes. + Returns: + The length of this string in bytes. + Notes: + This does not include the trailing null terminator in the count. + """ + ... + + fn unsafe_ptr(self) -> UnsafePointer[Byte]: + """Get raw pointer to the underlying data. + Returns: + The raw pointer to the data. + """ + ... + + fn find[T: Stringlike, //](self, substr: T, start: Int = 0) -> Int: + """Finds the offset of the first occurrence of `substr` starting at + `start`. If not found, returns -1. + + Parameters: + T: The type of the substring. + + Args: + substr: The substring to find. + start: The offset from which to find. + + Returns: + The offset of `substr` relative to the beginning of the string. + """ + ... + + fn _to_string_list[ T: CollectionElement, # TODO(MOCO-1446): Make `T` parameter inferred len_fn: fn (T) -> Int,