[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

434

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

481

(NumberOfLoops > 0 || NumberOfSwitch = 0)

498

NestedLoops would be clearer I think

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

521

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
 // - Void return type
 // - Cannot extract declarations that will be needed in the original function
 //   after extraction.
-// - Doesn't check for broken control flow (break/continue without loop/switch)
+// - Checks for broken control flow (break/continue without loop/switch)
 //
 // 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.
struct ExtractionZone {
   SourceLocation getInsertionPoint() const {
     return EnclosingFuncRange.getBegin();
   }
   bool isRootStmt(const Stmt *S) const;
   // The last root statement is important to decide where we need to insert a
   // semicolon after the extraction.
   const Node *getLastRootStmt() const { return Parent->Children.back(); }
   void generateRootStmts();
+
 private:
   llvm::DenseSet<const Stmt *> RootStmts;
 };
 
 bool ExtractionZone::isRootStmt(const Stmt *S) const {
   return RootStmts.find(S) != RootStmts.end();
 }
 
 // Generate RootStmts set
 void ExtractionZone::generateRootStmts() {
-  for(const Node *Child : Parent->Children)
+  for (const Node *Child : Parent->Children)
     RootStmts.insert(Child->ASTNode.get<Stmt>());
 }
 
 // Finds the function in which the zone lies.
 const FunctionDecl *findEnclosingFunction(const Node *CommonAnc) {
   // Walk up the SelectionTree until we find a function Decl
   for (const Node *CurNode = CommonAnc; CurNode; CurNode = CurNode->Parent) {
     // Don't extract from lambdas
     if (CurNode->ASTNode.get<LambdaExpr>())
       return nullptr;
     if (const FunctionDecl *Func = CurNode->ASTNode.get<FunctionDecl>()) {
       // FIXME: Support extraction from methods.
       if (isa<CXXMethodDecl>(Func))
         return nullptr;
       // FIXME: Support extraction from templated functions.
-      if(Func->isTemplated())
+      if (Func->isTemplated())
         return nullptr;
       return Func;
     }
   }
   return nullptr;
 }
 
 // Zone Range is the union of SourceRanges of all child Nodes in Parent since
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;
+  // Control flow is broken if we are extracting a break/continue without a
+  // corresponding parent loop/switch
+  bool BrokenControlFlow = false;
   // FIXME: capture TypeAliasDecl and UsingDirectiveDecl
   // FIXME: Capture type information as well.
   DeclInformation *createDeclInfo(const Decl *D, ZoneRelative RelativeLoc);
   DeclInformation *getDeclInfoFor(const Decl *D);
 };
 
 CapturedZoneInfo::DeclInformation *
 CapturedZoneInfo::createDeclInfo(const Decl *D, ZoneRelative RelativeLoc) {
void CapturedZoneInfo::DeclInformation::markOccurence(
   case ZoneRelative::After:
     IsReferencedInPostZone = true;
     break;
   default:
     break;
   }
 }
 
+bool isLoop(const Stmt *S) {
+  return isa<ForStmt>(S) || isa<DoStmt>(S) || isa<WhileStmt>(S) ||
+         isa<CXXForRangeStmt>(S);
+}
+
 // 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) {
       TraverseDecl(const_cast<FunctionDecl *>(ExtZone.EnclosingFunction));
     }
+
     bool TraverseStmt(Stmt *S) {
+      if (!S)
+        return true;
       bool IsRootStmt = ExtZone.isRootStmt(const_cast<const Stmt *>(S));
       // If we are starting traversal of a RootStmt, we are somewhere inside
       // ExtractionZone
       if (IsRootStmt)
         CurrentLocation = ZoneRelative::Inside;
+      addToLoopSwitchCounters(S, 1);
       // Traverse using base class's TraverseStmt
       RecursiveASTVisitor::TraverseStmt(S);
+      addToLoopSwitchCounters(S, -1);
       // We set the current location as after since next stmt will either be a
       // RootStmt (handled at the beginning) or after extractionZone
       if (IsRootStmt)
         CurrentLocation = ZoneRelative::After;
       return true;
     }
+
+    // Add Increment to CurNumberOf{Loops,Switch} if statement is
+    // {Loop,Switch} and inside Extraction Zone.
+    void addToLoopSwitchCounters(Stmt *S, int Increment) {
+      if (CurrentLocation != ZoneRelative::Inside)
+        return;
+      if (isLoop(S))
+        CurNumberOfNestedLoops += Increment;
+      else if (isa<SwitchStmt>(S))
+        CurNumberOfSwitch += Increment;
+    }
+
+    // Decrement CurNumberOf{NestedLoops,Switch} if statement is {Loop,Switch}
+    // and inside Extraction Zone.
+    void decrementLoopSwitchCounters(Stmt *S) {
+      if (CurrentLocation != ZoneRelative::Inside)
+        return;
+      if (isLoop(S))
+        CurNumberOfNestedLoops--;
+      else if (isa<SwitchStmt>(S))
+        CurNumberOfSwitch--;
+    }
+
     bool VisitDecl(Decl *D) {
       Info.createDeclInfo(D, CurrentLocation);
       return true;
     }
