SERVER-63192 Add slice, split, and sliceAndAdvance APIs to CDRC

5 files changed, 237 insertions, 1 deletions

diff --git a/src/mongo/base/data_range.h b/src/mongo/base/data_range.h
index d8b50807fbd..a366513ce3d 100644
--- a/src/mongo/base/data_range.h
+++ b/src/mongo/base/data_range.h
@@ -32,6 +32,7 @@
 #include <cstring>
 #include <tuple>
 #include <type_traits>
+#include <utility>
 
 #include "mongo/base/data_type.h"
 #include "mongo/base/error_codes.h"
@@ -162,6 +163,24 @@ public:
         return uassertStatusOK(readNoThrow<T>(offset));
     }
 
+    /**
+     * Split this ConstDataRange into two parts at `splitPoint`.
+     * May provide either a pointer within the range or an offset from the beginning.
+     */
+    template <typename T>
+    auto split(const T& splitPoint) const {
+        return doSplit<ConstDataRange>(splitPoint);
+    }
+
+    /**
+     * Create a smaller chunk of the original ConstDataRange.
+     * May provide either a pointer within the range or an offset from the beginning.
+     */
+    template <typename T>
+    auto slice(const T& splitPoint) const {
+        return doSlice<ConstDataRange>(splitPoint);
+    }
+
     friend bool operator==(const ConstDataRange& lhs, const ConstDataRange& rhs) {
         return std::tie(lhs._begin, lhs._end) == std::tie(rhs._begin, rhs._end);
     }
@@ -170,6 +189,34 @@ public:
         return !(lhs == rhs);
     }
 
+protected:
+    // Shared implementation of split() logic between DataRange and ConstDataRange.
+    template <typename RangeT,
+              typename ByteLike,
+              typename std::enable_if_t<isByteV<ByteLike>, int> = 0>
+    std::pair<RangeT, RangeT> doSplit(const ByteLike* splitPoint) const {
+        const auto* typedPoint = reinterpret_cast<const byte_type*>(splitPoint);
+        uassert(ErrorCodes::BadValue,
+                "Invalid split point",
+                (typedPoint >= _begin) && (typedPoint <= _end));
+        // RangeT will enforce constness, so use common-denominator for args to ctor.
+        auto* begin = const_cast<byte_type*>(_begin);
+        auto* split = const_cast<byte_type*>(typedPoint);
+        auto* end = const_cast<byte_type*>(_end);
+        return {{begin, split}, {split, end}};
+    }
+
+    template <typename RangeT>
+    auto doSplit(std::size_t splitPoint) const {
+        return doSplit<RangeT>(data() + splitPoint);
+    }
+
+    // Convenience wrapper to just grab the first half of a split.
+    template <typename RangeT, typename T>
+    RangeT doSlice(const T& splitPoint) const {
+        auto parts = doSplit<RangeT>(splitPoint);
+        return parts.first;
+    }
 
 protected:
     const byte_type* _begin;
@@ -246,6 +293,24 @@ public:
     void write(const T& value, std::size_t offset = 0) {
         uassertStatusOK(writeNoThrow(value, offset));
     }
+
+    using ConstDataRange::data;
+    template <typename ByteLike = byte_type>
+    ByteLike* data() noexcept {
+        return reinterpret_cast<ByteLike*>(const_cast<byte_type*>(_begin));
+    }
+
+    using ConstDataRange::split;
+    template <typename T>
+    auto split(const T& splitPoint) {
+        return doSplit<DataRange>(splitPoint);
+    }
+
+    using ConstDataRange::slice;
+    template <typename T>
+    auto slice(const T& splitPoint) {
+        return doSlice<DataRange>(splitPoint);
+    }
 };
 
 struct DataRangeTypeHelper {
diff --git a/src/mongo/base/data_range_cursor.h b/src/mongo/base/data_range_cursor.h
index 89dff93eb43..b11797c1a73 100644
--- a/src/mongo/base/data_range_cursor.h
+++ b/src/mongo/base/data_range_cursor.h
@@ -117,6 +117,17 @@ public:
         return uassertStatusOK(readAndAdvanceNoThrow<T>());
     }
 
+    /**
+     * Return a ConstDataRange based on `splitPoint`
+     * and advance the cursor past there.
+     */
+    template <typename ByteLike>
+    ConstDataRange sliceAndAdvance(const ByteLike& splitPoint) {
+        auto ret = slice(splitPoint);
+        advance(ret.length());
+        return ret;
+    }
+
 private:
     Status makeAdvanceStatus(size_t advance) const;
 };
@@ -220,6 +231,17 @@ public:
         uassertStatusOK(writeAndAdvanceNoThrow(value));
     }
 
