Fix use-after-free in Clang's ASTSelection unittest

llvm-svn: 311664

1 files changed, 271 insertions, 258 deletions

diff --git a/clang/unittests/Tooling/ASTSelectionTest.cpp b/clang/unittests/Tooling/ASTSelectionTest.cpp
index b4e8bc042b8..95326dded88 100644
--- a/clang/unittests/Tooling/ASTSelectionTest.cpp
+++ b/clang/unittests/Tooling/ASTSelectionTest.cpp
@@ -29,13 +29,14 @@ using FileRange = std::pair<FileLocation, FileLocation>;
 class SelectionFinderVisitor : public TestVisitor<SelectionFinderVisitor> {
   FileLocation Location;
   Optional<FileRange> SelectionRange;
+  llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer;
 
 public:
-  Optional<SelectedASTNode> Selection;
-
-  SelectionFinderVisitor(FileLocation Location,
-                         Optional<FileRange> SelectionRange)
-      : Location(Location), SelectionRange(SelectionRange) {}
+  SelectionFinderVisitor(
+      FileLocation Location, Optional<FileRange> SelectionRange,
+      llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer)
+      : Location(Location), SelectionRange(SelectionRange), Consumer(Consumer) {
+  }
 
   bool VisitTranslationUnitDecl(const TranslationUnitDecl *TU) {
     const ASTContext &Context = TU->getASTContext();
@@ -49,19 +50,18 @@ public:
       SourceLocation Loc = Location.translate(SM);
       SelRange = SourceRange(Loc, Loc);
     }
-    Selection = findSelectedASTNodes(Context, SelRange);
+    Consumer(findSelectedASTNodes(Context, SelRange));
     return false;
   }
 };
 
