Skip to content

Commit da8660b

Browse files
maskri17copybara-github
maskri17
authored and
copybara-github
committed
adding find and substring for byte_string
PiperOrigin-RevId: 805592111
1 parent 2abc4a6 commit da8660b

File tree

3 files changed

+233
-1
lines changed

3 files changed

+233
-1
lines changed

‎common/internal/byte_string.cc‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,106 @@ bool ByteString::EndsWith(const absl::Cord& rhs) const {
286286
[&rhs](const absl::Cord& lhs) -> bool { return lhs.EndsWith(rhs); }));
287287
}
288288

289+
absl::optional<size_t> ByteString::Find(absl::string_view needle,
290+
size_t pos) const {
291+
ABSL_DCHECK_LE(pos, size());
292+
293+
return Visit(absl::Overload(
294+
[&needle, pos](absl::string_view lhs) -> absl::optional<size_t> {
295+
absl::string_view::size_type i = lhs.find(needle, pos);
296+
if (i == absl::string_view::npos) {
297+
return absl::nullopt;
298+
}
299+
return i;
300+
},
301+
[&needle, pos](const absl::Cord& lhs) -> absl::optional<size_t> {
302+
absl::Cord cord = lhs.Subcord(pos, lhs.size() - pos);
303+
absl::Cord::CharIterator it = cord.Find(needle);
304+
if (it == cord.char_end()) {
305+
return absl::nullopt;
306+
}
307+
return pos +
308+
static_cast<size_t>(absl::Cord::Distance(cord.char_begin(), it));
309+
}));
310+
}
311+
312+
absl::optional<size_t> ByteString::Find(const absl::Cord& needle,
313+
size_t pos) const {
314+
ABSL_DCHECK_LE(pos, size());
315+
316+
return Visit(absl::Overload(
317+
[&needle, pos](absl::string_view lhs) -> absl::optional<size_t> {
318+
if (auto flat_needle = needle.TryFlat(); flat_needle) {
319+
absl::string_view::size_type i = lhs.find(*flat_needle, pos);
320+
if (i == absl::string_view::npos) {
321+
return absl::nullopt;
322+
}
323+
return i;
324+
}
325+
// Needle is fragmented, we have to do a linear scan.
326+
const size_t needle_size = needle.size();
327+
if (pos + needle_size > lhs.size()) {
328+
return absl::nullopt;
329+
}
330+
if (ABSL_PREDICT_FALSE(needle_size == 0)) {
331+
return pos;
332+
}
333+
// Optimization: find the first chunk of the needle, then compare the
334+
// rest. If the first chunk is empty, `lhs.find` will return
335+
// `current_pos`, which correctly degrades to a linear scan.
336+
absl::string_view first_chunk = *needle.Chunks().begin();
337+
absl::Cord rest_of_needle = needle.Subcord(
338+
first_chunk.size(), needle_size - first_chunk.size());
339+
size_t current_pos = pos;
340+
while (true) {
341+
size_t found_pos = lhs.find(first_chunk, current_pos);
342+
if (found_pos == absl::string_view::npos ||
343+
found_pos > lhs.size() - needle_size) {
344+
return absl::nullopt;
345+
}
346+
if (lhs.substr(found_pos + first_chunk.size(),
347+
rest_of_needle.size()) == rest_of_needle) {
348+
return found_pos;
349+
}
350+
current_pos = found_pos + 1;
351+
}
352+
},
353+
[&needle, pos](const absl::Cord& lhs) -> absl::optional<size_t> {
354+
absl::Cord cord = lhs.Subcord(pos, lhs.size() - pos);
355+
absl::Cord::CharIterator it = cord.Find(needle);
356+
if (it == cord.char_end()) {
357+
return absl::nullopt;
358+
}
359+
return pos +
360+
static_cast<size_t>(absl::Cord::Distance(cord.char_begin(), it));
361+
}));
362+
}
363+
364+
ByteString ByteString::Substring(size_t pos, size_t npos) const {
365+
ABSL_DCHECK_LE(npos, size());
366+
ABSL_DCHECK_LE(pos, npos);
367+
368+
switch (GetKind()) {
369+
case ByteStringKind::kSmall: {
370+
ByteString result;
371+
result.rep_.header.kind = ByteStringKind::kSmall;
372+
result.rep_.small.size = npos - pos;
373+
std::memcpy(result.rep_.small.data, rep_.small.data + pos,
374+
result.rep_.small.size);
375+
result.rep_.small.arena = GetSmallArena();
376+
return result;
377+
}
378+
case ByteStringKind::kMedium: {
379+
ByteString result(*this);
380+
result.rep_.medium.data += pos;
381+
result.rep_.medium.size = npos - pos;
382+
return result;
383+
}
384+
case ByteStringKind::kLarge:
385+
return ByteString(GetLarge().Subcord(pos, npos - pos));
386+
}
387+
}
388+
289389
void ByteString::RemovePrefix(size_t n) {
290390
ABSL_DCHECK_LE(n, size());
291391
if (n == 0) {

‎common/internal/byte_string.h‎

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,7 @@ union CEL_COMMON_INTERNAL_BYTE_STRING_TRIVIAL_ABI ByteStringRep final {
159159
absl::string_view LegacyByteString(const ByteString& string, bool stable,
160160
google::protobuf::Arena* absl_nonnull arena);
161161

162-
// `ByteString` is an vocabulary type capable of representing copy-on-write
162+
// `ByteString` is a vocabulary type capable of representing copy-on-write
163163
// strings efficiently for arenas and reference counting. The contents of the
164164
// byte string are owned by an arena or managed by a reference count. All byte
165165
// strings have an associated allocator specified at construction, once the byte
@@ -275,6 +275,24 @@ ByteString final {
275275
bool EndsWith(const absl::Cord& rhs) const;
276276
bool EndsWith(const ByteString& rhs) const;
277277

278+
// Finds the first occurrence of `needle` in this object, starting at byte
279+
// position `pos`. Returns `absl::nullopt` if `needle` is not found.
280+
// Note: Positions are byte-based, not code point based as in
281+
// `cel::StringValue`.
282+
absl::optional<size_t> Find(absl::string_view needle, size_t pos = 0) const;
283+
absl::optional<size_t> Find(const absl::Cord& needle, size_t pos = 0) const;
284+
absl::optional<size_t> Find(const ByteString& needle, size_t pos = 0) const;
285+
286+
// Returns a new `ByteString` that is a substring of this object, starting at
287+
// byte position `pos` and with a length of `npos` bytes.
288+
// Note: Positions are byte-based, not code point based as in
289+
// `cel::StringValue`.
290+
ByteString Substring(size_t pos, size_t npos) const;
291+
ByteString Substring(size_t pos) const {
292+
ABSL_DCHECK_LE(pos, size());
293+
return Substring(pos, size());
294+
}
295+
278296
void RemovePrefix(size_t n);
279297

280298
void RemoveSuffix(size_t n);
@@ -501,6 +519,17 @@ inline bool ByteString::EndsWith(const ByteString& rhs) const {
501519
[this](const absl::Cord& rhs) -> bool { return EndsWith(rhs); }));
502520
}
503521

522+
inline absl::optional<size_t> ByteString::Find(const ByteString& needle,
523+
size_t pos) const {
524+
return needle.Visit(absl::Overload(
525+
[this, pos](absl::string_view rhs) -> absl::optional<size_t> {
526+
return Find(rhs, pos);
527+
},
528+
[this, pos](const absl::Cord& rhs) -> absl::optional<size_t> {
529+
return Find(rhs, pos);
530+
}));
531+
}
532+
504533
inline bool operator==(const ByteString& lhs, const ByteString& rhs) {
505534
return lhs.Equals(rhs);
506535
}

‎common/internal/byte_string_test.cc‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -747,6 +747,109 @@ TEST_P(ByteStringTest, EndsWith) {
747747
GetMediumOrLargeCord().size() - kSmallByteStringCapacity)));
748748
}
749749

