[mlir][pdll] Add extra-dirs for LSP includes.

Enable specifying additional include directories to search. This is
consistent with what one can do with clangd (although there it is more
general compilation options) and Python LSP. We would in general expect
these to be provided by compilation database equivalent.

Differential Revision: https://reviews.llvm.org/D123474

diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/MlirPdllLspServerMain.cpp b/mlir/lib/Tools/mlir-pdll-lsp-server/MlirPdllLspServerMain.cpp
index b21786de918b76920bc959db57a08262a8c5d430..c83be8e5f8f137a1b1ff4e8914ed2eb214db1052 100644 (file)
--- a/mlir/lib/Tools/mlir-pdll-lsp-server/MlirPdllLspServerMain.cpp
+++ b/mlir/lib/Tools/mlir-pdll-lsp-server/MlirPdllLspServerMain.cpp
@@ -51,7 +51,11 @@ LogicalResult mlir::MlirPdllLspServerMain(int argc, char **argv) {
       llvm::cl::desc("Pretty-print JSON output"),
       llvm::cl::init(false),
   };
-  llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR LSP Language Server");
+  llvm::cl::list<std::string> extraIncludeDirs(
+      "I", llvm::cl::desc("Extra directory of include files"),
+      llvm::cl::value_desc("directory"), llvm::cl::Prefix);
+
+  llvm::cl::ParseCommandLineOptions(argc, argv, "PDLL LSP Language Server");
 
   if (litTest) {
     inputStyle = JSONStreamStyle::Delimited;
@@ -67,6 +71,7 @@ LogicalResult mlir::MlirPdllLspServerMain(int argc, char **argv) {
   JSONTransport transport(stdin, llvm::outs(), inputStyle, prettyPrint);
 
   // Configure the servers and start the main language server.
-  PDLLServer server;
+  PDLLServer::Options options(extraIncludeDirs);
+  PDLLServer server(options);
   return runPdllLSPServer(server, transport);
 }

diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp
index 4a8f85a6e02f299def3d404ad43c78d560762c85..8ed0bc2f87f838b376c24c01c789d3e6d469487f 100644 (file)
--- a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp
+++ b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp
@@ -238,6 +238,7 @@ namespace {
 /// document.
 struct PDLDocument {
   PDLDocument(const lsp::URIForFile &uri, StringRef contents,
+              const std::vector<std::string> &extraDirs,
               std::vector<lsp::Diagnostic> &diagnostics);
   PDLDocument(const PDLDocument &) = delete;
   PDLDocument &operator=(const PDLDocument &) = delete;
@@ -315,6 +316,7 @@ struct PDLDocument {
 } // namespace
 
 PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
+                         const std::vector<std::string> &extraDirs,
                          std::vector<lsp::Diagnostic> &diagnostics)
     : astContext(odsContext) {
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
@@ -327,6 +329,7 @@ PDLDocument::PDLDocument(const lsp::URIForFile &uri, StringRef contents,
   llvm::SmallString<32> uriDirectory(uri.file());
   llvm::sys::path::remove_filename(uriDirectory);
   includeDirs.push_back(uriDirectory.str().str());
+  includeDirs.insert(includeDirs.end(), extraDirs.begin(), extraDirs.end());
 
   sourceMgr.setIncludeDirs(includeDirs);
   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
@@ -991,8 +994,10 @@ namespace {
 struct PDLTextFileChunk {
   PDLTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
                    StringRef contents,
+                   const std::vector<std::string> &extraDirs,
                    std::vector<lsp::Diagnostic> &diagnostics)
-      : lineOffset(lineOffset), document(uri, contents, diagnostics) {}
+      : lineOffset(lineOffset),
+        document(uri, contents, extraDirs, diagnostics) {}
 
   /// Adjust the line number of the given range to anchor at the beginning of
   /// the file, instead of the beginning of this chunk.
@@ -1020,7 +1025,8 @@ namespace {
 class PDLTextFile {
 public:
   PDLTextFile(const lsp::URIForFile &uri, StringRef fileContents,
-              int64_t version, std::vector<lsp::Diagnostic> &diagnostics);
+              int64_t version, const std::vector<std::string> &extraDirs,
+              std::vector<lsp::Diagnostic> &diagnostics);
 
   /// Return the current version of this text file.
   int64_t getVersion() const { return version; }
@@ -1064,6 +1070,7 @@ private:
 
 PDLTextFile::PDLTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                          int64_t version,
+                         const std::vector<std::string> &extraDirs,
                          std::vector<lsp::Diagnostic> &diagnostics)
     : contents(fileContents.str()), version(version), totalNumLines(0) {
   // Split the file into separate PDL documents.
@@ -1073,13 +1080,13 @@ PDLTextFile::PDLTextFile(const lsp::URIForFile &uri, StringRef fileContents,
   SmallVector<StringRef, 8> subContents;
   StringRef(contents).split(subContents, "// -----");
   chunks.emplace_back(std::make_unique<PDLTextFileChunk>(
-      /*lineOffset=*/0, uri, subContents.front(), diagnostics));
+      /*lineOffset=*/0, uri, subContents.front(), extraDirs, diagnostics));
 
   uint64_t lineOffset = subContents.front().count('\n');
   for (StringRef docContents : llvm::drop_begin(subContents)) {
     unsigned currentNumDiags = diagnostics.size();
-    auto chunk = std::make_unique<PDLTextFileChunk>(lineOffset, uri,
-                                                    docContents, diagnostics);
+    auto chunk = std::make_unique<PDLTextFileChunk>(
+        lineOffset, uri, docContents, extraDirs, diagnostics);
     lineOffset += docContents.count('\n');
 
     // Adjust locations used in diagnostics to account for the offset from the
@@ -1218,6 +1225,11 @@ PDLTextFileChunk &PDLTextFile::getChunkFor(lsp::Position &pos) {
 //===----------------------------------------------------------------------===//
 
 struct lsp::PDLLServer::Impl {
+  explicit Impl(const Options &options) : options(options) {}
+
+  /// PDLL LSP options.
+  const Options &options;
+
   /// The files held by the server, mapped by their URI file name.
   llvm::StringMap<std::unique_ptr<PDLTextFile>> files;
 };
@@ -1226,14 +1238,15 @@ struct lsp::PDLLServer::Impl {
 // PDLLServer
 //===----------------------------------------------------------------------===//
 
-lsp::PDLLServer::PDLLServer() : impl(std::make_unique<Impl>()) {}
+lsp::PDLLServer::PDLLServer(const Options &options)
+    : impl(std::make_unique<Impl>(options)) {}
 lsp::PDLLServer::~PDLLServer() = default;
 
 void lsp::PDLLServer::addOrUpdateDocument(
     const URIForFile &uri, StringRef contents, int64_t version,
     std::vector<Diagnostic> &diagnostics) {
-  impl->files[uri.file()] =
-      std::make_unique<PDLTextFile>(uri, contents, version, diagnostics);
+  impl->files[uri.file()] = std::make_unique<PDLTextFile>(
+      uri, contents, version, impl->options.extraDirs, diagnostics);
 }
 
 Optional<int64_t> lsp::PDLLServer::removeDocument(const URIForFile &uri) {

diff --git a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h
index ab0ad48615eb14b4a3ed081ec285c3c1f1ab774b..947a3f53cb65792e3317b8b1fdf8f023198900ab 100644 (file)
--- a/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h
+++ b/mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.h
@@ -28,7 +28,14 @@ class URIForFile;
 /// separate from the logic that involves LSP server/client communication.
 class PDLLServer {
 public:
-  PDLLServer();
+  struct Options {
+    Options(const std::vector<std::string> &extraDirs) : extraDirs(extraDirs){};
+
+    /// Additional list of include directories to search.
+    const std::vector<std::string> &extraDirs;
+  };
+
+  PDLLServer(const Options &options);
   ~PDLLServer();
 
   /// Add or update the document, with the provided `version`, at the given URI.
@@ -69,7 +76,6 @@ public:
 
 private:
   struct Impl;
-
   std::unique_ptr<Impl> impl;
 };