[stdlib] Move StringRef find() implementation to Span

stdlib/src/builtin/string_literal.mojo

@@ -394,7 +394,7 @@ 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.as_string_slice(), start)
 
     fn rfind(self, substr: StringLiteral, start: Int = 0) -> Int:
         """Finds the offset of the last occurrence of `substr` starting at
@@ -407,7 +407,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.as_string_slice(), start)
 
     fn replace(self, old: StringLiteral, new: StringLiteral) -> StringLiteral:
         """Return a copy of the string with all occurrences of substring `old`

stdlib/src/collections/string.mojo

@@ -1511,7 +1511,6 @@ struct String(
         Returns:
           The offset of `substr` relative to the beginning of the string.
         """
-
         return self.as_string_slice().find(substr.as_string_slice(), start)
 
     fn rfind(self, substr: String, start: Int = 0) -> Int:
@@ -1525,10 +1524,7 @@ struct String(
         Returns:
           The offset of `substr` relative to the beginning of the string.
         """
-
-        return self._strref_dangerous().rfind(
-            substr._strref_dangerous(), start=start
-        )
+        return self.as_string_slice().rfind(substr.as_string_slice(), start)
 
     fn isspace(self) -> Bool:
         """Determines whether every character in the given String is a

stdlib/src/utils/span.mojo

@@ -21,8 +21,11 @@ from utils import Span
 """
 
 from collections import InlineArray
-from memory import Pointer, UnsafePointer
+from memory import Pointer, UnsafePointer, bitcast, memcmp
 from builtin.builtin_list import _lit_mut_cast
+from sys import simdwidthof
+from bit import count_trailing_zeros
+from builtin.dtype import _uint_type_of_width
 
 
 @value
@@ -335,3 +338,194 @@ struct Span[
         return Span[T, _lit_mut_cast[lifetime, False].result](
             unsafe_ptr=self._data, len=self._len
         )
+
+    fn find[
+        D: DType, //, from_left: Bool = True
+    ](self: Span[Scalar[D]], subseq: Span[Scalar[D]], start: Int = 0) -> Int:
+        """Finds the offset of the first occurrence of `subseq` starting at
+        `start`. If not found, returns -1.
+
+        Parameters:
+            D: The `DType` of the Scalar.
+            from_left: Whether to search the first occurrence from the left.
+
+        Args:
+            subseq: The sub sequence to find.
+            start: The offset from which to find.
+
+        Returns:
+            The offset of `subseq` relative to the beginning of the `Span`.
+        """
+        var _len = len(self)
+
+        if not subseq:
+
+            @parameter
+            if from_left:
+                return 0
+            else:
+                return _len
+
+        if _len < len(subseq) + start:
+            return -1
+
+        var start_norm = max(_len + start, 0) if start < 0 else min(_len, start)
+        var haystack = __type_of(self)(
+            unsafe_ptr=self.unsafe_ptr() + start_norm, len=_len - start_norm
+        )
+        var loc: UnsafePointer[Scalar[D]]
+
+        @parameter
+        if from_left:
+            loc = _memmem(haystack, subseq)
+        else:
+            loc = _memrmem(haystack, subseq)
+
+        return int(loc) - int(self.unsafe_ptr()) if loc else -1
+
+    @always_inline
+    fn rfind[
+        D: DType, //
+    ](self: Span[Scalar[D]], subseq: Span[Scalar[D]], start: Int = 0) -> Int:
+        """Finds the offset of the last occurrence of `subseq` starting at
+        `start`. If not found, returns -1.
+
+        Parameters:
+            D: The `DType` of the Scalar.
+
+        Args:
+            subseq: The sub sequence to find.
+            start: The offset from which to find.
+
+        Returns:
+            The offset of `subseq` relative to the beginning of the `Span`.
+        """
+        return self.find[from_left=False](subseq, start)
+
+
+# ===----------------------------------------------------------------------===#
+# Utilities
+# ===----------------------------------------------------------------------===#
+
+
+@always_inline
+fn _align_down(value: Int, alignment: Int) -> Int:
+    return value._positive_div(alignment) * alignment
+
+
+@always_inline
+fn _memchr[
+    type: DType
+](
+    source: UnsafePointer[Scalar[type]], char: Scalar[type], len: Int
+) -> UnsafePointer[Scalar[type]]:
+    if not len:
+        return UnsafePointer[Scalar[type]]()
+    alias bool_mask_width = simdwidthof[DType.bool]()
+    var first_needle = SIMD[type, bool_mask_width](char)
+    var vectorized_end = _align_down(len, bool_mask_width)
+
+    for i in range(0, vectorized_end, bool_mask_width):
+        var bool_mask = source.load[width=bool_mask_width](i) == first_needle
+        var mask = bitcast[_uint_type_of_width[bool_mask_width]()](bool_mask)
+        if mask:
+            return source + int(i + count_trailing_zeros(mask))
+
+    for i in range(vectorized_end, len):
+        if source[i] == char:
+            return source + i
+    return UnsafePointer[Scalar[type]]()
+
+
+@always_inline
+fn _memmem[
+    type: DType
+](
+    haystack_span: Span[Scalar[type]], needle_span: Span[Scalar[type]]
+) -> UnsafePointer[Scalar[type]]:
+    var haystack = haystack_span.unsafe_ptr()
+    var haystack_len = len(haystack_span)
+    var needle = needle_span.unsafe_ptr()
+    var needle_len = len(needle_span)
+    if not needle_len:
+        return haystack
+    if needle_len > haystack_len:
+        return UnsafePointer[Scalar[type]]()
+    if needle_len == 1:
+        return _memchr[type](haystack, needle[0], haystack_len)
+
+    alias bool_mask_width = simdwidthof[DType.bool]()
+    var vectorized_end = _align_down(
+        haystack_len - needle_len + 1, bool_mask_width
+    )
+
+    var first_needle = SIMD[type, bool_mask_width](needle[0])
+    var last_needle = SIMD[type, bool_mask_width](needle[needle_len - 1])
+
+    for i in range(0, vectorized_end, bool_mask_width):
+        var first_block = haystack.load[width=bool_mask_width](i)
+        var last_block = haystack.load[width=bool_mask_width](
+            i + needle_len - 1
+        )
+
+        var eq_first = first_needle == first_block
+        var eq_last = last_needle == last_block
+
+        var bool_mask = eq_first & eq_last
+        var mask = bitcast[_uint_type_of_width[bool_mask_width]()](bool_mask)
+
+        while mask:
+            var offset = int(i + count_trailing_zeros(mask))
+            if memcmp(haystack + offset + 1, needle + 1, needle_len - 1) == 0:
+                return haystack + offset
+            mask = mask & (mask - 1)
+
+    # remaining partial block compare using byte-by-byte
+    #
+    for i in range(vectorized_end, haystack_len - needle_len + 1):
+        if haystack[i] != needle[0]:
+            continue
+
+        if memcmp(haystack + i + 1, needle + 1, needle_len - 1) == 0:
+            return haystack + i
+    _ = haystack_span, needle_span
+    return UnsafePointer[Scalar[type]]()
+
+
+@always_inline
+fn _memrchr[
+    type: DType
+](
+    source: UnsafePointer[Scalar[type]], char: Scalar[type], len: Int
+) -> UnsafePointer[Scalar[type]]:
+    if not len:
+        return UnsafePointer[Scalar[type]]()
+    for i in reversed(range(len)):
+        if source[i] == char:
+            return source + i
+    return UnsafePointer[Scalar[type]]()
+
+
+@always_inline
+fn _memrmem[
+    type: DType
+](
+    haystack_span: Span[Scalar[type]], needle_span: Span[Scalar[type]]
+) -> UnsafePointer[Scalar[type]]:
+    var haystack = haystack_span.unsafe_ptr()
+    var haystack_len = len(haystack_span)
+    var needle = needle_span.unsafe_ptr()
+    var needle_len = len(needle_span)
+    if not needle_len:
+        return haystack
+    if needle_len > haystack_len:
+        return UnsafePointer[Scalar[type]]()
+    if needle_len == 1:
+        return _memrchr[type](haystack, needle[0], haystack_len)
+    for i in reversed(range(haystack_len - needle_len + 1)):
+        if haystack[i] != needle[0]:
+            continue
+        if memcmp(haystack + i + 1, needle + 1, needle_len - 1) == 0:
+            return haystack + i
+    _ = haystack_span, needle_span
+    return UnsafePointer[Scalar[type]]()

stdlib/src/utils/string_slice.mojo

@@ -561,27 +561,22 @@ struct StringSlice[
         Returns:
           The offset of `substr` relative to the beginning of the string.
         """
-        if not substr:
-            return 0
+        # FIXME(#3526): this should return unicode codepoint offsets
+        return self.as_bytes_span().find(substr.as_bytes_span(), start)
 
-        if self.byte_length() < substr.byte_length() + start:
-            return -1
-
-        # The substring to search within, offset from the beginning if `start`
-        # is positive, and offset from the end if `start` is negative.
-        var haystack_str = self._from_start(start)
-
-        var loc = stringref._memmem(
-            haystack_str.unsafe_ptr(),
-            haystack_str.byte_length(),
-            substr.unsafe_ptr(),
-            substr.byte_length(),
-        )
+    fn rfind(self, substr: StringSlice, start: Int = 0) -> Int:
+        """Finds the offset of the last occurrence of `substr` starting at
+        `start`. If not found, returns -1.
 
-        if not loc:
-            return -1
+        Args:
+          substr: The substring to find.
+          start: The offset from which to find.
 
-        return int(loc) - int(self.unsafe_ptr())
+        Returns:
+          The offset of `substr` relative to the beginning of the string.
+        """
+        # FIXME(#3526): this should return unicode codepoint offsets
+        return self.as_bytes_span().rfind(substr.as_bytes_span(), start)
 
     fn isspace(self) -> Bool:
         """Determines whether every character in the given StringSlice is a