Skip to content

Commit

Permalink
Reitroduce Stringlike trait and use it for Stringlike.find()
Browse files Browse the repository at this point in the history
Signed-off-by: martinvuyk <[email protected]>
  • Loading branch information
martinvuyk committed Dec 11, 2024
1 parent 5655490 commit 2fde772
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 16 deletions.
14 changes: 10 additions & 4 deletions stdlib/src/builtin/string_literal.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -591,31 +591,37 @@ 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.
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.
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`
Expand Down
18 changes: 10 additions & 8 deletions stdlib/src/collections/string.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -1707,35 +1707,37 @@ 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.
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.
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
Expand Down
53 changes: 49 additions & 4 deletions stdlib/src/utils/string_slice.mojo
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand All @@ -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`
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit 2fde772

Please sign in to comment.