+
     bool VisitDeclRefExpr(DeclRefExpr *DRE) {
       // Find the corresponding Decl and mark it's occurence.
       const Decl *D = DRE->getDecl();
       auto *DeclInfo = Info.getDeclInfoFor(D);
       // If no Decl was found, the Decl must be outside the enclosingFunc.
       if (!DeclInfo)
         DeclInfo = Info.createDeclInfo(D, ZoneRelative::OutsideFunc);
       DeclInfo->markOccurence(CurrentLocation);
       // FIXME: check if reference mutates the Decl being referred.
       return true;
     }
+
     bool VisitReturnStmt(ReturnStmt *Return) {
       if (CurrentLocation == ZoneRelative::Inside)
         Info.HasReturnStmt = true;
       return true;
     }
 
-    // FIXME: check for broken break/continue only.
     bool VisitBreakStmt(BreakStmt *Break) {
-      if (CurrentLocation == ZoneRelative::Inside)
-        Info.HasBreakOrContinue = true;
+      // Control flow is broken if break statement is selected without any
+      // parent loop or switch statement.
+      if (CurrentLocation == ZoneRelative::Inside &&
+          !(CurNumberOfNestedLoops || CurNumberOfSwitch))
+        Info.BrokenControlFlow = true;
       return true;
     }
+
     bool VisitContinueStmt(ContinueStmt *Continue) {
-      if (CurrentLocation == ZoneRelative::Inside)
-        Info.HasBreakOrContinue = true;
+      // Control flow is broken if Continue statement is selected without any
+      // parent loop
+      if (CurrentLocation == ZoneRelative::Inside && !CurNumberOfNestedLoops)
+        Info.BrokenControlFlow = true;
       return true;
     }
     CapturedZoneInfo Info;
     const ExtractionZone &ExtZone;
     ZoneRelative CurrentLocation = ZoneRelative::Before;
+    // Number of {loop,switch} statements that are currently in the traversal
+    // stack inside Extraction Zone. Used to check for broken control flow.
+    unsigned CurNumberOfNestedLoops = 0;
+    unsigned CurNumberOfSwitch = 0;
   };
   ExtractionZoneVisitor Visitor(ExtZone);
   return std::move(Visitor.Info);
 }
 
 // Adds parameters to ExtractedFunc.
 // Returns true if able to find the parameters successfully and no hoisting
 // needed.

 
 // FIXME: add support for adding other function return types besides void.
 // FIXME: assign the value returned by non void extracted function.
 llvm::Expected<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)
+  if (CapturedInfo.BrokenControlFlow)
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
-                                   +"Cannot extract break or continue.");
+                                   +"Cannot extract break/continue without "
+                                    "corresponding loop/switch statement.");
   NewFunction ExtractedFunc(getSemicolonPolicy(ExtZone, SM, LangOpts));
   ExtractedFunc.BodyRange = ExtZone.ZoneRange;
   ExtractedFunc.InsertionPoint = ExtZone.getInsertionPoint();
   ExtractedFunc.EnclosingFuncContext =
       ExtZone.EnclosingFunction->getDeclContext();
   if (!createParameters(ExtractedFunc, CapturedInfo) ||
       !generateReturnProperties(ExtractedFunc, CapturedInfo))
     return llvm::createStringError(llvm::inconvertibleErrorCode(),

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

TEST_F(ExtractFunctionTest, FunctionTest) {
   // lead to break being included in the extraction zone.
   EXPECT_THAT(apply("for(;;) { [[int x;]]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"));
   // 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"));
+  
 }
 
 TEST_F(ExtractFunctionTest, FileTest) {
   // Check all parameters are in order
   std::string ParameterCheckInput = R"cpp(
 struct Foo {
   int x;
 };
void f(const int c) {
   // macros (currently it doesn't select anything for the following case)
   std::string MacroFailInput = R"cpp(
     #define F(BODY) void f() { BODY }
     F ([[int x = 0;]])
   )cpp";
   EXPECT_EQ(apply(MacroFailInput), "unavailable");
 }
 
+TEST_F(ExtractFunctionTest, ControlFlow) {
+  Context = Function;
+  // 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 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"));
+}
 } // namespace
 } // namespace clangd
 } // namespace clang