[clangd] Allow extract-to-function on regions that always return.

[sammccall]

Summary

We only do a trivial check whether the region always returns - it has to end
with a return statement.

[merge_guards_bot]

Build result: pass - 60231 tests passed, 0 failed and 732 were skipped.
Log files: console-log.txt, CMakeCache.txt

[merge_guards_bot]

Build result: pass - 60231 tests passed, 0 failed and 732 were skipped.

Log files: console-log.txt, CMakeCache.txt

[kadircet]

LGTM, thanks!

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

173

EZ.getLastRootStmt() instead of EZ.Parent->Children.back()

175

nit: braces

180

what if there's a goto in the selection :P

612

nit: early exits

[kadircet]

btw forgot to add a comment, it might be worth mentioning in change description that, this works in its current state because we don't allow extractions inside lambda expressions.

[sammccall]

(Thanks! Will address comments, just wanted to ask for a clarification)

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

180

If the goto target is in the selection, that's fine. Otherwise it's broken control flow like break/continue. It doesn't seem to be detected yet indeed.

612

I'm not sure what this comment means, can you elaborate?
(I do mean "sometimes returns and sometimes doesn't". Early-exits aren't a problem, e.g. if (x) return 1; return 2; can be extracted.)

[kadircet]

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

612

ah sorry, i was trying to imply usage of early exits in the code, i.e:

ExtractedFunc.ReturnType = EnclosingFunc.getParentASTContext().VoidTy;
if not HasReturnStmt
  return true;
if not AlwaysReturn
  return false;
.
.
.

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

struct ExtractionZone {
   // semicolon after the extraction.
   const Node *getLastRootStmt() const { return Parent->Children.back(); }
   void generateRootStmts();
 
 private:
   llvm::DenseSet<const Stmt *> RootStmts;
 };
 
+// Whether the code in the extraction zone is guaranteed to return, assuming
+// no broken control flow (unbound break/continue).
+// This is a very naive check (does it end with a return stmt).
+// Doing some rudimentary control flow analysis would cover more cases.
+bool alwaysReturns(const ExtractionZone &EZ) {
+  const Stmt *Last = EZ.getLastRootStmt()->ASTNode.get<Stmt>();
+  // Unwrap enclosing (unconditional) compound statement.
+  while (const auto *CS = llvm::dyn_cast<CompoundStmt>(Last)) {
+    if (CS->body_empty())
+      return false;
+    else
+      Last = CS->body_back();
+  }
+  return llvm::isa<ReturnStmt>(Last);
+}
+
 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)
     RootStmts.insert(Child->ASTNode.get<Stmt>());
struct Parameter {
     bool PassByReference;
     unsigned OrderPriority; // Lower value parameters are preferred first.
     std::string render(const DeclContext *Context) const;
     bool operator<(const Parameter &Other) const {
       return OrderPriority < Other.OrderPriority;
     }
   };
   std::string Name = "extracted";
-  std::string ReturnType;
+  QualType ReturnType;
   std::vector<Parameter> Parameters;
   SourceRange BodyRange;
   SourceLocation InsertionPoint;
   const DeclContext *EnclosingFuncContext;
+  bool CallerReturnsValue = false;
   // Decides whether the extracted function body and the function call need a
   // semicolon after extraction.
   tooling::ExtractionSemicolonPolicy SemicolonPolicy;
   NewFunction(tooling::ExtractionSemicolonPolicy SemicolonPolicy)
       : SemicolonPolicy(SemicolonPolicy) {}
   // Render the call for this function.
   std::string renderCall() const;
   // Render the definition for this function.
if (NeedCommaBefore)
       Result += ", ";
     NeedCommaBefore = true;
     Result += P.Name;
   }
   return Result;
 }
 
 std::string NewFunction::renderCall() const {
-  return Name + "(" + renderParametersForCall() + ")" +
-         (SemicolonPolicy.isNeededInOriginalFunction() ? ";" : "");
+  return llvm::formatv(
+      "{0}{1}({2}){3}", CallerReturnsValue ? "return " : "", Name,
+      renderParametersForCall(),
+      (SemicolonPolicy.isNeededInOriginalFunction() ? ";" : ""));
 }
 
 std::string NewFunction::renderDefinition(const SourceManager &SM) const {
-  return ReturnType + " " + Name + "(" + renderParametersForDefinition() + ")" +
-         " {\n" + getFuncBody(SM) + "\n}\n";
+  return llvm::formatv("{0} {1}({2}) {\n{3}\n}\n",
+                       printType(ReturnType, *EnclosingFuncContext), Name,
+                       renderParametersForDefinition(), getFuncBody(SM));
 }
 
 std::string NewFunction::getFuncBody(const SourceManager &SM) const {
   // FIXME: Generate tooling::Replacements instead of std::string to
   // - hoist decls
   // - add return statement
   // - Add semicolon
   return toSourceCode(SM, BodyRange).str() +
struct DeclInformation {
     DeclInformation(const Decl *TheDecl, ZoneRelative DeclaredIn,
                     unsigned DeclIndex)
         : 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;
+  bool HasReturnStmt = false; // Are there any return statements in the zone?
+  bool AlwaysReturns = false; // Does the zone always return?
   // 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);
 };
public:
     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);
