1 files changed, 23 insertions, 12 deletions
diff --git a/src/google/protobuf/compiler/plugin.cc b/src/google/protobuf/compiler/plugin.cc
index a4aedaf..727f942 100644
--- a/src/google/protobuf/compiler/plugin.cc
+++ b/src/google/protobuf/compiler/plugin.cc
@@ -59,13 +59,15 @@ namespace google {
 namespace protobuf {
 namespace compiler {
 
-class GeneratorResponseOutputDirectory : public OutputDirectory {
+class GeneratorResponseContext : public GeneratorContext {
  public:
-  GeneratorResponseOutputDirectory(CodeGeneratorResponse* response)
-      : response_(response) {}
-  virtual ~GeneratorResponseOutputDirectory() {}
+  GeneratorResponseContext(CodeGeneratorResponse* response,
+                           const vector<const FileDescriptor*>& parsed_files)
+      : response_(response),
+        parsed_files_(parsed_files) {}
+  virtual ~GeneratorResponseContext() {}
 
-  // implements OutputDirectory --------------------------------------
+  // implements GeneratorContext --------------------------------------
 
   virtual io::ZeroCopyOutputStream* Open(const string& filename) {
     CodeGeneratorResponse::File* file = response_->add_file();
@@ -81,8 +83,13 @@ class GeneratorResponseOutputDirectory : public OutputDirectory {
     return new io::StringOutputStream(file->mutable_content());
   }
 
+  void ListParsedFiles(vector<const FileDescriptor*>* output) {
+    *output = parsed_files_;
+  }
+
  private:
   CodeGeneratorResponse* response_;
+  const vector<const FileDescriptor*>& parsed_files_;
 };
 
 int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
@@ -112,22 +119,26 @@ int PluginMain(int argc, char* argv[], const CodeGenerator* generator) {
     }
   }
 
-  CodeGeneratorResponse response;
-  GeneratorResponseOutputDirectory output_directory(&response);
-
+  vector<const FileDescriptor*> parsed_files;
   for (int i = 0; i < request.file_to_generate_size(); i++) {
-    const FileDescriptor* file =
-        pool.FindFileByName(request.file_to_generate(i));
-    if (file == NULL) {
+    parsed_files.push_back(pool.FindFileByName(request.file_to_generate(i)));
+    if (parsed_files.back() == NULL) {
       cerr << argv[0] << ": protoc asked plugin to generate a file but "
               "did not provide a descriptor for the file: "
            << request.file_to_generate(i) << endl;
       return 1;
     }
+  }
+
+  CodeGeneratorResponse response;
+  GeneratorResponseContext context(&response, parsed_files);
+
+  for (int i = 0; i < parsed_files.size(); i++) {
+    const FileDescriptor* file = parsed_files[i];
 
     string error;
     bool succeeded = generator->Generate(
-        file, request.parameter(), &output_directory, &error);
+        file, request.parameter(), &context, &error);
 
     if (!succeeded && error.empty()) {
       error = "Code generator returned false but provided no error "