Improve performance of loplugin:flatten

...by avoiding calls to parentStmt.

Change-Id: I4f3d66a0529e9c3abf5c963bcf70db7a2afa1bf9

2 files changed, 56 insertions, 28 deletions

diff --git a/compilerplugins/clang/flatten.cxx b/compilerplugins/clang/flatten.cxx
index 7fa408ea1731..ce16bc2384ee 100644
--- a/compilerplugins/clang/flatten.cxx
+++ b/compilerplugins/clang/flatten.cxx
@@ -7,12 +7,13 @@
  * file, You can obtain one at http://mozilla.org/MPL/2.0/.
  */
 
+#include "plugin.hxx"
 #include <cassert>
 #include <string>
 #include <iostream>
 #include <fstream>
 #include <set>
-#include "plugin.hxx"
+#include <stack>
 
 /**
   Look for places where we can flatten the control flow in a method by returning early.
@@ -30,7 +31,9 @@ public:
         TraverseDecl(compiler.getASTContext().getTranslationUnitDecl());
     }
 
+    bool TraverseIfStmt(IfStmt *);
     bool TraverseCXXCatchStmt(CXXCatchStmt * );
+    bool TraverseCompoundStmt(CompoundStmt *);
     bool VisitIfStmt(IfStmt const * );
 private:
     bool rewrite1(IfStmt const * );
@@ -40,8 +43,8 @@ private:
     std::string getSourceAsString(SourceRange range);
     std::string invertCondition(Expr const * condExpr, SourceRange conditionRange);
     bool checkOverlap(SourceRange);
-    bool lastStmtInParent(Stmt const * stmt);
 
+    std::stack<bool> mLastStmtInParentStack;
     std::vector<std::pair<char const *, char const *>> mvModifiedRanges;
 };
 
@@ -72,23 +75,30 @@ static bool containsVarDecl(Stmt const * stmt)
     return declStmt && isa<VarDecl>(*declStmt->decl_begin());
 }
 
-bool Flatten::lastStmtInParent(Stmt const * stmt)
+bool Flatten::TraverseCXXCatchStmt(CXXCatchStmt* )
 {
-    auto parent = parentStmt(stmt);
-    if (!parent) {
-        return true;
-    }
-    auto parentCompound = dyn_cast<CompoundStmt>(parent);
-    if (!parentCompound) {
+    // ignore stuff inside catch statements, where doing a "if...else..throw" is more natural
+    return true;
+}
+
+bool Flatten::TraverseIfStmt(IfStmt * ifStmt)
+{
+    // ignore if we are part of an if/then/else/if chain
+    if (ifStmt->getElse() && isa<IfStmt>(ifStmt->getElse()))
         return true;
-    }
-    return parentCompound->body_back() == stmt;
+    return RecursiveASTVisitor<Flatten>::TraverseIfStmt(ifStmt);
 }
 
-bool Flatten::TraverseCXXCatchStmt(CXXCatchStmt* )
+bool Flatten::TraverseCompoundStmt(CompoundStmt * compoundStmt)
 {
-    // ignore stuff inside catch statements, where doing a "if...else..throw" is more natural
-    return true;
+    // if the "if" statement is not the last statement in its block, and it contains
+    // var decls in its then block, we cannot de-indent the then block without
+    // extending the lifetime of some variables, which may be problematic
+    // ignore if we are part of an if/then/else/if chain
+    mLastStmtInParentStack.push(compoundStmt->size() > 0 && isa<IfStmt>(*compoundStmt->body_back()));
+    bool rv = RecursiveASTVisitor<Flatten>::TraverseCompoundStmt(compoundStmt);
+    mLastStmtInParentStack.pop();
+    return rv;
 }
 
 bool Flatten::VisitIfStmt(IfStmt const * ifStmt)
@@ -99,15 +109,6 @@ bool Flatten::VisitIfStmt(IfStmt const * ifStmt)
     if (!ifStmt->getElse())
         return true;
 
-    // ignore if/then/else/if chains for now
-    if (isa<IfStmt>(ifStmt->getElse()))
-        return true;
-
-    // ignore if we are part of an if/then/else/if chain
-    auto parentIfStmt = dyn_cast<IfStmt>(parentStmt(ifStmt));
-    if (parentIfStmt && parentIfStmt->getElse() == ifStmt)
-        return true;
-
     if (containsPreprocessingConditionalInclusion(ifStmt->getSourceRange())) {
         return true;
     }
@@ -118,10 +119,10 @@ bool Flatten::VisitIfStmt(IfStmt const * ifStmt)
         // if both the "if" and the "else" contain throws, no improvement
         if (containsSingleThrowExpr(ifStmt->getThen()))
             return true;
-        // if the "if" statement is not the last statement in it's block, and it contains
-        // var decls in it's then block, we cannot de-indent the then block without
+        // if the "if" statement is not the last statement in its block, and it contains
+        // var decls in its then block, we cannot de-indent the then block without
         // extending the lifetime of some variables, which may be problematic
-        if (!lastStmtInParent(ifStmt) && containsVarDecl(ifStmt->getThen()))
+        if (!mLastStmtInParentStack.top() || containsVarDecl(ifStmt->getThen()))
             return true;
 
         if (!rewrite1(ifStmt))
@@ -147,7 +148,7 @@ bool Flatten::VisitIfStmt(IfStmt const * ifStmt)
         // if the "if" statement is not the last statement in it's block, and it contains
         // var decls in it's else block, we cannot de-indent the else block without
         // extending the lifetime of some variables, which may be problematic
-        if (!lastStmtInParent(ifStmt) && containsVarDecl(ifStmt->getElse()))
+        if (!mLastStmtInParentStack.top() || containsVarDecl(ifStmt->getElse()))
             return true;
 
         if (!rewrite2(ifStmt))
diff --git a/compilerplugins/clang/test/flatten.cxx b/compilerplugins/clang/test/flatten.cxx
index a901d273b9a1..dcc7cb3b81c9 100644
--- a/compilerplugins/clang/test/flatten.cxx
+++ b/compilerplugins/clang/test/flatten.cxx
@@ -15,10 +15,16 @@ class Class {};
 
 void top1() {
     if (foo() == 1) { // expected-note {{if condition here [loplugin:flatten]}}
+        foo();
+    } else {
+        throw std::exception(); // expected-error {{unconditional throw in else branch, rather invert the condition, throw early, and flatten the normal case [loplugin:flatten]}}
+    }
+    // no warning expected
+    if (foo() == 1) {
         Class aClass;
         (void)aClass;
     } else {
-        throw std::exception(); // expected-error {{unconditional throw in else branch, rather invert the condition, throw early, and flatten the normal case [loplugin:flatten]}}
+        throw std::exception();
     }
 }
 
@@ -26,6 +32,12 @@ void top2() {
     if (foo() == 2) {
         throw std::exception(); // expected-error {{unconditional throw in then branch, just flatten the else [loplugin:flatten]}}
     } else {
+        foo();
+    }
+    // no warning expected
+    if (foo() == 2) {
+        throw std::exception();
+    } else {
         Class aClass;
         (void)aClass;
     }
@@ -77,4 +89,19 @@ int main() {
     }
 }
 
+void top6() {
+    // no warning expected
+    if (foo() == 2) {
+        Class aClass;
+        (void)aClass;
+    } else if (foo() == 2) {
+        Class aClass;
+        (void)aClass;
+    } else {
+        throw std::exception();
+    }
+    int x = 1;
+    (void)x;
+}
+
 /* vim:set shiftwidth=4 softtabstop=4 expandtab cinoptions=b1,g0,N-s cinkeys+=0=break: */