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, 4 deletions

diff --git a/src/proto/protoio.cc b/src/proto/protoio.cc
index 852ba56fa..ae1e2d8c2 100644
--- a/src/proto/protoio.cc
+++ b/src/proto/protoio.cc
@@ -62,10 +62,8 @@ ProtoOutputStream::ProtoOutputStream(const string& filename) :
         codedStream = new io::CodedOutputStream(zeroCopyStream);
     }
 
-    // Use the ASCII characters gem5 as our magic number and write it
-    // to the file
-    const uint32_t magic_number = 0x356d6567;
-    codedStream->WriteLittleEndian32(magic_number);
+    // Write the magic number to the file
+    codedStream->WriteLittleEndian32(magicNumber);
 
     // Note that each type of stream (packet, instruction etc) should
     // add its own header and perform the appropriate checks
@@ -91,3 +89,102 @@ ProtoOutputStream::write(const Message& msg)
     if (!msg.SerializeToCodedStream(codedStream))
         panic("Unable to write message to coded stream\n");
 }
+
+ProtoInputStream::ProtoInputStream(const string& filename) :
+    fileStream(filename, ios::in | ios::binary), fileName(filename),
+    useGzip(false),
+    zeroCopyStream(NULL), gzipStream(NULL), codedStream(NULL)
+{
+    if (!fileStream.good())
+        panic("Could not open %s for reading\n", filename);
+
+    // check the magic number to see if this is a gzip stream
+    unsigned char bytes[2];
+    fileStream.read((char*) bytes, 2);
+    useGzip = fileStream.good() && bytes[0] == 0x1f && bytes[1] == 0x8b;
+
+    // seek to the start of the input file and clear any flags
+    fileStream.clear();
+    fileStream.seekg(0, ifstream::beg);
+
+    createStreams();
+}
+
+void
+ProtoInputStream::createStreams()
+{
+    // All streams should be NULL at this point
+    assert(zeroCopyStream == NULL && gzipStream == NULL &&
+           codedStream == NULL);
+
+    // Wrap the input file in a zero copy stream, that in turn is
+    // wrapped in a gzip stream if the filename ends with .gz. The
+    // latter stream is in turn wrapped in a coded stream
+    zeroCopyStream = new io::IstreamInputStream(&fileStream);
+    if (useGzip) {
+        gzipStream = new io::GzipInputStream(zeroCopyStream);
+        codedStream = new io::CodedInputStream(gzipStream);
+    } else {
+        codedStream = new io::CodedInputStream(zeroCopyStream);
+    }
+
+    uint32_t magic_check;
+    if (!codedStream->ReadLittleEndian32(&magic_check) ||
+        magic_check != magicNumber)
+        panic("Input file %s is not a valid gem5 proto format.\n",
+              fileName);
+}
+
+void
+ProtoInputStream::destroyStreams()
+{
+    delete codedStream;
+    codedStream = NULL;
+    // As the compression is optional, see if the stream exists
+    if (gzipStream != NULL) {
+        delete gzipStream;
+        gzipStream = NULL;
+    }
+    delete zeroCopyStream;
+    zeroCopyStream = NULL;
+}
+
+
+ProtoInputStream::~ProtoInputStream()
+{
+    destroyStreams();
+    fileStream.close();
+}
+
+
+void
+ProtoInputStream::reset()
+{
+    destroyStreams();
+    // seek to the start of the input file and clear any flags
+    fileStream.clear();
+    fileStream.seekg(0, ifstream::beg);
+    createStreams();
+}
+
+bool
+ProtoInputStream::read(Message& msg)
+{
+    // Read a message from the stream by getting the size, using it as
+    // a limit when parsing the message, then popping the limit again
+    uint32_t size;
+    if (codedStream->ReadVarint32(&size)) {
+        io::CodedInputStream::Limit limit = codedStream->PushLimit(size);
+        if (msg.ParseFromCodedStream(codedStream)) {
+            codedStream->PopLimit(limit);
+            // All went well, the message is parsed and the limit is
+            // popped again
+            return true;
+        } else {
+            panic("Unable to read message from coded stream %s\n",
+                  fileName);
+        }
+    }
+
+    return false;
+}