-Optional<SelectedASTNode>
-findSelectedASTNodes(StringRef Source, FileLocation Location,
-                     Optional<FileRange> SelectionRange,
-                     SelectionFinderVisitor::Language Language =
-                         SelectionFinderVisitor::Lang_CXX11) {
-  SelectionFinderVisitor Visitor(Location, SelectionRange);
+void findSelectedASTNodes(
+    StringRef Source, FileLocation Location, Optional<FileRange> SelectionRange,
+    llvm::function_ref<void(Optional<SelectedASTNode>)> Consumer,
+    SelectionFinderVisitor::Language Language =
+        SelectionFinderVisitor::Lang_CXX11) {
+  SelectionFinderVisitor Visitor(Location, SelectionRange, Consumer);
   EXPECT_TRUE(Visitor.runOver(Source, Language));
-  return std::move(Visitor.Selection);
 }
 
 void checkNodeImpl(bool IsTypeMatched, const SelectedASTNode &Node,
@@ -125,123 +125,126 @@ ForAllChildrenOf allChildrenOf(const SelectedASTNode &Node) {
 }
 
 TEST(ASTSelectionFinder, CursorNoSelection) {
-  Optional<SelectedASTNode> Node =
-      findSelectedASTNodes(" void f() { }", {1, 1}, None);
-  EXPECT_FALSE(Node);
+  findSelectedASTNodes(
+      " void f() { }", {1, 1}, None,
+      [](Optional<SelectedASTNode> Node) { EXPECT_FALSE(Node); });
 }
 
 TEST(ASTSelectionFinder, CursorAtStartOfFunction) {
-  Optional<SelectedASTNode> Node =
-      findSelectedASTNodes("void f() { }", {1, 1}, None);
-  EXPECT_TRUE(Node);
-  checkNode<TranslationUnitDecl>(*Node, SourceSelectionKind::None,
-                                 /*NumChildren=*/1);
-  checkNode<FunctionDecl>(Node->Children[0],
-                          SourceSelectionKind::ContainsSelection,
-                          /*NumChildren=*/0, /*Name=*/"f");
-
-  // Check that the dumping works.
-  std::string DumpValue;
-  llvm::raw_string_ostream OS(DumpValue);
-  Node->Children[0].dump(OS);
-  ASSERT_EQ(OS.str(), "FunctionDecl \"f\" contains-selection\n");
+  findSelectedASTNodes(
+      "void f() { }", {1, 1}, None, [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        checkNode<TranslationUnitDecl>(*Node, SourceSelectionKind::None,
+                                       /*NumChildren=*/1);
+        checkNode<FunctionDecl>(Node->Children[0],
+                                SourceSelectionKind::ContainsSelection,
+                                /*NumChildren=*/0, /*Name=*/"f");
+
+        // Check that the dumping works.
+        std::string DumpValue;
+        llvm::raw_string_ostream OS(DumpValue);
+        Node->Children[0].dump(OS);
+        ASSERT_EQ(OS.str(), "FunctionDecl \"f\" contains-selection\n");
+      });
 }
 
 TEST(ASTSelectionFinder, RangeNoSelection) {
-  {
-    Optional<SelectedASTNode> Node = findSelectedASTNodes(
-        " void f() { }", {1, 1}, FileRange{{1, 1}, {1, 1}});
-    EXPECT_FALSE(Node);
-  }
-  {
-    Optional<SelectedASTNode> Node = findSelectedASTNodes(
-        "  void f() { }", {1, 1}, FileRange{{1, 1}, {1, 2}});
-    EXPECT_FALSE(Node);
-  }
+  findSelectedASTNodes(
+      " void f() { }", {1, 1}, FileRange{{1, 1}, {1, 1}},
+      [](Optional<SelectedASTNode> Node) { EXPECT_FALSE(Node); });
+  findSelectedASTNodes(
+      "  void f() { }", {1, 1}, FileRange{{1, 1}, {1, 2}},
+      [](Optional<SelectedASTNode> Node) { EXPECT_FALSE(Node); });
 }
 
 TEST(ASTSelectionFinder, EmptyRangeFallbackToCursor) {
-  Optional<SelectedASTNode> Node =
-      findSelectedASTNodes("void f() { }", {1, 1}, FileRange{{1, 1}, {1, 1}});
-  EXPECT_TRUE(Node);
-  checkNode<FunctionDecl>(Node->Children[0],
-                          SourceSelectionKind::ContainsSelection,
-                          /*NumChildren=*/0, /*Name=*/"f");
+  findSelectedASTNodes("void f() { }", {1, 1}, FileRange{{1, 1}, {1, 1}},
+                       [](Optional<SelectedASTNode> Node) {
+                         EXPECT_TRUE(Node);
+                         checkNode<FunctionDecl>(
+                             Node->Children[0],
+                             SourceSelectionKind::ContainsSelection,
+                             /*NumChildren=*/0, /*Name=*/"f");
+                       });
 }
 
 TEST(ASTSelectionFinder, WholeFunctionSelection) {
   StringRef Source = "int f(int x) { return x;\n}\nvoid f2() { }";
   // From 'int' until just after '}':
-  {
-    auto Node = findSelectedASTNodes(Source, {1, 1}, FileRange{{1, 1}, {2, 2}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/2, /*Name=*/"f");
-    checkNode<ParmVarDecl>(Fn.Children[0],
-                           SourceSelectionKind::InsideSelection);
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[1], SourceSelectionKind::InsideSelection,
-        /*NumChildren=*/1);
-    const auto &Return = checkNode<ReturnStmt>(
-        Body.Children[0], SourceSelectionKind::InsideSelection,
-        /*NumChildren=*/1);
-    checkNode<ImplicitCastExpr>(Return.Children[0],
-                                SourceSelectionKind::InsideSelection,
-                                /*NumChildren=*/1);
-    checkNode<DeclRefExpr>(Return.Children[0].Children[0],
-                           SourceSelectionKind::InsideSelection);
-  }
+
+  findSelectedASTNodes(
+      Source, {1, 1}, FileRange{{1, 1}, {2, 2}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/2, /*Name=*/"f");
+        checkNode<ParmVarDecl>(Fn.Children[0],
+                               SourceSelectionKind::InsideSelection);
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[1], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/1);
+        const auto &Return = checkNode<ReturnStmt>(
+            Body.Children[0], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/1);
+        checkNode<ImplicitCastExpr>(Return.Children[0],
+                                    SourceSelectionKind::InsideSelection,
+                                    /*NumChildren=*/1);
+        checkNode<DeclRefExpr>(Return.Children[0].Children[0],
+                               SourceSelectionKind::InsideSelection);
+      });
+
   // From 'int' until just before '}':
-  {
-    auto Node = findSelectedASTNodes(Source, {2, 1}, FileRange{{1, 1}, {2, 1}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/2, /*Name=*/"f");
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[1], SourceSelectionKind::ContainsSelectionEnd,
-        /*NumChildren=*/1);
-    checkNode<ReturnStmt>(Body.Children[0],
-                          SourceSelectionKind::InsideSelection,
-                          /*NumChildren=*/1);
-  }
+  findSelectedASTNodes(
+      Source, {2, 1}, FileRange{{1, 1}, {2, 1}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/2, /*Name=*/"f");
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[1], SourceSelectionKind::ContainsSelectionEnd,
+            /*NumChildren=*/1);
+        checkNode<ReturnStmt>(Body.Children[0],
+                              SourceSelectionKind::InsideSelection,
+                              /*NumChildren=*/1);
+      });
   // From '{' until just after '}':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {1, 14}, FileRange{{1, 14}, {2, 2}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"f");
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1);
-    checkNode<ReturnStmt>(Body.Children[0],
-                          SourceSelectionKind::InsideSelection,
-                          /*NumChildren=*/1);
-  }
+  findSelectedASTNodes(
+      Source, {1, 14}, FileRange{{1, 14}, {2, 2}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"f");
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1);
+        checkNode<ReturnStmt>(Body.Children[0],
+                              SourceSelectionKind::InsideSelection,
+                              /*NumChildren=*/1);
+      });
   // From 'x' until just after '}':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {2, 2}, FileRange{{1, 11}, {2, 2}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/2, /*Name=*/"f");
