Skip to content

Commit

Permalink
Add 'ReadCord and 'WriteCord functions to CodedStream
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 493620002
  • Loading branch information
martijnvels authored and copybara-github committed Dec 7, 2022
1 parent e173fb6 commit e5e2ad8
Show file tree
Hide file tree
Showing 3 changed files with 297 additions and 0 deletions.
93 changes: 93 additions & 0 deletions src/google/protobuf/io/coded_stream.cc
Original file line number Diff line number Diff line change
Expand Up @@ -43,10 +43,13 @@
#include <limits.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <utility>

#include "google/protobuf/stubs/logging.h"
#include "absl/strings/cord.h"
#include "absl/strings/internal/resize_uninitialized.h"
#include "absl/strings/string_view.h"
#include "google/protobuf/arena.h"
Expand All @@ -67,6 +70,9 @@ namespace {
static const int kMaxVarintBytes = 10;
static const int kMaxVarint32Bytes = 5;

// When reading / writing Cords, if we have fewer than this many bytes we
// won't bother trying to avoid coping the contents.
static const int kMaxCordBytesToCopy = 512;

inline bool NextNonEmpty(ZeroCopyInputStream* input, const void** data,
int* size) {
Expand All @@ -77,6 +83,14 @@ inline bool NextNonEmpty(ZeroCopyInputStream* input, const void** data,
return success;
}

inline uint8_t* CopyCordToArray(const absl::Cord& cord, uint8_t* target) {
for (absl::string_view sv : cord.Chunks()) {
memcpy(target, sv.data(), sv.size());
target += sv.size();
}
return target;
}

} // namespace

// CodedInputStream ==================================================
Expand Down Expand Up @@ -309,6 +323,46 @@ bool CodedInputStream::ReadStringFallback(std::string* buffer, int size) {
return true;
}

bool CodedInputStream::ReadCord(absl::Cord* output, int size) {
GOOGLE_DCHECK_NE(output, nullptr);

// security: size is often user-supplied
if (size < 0) {
output->Clear();
return false; // security: size is often user-supplied
}

// Grab whatever is in the current output if `size` is relatively small,
// or if we are not sourcing data from an input stream.
if (input_ == nullptr || size < kMaxCordBytesToCopy) {
// Just copy the current buffer into the output rather than backing up.
absl::string_view buffer(reinterpret_cast<const char*>(buffer_),
static_cast<size_t>(std::min(size, BufferSize())));
*output = buffer;
Advance(static_cast<int>(buffer.size()));
size -= static_cast<int>(buffer.size());
if (size == 0) return true;
if (input_ == nullptr || buffer_size_after_limit_ + overflow_bytes_ > 0) {
// We hit a limit.
return false;
}
} else {
output->Clear();
BackUpInputToCurrentPosition();
}

// Make sure to not cross a limit set by PushLimit() or SetTotalBytesLimit().
const int closest_limit = std::min(current_limit_, total_bytes_limit_);
const int available = closest_limit - total_bytes_read_;
if (PROTOBUF_PREDICT_FALSE(size > available)) {
total_bytes_read_ = closest_limit;
input_->ReadCord(output, available);
return false;
}
total_bytes_read_ += size;
return input_->ReadCord(output, size);
}


bool CodedInputStream::ReadLittleEndian32Fallback(uint32_t* value) {
uint8_t bytes[sizeof(*value)];
Expand Down Expand Up @@ -910,6 +964,26 @@ uint8_t* EpsCopyOutputStream::WriteRawLittleEndian64(const void* data, int size,
}
#endif

uint8_t* EpsCopyOutputStream::WriteCord(const absl::Cord& cord, uint8_t* ptr) {
int s = GetSize(ptr);
if (stream_ == nullptr) {
if (static_cast<int64_t>(cord.size()) <= s) {
// Just copy it to the current buffer.
return CopyCordToArray(cord, ptr);
} else {
return Error();
}
} else if (static_cast<int64_t>(cord.size()) <= s &&
static_cast<int64_t>(cord.size()) < kMaxCordBytesToCopy) {
// Just copy it to the current buffer.
return CopyCordToArray(cord, ptr);
} else {
// Back up to the position where the Cord should start.
ptr = Trim(ptr);
if (!stream_->WriteCord(cord)) return Error();
return ptr;
}
}

uint8_t* EpsCopyOutputStream::WriteStringMaybeAliasedOutline(uint32_t num,
const std::string& s,
Expand All @@ -928,11 +1002,30 @@ uint8_t* EpsCopyOutputStream::WriteStringOutline(uint32_t num, const std::string
return WriteRaw(s.data(), size, ptr);
}

uint8_t* EpsCopyOutputStream::WriteStringOutline(uint32_t num, absl::string_view s,
uint8_t* ptr) {
ptr = EnsureSpace(ptr);
uint32_t size = s.size();
ptr = WriteLengthDelim(num, size, ptr);
return WriteRaw(s.data(), size, ptr);
}

uint8_t* EpsCopyOutputStream::WriteCordOutline(const absl::Cord& c, uint8_t* ptr) {
uint32_t size = c.size();
ptr = UnsafeWriteSize(size, ptr);
return WriteCord(c, ptr);
}

std::atomic<bool> CodedOutputStream::default_serialization_deterministic_{
false};

CodedOutputStream::~CodedOutputStream() { Trim(); }

uint8_t* CodedOutputStream::WriteCordToArray(const absl::Cord& cord,
uint8_t* target) {
return CopyCordToArray(cord, target);
}


uint8_t* CodedOutputStream::WriteStringWithSizeToArray(const std::string& str,
uint8_t* target) {
Expand Down
21 changes: 21 additions & 0 deletions src/google/protobuf/io/coded_stream.h
Original file line number Diff line number Diff line change
Expand Up @@ -129,7 +129,9 @@

#include "google/protobuf/stubs/common.h"
#include "google/protobuf/stubs/logging.h"
#include "absl/base/attributes.h"
#include "absl/numeric/bits.h"
#include "absl/strings/cord.h"
#include "absl/strings/string_view.h"
#include "google/protobuf/port.h"

Expand Down Expand Up @@ -215,6 +217,9 @@ class PROTOBUF_EXPORT CodedInputStream {
// Like ReadRaw, but reads into a string.
bool ReadString(std::string* buffer, int size);

// Like ReadString(), but reads to a Cord.
bool ReadCord(absl::Cord* output, int size);


// Read a 32-bit little-endian integer.
bool ReadLittleEndian32(uint32_t* value);
Expand Down Expand Up @@ -698,6 +703,7 @@ class PROTOBUF_EXPORT EpsCopyOutputStream {
}
}

uint8_t* WriteCord(const absl::Cord& cord, uint8_t* ptr);

#ifndef NDEBUG
PROTOBUF_NOINLINE
Expand Down Expand Up @@ -732,6 +738,13 @@ class PROTOBUF_EXPORT EpsCopyOutputStream {
std::memcpy(ptr, s.data(), size);
return ptr + size;
}

uint8_t* WriteString(uint32_t num, const absl::Cord& s, uint8_t* ptr) {
ptr = EnsureSpace(ptr);
ptr = WriteTag(num, 2, ptr);
return WriteCordOutline(s, ptr);
}

template <typename T>
#ifndef NDEBUG
PROTOBUF_NOINLINE
Expand Down Expand Up @@ -869,6 +882,8 @@ class PROTOBUF_EXPORT EpsCopyOutputStream {
uint8_t* WriteStringMaybeAliasedOutline(uint32_t num, const std::string& s,
uint8_t* ptr);
uint8_t* WriteStringOutline(uint32_t num, const std::string& s, uint8_t* ptr);
uint8_t* WriteStringOutline(uint32_t num, absl::string_view s, uint8_t* ptr);
uint8_t* WriteCordOutline(const absl::Cord& c, uint8_t* ptr);

template <typename T, typename E>
PROTOBUF_ALWAYS_INLINE uint8_t* WriteVarintPacked(int num, const T& r,
Expand Down Expand Up @@ -1124,6 +1139,12 @@ class PROTOBUF_EXPORT CodedOutputStream {
static uint8_t* WriteStringWithSizeToArray(const std::string& str,
uint8_t* target);

// Like WriteString() but writes a Cord.
void WriteCord(const absl::Cord& cord) { cur_ = impl_.WriteCord(cord, cur_); }

// Like WriteCord() but writing directly to the target array.
static uint8_t* WriteCordToArray(const absl::Cord& cord, uint8_t* target);


// Write a 32-bit little-endian integer.
void WriteLittleEndian32(uint32_t value) {
Expand Down
Loading

0 comments on commit e5e2ad8

Please sign in to comment.