1 files changed, 69 insertions, 23 deletions

diff --git a/include/llvm/Support/BinaryStreamRef.h b/include/llvm/Support/BinaryStreamRef.h
index 6d5135cb258dd..5cf355be6fe9c 100644
--- a/include/llvm/Support/BinaryStreamRef.h
+++ b/include/llvm/Support/BinaryStreamRef.h
@@ -11,6 +11,7 @@
 #define LLVM_SUPPORT_BINARYSTREAMREF_H
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/Optional.h"
 #include "llvm/Support/BinaryStream.h"
 #include "llvm/Support/BinaryStreamError.h"
 #include "llvm/Support/Error.h"
@@ -24,47 +25,74 @@ namespace llvm {
 template <class RefType, class StreamType> class BinaryStreamRefBase {
 protected:
   BinaryStreamRefBase() = default;
+  explicit BinaryStreamRefBase(StreamType &BorrowedImpl)
+      : BorrowedImpl(&BorrowedImpl), ViewOffset(0) {
+    if (!(BorrowedImpl.getFlags() & BSF_Append))
+      Length = BorrowedImpl.getLength();
+  }
+
   BinaryStreamRefBase(std::shared_ptr<StreamType> SharedImpl, uint32_t Offset,
-                      uint32_t Length)
+                      Optional<uint32_t> Length)
       : SharedImpl(SharedImpl), BorrowedImpl(SharedImpl.get()),
         ViewOffset(Offset), Length(Length) {}
   BinaryStreamRefBase(StreamType &BorrowedImpl, uint32_t Offset,
-                      uint32_t Length)
+                      Optional<uint32_t> Length)
       : BorrowedImpl(&BorrowedImpl), ViewOffset(Offset), Length(Length) {}
-  BinaryStreamRefBase(const BinaryStreamRefBase &Other) {
-    SharedImpl = Other.SharedImpl;
-    BorrowedImpl = Other.BorrowedImpl;
-    ViewOffset = Other.ViewOffset;
-    Length = Other.Length;
-  }
+  BinaryStreamRefBase(const BinaryStreamRefBase &Other) = default;
+  BinaryStreamRefBase &operator=(const BinaryStreamRefBase &Other) = default;
+
+  BinaryStreamRefBase &operator=(BinaryStreamRefBase &&Other) = default;
+  BinaryStreamRefBase(BinaryStreamRefBase &&Other) = default;
 
 public:
   llvm::support::endianness getEndian() const {
     return BorrowedImpl->getEndian();
   }
 
-  uint32_t getLength() const { return Length; }
+  uint32_t getLength() const {
+    if (Length.hasValue())
+      return *Length;
 
-  /// Return a new BinaryStreamRef with the first \p N elements removed.
+    return BorrowedImpl ? (BorrowedImpl->getLength() - ViewOffset) : 0;
+  }
+
+  /// Return a new BinaryStreamRef with the first \p N elements removed.  If
+  /// this BinaryStreamRef is length-tracking, then the resulting one will be
+  /// too.
   RefType drop_front(uint32_t N) const {
     if (!BorrowedImpl)
       return RefType();
 
-    N = std::min(N, Length);
+    N = std::min(N, getLength());
     RefType Result(static_cast<const RefType &>(*this));
+    if (N == 0)
+      return Result;
+
     Result.ViewOffset += N;
-    Result.Length -= N;
+    if (Result.Length.hasValue())
+      *Result.Length -= N;
     return Result;
   }
 
-  /// Return a new BinaryStreamRef with the first \p N elements removed.
+  /// Return a new BinaryStreamRef with the last \p N elements removed.  If
+  /// this BinaryStreamRef is length-tracking and \p N is greater than 0, then
+  /// this BinaryStreamRef will no longer length-track.
   RefType drop_back(uint32_t N) const {
     if (!BorrowedImpl)
       return RefType();
 
-    N = std::min(N, Length);
     RefType Result(static_cast<const RefType &>(*this));
-    Result.Length -= N;
+    N = std::min(N, getLength());
+
+    if (N == 0)
+      return Result;
+
+    // Since we're dropping non-zero bytes from the end, stop length-tracking
+    // by setting the length of the resulting StreamRef to an explicit value.
+    if (!Result.Length.hasValue())
+      Result.Length = getLength();
+
+    *Result.Length -= N;
     return Result;
   }
 
