[Clangd] ExtractFunction Added checks for broken control flow

[SureYeaah]

Summary

• Added checks for broken control flow
• Added unittests

[SureYeaah]

Removed selectionTree Visitor

[SureYeaah]

Added null statement check in TraverseStmt

[sammccall]

clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp

415

rather than writing this twice, take an int Increment parameter?

431

(NumberOfLoops > 0 || NumberOfSwitch = 0)

433

NestedLoops would be clearer I think

clang-tools-extra/clangd/unittests/TweakTests.cpp

628

please add these to a new test case e.g. TEST_F(ExtractFunctionTest, ControlFlow) to avoid each getting too long

[SureYeaah]

Addressed review comments

clang-tools-extra/clangd/refactor/tweaks/ExtractFunction.cpp

 //   }
 //
 // - Only extract statements
 // - Extracts from non-templated free functions only.
 // - Parameters are const only if the declaration was const
 //   - Always passed by l-value reference
 // - No return
 // - Cannot move declarations before extracting
-// - Doesn't check for broken control flow
 //
 // 1. ExtractFunction is the tweak subclass
 //    - Prepare does basic analysis of the selection and is therefore fast.
 //      Successful prepare doesn't always mean we can apply the tweak.
 //    - Apply does a more detailed analysis and can be slower. In case of
 //      failure, we let the user know that we are unable to perform extraction.
 // 2. ExtractionZone store information about the range being extracted and the
 //    enclosing function.
 // 3. NewFunction stores properties of the extracted function and provides
 //    methods for rendering it.
-// 4. CapturedZoneInfo uses a RAV to capture information about the extraction
-//    like declarations, existing return statements, broken control flow, etc.
+// 4. CapturedZoneInfo uses a RAV and SelectionTreeVisitor to capture
+//    information about the extraction like declarations, existing return
+//    statements, broken control flow, etc.
 // 5. getExtractedFunction is responsible for analyzing the CapturedZoneInfo and
 //    creating a NewFunction.
 // Design Doc at <TODO: Share design doc>
 //===----------------------------------------------------------------------===//
 
 #include "ClangdUnit.h"
 #include "Logger.h"
 #include "Selection.h"
 #include "SourceCode.h"
 #include "refactor/Tweak.h"
 #include "clang/AST/ASTContext.h"
 #include "clang/AST/DeclTemplate.h"
 #include "clang/AST/RecursiveASTVisitor.h"
 #include "clang/AST/Stmt.h"
+#include "clang/AST/StmtCXX.h"
 #include "clang/Basic/LangOptions.h"
 #include "clang/Basic/SourceLocation.h"
 #include "clang/Basic/SourceManager.h"
 #include "clang/Lex/Lexer.h"
 #include "clang/Tooling/Core/Replacement.h"
 #include "clang/Tooling/Refactoring/Extract/SourceExtraction.h"
 #include "llvm/ADT/None.h"
 #include "llvm/ADT/Optional.h"
DeclInformation(const Decl *TheDecl, ZoneRelative DeclaredIn,
         : TheDecl(TheDecl), DeclaredIn(DeclaredIn), DeclIndex(DeclIndex){};
     // Marks the occurence of a reference for this declaration
     void markOccurence(ZoneRelative ReferenceLoc);
   };
   // Maps Decls to their DeclInfo
   llvm::DenseMap<const Decl *, DeclInformation> DeclInfoMap;
   // True if there is a return statement in zone.
   bool HasReturnStmt = false;
