diff --git a/include/tvm/ffi/string.h b/include/tvm/ffi/string.h index b18e690e..b93c3a99 100644 --- a/include/tvm/ffi/string.h +++ b/include/tvm/ffi/string.h @@ -611,6 +611,50 @@ class String { } } + /*! \brief Value returned by find() when no match is found */ + static constexpr size_t npos = static_cast(-1); + + /*! + * \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 npos 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 npos 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 npos if not found + */ + size_t find(const char* str, size_t pos, size_t count) const { + return std::string_view(data(), size()).find(std::string_view(str, count), pos); + } + + /*! + * \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 * diff --git a/tests/cpp/test_string.cc b/tests/cpp/test_string.cc index 072c5cd3..afcb6516 100644 --- a/tests/cpp/test_string.cc +++ b/tests/cpp/test_string.cc @@ -445,4 +445,41 @@ TEST(String, StdHash) { EXPECT_EQ(std::hash()(s3), std::hash()(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"), String::npos); + EXPECT_EQ(s.find(""), 0); + EXPECT_EQ(s.find("", 5), 5); + EXPECT_EQ(s.find("", 11), 11); + EXPECT_EQ(s.find("", 20), String::npos); + + String pattern{"world"}; + EXPECT_EQ(s.find(pattern), 6); + + String empty{""}; + EXPECT_EQ(empty.find("x"), String::npos); + 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 diff --git a/tests/python/test_string.py b/tests/python/test_string.py index 5dd06f4a..499270bb 100644 --- a/tests/python/test_string.py +++ b/tests/python/test_string.py @@ -53,3 +53,21 @@ 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.find("", 5) == 5 + assert s.find("", 11) == 11 + assert s.find("", 20) == -1 + + assert s[6:11] == "world" + assert s[0:5] == "hello" + assert s[6:] == "world" + assert s[:5] == "hello"