raw_istream enhancements.

wpiutil/src/main/native/include/wpi/raw_istream.h

@@ -8,10 +8,15 @@
 #ifndef WPIUTIL_WPI_RAW_ISTREAM_H_
 #define WPIUTIL_WPI_RAW_ISTREAM_H_
 
+#include <stdint.h>
+
 #include <algorithm>
 #include <cstddef>
+#include <string>
 #include <system_error>
+#include <vector>
 
+#include "wpi/ArrayRef.h"
 #include "wpi/SmallVector.h"
 #include "wpi/StringRef.h"
 #include "wpi/Twine.h"
@@ -47,7 +52,47 @@ class raw_istream {
     size_t readlen = std::min(in_avail(), len);
     if (readlen == 0) return 0;
     read_impl(data, readlen);
-    return readlen;
+    return m_read_count;
+  }
+
+  raw_istream& readinto(SmallVectorImpl<char>& buf, size_t len) {
+    size_t old_size = buf.size();
+    buf.append(len, 0);
+    read_impl(&buf[old_size], len);
+    buf.resize(old_size + m_read_count);
+    return *this;
+  }
+
+  raw_istream& readinto(SmallVectorImpl<uint8_t>& buf, size_t len) {
+    size_t old_size = buf.size();
+    buf.append(len, 0);
+    read_impl(&buf[old_size], len);
+    buf.resize(old_size + m_read_count);
+    return *this;
+  }
+
+  raw_istream& readinto(std::vector<char>& buf, size_t len) {
+    size_t old_size = buf.size();
+    buf.insert(buf.end(), len, 0);
+    read_impl(&buf[old_size], len);
+    buf.resize(old_size + m_read_count);
+    return *this;
+  }
+
+  raw_istream& readinto(std::vector<uint8_t>& buf, size_t len) {
+    size_t old_size = buf.size();
+    buf.insert(buf.end(), len, 0);
+    read_impl(&buf[old_size], len);
+    buf.resize(old_size + m_read_count);
+    return *this;
+  }
+
+  raw_istream& readinto(std::string& buf, size_t len) {
+    size_t old_size = buf.size();
+    buf.insert(buf.end(), len, 0);
+    read_impl(&buf[old_size], len);
+    buf.resize(old_size + m_read_count);
+    return *this;
   }
 
   // Read a line from an input stream (up to a maximum length).
@@ -59,8 +104,15 @@ class raw_istream {
   StringRef getline(SmallVectorImpl<char>& buf, int maxLen);
 
   virtual void close() = 0;
+
+  // Number of bytes available to read without potentially blocking.
+  // Note this can return zero even if there are bytes actually available to
+  // read.
   virtual size_t in_avail() const = 0;
 
+  // Return the number of bytes read by the last read operation.
+  size_t read_count() const { return m_read_count; }
+
   bool has_error() const { return m_error; }
   void clear_error() { m_error = false; }
 
@@ -69,16 +121,27 @@ class raw_istream {
 
  protected:
   void error_detected() { m_error = true; }
+  void set_read_count(size_t count) { m_read_count = count; }
 
  private:
   virtual void read_impl(void* data, size_t len) = 0;
 
   bool m_error = false;
+  size_t m_read_count = 0;
 };
 
 class raw_mem_istream : public raw_istream {
  public:
-  explicit raw_mem_istream(StringRef mem);
+  // not const as we don't want to allow temporaries
+  explicit raw_mem_istream(std::string& str)
+      : raw_mem_istream(str.data(), str.size()) {}
+  explicit raw_mem_istream(ArrayRef<char> mem)
+      : raw_mem_istream(mem.data(), mem.size()) {}
+  explicit raw_mem_istream(ArrayRef<uint8_t> mem)
+      : raw_mem_istream(reinterpret_cast<const char*>(mem.data()), mem.size()) {
+  }
+  explicit raw_mem_istream(const char* str)
+      : m_cur(str), m_left(std::strlen(str)) {}
   raw_mem_istream(const char* mem, size_t len) : m_cur(mem), m_left(len) {}
   void close() override;
   size_t in_avail() const override;