[lldb-instr] Wrap returns of struct/classes in LLDB_RECORD_RESULT

The instrumentation framework requires return values of custom classes
and structs to be wrapped in the LLDB_RECORD_RESULT macro.

llvm-svn: 354301

5 files changed, 50 insertions, 4 deletions

diff --git a/lldb/lit/tools/lldb-instr/Inputs/foo.cpp b/lldb/lit/tools/lldb-instr/Inputs/foo.cpp
index 7db0452b9e5..51c89758cf9 100644
--- a/lldb/lit/tools/lldb-instr/Inputs/foo.cpp
+++ b/lldb/lit/tools/lldb-instr/Inputs/foo.cpp
@@ -7,3 +7,5 @@ int Foo::C(int i) { return i; }
 int Foo::D(bool b) const { return 1; }
 void Foo::E() {}
 int Foo::F(int i) { return i; }
+void Foo::G(const char *fmt...) {}
+Foo Foo::H() { return Foo(); }
diff --git a/lldb/lit/tools/lldb-instr/Inputs/foo.h b/lldb/lit/tools/lldb-instr/Inputs/foo.h
index 63b6cab828c..913960440c4 100644
--- a/lldb/lit/tools/lldb-instr/Inputs/foo.h
+++ b/lldb/lit/tools/lldb-instr/Inputs/foo.h
@@ -8,5 +8,6 @@ struct Foo {
   int D(bool b) const;
   static void E();
   static int F(int i);
-  int G() { return 0; }
+  void G(const char* fmt...);
+  static Foo H();
 };
diff --git a/lldb/lit/tools/lldb-instr/TestInstrumentationRecord.test b/lldb/lit/tools/lldb-instr/TestInstrumentationRecord.test
index 6f940b64c56..6921a040ecc 100644
--- a/lldb/lit/tools/lldb-instr/TestInstrumentationRecord.test
+++ b/lldb/lit/tools/lldb-instr/TestInstrumentationRecord.test
@@ -12,3 +12,6 @@
 # CHECK: LLDB_RECORD_METHOD_CONST(int, Foo, D, (bool), b);
 # CHECK: LLDB_RECORD_STATIC_METHOD_NO_ARGS(void, Foo, E);
 # CHECK: LLDB_RECORD_STATIC_METHOD(int, Foo, F, (int), i);
+# CHECK-NOT: LLDB_RECORD_STATIC_METHOD(void, Foo, G
+# CHECK: LLDB_RECORD_STATIC_METHOD_NO_ARGS(Foo, Foo, H);
+# CHECK: LLDB_RECORD_RESULT(Foo())
diff --git a/lldb/lit/tools/lldb-instr/TestInstrumentationRegister.test b/lldb/lit/tools/lldb-instr/TestInstrumentationRegister.test
index 300158c7220..baf737be22b 100644
--- a/lldb/lit/tools/lldb-instr/TestInstrumentationRegister.test
+++ b/lldb/lit/tools/lldb-instr/TestInstrumentationRegister.test
@@ -10,3 +10,4 @@
 # CHECK: LLDB_REGISTER_METHOD_CONST(int, Foo, D, (bool));
 # CHECK: LLDB_REGISTER_STATIC_METHOD(void, Foo, E, ());
 # CHECK: LLDB_REGISTER_STATIC_METHOD(int, Foo, F, (int));
+# CHECK-NOT: LLDB_REGISTER_STATIC_METHOD(void, Foo, G
diff --git a/lldb/tools/lldb-instr/Instrument.cpp b/lldb/tools/lldb-instr/Instrument.cpp
index cbc28423375..06a34cee556 100644
--- a/lldb/tools/lldb-instr/Instrument.cpp
+++ b/lldb/tools/lldb-instr/Instrument.cpp
@@ -120,6 +120,33 @@ static std::string GetRegisterMethodMacro(StringRef Result, StringRef Class,
   return OS.str();
 }
 
+class SBReturnVisitor : public RecursiveASTVisitor<SBReturnVisitor> {
+public:
+  SBReturnVisitor(Rewriter &R) : MyRewriter(R) {}
+
+  bool VisitReturnStmt(ReturnStmt *Stmt) {
+    Expr *E = Stmt->getRetValue();
+
+    if (E->getBeginLoc().isMacroID())
+      return false;
+
+    SourceRange R(E->getBeginLoc(), E->getEndLoc());
+
+    StringRef WrittenExpr = Lexer::getSourceText(
+        CharSourceRange::getTokenRange(R), MyRewriter.getSourceMgr(),
+        MyRewriter.getLangOpts());
+
+    std::string ReplacementText =
+        "LLDB_RECORD_RESULT(" + WrittenExpr.str() + ")";
+    MyRewriter.ReplaceText(R, ReplacementText);
+
+    return true;
+  }
+
+private:
+  Rewriter &MyRewriter;
+};
+
 class SBVisitor : public RecursiveASTVisitor<SBVisitor> {
 public:
   SBVisitor(Rewriter &R, ASTContext &Context)
@@ -200,6 +227,13 @@ public:
         MyRewriter.getLangOpts());
     MyRewriter.InsertTextAfter(InsertLoc, Macro);
 
+    // If the function returns a class or struct, we need to wrap its return
+    // statement(s).
+    if (ReturnType->isStructureOrClassType()) {
+      SBReturnVisitor Visitor(MyRewriter);
+      Visitor.TraverseDecl(Decl);
+    }
+
     return true;
   }
 
@@ -210,8 +244,9 @@ private:
   ///  1. Decls outside the main source file,
   ///  2. Decls that are only present in the source file,
   ///  3. Decls that are not definitions,
-  ///  4. Non-public decls,
-  ///  5. Destructors.
+  ///  4. Non-public methods,
+  ///  5. Variadic methods.
+  ///  6. Destructors.
   bool ShouldSkip(CXXMethodDecl *Decl) {
     // Skip anything outside the main file.
     if (!MyRewriter.getSourceMgr().isInMainFile(Decl->getBeginLoc()))
@@ -228,11 +263,15 @@ private:
     if (!Body)
       return true;
 
-    // Skip non-public decls.
+    // Skip non-public methods.
     AccessSpecifier AS = Decl->getAccess();
     if (AS != AccessSpecifier::AS_public)
       return true;
 
+    // Skip variadic methods.
+    if (Decl->isVariadic())
+      return true;
+
     // Skip destructors.
     if (isa<CXXDestructorDecl>(Decl))
       return true;