base: Add wrapped protobuf input stream

This patch adds support for inputting protobuf messages through a
ProtoInputStream which hides the internal streams used by the
library. The stream is created based on the name of an input file and
optionally includes decompression using gzip.

The input stream will start by getting a magic number from the file,
and also verify that it matches with the expected value. Once opened,
messages can be read incrementally from the stream, returning
true/false until an error occurs or the end of the file is reached.

1 files changed, 101 insertions, 2 deletions

diff --git a/src/proto/protoio.hh b/src/proto/protoio.hh
index bb11bc7e9..d5c6a4bac 100644
--- a/src/proto/protoio.hh
+++ b/src/proto/protoio.hh
@@ -40,7 +40,7 @@
 
 /**
  * @file
- * Declaration of a wrapper for protobuf output streams.
+ * Declaration of a wrapper for protobuf output streams and input streams.
  */
 
 #ifndef __PROTO_PROTOIO_HH__
@@ -54,6 +54,34 @@
 #include <fstream>
 
 /**
+ * A ProtoStream provides the shared functionality of the input and
+ * output streams. At the moment this is limited to magic number.
+ */
+class ProtoStream
+{
+
+  protected:
+
+    /// Use the ASCII characters gem5 as our magic number
+    static const uint32_t magicNumber = 0x356d6567;
+
+    /**
+     * Create a ProtoStream.
+     */
+    ProtoStream() {}
+
+  private:
+
+    /**
+     * Hide the copy constructor and assignment operator.
+     * @{
+     */
+    ProtoStream(const ProtoStream&);
+    ProtoStream& operator=(const ProtoStream&);
+    /** @} */
+};
+
+/**
  * A ProtoOutputStream wraps a coded stream, potentially with
  * compression, based on looking at the file name. Writing to the
  * stream is done to enable interaction with the file on a per-message
@@ -61,7 +89,7 @@
  * is made possible by encoding the length of each message in the
  * stream.
  */
-class ProtoOutputStream
+class ProtoOutputStream : public ProtoStream
 {
 
   public:
@@ -104,4 +132,75 @@ class ProtoOutputStream
 
 };
 
+/**
+ * A ProtoInputStream wraps a coded stream, potentially with
+ * decompression, based on looking at the file name. Reading from the
+ * stream is done on a per-message basis to avoid having to deal with
+ * huge data structures. The latter assumes the length of each message
+ * is encoded in the stream when it is written.
+ */
+class ProtoInputStream : public ProtoStream
+{
+
+  public:
+
+    /**
+     * Create an input stream for a given file name. If the filename
+     * ends with .gz then the file will be decompressed accordingly.
+     *
+     * @param filename Path to the file to read from
+     */
+    ProtoInputStream(const std::string& filename);
+
+    /**
+     * Destruct the input stream, and also close the underlying file
+     * streams and coded streams.
+     */
+    ~ProtoInputStream();
+
+    /**
+     * Read a message from the stream.
+     *
+     * @param msg Message read from the stream
+     * @param return True if a message was read, false if reading fails
+     */
+    bool read(google::protobuf::Message& msg);
+
+    /**
+     * Reset the input stream and seek to the beginning of the file.
+     */
+    void reset();
+
+  private:
+
+    /**
+     * Create the internal streams that are wrapping the input file.
+     */
+    void createStreams();
+
+    /**
+     * Destroy the internal streams that are wrapping the input file.
+     */
+    void destroyStreams();
+
+    /// Underlying file input stream
+    std::ifstream fileStream;
+
+    /// Hold on to the file name for debug messages
+    const std::string fileName;
+
+    /// Boolean flag to remember whether we use gzip or not
+    bool useGzip;
+
+    /// Zero Copy stream wrapping the STL input stream
+    google::protobuf::io::IstreamInputStream* zeroCopyStream;
+
+    /// Optional Gzip stream to wrap the Zero Copy stream
+    google::protobuf::io::GzipInputStream* gzipStream;
+
+    /// Top-level coded stream that messages are read from
+    google::protobuf::io::CodedInputStream* codedStream;
+
+};
+
 #endif //__PROTO_PROTOIO_HH