-    checkNode<ParmVarDecl>(Fn.Children[0],
-                           SourceSelectionKind::ContainsSelectionStart);
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[1], SourceSelectionKind::InsideSelection,
-        /*NumChildren=*/1);
-    checkNode<ReturnStmt>(Body.Children[0],
-                          SourceSelectionKind::InsideSelection,
-                          /*NumChildren=*/1);
-  }
+  findSelectedASTNodes(
+      Source, {2, 2}, FileRange{{1, 11}, {2, 2}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/2, /*Name=*/"f");
+        checkNode<ParmVarDecl>(Fn.Children[0],
+                               SourceSelectionKind::ContainsSelectionStart);
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[1], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/1);
+        checkNode<ReturnStmt>(Body.Children[0],
+                              SourceSelectionKind::InsideSelection,
+                              /*NumChildren=*/1);
+      });
 }
 
 TEST(ASTSelectionFinder, MultipleFunctionSelection) {
@@ -262,10 +265,10 @@ void f3() { }
                             /*NumChildren=*/1, /*Name=*/"f2");
   };
   // Just after '}' of f0 and just before 'void' of f3:
-  SelectedF1F2(findSelectedASTNodes(Source, {2, 2}, FileRange{{2, 2}, {5, 1}}));
+  findSelectedASTNodes(Source, {2, 2}, FileRange{{2, 2}, {5, 1}}, SelectedF1F2);
   // Just before 'void' of f1 and just after '}' of f2:
-  SelectedF1F2(
-      findSelectedASTNodes(Source, {3, 1}, FileRange{{3, 1}, {4, 14}}));
+  findSelectedASTNodes(Source, {3, 1}, FileRange{{3, 1}, {4, 14}},
+                       SelectedF1F2);
 }
 
 TEST(ASTSelectionFinder, MultipleStatementSelection) {
@@ -279,67 +282,73 @@ TEST(ASTSelectionFinder, MultipleStatementSelection) {
   return;
 })";
   // From 'f(2,3)' until just before 'x = 1;':
