feat: add find and substr methods to String

include/tvm/ffi/string.h

@@ -611,6 +611,58 @@ class String {
     }
   }
 
+  /*!
+   * \brief Find the first occurrence of a substring
+   * \param str The substring to search for
+   * \param pos The position at which to start the search
+   * \return The position of the first character of the first match, or size_t(-1) if not found
+   */
+  size_t find(const String& str, size_t pos = 0) const { return find(str.data(), pos, str.size()); }
+
+  /*!
+   * \brief Find the first occurrence of a substring
+   * \param str The substring to search for
+   * \param pos The position at which to start the search
+   * \return The position of the first character of the first match, or size_t(-1) if not found
+   */
+  size_t find(const char* str, size_t pos = 0) const { return find(str, pos, std::strlen(str)); }
+
+  /*!
+   * \brief Find the first occurrence of a substring
+   * \param str The substring to search for
+   * \param pos The position at which to start the search
+   * \param count The length of the substring
+   * \return The position of the first character of the first match, or size_t(-1) if not found
+   */
+  size_t find(const char* str, size_t pos, size_t count) const {
+    if (count == 0) return pos <= size() ? pos : size_t(-1);
+    if (pos >= size() || count > size() - pos) return size_t(-1);
+
+    const char* this_data = data();
+    size_t this_size = size();
+
+    for (size_t i = pos; i <= this_size - count; ++i) {
+      if (std::memcmp(this_data + i, str, count) == 0) {
+        return i;
+      }
+    }
+    return size_t(-1);
+  }
+
+  /*!
+   * \brief Returns a substring [pos, pos+count)
+   * \param pos The position of the first character to include
+   * \param count The length of the substring (default: until end of string)
+   * \return A string containing the substring
+   */
+  String substr(size_t pos = 0, size_t count = size_t(-1)) const {
+    if (pos > size()) {
+      throw std::out_of_range("tvm::String substr index out of bounds");
+    }
+    size_t rcount = std::min(count, size() - pos);
+    return String(data() + pos, rcount);
+  }
+
   /*!
    * \brief Convert String to an std::string object
    *

tests/cpp/test_string.cc

@@ -445,4 +445,40 @@ TEST(String, StdHash) {
   EXPECT_EQ(std::hash<Bytes>()(s3), std::hash<Bytes>()(s4));
 }
 
+TEST(String, Find) {
+  String s{"hello world"};
+  EXPECT_EQ(s.find("world"), 6);
+  EXPECT_EQ(s.find("hello"), 0);
+  EXPECT_EQ(s.find("o"), 4);
+  EXPECT_EQ(s.find("o", 5), 7);
+  EXPECT_EQ(s.find("notfound"), size_t(-1));
+  EXPECT_EQ(s.find(""), 0);
+  EXPECT_EQ(s.find("", 5), 5);
+  EXPECT_EQ(s.find("", 20), size_t(-1));
+
+  String pattern{"world"};
+  EXPECT_EQ(s.find(pattern), 6);
+
+  String empty{""};
+  EXPECT_EQ(empty.find("x"), size_t(-1));
+  EXPECT_EQ(empty.find(""), 0);
+}
+
+TEST(String, Substr) {
+  String s{"hello world"};
+  EXPECT_EQ(s.substr(0, 5), "hello");
+  EXPECT_EQ(s.substr(6, 5), "world");
+  EXPECT_EQ(s.substr(6), "world");
+  EXPECT_EQ(s.substr(0), "hello world");
+  EXPECT_EQ(s.substr(11), "");
+  EXPECT_EQ(s.substr(0, 0), "");
+
+  EXPECT_THROW(s.substr(12), std::out_of_range);
+  EXPECT_THROW(s.substr(100), std::out_of_range);
+
+  String empty{""};
+  EXPECT_EQ(empty.substr(0), "");
+  EXPECT_THROW(empty.substr(1), std::out_of_range);
+}
+
 }  // namespace

tests/python/test_string.py

@@ -53,3 +53,18 @@ def test_bytes() -> None:
     b5 = pickle.loads(pickle.dumps(b))
     assert b5 == b"hello"
     assert isinstance(b5, tvm_ffi.core.Bytes)
+
+
+def test_string_find_substr() -> None:
+    s = tvm_ffi.core.String("hello world")
+    assert s.find("world") == 6
+    assert s.find("hello") == 0
+    assert s.find("o") == 4
+    assert s.find("o", 5) == 7
+    assert s.find("notfound") == -1
+    assert s.find("") == 0
+
+    assert s[6:11] == "world"
+    assert s[0:5] == "hello"
+    assert s[6:] == "world"
+    assert s[:5] == "hello"