+  CapturedZoneInfo Result = std::move(Visitor.Info);
+  Result.AlwaysReturns = alwaysReturns(ExtZone);
+  return Result;
 }
 
 // Adds parameters to ExtractedFunc.
 // Returns true if able to find the parameters successfully and no hoisting
 // needed.
 // FIXME: Check if the declaration has a local/anonymous type
 bool createParameters(NewFunction &ExtractedFunc,
                       const CapturedZoneInfo &CapturedInfo) {
auto SemicolonPolicy = tooling::ExtractionSemicolonPolicy::compute(
       LangOpts);
   // Update ZoneRange.
   ExtZone.ZoneRange.setEnd(FuncBodyRange.getEnd().getLocWithOffset(1));
   return SemicolonPolicy;
 }
 
 // Generate return type for ExtractedFunc. Return false if unable to do so.
 bool generateReturnProperties(NewFunction &ExtractedFunc,
+                              const FunctionDecl &EnclosingFunc,
                               const CapturedZoneInfo &CapturedInfo) {
-
-  // FIXME: Use Existing Return statements (if present)
-  // FIXME: Generate new return statement if needed.
-  if (CapturedInfo.HasReturnStmt)
+  // If the selected code always returns, we preserve those return statements.
+  // The return type should be the same as the enclosing function.
+  // (Others are possible if there are conversions, but this seems clearest).
+  if (CapturedInfo.HasReturnStmt) {
+    // If the return is conditional, neither replacing the code with
+    // `extracted()` nor `return extracted()` is correct.
+    if (!CapturedInfo.AlwaysReturns)
+      return false;
+    QualType Ret = EnclosingFunc.getReturnType();
+    // Once we support members, it'd be nice to support e.g. extracting a method
+    // of Foo<T> that returns T. But it's not clear when that's safe.
+    if (Ret->isDependentType())
       return false;
-  ExtractedFunc.ReturnType = "void";
+    ExtractedFunc.ReturnType = Ret;
+    return true;
+  }
+  // FIXME: Generate new return statement if needed.
+  ExtractedFunc.ReturnType = EnclosingFunc.getParentASTContext().VoidTy;
   return true;
 }
 
 // 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
   if (CapturedInfo.BrokenControlFlow)
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    +"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();
+  ExtractedFunc.CallerReturnsValue = CapturedInfo.AlwaysReturns;
   if (!createParameters(ExtractedFunc, CapturedInfo) ||
-      !generateReturnProperties(ExtractedFunc, CapturedInfo))
+      !generateReturnProperties(ExtractedFunc, *ExtZone.EnclosingFunction,
+                                CapturedInfo))
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    +"Too complex to extract.");
   return ExtractedFunc;
 }
 
 class ExtractFunction : public Tweak {
 public:
   const char *id() const override final;

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

TEST_F(ExtractFunctionTest, FunctionTest) {
   // Ensure that end of Zone and Beginning of PostZone being adjacent doesn't
   // lead to break being included in the extraction zone.
   EXPECT_THAT(apply("for(;;) { [[int x;]]break; }"), 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"));
+  // Extract certain return
+  EXPECT_THAT(apply(" if(true) [[{ return; }]] "), HasSubstr("extracted"));
+  // Don't extract uncertain return
+  EXPECT_THAT(apply(" if(true) [[if (false) return;]] "), StartsWith("fail"));
 }
 
 TEST_F(ExtractFunctionTest, FileTest) {
   // Check all parameters are in order
   std::string ParameterCheckInput = R"cpp(
 struct Foo {
   int x;
 };
TEST_F(ExtractFunctionTest, ControlFlow) {
   // 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, ExistingReturnStatement) {
+  Context = File;
+  const char* Before = R"cpp(
+    bool lucky(int N);
+    int getNum(bool Superstitious, int Min, int Max) {
+      if (Superstitious) [[{
+        for (int I = Min; I <= Max; ++I)
+          if (lucky(I))
+            return I;
+        return -1;
+      }]] else {
+        return (Min + Max) / 2;
+      }
+    }
+  )cpp";
+  // FIXME: min/max should be by value.
+  // FIXME: avoid emitting redundant braces
+  const char* After = R"cpp(
+    bool lucky(int N);
+    int extracted(int &Min, int &Max) {
+{
+        for (int I = Min; I <= Max; ++I)
+          if (lucky(I))
+            return I;
+        return -1;
+      }
+}
+int getNum(bool Superstitious, int Min, int Max) {
+      if (Superstitious) return extracted(Min, Max); else {
+        return (Min + Max) / 2;
+      }
+    }
+  )cpp";
+  EXPECT_EQ(apply(Before), After);
+}
+
 TWEAK_TEST(RemoveUsingNamespace);
 TEST_F(RemoveUsingNamespaceTest, All) {
   std::pair<llvm::StringRef /*Input*/, llvm::StringRef /*Expected*/> Cases[] = {
       {// Remove all occurrences of ns. Qualify only unqualified.
        R"cpp(
       namespace ns1 { struct vector {}; }
       namespace ns2 { struct map {}; }
       using namespace n^s1;