750+
TEST_P(ByteStringTest, Find) {
751+
ByteString byte_string = ByteString(GetAllocator(), GetMediumStringView());
752+
// Find string_view
753+
EXPECT_THAT(byte_string.Find("A string"), Optional(0));
754+
EXPECT_THAT(
755+
byte_string.Find("small string optimization!"),
756+
Optional(GetMediumStringView().find("small string optimization!")));
757+
EXPECT_THAT(byte_string.Find("not found"), Eq(absl::nullopt));
758+
EXPECT_THAT(byte_string.Find(""), Optional(0));
759+
EXPECT_THAT(byte_string.Find("", 3), Optional(3));
760+
EXPECT_THAT(byte_string.Find("A string", 1), Eq(absl::nullopt));
761+
// Find cord
762+
EXPECT_THAT(byte_string.Find(absl::Cord("A string")), Optional(0));
763+
EXPECT_THAT(
764+
byte_string.Find(absl::Cord("small string optimization!")),
765+
Optional(GetMediumStringView().find("small string optimization!")));
766+
EXPECT_THAT(
767+
byte_string.Find(absl::MakeFragmentedCord(
768+
{"A string", " that is too large for the small string optimization!",
769+
" extra"})),
770+
Eq(absl::nullopt));
771+
EXPECT_THAT(byte_string.Find(GetMediumOrLargeFragmentedCord()), Optional(0));
772+
EXPECT_THAT(byte_string.Find(absl::Cord("not found")), Eq(absl::nullopt));
773+
}
774+
775+
TEST_P(ByteStringTest, Substring) {
776+
// small byte_string substring
777+
ByteString small_byte_string =
778+
ByteString(GetAllocator(), GetSmallStringView());
779+
EXPECT_EQ(small_byte_string.Substring(1, 5),
780+
GetSmallStringView().substr(1, 4));
781+
EXPECT_EQ(small_byte_string.Substring(0, small_byte_string.size()),
782+
GetSmallStringView());
783+
EXPECT_EQ(small_byte_string.Substring(1, 1), "");
784+
// medium byte_string substring
785+
ByteString medium_byte_string =
786+
ByteString(GetAllocator(), GetMediumStringView());
787+
EXPECT_EQ(medium_byte_string.Substring(2, 12),
788+
GetMediumStringView().substr(2, 10));
789+
EXPECT_EQ(medium_byte_string.Substring(0, medium_byte_string.size()),
790+
GetMediumStringView());
791+
// large byte_string substring
792+
ByteString large_byte_string =
793+
ByteString(GetAllocator(), GetMediumOrLargeCord());
794+
EXPECT_EQ(large_byte_string.Substring(3, 15),
795+
GetMediumOrLargeCord().Subcord(3, 12));
796+
EXPECT_EQ(large_byte_string.Substring(0, large_byte_string.size()),
797+
GetMediumOrLargeCord());
798+
// substring with one parameter
799+
ByteString tacocat_byte_string = ByteString(GetAllocator(), "tacocat");
800+
EXPECT_EQ(tacocat_byte_string.Substring(4), "cat");
801+
}
802+
803+
TEST_P(ByteStringTest, FindEdgeCases) {
804+
ByteString empty_byte_string(GetAllocator(), "");
805+
EXPECT_THAT(empty_byte_string.Find("a"), Eq(absl::nullopt));
806+
EXPECT_THAT(empty_byte_string.Find(""), Optional(0));
807+
ByteString byte_string = ByteString(GetAllocator(), GetMediumStringView());
808+
// Needle longer than haystack.
809+
EXPECT_THAT(byte_string.Find(std::string(byte_string.size() + 1, 'a')),
810+
Eq(absl::nullopt));
811+
// Needle at the end.
812+
absl::string_view suffix = "optimization!";
813+
EXPECT_THAT(byte_string.Find(suffix),
814+
Optional(byte_string.size() - suffix.size()));
815+
// pos at the end.
816+
EXPECT_THAT(byte_string.Find("a", byte_string.size()), Eq(absl::nullopt));
817+
EXPECT_THAT(byte_string.Find("", byte_string.size()),
818+
Optional(byte_string.size()));
819+
// Search in a cord-backed ByteString.
820+
ByteString cord_byte_string(GetAllocator(), GetMediumOrLargeFragmentedCord());
821+
EXPECT_THAT(cord_byte_string.Find("string that is too large"),
822+
Optional(GetMediumStringView().find("string that is too large")));
823+
EXPECT_THAT(cord_byte_string.Find(
824+
absl::MakeFragmentedCord({"string that ", "is too large"})),
825+
Optional(GetMediumStringView().find("string that is too large")));
826+
// Fragmented needle with empty first chunk.
827+
absl::Cord fragmented_with_empty_chunk;
828+
fragmented_with_empty_chunk.Append("");
829+
fragmented_with_empty_chunk.Append("A string");
830+
EXPECT_THAT(byte_string.Find(fragmented_with_empty_chunk), Optional(0));
831+
}
832+
833+
TEST_P(ByteStringTest, SubstringEdgeCases) {
834+
ByteString byte_string = ByteString(GetAllocator(), GetSmallStringView());
835+
EXPECT_EQ(byte_string.Substring(byte_string.size(), byte_string.size()), "");
836+
EXPECT_EQ(byte_string.Substring(0, 0), "");
837+
}
838+
839+
#ifndef NDEBUG
840+
TEST_P(ByteStringTest, FindOutOfBounds) {
841+
ByteString byte_string = ByteString(GetAllocator(), "test");
842+
EXPECT_DEATH(byte_string.Find("t", 5), _);
843+
}
844+
845+
TEST_P(ByteStringTest, SubstringOutOfBounds) {
846+
ByteString byte_string = ByteString(GetAllocator(), "test");
847+
EXPECT_DEATH(static_cast<void>(byte_string.Substring(5, 5)), _);
848+
EXPECT_DEATH(static_cast<void>(byte_string.Substring(0, 5)), _);
849+
EXPECT_DEATH(static_cast<void>(byte_string.Substring(3, 2)), _);
850+
}
851+
#endif
852+
750853
TEST_P(ByteStringTest, RemovePrefixSmall) {
751854
ByteString byte_string = ByteString(GetAllocator(), GetSmallStringView());
752855
byte_string.RemovePrefix(1);

0 commit comments

Comments
 (0)