+    /**
+     * Return a DataRange based on `splitPoint`
+     * and advance the cursor past there.
+     */
+    template <typename ByteLike>
+    DataRange sliceAndAdvance(const ByteLike& splitPoint) {
+        auto ret = slice(splitPoint);
+        advance(ret.length());
+        return ret;
+    }
+
 private:
     Status makeAdvanceStatus(size_t advance) const;
 };
diff --git a/src/mongo/base/data_range_cursor_test.cpp b/src/mongo/base/data_range_cursor_test.cpp
index 60543c6af6a..a91bf708d28 100644
--- a/src/mongo/base/data_range_cursor_test.cpp
+++ b/src/mongo/base/data_range_cursor_test.cpp
@@ -34,6 +34,26 @@
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
+namespace {
+// ConstDataRange::operator==() requires that the pointers
+// refer to the same memory addresses.
+// So just promote to a string that we can do direct comparisons on.
+std::string toString(ConstDataRange cdr) {
+    return std::string(cdr.data(), cdr.length());
+}
+
+// The ASSERT macro can't handle template specialization,
+// so work out the value external to the macro call.
+template <typename T>
+bool isConstDataRange(const T& val) {
+    return std::is_same_v<T, ConstDataRange>;
+}
+
+template <typename T>
+bool isDataRange(const T& val) {
+    return std::is_same_v<T, DataRange>;
+}
+}  // namespace
 
 TEST(DataRangeCursor, ConstDataRangeCursor) {
     char buf[14];
@@ -102,4 +122,41 @@ TEST(DataRangeCursor, DataRangeCursorType) {
 
     ASSERT_EQUALS(std::string("fooZ"), buf2);
 }
+
+TEST(DataRangeCursor, sliceAndAdvance) {
+    std::string buffer = "Hello World";
+    ConstDataRangeCursor bufferCDRC(buffer.c_str(), buffer.size());
+
+    // Split by position in range [0..length)
+    auto helloCDR = bufferCDRC.sliceAndAdvance(5);
+    ASSERT_EQ(toString(helloCDR), "Hello");
+    ASSERT_EQ(toString(bufferCDRC), " World");
+
+    // Split by pointer
+    auto spaceCDR = bufferCDRC.sliceAndAdvance(bufferCDRC.data() + 1);
+    ASSERT_EQ(toString(spaceCDR), " ");
+    ASSERT_EQ(toString(bufferCDRC), "World");
+
+    // Get DataRange from a DataRangeCursor if original was non-const.
+    DataRangeCursor mutableDRC(const_cast<char*>(buffer.c_str()), buffer.size());
+    auto mutableByLen = mutableDRC.sliceAndAdvance(1);
+    ASSERT_TRUE(isDataRange(mutableByLen));
+
+    // Get ConstDataRange from a DataRangeCursor if original was const.
+    const DataRange nonmutableDRC = mutableDRC;
+    auto nonmutableCDR = nonmutableDRC.slice(2);
+    ASSERT_TRUE(isConstDataRange(nonmutableCDR));
+}
+
+TEST(DataRange, sliceAndAdvanceThrow) {
+    std::string buffer("Hello World");
+    ConstDataRangeCursor bufferCDRC(buffer.c_str(), buffer.size());
+
+    // Split point is out of range.
+    ASSERT_THROWS(bufferCDRC.sliceAndAdvance(bufferCDRC.length() + 1), AssertionException);
+    ASSERT_THROWS(bufferCDRC.sliceAndAdvance(bufferCDRC.data() + bufferCDRC.length() + 1),
+                  AssertionException);
+    ASSERT_THROWS(bufferCDRC.sliceAndAdvance(bufferCDRC.data() - 1), AssertionException);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/base/data_range_test.cpp b/src/mongo/base/data_range_test.cpp
index f45c6e26bbf..9d121406ed8 100644
--- a/src/mongo/base/data_range_test.cpp
+++ b/src/mongo/base/data_range_test.cpp
@@ -36,6 +36,26 @@
 #include "mongo/unittest/unittest.h"
 
 namespace mongo {
+namespace {
+// ConstDataRange::operator==() requires that the pointers
+// refer to the same memory addresses.
+// So just promote to a string that we can do direct comparisons on.
+std::string toString(ConstDataRange cdr) {
+    return std::string(cdr.data(), cdr.length());
+}
+
+// The ASSERT macro can't handle template specialization,
+// so work out the value external to the macro call.
+template <typename T>
+bool isConstDataRange(const T& val) {
+    return std::is_same_v<T, ConstDataRange>;
+}
+
+template <typename T>
+bool isDataRange(const T& val) {
+    return std::is_same_v<T, DataRange>;
+}
+}  // namespace
 
 TEST(DataRange, ConstDataRange) {
     unsigned char buf[sizeof(uint32_t) * 3];
@@ -128,4 +148,76 @@ TEST(DataRange, InitFromContainer) {
     ASSERT_EQUALS(status, ErrorCodes::Overflow);
 }
 
+TEST(DataRange, slice) {
+    std::string buffer("Hello World");
+    ConstDataRange bufferCDR(buffer.c_str(), buffer.size());
+
+    // Split by position in range [0..length)
+    auto helloByLen = bufferCDR.slice(5);
+    ASSERT_EQ(helloByLen.length(), 5);
+    ASSERT_EQ(toString(helloByLen), "Hello");
+
+    // Split by pointer within range
+    auto helloByPtr = bufferCDR.slice(bufferCDR.data() + 4);
+    ASSERT_EQ(helloByPtr.length(), 4);
+    ASSERT_EQ(toString(helloByPtr), "Hell");
+
+    // Get DataRange from a DataRange if original was non-const.
+    DataRange mutableDR(const_cast<char*>(buffer.c_str()), buffer.size());
+    auto mutableByLen = mutableDR.slice(1);
+    ASSERT_TRUE(isDataRange(mutableByLen));
+
+    // Get ConstDataRange from a DataRange if original was const.
+    const DataRange nonmutableDR = mutableDR;
+    auto nonmutableCDR = nonmutableDR.slice(2);
+    ASSERT_TRUE(isConstDataRange(nonmutableCDR));
+}
+
+TEST(DataRange, sliceThrow) {
+    std::string buffer("Hello World");
+    ConstDataRange bufferCDR(buffer.c_str(), buffer.size());
+
+    // Split point is out of range.
+    ASSERT_THROWS(bufferCDR.slice(bufferCDR.length() + 1), AssertionException);
+    ASSERT_THROWS(bufferCDR.slice(bufferCDR.data() + bufferCDR.length() + 1), AssertionException);
+    ASSERT_THROWS(bufferCDR.slice(bufferCDR.data() - 1), AssertionException);
+}
+
+TEST(DataRange, split) {
+    std::string buffer("Hello World");
+    ConstDataRange bufferCDR(buffer.c_str(), buffer.size());
+
+    // Split by position in range [0..length)
+    auto [hello, world] = bufferCDR.split(6);
+    ASSERT_EQ(toString(hello), "Hello ");
+    ASSERT_EQ(toString(world), "World");
+
+    // Split by pointer within range
+    auto [hell, oWorld] = bufferCDR.split(bufferCDR.data() + 4);
+    ASSERT_EQ(toString(hell), "Hell");
+    ASSERT_EQ(toString(oWorld), "o World");
+
+    // Get DataRange from a DataRange if original was non-const.
+    DataRange bufferDR(const_cast<char*>(buffer.c_str()), buffer.size());
+    auto [dr1, dr2] = bufferDR.split(6);
+    ASSERT_TRUE(isDataRange(dr1));
+    ASSERT_TRUE(isDataRange(dr2));
+
+    // Get ConstDataRange from a DataRange if original was const.
+    const DataRange constBufferDR = bufferDR;
+    auto [cdr1, cdr2] = constBufferDR.split(6);
+    ASSERT_TRUE(isConstDataRange(cdr1));
+    ASSERT_TRUE(isConstDataRange(cdr2));
+}
+
+TEST(DataRange, splitThrow) {
+    std::string buffer("Hello World");
+    ConstDataRange bufferCDR(buffer.c_str(), buffer.size());
+
+    // Split point is out of range.
+    ASSERT_THROWS(bufferCDR.split(bufferCDR.length() + 1), AssertionException);
+    ASSERT_THROWS(bufferCDR.split(bufferCDR.data() + bufferCDR.length() + 1), AssertionException);
+    ASSERT_THROWS(bufferCDR.split(bufferCDR.data() - 1), AssertionException);
+}
+
 }  // namespace mongo
diff --git a/src/mongo/base/data_type_terminated_test.cpp b/src/mongo/base/data_type_terminated_test.cpp
index d516626e3b9..fdc5d1a6ef7 100644
--- a/src/mongo/base/data_type_terminated_test.cpp
+++ b/src/mongo/base/data_type_terminated_test.cpp
@@ -230,7 +230,7 @@ TEST(DataTypeTerminated, ThroughDataRangeCursor) {
             Terminated<'\0', ConstDataRange> tcdr(ConstDataRange(s.data(), s.data() + s.size()));
             ASSERT_OK(buf_writer.writeAndAdvanceNoThrow(tcdr));
         }
-        const auto written = std::string(static_cast<const char*>(buf), buf_writer.data());
+        const auto written = std::string(static_cast<char*>(buf), buf_writer.data());
         ASSERT_EQUALS(written, serialized);
     }
     {