diff --git a/lit/tools/lldb-instr/TestInstrumentationRegister.test b/lit/tools/lldb-instr/TestInstrumentationRegister.test
index 4767221..651880c 100644
--- a/lit/tools/lldb-instr/TestInstrumentationRegister.test
+++ b/lit/tools/lldb-instr/TestInstrumentationRegister.test
@@ -4,6 +4,7 @@
 
 # RUN: lldb-instr %t.dir/foo.cpp | FileCheck %s
 
+# CHECK: {
 # CHECK: LLDB_REGISTER_METHOD(void, Foo, A, ());
 # CHECK: LLDB_REGISTER_METHOD(void, Foo, B, (int));
 # CHECK: LLDB_REGISTER_METHOD(int, Foo, C, (int));
@@ -12,3 +13,4 @@
 # CHECK: LLDB_REGISTER_STATIC_METHOD(int, Foo, F, (int));
 # CHECK-NOT: LLDB_REGISTER_STATIC_METHOD(void, Foo, G
 # CHECK-NOT: LLDB_REGISTER_METHOD_CONST(void, Foo, I, ());
+# CHECK: }
diff --git a/tools/lldb-instr/Instrument.cpp b/tools/lldb-instr/Instrument.cpp
index 3a5df4b..476b5e6 100644
--- a/tools/lldb-instr/Instrument.cpp
+++ b/tools/lldb-instr/Instrument.cpp
@@ -104,8 +104,7 @@
                                                StringRef Signature) {
   std::string Macro;
   llvm::raw_string_ostream OS(Macro);
-  OS << "LLDB_REGISTER_CONSTRUCTOR(" << Class << ", (" << Signature
-     << "));\n\n";
+  OS << "LLDB_REGISTER_CONSTRUCTOR(" << Class << ", (" << Signature << "));\n";
   return OS.str();
 }
 
@@ -305,10 +304,18 @@
 public:
   SBAction() = default;
 
-  void EndSourceFileAction() override { MyRewriter.overwriteChangedFiles(); }
+  bool BeginSourceFileAction(CompilerInstance &CI) override {
+    llvm::outs() << "{\n";
+    return true;
+  }
+
+  void EndSourceFileAction() override {
+    llvm::outs() << "}\n";
+    MyRewriter.overwriteChangedFiles();
+  }
 
   std::unique_ptr<ASTConsumer> CreateASTConsumer(CompilerInstance &CI,
-                                                 StringRef file) override {
+                                                 StringRef File) override {
     MyRewriter.setSourceMgr(CI.getSourceManager(), CI.getLangOpts());
     return llvm::make_unique<SBConsumer>(MyRewriter, CI.getASTContext());
   }