-  // For now we just care whether there exists a break/continue in zone.
-  bool HasBreakOrContinue = false;
+  // True if extracting break or continue without corresponding loop/switch.
+  bool BrokenControlFlow = false;
   // FIXME: capture TypeAliasDecl and UsingDirectiveDecl
   // FIXME: Capture type information as well.
   // Return reference for a Decl, adding it if not already present.
   DeclInformation &getDeclInformationFor(const Decl *D,
                                          const ExtractionZone &ExtZone);
 };
 
 CapturedZoneInfo::DeclInformation &

 // Captures information from Extraction Zone
 CapturedZoneInfo captureZoneInfo(const ExtractionZone &ExtZone) {
   // We use the ASTVisitor instead of using the selection tree since we need to
   // find references in the PostZone as well.
   // FIXME: Check which statements we don't allow to extract.
   class ExtractionZoneVisitor
       : public clang::RecursiveASTVisitor<ExtractionZoneVisitor> {
   public:
-    ExtractionZoneVisitor(const ExtractionZone &ExtZone) : ExtZone(ExtZone) {
+    ExtractionZoneVisitor(CapturedZoneInfo &Info, const ExtractionZone &ExtZone)
+        : Info(Info), ExtZone(ExtZone) {
       TraverseDecl(const_cast<FunctionDecl *>(
           ExtZone.EnclosingFunction->ASTNode.get<FunctionDecl>()));
     }
     bool VisitDecl(Decl *D) { // NOLINT
       // Create new DeclInfo
       Info.getDeclInformationFor(D, ExtZone);
       return true;
     }
     bool VisitDeclRefExpr(DeclRefExpr *DRE) { // NOLINT
       // Find the corresponding Decl and mark it's occurence.
       auto &DeclInfo = Info.getDeclInformationFor(DRE->getDecl(), ExtZone);
       DeclInfo.markOccurence(ExtZone.getRelationToZone(DRE->getLocation()));
       // FIXME: check if reference mutates the Decl being referred.
       return true;
     }
     bool VisitReturnStmt(ReturnStmt *Return) { // NOLINT
       if (isStmtInExtractionZone(Return, ExtZone))
         Info.HasReturnStmt = true;
       return true;
     }
-
-    CapturedZoneInfo Info;
+    CapturedZoneInfo &Info;
     const ExtractionZone &ExtZone;
   };
-  ExtractionZoneVisitor Visitor(ExtZone);
-  return std::move(Visitor.Info);
+
+  class SelectionTreeVisitor {
+    // Returns tuple denoting whether we can Break/Continue out of a node.
+    std::tuple<bool, bool> canHaveBreakContinueAsChild(const Node *N) {
+      const Stmt *S = N->ASTNode.get<Stmt>();
+      bool Break = false;
+      bool Continue = false;
+      // We can break/continue out of a loop
+      if (dyn_cast_or_null<WhileStmt>(S) || dyn_cast_or_null<ForStmt>(S) ||
+          dyn_cast_or_null<CXXForRangeStmt>(S) || dyn_cast_or_null<DoStmt>(S))
+        Break = Continue = true;
+      // We can only break out of a switch
+      else if (dyn_cast_or_null<SwitchStmt>(S))
+        Break = true;
+      return {Break, Continue};
+    }
+    // Perform a depth first traversal of selection tree to look for broken
+    // break/continue in Extraction zone.
+    void traverse(const Node *CurNode) {
+      bool CanHaveBreak, CanHaveContinue;
+      std::tie(CanHaveBreak, CanHaveContinue) =
+          canHaveBreakContinueAsChild(CurNode);
+      BreakParents += CanHaveBreak;
+      ContinueParents += CanHaveContinue;
+      // If we find a break or continue without a corresponding parent that can
+      // have a break/continue inside it, we have broken control flow.
+      if ((CurNode->ASTNode.get<BreakStmt>() && !BreakParents) ||
+          (CurNode->ASTNode.get<ContinueStmt>() && !ContinueParents))
+        Info.BrokenControlFlow = true;
+      for (const Node *ChildNode : CurNode->Children) {
+        traverse(ChildNode);
+      }
+      BreakParents -= CanHaveBreak;
+      ContinueParents -= CanHaveContinue;
+    }
+    // Current number of nodes that can we can break/continue out of.
+    unsigned BreakParents = 0;
+    unsigned ContinueParents = 0;
+    CapturedZoneInfo &Info;
+
+  public:
+    SelectionTreeVisitor(CapturedZoneInfo &Info, const Node *ZoneParent)
+        : Info(Info) {
+      for (const Node *RootStmt : ZoneParent->Children) {
+        traverse(RootStmt);
+      }
+    }
+  };
+  CapturedZoneInfo Info;
+  ExtractionZoneVisitor Visitor(Info, ExtZone);
+  SelectionTreeVisitor SelectionVisitor(Info, ExtZone.Parent);
+  return Info;
 }
 
 // Adds parameters to ExtractedFunc.
 // Returns true if able to find the parameters successfully and no hoisting
 // needed.
 bool createParameters(NewFunction &ExtractedFunc,
                       const CapturedZoneInfo &CapturedInfo) {
 

 }
 
 // FIXME: add support for adding other function return types besides void.
 // FIXME: assign the value returned by non void extracted function.
 llvm::Optional<NewFunction> getExtractedFunction(ExtractionZone &ExtZone,
                                                  const SourceManager &SM,
                                                  const LangOptions &LangOpts) {
   CapturedZoneInfo CapturedInfo = captureZoneInfo(ExtZone);
-  // Bail out if any break of continue exists
-  // FIXME: check for broken control flow only
-  if (CapturedInfo.HasBreakOrContinue)
+  // Bail out if trying to extract broken control flow
+  if (CapturedInfo.BrokenControlFlow)
     return llvm::None;
   auto SemicolonPolicy = getSemicolonPolicy(ExtZone, SM, LangOpts);
   NewFunction ExtractedFunc(ExtZone.ZoneRange, ExtZone.getInsertionPoint(),
                             std::move(SemicolonPolicy), SM);
   if (!createParameters(ExtractedFunc, CapturedInfo))
     return llvm::None;
   if (!generateReturnProperties(ExtractedFunc, CapturedInfo))
     return llvm::None;

clang-tools-extra/clangd/unittests/TweakTests.cpp

TEST_F(ExtractFunctionTest, FunctionTest) {
   // lead to break being included in the extraction zone.
   EXPECT_THAT(apply("for(;;) { [[{}]]break; }"), HasSubstr("extracted"));
   // FIXME: This should be unavailable since partially selected but
   // selectionTree doesn't always work correctly for VarDecls.
   EXPECT_THAT(apply("int [[x = 0]];"), HasSubstr("extracted"));
   // FIXME: ExtractFunction should be unavailable inside loop construct
   // initalizer/condition.
   EXPECT_THAT(apply(" for([[int i = 0;]];);"), HasSubstr("extracted"));
+  // We should be able to extract break/continue with a parent loop/switch.
+  EXPECT_THAT(apply(" [[for(;;) if(1) break;]] "), HasSubstr("extracted"));
+  EXPECT_THAT(apply(" for(;;) [[while(1) break;]] "), HasSubstr("extracted"));
+  EXPECT_THAT(apply(" [[switch(1) { break; }]]"), HasSubstr("extracted"));
+  EXPECT_THAT(apply(" [[while(1) switch(1) { continue; }]]"),
+              HasSubstr("extracted"));
   // Don't extract because needs hoisting.
   EXPECT_THAT(apply(" [[int a = 5;]] a++; "), StartsWith("fail"));
   // Don't extract return
   EXPECT_THAT(apply(" if(true) [[return;]] "), StartsWith("fail"));
-  // Don't extract break and continue.
-  // FIXME: We should be able to extract this since it's non broken.
-  EXPECT_THAT(apply(" [[for(;;) break;]] "), StartsWith("fail"));
-  EXPECT_THAT(apply(" for(;;) [[continue;]] "), StartsWith("fail"));
+  // Don't extract break and continue without a loop/switch parent.
+  EXPECT_THAT(apply(" for(;;) [[if(1) continue;]] "), StartsWith("fail"));
+  EXPECT_THAT(apply(" while(1) [[if(1) break;]] "), StartsWith("fail"));
+  EXPECT_THAT(apply(" switch(1) { [[break;]] }"), StartsWith("fail"));
+  EXPECT_THAT(apply(" for(;;) { [[while(1) break; break;]] }"),
+              StartsWith("fail"));
 }
 
 TEST_F(ExtractFunctionTest, FileTest) {
   // Check all parameters are in order
   std::string ParameterCheckInput = R"cpp(
 struct FOO {
   int x;
 };