diff --git a/sandboxed_api/tools/clang_generator/generator.cc b/sandboxed_api/tools/clang_generator/generator.cc
index de0052a..74c6044 100644
--- a/sandboxed_api/tools/clang_generator/generator.cc
+++ b/sandboxed_api/tools/clang_generator/generator.cc
@@ -90,9 +90,9 @@ void GeneratorASTConsumer::HandleTranslationUnit(clang::ASTContext& context) {
     return;
   }
 
-  emitter_.AddTypeDeclarations(visitor_.collector_.GetTypeDeclarations());
-
-  for (clang::FunctionDecl* func : visitor_.functions_) {
+  // TODO(cblichmann): Move below to emit all functions after traversing TUs.
+  emitter_.AddTypeDeclarations(visitor_.collector().GetTypeDeclarations());
+  for (clang::FunctionDecl* func : visitor_.functions()) {
     absl::Status status = emitter_.AddFunction(func);
     if (!status.ok()) {
       clang::SourceLocation loc =
diff --git a/sandboxed_api/tools/clang_generator/generator.h b/sandboxed_api/tools/clang_generator/generator.h
index 9e4f9a4..e1b5962 100644
--- a/sandboxed_api/tools/clang_generator/generator.h
+++ b/sandboxed_api/tools/clang_generator/generator.h
@@ -64,13 +64,15 @@ class GeneratorASTVisitor
   bool VisitTypeDecl(clang::TypeDecl* decl);
   bool VisitFunctionDecl(clang::FunctionDecl* decl);
 
+  TypeCollector& collector() { return collector_; }
+
+  const std::vector<clang::FunctionDecl*>& functions() const {
+    return functions_;
+  }
+
  private:
-  friend class GeneratorASTConsumer;
-
   TypeCollector collector_;
-
   std::vector<clang::FunctionDecl*> functions_;
-
   const GeneratorOptions& options_;
 };
 
@@ -85,7 +87,6 @@ class GeneratorASTConsumer : public clang::ASTConsumer {
 
   std::string in_file_;
   GeneratorASTVisitor visitor_;
-
   Emitter& emitter_;
 };