-  {
-    auto Node = findSelectedASTNodes(Source, {3, 2}, FileRange{{3, 2}, {7, 1}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"f");
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/2);
-    allChildrenOf(checkNode<CallExpr>(Body.Children[0],
-                                      SourceSelectionKind::InsideSelection,
-                                      /*NumChildren=*/3))
-        .shouldHaveSelectionKind(SourceSelectionKind::InsideSelection);
-    allChildrenOf(checkNode<IfStmt>(Body.Children[1],
-                                    SourceSelectionKind::InsideSelection,
-                                    /*NumChildren=*/2))
-        .shouldHaveSelectionKind(SourceSelectionKind::InsideSelection);
-  }
+  findSelectedASTNodes(
+      Source, {3, 2}, FileRange{{3, 2}, {7, 1}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"f");
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/2);
+        allChildrenOf(checkNode<CallExpr>(Body.Children[0],
+                                          SourceSelectionKind::InsideSelection,
+                                          /*NumChildren=*/3))
+            .shouldHaveSelectionKind(SourceSelectionKind::InsideSelection);
+        allChildrenOf(checkNode<IfStmt>(Body.Children[1],
+                                        SourceSelectionKind::InsideSelection,
+                                        /*NumChildren=*/2))
+            .shouldHaveSelectionKind(SourceSelectionKind::InsideSelection);
+      });
   // From 'f(2,3)' until just before ';' in 'x = 1;':
-  {
-    auto Node = findSelectedASTNodes(Source, {3, 2}, FileRange{{3, 2}, {7, 8}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"f");
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/3);
-    checkNode<CallExpr>(Body.Children[0], SourceSelectionKind::InsideSelection,
-                        /*NumChildren=*/3);
-    checkNode<IfStmt>(Body.Children[1], SourceSelectionKind::InsideSelection,
-                      /*NumChildren=*/2);
-    checkNode<BinaryOperator>(Body.Children[2],
-                              SourceSelectionKind::InsideSelection,
-                              /*NumChildren=*/2);
-  }
+  findSelectedASTNodes(
+      Source, {3, 2}, FileRange{{3, 2}, {7, 8}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"f");
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/3);
+        checkNode<CallExpr>(Body.Children[0],
+                            SourceSelectionKind::InsideSelection,
+                            /*NumChildren=*/3);
+        checkNode<IfStmt>(Body.Children[1],
+                          SourceSelectionKind::InsideSelection,
+                          /*NumChildren=*/2);
+        checkNode<BinaryOperator>(Body.Children[2],
+                                  SourceSelectionKind::InsideSelection,
+                                  /*NumChildren=*/2);
+      });
   // From the middle of 'int z = 3' until the middle of 'x = 1;':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {2, 10}, FileRange{{2, 10}, {7, 5}});
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Fn = checkNode<FunctionDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"f");
-    const auto &Body = checkNode<CompoundStmt>(
-        Fn.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/4);
-    checkNode<DeclStmt>(Body.Children[0],
-                        SourceSelectionKind::ContainsSelectionStart,
-                        /*NumChildren=*/1);
-    checkNode<CallExpr>(Body.Children[1], SourceSelectionKind::InsideSelection,
-                        /*NumChildren=*/3);
-    checkNode<IfStmt>(Body.Children[2], SourceSelectionKind::InsideSelection,
-                      /*NumChildren=*/2);
-    checkNode<BinaryOperator>(Body.Children[3],
-                              SourceSelectionKind::ContainsSelectionEnd,
-                              /*NumChildren=*/1);
-  }
+  findSelectedASTNodes(
+      Source, {2, 10}, FileRange{{2, 10}, {7, 5}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Fn = checkNode<FunctionDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"f");
+        const auto &Body = checkNode<CompoundStmt>(
+            Fn.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/4);
+        checkNode<DeclStmt>(Body.Children[0],
+                            SourceSelectionKind::ContainsSelectionStart,
+                            /*NumChildren=*/1);
+        checkNode<CallExpr>(Body.Children[1],
+                            SourceSelectionKind::InsideSelection,
+                            /*NumChildren=*/3);
+        checkNode<IfStmt>(Body.Children[2],
+                          SourceSelectionKind::InsideSelection,
+                          /*NumChildren=*/2);
+        checkNode<BinaryOperator>(Body.Children[3],
+                                  SourceSelectionKind::ContainsSelectionEnd,
+                                  /*NumChildren=*/1);
+      });
 }
 
 TEST(ASTSelectionFinder, SelectionInFunctionInObjCImplementation) {
@@ -364,72 +373,72 @@ void catF() { }
 void outerFunction() { }
 )";
   // Just the 'x' expression in 'selected':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {9, 10}, FileRange{{9, 10}, {9, 11}},
