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
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;
+}
/**
* @file
- * Declaration of a wrapper for protobuf output streams.
+ * Declaration of a wrapper for protobuf output streams and input streams.
*/
#ifndef __PROTO_PROTOIO_HH__
#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
* is made possible by encoding the length of each message in the
* stream.
*/
-class ProtoOutputStream
+class ProtoOutputStream : public ProtoStream
{
public:
};
+/**
+ * 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
projects / gem5.git / commitdiff