5 files changed, 57 insertions, 0 deletions
diff --git a/clang-tools-extra/clangd/ClangdServer.cpp b/clang-tools-extra/clangd/ClangdServer.cpp
index dbb3ce49306..e73834be32d 100644
--- a/clang-tools-extra/clangd/ClangdServer.cpp
+++ b/clang-tools-extra/clangd/ClangdServer.cpp
@@ -286,6 +286,13 @@ static llvm::Expected<HeaderFile> toHeaderFile(StringRef Header,
   auto U = URI::parse(Header);
   if (!U)
     return U.takeError();
+
+  auto IncludePath = URI::includeSpelling(*U);
+  if (!IncludePath)
+    return IncludePath.takeError();
+  if (!IncludePath->empty())
+    return HeaderFile{std::move(*IncludePath), /*Verbatim=*/true};
+
   auto Resolved = URI::resolve(*U, HintPath);
   if (!Resolved)
     return Resolved.takeError();
diff --git a/clang-tools-extra/clangd/URI.cpp b/clang-tools-extra/clangd/URI.cpp
index c0e10cef173..a722640890d 100644
--- a/clang-tools-extra/clangd/URI.cpp
+++ b/clang-tools-extra/clangd/URI.cpp
@@ -196,5 +196,12 @@ llvm::Expected<std::string> URI::resolve(const URI &Uri,
   return S->get()->getAbsolutePath(Uri.Authority, Uri.Body, HintPath);
 }
 
+llvm::Expected<std::string> URI::includeSpelling(const URI &Uri) {
+  auto S = findSchemeByName(Uri.Scheme);
+  if (!S)
+    return S.takeError();
+  return S->get()->getIncludeSpelling(Uri);
+}
+
 } // namespace clangd
 } // namespace clang
diff --git a/clang-tools-extra/clangd/URI.h b/clang-tools-extra/clangd/URI.h
index 51b23e37da0..ff4bc2d31fb 100644
--- a/clang-tools-extra/clangd/URI.h
+++ b/clang-tools-extra/clangd/URI.h
@@ -60,6 +60,16 @@ public:
   static llvm::Expected<std::string> resolve(const URI &U,
                                              llvm::StringRef HintPath = "");
 
+  /// Gets the preferred spelling of this file for #include, if there is one,
+  /// e.g. <system_header.h>, "path/to/x.h".
+  ///
+  /// This allows URI schemas to provide their customized include paths.
+  ///
+  /// Returns an empty string if normal include-shortening based on the absolute
+  /// path should be used.
+  /// Fails if the URI is not valid in the schema.
+  static llvm::Expected<std::string> includeSpelling(const URI &U);
+
   friend bool operator==(const URI &LHS, const URI &RHS) {
     return std::tie(LHS.Scheme, LHS.Authority, LHS.Body) ==
            std::tie(RHS.Scheme, RHS.Authority, RHS.Body);
@@ -94,6 +104,13 @@ public:
 
   virtual llvm::Expected<URI>
   uriFromAbsolutePath(llvm::StringRef AbsolutePath) const = 0;
+
+  /// Returns the include path of the file (e.g. <path>, "path"), which can be
+  /// #included directly. See URI::includeSpelling for details.
+  virtual llvm::Expected<std::string>
+  getIncludeSpelling(const URI& U) const {
+    return "";  // no customized include path for this scheme.
+  }
 };
 
 /// By default, a "file" scheme is supported where URI paths are always absolute
diff --git a/clang-tools-extra/unittests/clangd/ClangdTests.cpp b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
index b0e56ba78e7..4704fad3a3d 100644
--- a/clang-tools-extra/unittests/clangd/ClangdTests.cpp
+++ b/clang-tools-extra/unittests/clangd/ClangdTests.cpp
@@ -152,6 +152,28 @@ protected:
   }
 };
 
+constexpr const char* ClangdTestScheme = "ClangdTests";
+class TestURIScheme : public URIScheme {
+public:
+  llvm::Expected<std::string>
+  getAbsolutePath(llvm::StringRef /*Authority*/, llvm::StringRef Body,
+                  llvm::StringRef /*HintPath*/) const override {
+    llvm_unreachable("ClangdTests never makes absolute path.");
+  }
+
+  llvm::Expected<URI>
+  uriFromAbsolutePath(llvm::StringRef AbsolutePath) const override {
+    llvm_unreachable("ClangdTest never creates a test URI.");
+  }
+
+  llvm::Expected<std::string> getIncludeSpelling(const URI &U) const override {
+    return ("\"" + U.body().trim("/") + "\"").str();
+  }
+};
+
+static URISchemeRegistry::Add<TestURIScheme>
+    X(ClangdTestScheme, "Test scheme for ClangdTests.");
+
 TEST_F(ClangdVFSTest, Parse) {
   // FIXME: figure out a stable format for AST dumps, so that we can check the
   // output of the dump itself is equal to the expected one, not just that it's
@@ -961,6 +983,10 @@ void f() {}
                        /*Preferred=*/"<Y.h>", "<Y.h>"));
   EXPECT_TRUE(Inserted(OriginalHeader, PreferredHeader, "\"Y.h\""));
   EXPECT_TRUE(Inserted("<y.h>", PreferredHeader, "\"Y.h\""));
+  auto TestURIHeader =
+      URI::parse(llvm::formatv("{0}:///x/y/z.h", ClangdTestScheme).str());
+  EXPECT_TRUE(static_cast<bool>(TestURIHeader));
+  EXPECT_TRUE(Inserted(TestURIHeader->toString(), "", "\"x/y/z.h\""));
 
   // Check that includes are sorted.
   const auto Expected = R"cpp(
diff --git a/clang-tools-extra/unittests/clangd/TestScheme.h b/clang-tools-extra/unittests/clangd/TestScheme.h
new file mode 100644
index 00000000000..e69de29bb2d
--- /dev/null
+++ b/clang-tools-extra/unittests/clangd/TestScheme.h