-                             SelectionFinderVisitor::Lang_OBJC);
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Impl = checkNode<ObjCImplementationDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"I");
-    const auto &Fn = checkNode<FunctionDecl>(
-        Impl.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"selected");
-    allChildrenOf(Fn).shouldHaveSelectionKind(
-        SourceSelectionKind::ContainsSelection);
-  }
+  findSelectedASTNodes(
+      Source, {9, 10}, FileRange{{9, 10}, {9, 11}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Impl = checkNode<ObjCImplementationDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"I");
+        const auto &Fn = checkNode<FunctionDecl>(
+            Impl.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"selected");
+        allChildrenOf(Fn).shouldHaveSelectionKind(
+            SourceSelectionKind::ContainsSelection);
+      },
+      SelectionFinderVisitor::Lang_OBJC);
   // The entire 'catF':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {15, 1}, FileRange{{15, 1}, {15, 16}},
-                             SelectionFinderVisitor::Lang_OBJC);
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Impl = checkNode<ObjCCategoryImplDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"Cat");
-    const auto &Fn = checkNode<FunctionDecl>(
-        Impl.Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"catF");
-    allChildrenOf(Fn).shouldHaveSelectionKind(
-        SourceSelectionKind::ContainsSelection);
-  }
+  findSelectedASTNodes(
+      Source, {15, 1}, FileRange{{15, 1}, {15, 16}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Impl = checkNode<ObjCCategoryImplDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"Cat");
+        const auto &Fn = checkNode<FunctionDecl>(
+            Impl.Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"catF");
+        allChildrenOf(Fn).shouldHaveSelectionKind(
+            SourceSelectionKind::ContainsSelection);
+      },
+      SelectionFinderVisitor::Lang_OBJC);
   // From the line before 'selected' to the line after 'catF':
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {16, 1}, FileRange{{7, 1}, {16, 1}},
-                             SelectionFinderVisitor::Lang_OBJC);
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 2u);
-    const auto &Impl = checkNode<ObjCImplementationDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelectionStart,
-        /*NumChildren=*/1, /*Name=*/"I");
-    const auto &Selected = checkNode<FunctionDecl>(
-        Impl.Children[0], SourceSelectionKind::InsideSelection,
-        /*NumChildren=*/2, /*Name=*/"selected");
-    allChildrenOf(Selected).shouldHaveSelectionKind(
-        SourceSelectionKind::InsideSelection);
-    const auto &Cat = checkNode<ObjCCategoryImplDecl>(
-        Node->Children[1], SourceSelectionKind::ContainsSelectionEnd,
-        /*NumChildren=*/1, /*Name=*/"Cat");
-    const auto &CatF = checkNode<FunctionDecl>(
-        Cat.Children[0], SourceSelectionKind::InsideSelection,
-        /*NumChildren=*/1, /*Name=*/"catF");
-    allChildrenOf(CatF).shouldHaveSelectionKind(
-        SourceSelectionKind::InsideSelection);
-  }
+  findSelectedASTNodes(
+      Source, {16, 1}, FileRange{{7, 1}, {16, 1}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 2u);
+        const auto &Impl = checkNode<ObjCImplementationDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelectionStart,
+            /*NumChildren=*/1, /*Name=*/"I");
+        const auto &Selected = checkNode<FunctionDecl>(
+            Impl.Children[0], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/2, /*Name=*/"selected");
+        allChildrenOf(Selected).shouldHaveSelectionKind(
+            SourceSelectionKind::InsideSelection);
+        const auto &Cat = checkNode<ObjCCategoryImplDecl>(
+            Node->Children[1], SourceSelectionKind::ContainsSelectionEnd,
+            /*NumChildren=*/1, /*Name=*/"Cat");
+        const auto &CatF = checkNode<FunctionDecl>(
+            Cat.Children[0], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/1, /*Name=*/"catF");
+        allChildrenOf(CatF).shouldHaveSelectionKind(
+            SourceSelectionKind::InsideSelection);
+      },
+      SelectionFinderVisitor::Lang_OBJC);
   // Just the 'outer' function:
-  {
-    auto Node =
-        findSelectedASTNodes(Source, {19, 1}, FileRange{{19, 1}, {19, 25}},
-                             SelectionFinderVisitor::Lang_OBJC);
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    checkNode<FunctionDecl>(Node->Children[0],
-                            SourceSelectionKind::ContainsSelection,
-                            /*NumChildren=*/1, /*Name=*/"outerFunction");
-  }
+  findSelectedASTNodes(Source, {19, 1}, FileRange{{19, 1}, {19, 25}},
+                       [](Optional<SelectedASTNode> Node) {
+                         EXPECT_TRUE(Node);
+                         EXPECT_EQ(Node->Children.size(), 1u);
+                         checkNode<FunctionDecl>(
+                             Node->Children[0],
+                             SourceSelectionKind::ContainsSelection,
+                             /*NumChildren=*/1, /*Name=*/"outerFunction");
+                       },
+                       SelectionFinderVisitor::Lang_OBJC);
 }
 
 TEST(ASTSelectionFinder, FunctionInObjCImplementationCarefulWithEarlyExit) {
@@ -446,18 +455,19 @@ void selected() {
 @end
 )";
   // Just 'selected'
-  {
-    auto Node = findSelectedASTNodes(Source, {6, 1}, FileRange{{6, 1}, {7, 2}},
-                                     SelectionFinderVisitor::Lang_OBJC);
-    EXPECT_TRUE(Node);
-    EXPECT_EQ(Node->Children.size(), 1u);
-    const auto &Impl = checkNode<ObjCImplementationDecl>(
-        Node->Children[0], SourceSelectionKind::ContainsSelection,
-        /*NumChildren=*/1, /*Name=*/"I");
-    checkNode<FunctionDecl>(Impl.Children[0],
-                            SourceSelectionKind::ContainsSelection,
-                            /*NumChildren=*/1, /*Name=*/"selected");
-  }
+  findSelectedASTNodes(
+      Source, {6, 1}, FileRange{{6, 1}, {7, 2}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Impl = checkNode<ObjCImplementationDecl>(
+            Node->Children[0], SourceSelectionKind::ContainsSelection,
+            /*NumChildren=*/1, /*Name=*/"I");
+        checkNode<FunctionDecl>(Impl.Children[0],
+                                SourceSelectionKind::ContainsSelection,
+                                /*NumChildren=*/1, /*Name=*/"selected");
+      },
+      SelectionFinderVisitor::Lang_OBJC);
 }
 
 TEST(ASTSelectionFinder, AvoidImplicitDeclarations) {
@@ -471,14 +481,17 @@ void foo() {
 }
 )";
   // The entire struct 'Copy':
-  auto Node = findSelectedASTNodes(Source, {2, 1}, FileRange{{2, 1}, {4, 3}});
-  EXPECT_TRUE(Node);
-  EXPECT_EQ(Node->Children.size(), 1u);
-  const auto &Record = checkNode<CXXRecordDecl>(
-      Node->Children[0], SourceSelectionKind::InsideSelection,
-      /*NumChildren=*/1, /*Name=*/"Copy");
-  checkNode<FieldDecl>(Record.Children[0],
-                       SourceSelectionKind::InsideSelection);
+  findSelectedASTNodes(
+      Source, {2, 1}, FileRange{{2, 1}, {4, 3}},
+      [](Optional<SelectedASTNode> Node) {
+        EXPECT_TRUE(Node);
+        EXPECT_EQ(Node->Children.size(), 1u);
+        const auto &Record = checkNode<CXXRecordDecl>(
+            Node->Children[0], SourceSelectionKind::InsideSelection,
+            /*NumChildren=*/1, /*Name=*/"Copy");
+        checkNode<FieldDecl>(Record.Children[0],
+                             SourceSelectionKind::InsideSelection);
+      });
 }
 
 } // end anonymous namespace