@@ -105,7 +133,7 @@ public:
   }
 
 protected:
-  Error checkOffset(uint32_t Offset, uint32_t DataSize) const {
+  Error checkOffsetForRead(uint32_t Offset, uint32_t DataSize) const {
     if (Offset > getLength())
       return make_error<BinaryStreamError>(stream_error_code::invalid_offset);
     if (getLength() < DataSize + Offset)
@@ -116,7 +144,7 @@ protected:
   std::shared_ptr<StreamType> SharedImpl;
   StreamType *BorrowedImpl = nullptr;
   uint32_t ViewOffset = 0;
-  uint32_t Length = 0;
+  Optional<uint32_t> Length;
 };
 
 /// \brief BinaryStreamRef is to BinaryStream what ArrayRef is to an Array.  It
@@ -131,18 +159,22 @@ class BinaryStreamRef
   friend BinaryStreamRefBase<BinaryStreamRef, BinaryStream>;
   friend class WritableBinaryStreamRef;
   BinaryStreamRef(std::shared_ptr<BinaryStream> Impl, uint32_t ViewOffset,
-                  uint32_t Length)
+                  Optional<uint32_t> Length)
       : BinaryStreamRefBase(Impl, ViewOffset, Length) {}
 
 public:
   BinaryStreamRef() = default;
   BinaryStreamRef(BinaryStream &Stream);
-  BinaryStreamRef(BinaryStream &Stream, uint32_t Offset, uint32_t Length);
+  BinaryStreamRef(BinaryStream &Stream, uint32_t Offset,
+                  Optional<uint32_t> Length);
   explicit BinaryStreamRef(ArrayRef<uint8_t> Data,
                            llvm::support::endianness Endian);
   explicit BinaryStreamRef(StringRef Data, llvm::support::endianness Endian);
 
-  BinaryStreamRef(const BinaryStreamRef &Other);
+  BinaryStreamRef(const BinaryStreamRef &Other) = default;
+  BinaryStreamRef &operator=(const BinaryStreamRef &Other) = default;
+  BinaryStreamRef(BinaryStreamRef &&Other) = default;
+  BinaryStreamRef &operator=(BinaryStreamRef &&Other) = default;
 
   // Use BinaryStreamRef.slice() instead.
   BinaryStreamRef(BinaryStreamRef &S, uint32_t Offset,
@@ -193,17 +225,31 @@ class WritableBinaryStreamRef
                                  WritableBinaryStream> {
   friend BinaryStreamRefBase<WritableBinaryStreamRef, WritableBinaryStream>;
   WritableBinaryStreamRef(std::shared_ptr<WritableBinaryStream> Impl,
-                          uint32_t ViewOffset, uint32_t Length)
+                          uint32_t ViewOffset, Optional<uint32_t> Length)
       : BinaryStreamRefBase(Impl, ViewOffset, Length) {}
 
+  Error checkOffsetForWrite(uint32_t Offset, uint32_t DataSize) const {
+    if (!(BorrowedImpl->getFlags() & BSF_Append))
+      return checkOffsetForRead(Offset, DataSize);
+
+    if (Offset > getLength())
+      return make_error<BinaryStreamError>(stream_error_code::invalid_offset);
+    return Error::success();
+  }
+
 public:
   WritableBinaryStreamRef() = default;
   WritableBinaryStreamRef(WritableBinaryStream &Stream);
   WritableBinaryStreamRef(WritableBinaryStream &Stream, uint32_t Offset,
-                          uint32_t Length);
+                          Optional<uint32_t> Length);
   explicit WritableBinaryStreamRef(MutableArrayRef<uint8_t> Data,
                                    llvm::support::endianness Endian);
-  WritableBinaryStreamRef(const WritableBinaryStreamRef &Other);
+  WritableBinaryStreamRef(const WritableBinaryStreamRef &Other) = default;
+  WritableBinaryStreamRef &
+  operator=(const WritableBinaryStreamRef &Other) = default;
+
+  WritableBinaryStreamRef(WritableBinaryStreamRef &&Other) = default;
+  WritableBinaryStreamRef &operator=(WritableBinaryStreamRef &&Other) = default;
 
   // Use WritableBinaryStreamRef.slice() instead.
   WritableBinaryStreamRef(WritableBinaryStreamRef &S, uint32_t Offset,