Index: clang-tools-extra/trunk/clangd/ClangdLSPServer.cpp
===================================================================
--- clang-tools-extra/trunk/clangd/ClangdLSPServer.cpp
+++ clang-tools-extra/trunk/clangd/ClangdLSPServer.cpp
@@ -1038,10 +1038,16 @@
 void ClangdLSPServer::onSwitchSourceHeader(
     const TextDocumentIdentifier &Params,
     Callback<llvm::Optional<URIForFile>> Reply) {
-  if (auto Result = Server->switchSourceHeader(Params.uri.file()))
-    Reply(URIForFile::canonicalize(*Result, Params.uri.file()));
-  else
-    Reply(llvm::None);
+  Server->switchSourceHeader(
+      Params.uri.file(),
+      [Reply = std::move(Reply),
+       Params](llvm::Expected<llvm::Optional<clangd::Path>> Path) mutable {
+        if (!Path)
+          return Reply(Path.takeError());
+        if (*Path)
+          Reply(URIForFile::canonicalize(**Path, Params.uri.file()));
+        return Reply(llvm::None);
+      });
 }
 
 void ClangdLSPServer::onDocumentHighlight(
Index: clang-tools-extra/trunk/clangd/ClangdServer.h
===================================================================
--- clang-tools-extra/trunk/clangd/ClangdServer.h
+++ clang-tools-extra/trunk/clangd/ClangdServer.h
@@ -192,9 +192,10 @@
   void locateSymbolAt(PathRef File, Position Pos,
                       Callback<std::vector<LocatedSymbol>> CB);
 
-  /// Helper function that returns a path to the corresponding source file when
-  /// given a header file and vice versa.
-  llvm::Optional<Path> switchSourceHeader(PathRef Path);
+  /// Switch to a corresponding source file when given a header file, and vice
+  /// versa.
+  void switchSourceHeader(PathRef Path,
+                          Callback<llvm::Optional<clangd::Path>> CB);
 
   /// Get document highlights for a given position.
   void findDocumentHighlights(PathRef File, Position Pos,
Index: clang-tools-extra/trunk/clangd/ClangdServer.cpp
===================================================================
--- clang-tools-extra/trunk/clangd/ClangdServer.cpp
+++ clang-tools-extra/trunk/clangd/ClangdServer.cpp
@@ -449,8 +449,24 @@
   WorkScheduler.runWithAST("Definitions", File, std::move(Action));
 }
 
-llvm::Optional<Path> ClangdServer::switchSourceHeader(PathRef Path) {
-  return getCorrespondingHeaderOrSource(Path, FSProvider.getFileSystem());
+void ClangdServer::switchSourceHeader(
+    PathRef Path, Callback<llvm::Optional<clangd::Path>> CB) {
+  // We want to return the result as fast as possible, stragety is:
+  //  1) use the file-only heuristic, it requires some IO but it is much
+  //     faster than building AST, but it only works when .h/.cc files are in
+  //     the same directory.
+  //  2) if 1) fails, we use the AST&Index approach, it is slower but supports
+  //     different code layout.
+  if (auto CorrespondingFile =
+          getCorrespondingHeaderOrSource(Path, FSProvider.getFileSystem()))
+    return CB(std::move(CorrespondingFile));
+  auto Action = [Path, CB = std::move(CB),
+                 this](llvm::Expected<InputsAndAST> InpAST) mutable {
+    if (!InpAST)
+      return CB(InpAST.takeError());
+    CB(getCorrespondingHeaderOrSource(Path, InpAST->AST, Index));
+  };
+  WorkScheduler.runWithAST("SwitchHeaderSource", Path, std::move(Action));
 }
 
 llvm::Expected<tooling::Replacements>
Index: clang-tools-extra/trunk/clangd/unittests/HeaderSourceSwitchTests.cpp
===================================================================
--- clang-tools-extra/trunk/clangd/unittests/HeaderSourceSwitchTests.cpp
+++ clang-tools-extra/trunk/clangd/unittests/HeaderSourceSwitchTests.cpp
@@ -8,6 +8,7 @@
 
 #include "HeaderSourceSwitch.h"
 
+#include "SyncAPI.h"
 #include "TestFS.h"
 #include "TestTU.h"
 #include "index/MemIndex.h"
@@ -240,6 +241,32 @@
   }
 }
 
+TEST(HeaderSourceSwitchTest, ClangdServerIntegration) {
+  class IgnoreDiagnostics : public DiagnosticsConsumer {
+    void onDiagnosticsReady(PathRef File,
+                            std::vector<Diag> Diagnostics) override {}
+  } DiagConsumer;
+  MockCompilationDatabase CDB;
+  CDB.ExtraClangFlags = {"-I" +
+                         testPath("src/include")}; // add search directory.
+  MockFSProvider FS;
+  // File heuristic fails here, we rely on the index to find the .h file.
+  std::string CppPath = testPath("src/lib/test.cpp");
+  std::string HeaderPath = testPath("src/include/test.h");
+  FS.Files[HeaderPath] = "void foo();";
+  const std::string FileContent = R"cpp(
+    #include "test.h"
+    void foo() {};
+  )cpp";
+  FS.Files[CppPath] = FileContent;
+  auto Options = ClangdServer::optsForTest();
+  Options.BuildDynamicSymbolIndex = true;
+  ClangdServer Server(CDB, FS, DiagConsumer, Options);
+  runAddDocument(Server, CppPath, FileContent);
+  EXPECT_EQ(HeaderPath,
+            *llvm::cantFail(runSwitchHeaderSource(Server, CppPath)));
+}
+
 } // namespace
 } // namespace clangd
 } // namespace clang
Index: clang-tools-extra/trunk/clangd/unittests/SyncAPI.h
===================================================================
--- clang-tools-extra/trunk/clangd/unittests/SyncAPI.h
+++ clang-tools-extra/trunk/clangd/unittests/SyncAPI.h
@@ -56,6 +56,9 @@
 llvm::Expected<std::vector<Range>>
 runSemanticRanges(ClangdServer &Server, PathRef File, Position Pos);
 
+llvm::Expected<llvm::Optional<clangd::Path>>
+runSwitchHeaderSource(ClangdServer &Server, PathRef File);
+
 } // namespace clangd
 } // namespace clang
 
Index: clang-tools-extra/trunk/clangd/unittests/SyncAPI.cpp
===================================================================
--- clang-tools-extra/trunk/clangd/unittests/SyncAPI.cpp
+++ clang-tools-extra/trunk/clangd/unittests/SyncAPI.cpp
@@ -152,5 +152,12 @@
   return std::move(*Result);
 }
 
+llvm::Expected<llvm::Optional<clangd::Path>>
+runSwitchHeaderSource(ClangdServer &Server, PathRef File) {
+  llvm::Optional<llvm::Expected<llvm::Optional<clangd::Path>>> Result;
+  Server.switchSourceHeader(File, capture(Result));
+  return std::move(*Result);
+}
+
 } // namespace clangd
 } // namespace clang