[mlir][PDLL-LSP] Add code completion for include file paths

[rriddle]

Summary

This allows for providing completion results for include directive
file paths by searching the set of include directories for the current
file.

[jpienaar]

mlir/lib/Tools/PDLL/Parser/Parser.cpp

685

Is string sufficient differentiator? Don't we need to be inside include too?

[rriddle]

mlir/lib/Tools/PDLL/Parser/Parser.cpp

685

Can you clarify what you mean here? If we hit a completion here inside of a string, that is for the include file. This method is only called from inside an include (Parser::parseInclude).

[jpienaar]

mlir/lib/Tools/PDLL/Parser/Parser.cpp

685

The token kind

/// Token signifying a code completion location within a string.
    code_complete_string

is more general, but I missed this part was only inside parseInclude.

mlir/include/mlir/Tools/PDLL/Parser/CodeComplete.h

public:
   virtual void codeCompleteDialectName() {}
 
   /// Signal code completion for an operation name in the given dialect.
   virtual void codeCompleteOperationName(StringRef dialectName) {}
 
   /// Signal code completion for Pattern metadata.
   virtual void codeCompletePatternMetadata() {}
 
+  /// Signal code completion for an include filename.
+  virtual void codeCompleteIncludeFilename(StringRef curPath) {}
+
   //===--------------------------------------------------------------------===//
   // Signature Hooks
   //===--------------------------------------------------------------------===//
 
   /// Signal code completion for the signature of a callable.
   virtual void codeCompleteCallSignature(const ast::CallableDecl *callable,
                                          unsigned currentNumArgs) {}
 

mlir/lib/Tools/PDLL/Parser/Lexer.h

 
 //===----------------------------------------------------------------------===//
 // Token
 //===----------------------------------------------------------------------===//
 
 class Token {
 public:
   enum Kind {
-    // Markers.
+    /// Markers.
     eof,
     error,
+    /// Token signifying a code completion location.
     code_complete,
+    /// Token signifying a code completion location within a string.
+    code_complete_string,
 
-    // Keywords.
+    /// Keywords.
     KW_BEGIN,
-    // Dependent keywords, i.e. those that are treated as keywords depending on
-    // the current parser context.
+    /// Dependent keywords, i.e. those that are treated as keywords depending on
+    /// the current parser context.
     KW_DEPENDENT_BEGIN,
     kw_attr,
     kw_op,
     kw_type,
     KW_DEPENDENT_END,
 
-    // General keywords.
+    /// General keywords.
     kw_Attr,
     kw_erase,
     kw_let,
     kw_Constraint,
     kw_Op,
     kw_OpName,
     kw_Pattern,
     kw_replace,
     kw_return,
     kw_rewrite,
     kw_Rewrite,
     kw_Type,
     kw_TypeRange,
     kw_Value,
     kw_ValueRange,
     kw_with,
     KW_END,
 
-    // Punctuation.
+    /// Punctuation.
     arrow,
     colon,
     comma,
     dot,
     equal,
     equal_arrow,
     semicolon,
-    // Paired punctuation.
+    /// Paired punctuation.
     less,
     greater,
     l_brace,
     r_brace,
     l_paren,
     r_paren,
     l_square,
     r_square,
     underscore,
 
-    // Tokens.
+    /// Tokens.
     directive,
     identifier,
     integer,
     string_block,
     string
   };
   Token(Kind kind, StringRef spelling) : kind(kind), spelling(spelling) {}
 

mlir/lib/Tools/PDLL/Parser/Lexer.cpp

 using namespace mlir;
 using namespace mlir::pdll;
 
 //===----------------------------------------------------------------------===//
 // Token
 //===----------------------------------------------------------------------===//
 
 std::string Token::getStringValue() const {
-  assert(getKind() == string || getKind() == string_block);
+  assert(getKind() == string || getKind() == string_block ||
+         getKind() == code_complete_string);
 
   // Start by dropping the quotes.
-  StringRef bytes = getSpelling().drop_front().drop_back();
-  if (is(string_block)) bytes = bytes.drop_front().drop_back();
+  StringRef bytes = getSpelling();
+  if (is(string))
+    bytes = bytes.drop_front().drop_back();
+  else if (is(string_block))
+    bytes = bytes.drop_front(2).drop_back(2);
 
   std::string result;
   result.reserve(bytes.size());
   for (unsigned i = 0, e = bytes.size(); i != e;) {
     auto c = bytes[i++];
     if (c != '\\') {
       result.push_back(c);
       continue;
Token Lexer::lexNumber(const char *tokStart) {
   // Handle the normal decimal case.
   while (isdigit(*curPtr)) ++curPtr;
 
   return formToken(Token::integer, tokStart);
 }
 
 Token Lexer::lexString(const char *tokStart, bool isStringBlock) {
   while (true) {
+    // Check to see if there is a code completion location within the string. In
+    // these cases we generate a completion location and place the currently
+    // lexed string within the token (without the quotes). This allows for the
+    // parser to use the partially lexed string when computing the completion
+    // results.
+    if (curPtr == codeCompletionLocation) {
+      return formToken(Token::code_complete_string,
+                       tokStart + (isStringBlock ? 2 : 1));
+    }
+
     switch (*curPtr++) {
       case '"':
         // If this is a string block, we only end the string when we encounter a
         // `}]`.
         if (!isStringBlock) return formToken(Token::string, tokStart);
         continue;
       case '}':
         // If this is a string block, we only end the string when we encounter a

mlir/lib/Tools/PDLL/Parser/Parser.cpp

private:
   LogicalResult codeCompleteMemberAccess(ast::Expr *parentExpr);
   LogicalResult codeCompleteAttributeName(Optional<StringRef> opName);
   LogicalResult codeCompleteConstraintName(ast::Type inferredType,
                                            bool allowNonCoreConstraints,
                                            bool allowInlineTypeConstraints);
   LogicalResult codeCompleteDialectName();
   LogicalResult codeCompleteOperationName(StringRef dialectName);
   LogicalResult codeCompletePatternMetadata();
+  LogicalResult codeCompleteIncludeFilename(StringRef curPath);
 
   void codeCompleteCallSignature(ast::Node *parent, unsigned currentNumArgs);
   void codeCompleteOperationOperandsSignature(Optional<StringRef> opName,
                                               unsigned currentNumOperands);
   void codeCompleteOperationResultsSignature(Optional<StringRef> opName,
                                              unsigned currentNumResults);
 
   //===--------------------------------------------------------------------===//
LogicalResult Parser::parseDirective(SmallVectorImpl<ast::Decl *> &decls) {
 
   return emitError("unknown directive `" + directive + "`");
 }
 
 LogicalResult Parser::parseInclude(SmallVectorImpl<ast::Decl *> &decls) {
   SMRange loc = curToken.getLoc();
   consumeToken(Token::directive);
 
+  // Handle code completion of the include file path.
+  if (curToken.is(Token::code_complete_string))
+    return codeCompleteIncludeFilename(curToken.getStringValue());
+
   // Parse the file being included.
   if (!curToken.isString())
     return emitError(loc,
                      "expected string file name after `include` directive");
   SMRange fileLoc = curToken.getLoc();
   std::string filenameStr = curToken.getStringValue();
   StringRef filename = filenameStr;
   consumeToken();
LogicalResult Parser::codeCompleteOperationName(StringRef dialectName) {
   return failure();
 }
 
 LogicalResult Parser::codeCompletePatternMetadata() {
   codeCompleteContext->codeCompletePatternMetadata();
   return failure();
 }
 
+LogicalResult Parser::codeCompleteIncludeFilename(StringRef curPath) {
+  codeCompleteContext->codeCompleteIncludeFilename(curPath);
+  return failure();
+}
+
 void Parser::codeCompleteCallSignature(ast::Node *parent,
                                        unsigned currentNumArgs) {
   ast::CallableDecl *callableDecl = tryExtractCallableDecl(parent);
   if (!callableDecl)
     return;
 
   codeCompleteContext->codeCompleteCallSignature(callableDecl, currentNumArgs);
 }

mlir/lib/Tools/mlir-pdll-lsp-server/LSPServer.cpp

llvm::json::Object serverCaps{
        }},
       {"completionProvider",
        llvm::json::Object{
            {"allCommitCharacters",
             {" ", "\t", "(", ")", "[", "]", "{",  "}", "<",
              ">", ":",  ";", ",", "+", "-", "/",  "*", "%",
              "^", "&",  "#", "?", ".", "=", "\"", "'", "|"}},
            {"resolveProvider", false},
-           {"triggerCharacters", {".", ">", "(", "{", ",", "<", ":", "[", " "}},
+           {"triggerCharacters",
+            {".", ">", "(", "{", ",", "<", ":", "[", " ", "\"", "/"}},
        }},
       {"signatureHelpProvider",
        llvm::json::Object{
            {"triggerCharacters", {"(", ","}},
        }},
       {"definitionProvider", true},
       {"referencesProvider", true},
       {"documentLinkProvider",

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

 #include "mlir/Tools/PDLL/ODS/Constraint.h"
 #include "mlir/Tools/PDLL/ODS/Context.h"
 #include "mlir/Tools/PDLL/ODS/Dialect.h"
 #include "mlir/Tools/PDLL/ODS/Operation.h"
 #include "mlir/Tools/PDLL/Parser/CodeComplete.h"
 #include "mlir/Tools/PDLL/Parser/Parser.h"
 #include "llvm/ADT/IntervalMap.h"
 #include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 
 using namespace mlir;
 using namespace mlir::pdll;
 
 /// Returns a language server uri for the given source location. `mainFileURI`
 /// corresponds to the uri for the main file of the source manager.
 static lsp::URIForFile getURIFromLoc(llvm::SourceMgr &mgr, SMRange loc,

 //===----------------------------------------------------------------------===//
 // PDLDocument: Code Completion
 //===----------------------------------------------------------------------===//
 
 namespace {
 class LSPCodeCompleteContext : public CodeCompleteContext {
 public:
   LSPCodeCompleteContext(SMLoc completeLoc, lsp::CompletionList &completionList,
-                         ods::Context &odsContext)
+                         ods::Context &odsContext,
+                         ArrayRef<std::string> includeDirs)
       : CodeCompleteContext(completeLoc), completionList(completionList),
-        odsContext(odsContext) {}
+        odsContext(odsContext), includeDirs(includeDirs) {}
 
   void codeCompleteTupleMemberAccess(ast::TupleType tupleType) final {
     ArrayRef<ast::Type> elementTypes = tupleType.getElementTypes();
     ArrayRef<StringRef> elementNames = tupleType.getElementNames();
     for (unsigned i = 0, e = tupleType.size(); i < e; ++i) {
       // Push back a completion item that uses the result index.
       lsp::CompletionItem item;
       item.label = llvm::formatv("{0} (field #{0})", i).str();
void codeCompletePatternMetadata() final {
     };
 
     addSimpleConstraint("benefit", "The `benefit` of matching the pattern.",
                         "benefit($1)");
     addSimpleConstraint("recursion",
                         "The pattern properly handles recursive application.");
   }
 
+  void codeCompleteIncludeFilename(StringRef curPath) final {
+    // Normalize the path to allow for interacting with the file system
+    // utilities.
+    SmallString<128> nativeRelDir(llvm::sys::path::convert_to_slash(curPath));
+    llvm::sys::path::native(nativeRelDir);
+
+    // Set of already included completion paths.
+    StringSet<> seenResults;
+
+    // Functor used to add a single include completion item.
+    auto addIncludeCompletion = [&](StringRef path, bool isDirectory) {
+      lsp::CompletionItem item;
+      item.label = (path + (isDirectory ? "/" : "")).str();
+      item.kind = isDirectory ? lsp::CompletionItemKind::Folder
+                              : lsp::CompletionItemKind::File;
+      if (seenResults.insert(item.label).second)
+        completionList.items.emplace_back(item);
+    };
+
+    // Process the include directories for this file, adding any potential
+    // nested include files or directories.
+    for (StringRef includeDir : includeDirs) {
+      llvm::SmallString<128> dir = includeDir;
+      if (!nativeRelDir.empty())
+        llvm::sys::path::append(dir, nativeRelDir);
+
+      std::error_code errorCode;
+      for (auto it = llvm::sys::fs::directory_iterator(dir, errorCode),
+                e = llvm::sys::fs::directory_iterator();
+           !errorCode && it != e; it.increment(errorCode)) {
+        StringRef filename = llvm::sys::path::filename(it->path());
+
+        // To know whether a symlink should be treated as file or a directory,
+        // we have to stat it. This should be cheap enough as there shouldn't be
+        // many symlinks.
+        llvm::sys::fs::file_type fileType = it->type();
+        if (fileType == llvm::sys::fs::file_type::symlink_file) {
+          if (auto fileStatus = it->status())
+            fileType = fileStatus->type();
+        }
+
+        switch (fileType) {
+        case llvm::sys::fs::file_type::directory_file:
+          addIncludeCompletion(filename, /*isDirectory=*/true);
+          break;
+        case llvm::sys::fs::file_type::regular_file: {
+          // Only consider concrete files that can actually be included by PDLL.
+          if (filename.endswith(".pdll") || filename.endswith(".td"))
+            addIncludeCompletion(filename, /*isDirectory=*/false);
+          break;
+        }
+        default:
+          break;
+        }
+      }
+    };
+  }
+
 private:
   lsp::CompletionList &completionList;
   ods::Context &odsContext;
+  ArrayRef<std::string> includeDirs;
 };
 } // namespace
 
 lsp::CompletionList
 PDLDocument::getCodeCompletion(const lsp::URIForFile &uri,
                                const lsp::Position &completePos) {
   SMLoc posLoc = completePos.getAsSMLoc(sourceMgr);
   if (!posLoc.isValid())
     return lsp::CompletionList();
 
   // Adjust the position one further to after the completion trigger token.
   posLoc = SMLoc::getFromPointer(posLoc.getPointer() + 1);
 
   // To perform code completion, we run another parse of the module with the
   // code completion context provided.
   ods::Context tmpODSContext;
   lsp::CompletionList completionList;
-  LSPCodeCompleteContext lspCompleteContext(posLoc, completionList,
-                                            tmpODSContext);
+  LSPCodeCompleteContext lspCompleteContext(
+      posLoc, completionList, tmpODSContext, sourceMgr.getIncludeDirs());
 
   ast::Context tmpContext(tmpODSContext);
   (void)parsePDLAST(tmpContext, sourceMgr, &lspCompleteContext);
 
   return completionList;
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/mlir-pdll-lsp-server/completion.test

-// RUN: mlir-pdll-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+// RUN: mlir-pdll-lsp-server -pdll-extra-dir %S -pdll-extra-dir %S/../../include  -lit-test < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"pdll","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.pdll",
   "languageId":"pdll",
   "version":1,
-  "text":"Constraint ValueCst(value: Value);\nConstraint Cst();\nPattern FooPattern with benefit(1) {\nlet tuple = (value1 = _: Op, _: Op<test.op>);\nerase tuple.value1;\n}"
+  "text":"#include \"include/included.pdll\"\nConstraint ValueCst(value: Value);\nConstraint Cst();\nPattern FooPattern with benefit(1) {\nlet tuple = (value1 = _: Op, _: Op<test.op>);\nerase tuple.value1;\n}"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":4,"character":12}
+  "position":{"line":5,"character":12}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "isIncomplete": false,
 // CHECK-NEXT:    "items": [
 // CHECK-NEXT:      {
 // CHECK-NEXT:        "detail": "0: Op",

 // CHECK-NEXT:        "label": "1 (field #1)",
 // CHECK-NEXT:        "sortText": "1"
 // CHECK-NEXT:      }
 // CHECK-NEXT:    ]
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":2,"character":23}
+  "position":{"line":3,"character":23}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "isIncomplete": false,
 // CHECK-NEXT:    "items": [
 // CHECK-NEXT:      {
 // CHECK-NEXT:        "detail": "pattern metadata",

 // CHECK-NEXT:        "kind": 7,
 // CHECK-NEXT:        "label": "recursion"
 // CHECK-NEXT:      }
 // CHECK-NEXT:    ]
 // CHECK-NEXT:  }
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
   "textDocument":{"uri":"test:///foo.pdll"},
-  "position":{"line":3,"character":24}
+  "position":{"line":4,"character":24}
 }}
 //      CHECK:  "id": 1
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": {
 // CHECK-NEXT:    "isIncomplete": false,
 // CHECK-NEXT:    "items": [
 // CHECK-NEXT:      {
 // CHECK-NEXT:        "detail": "Attr constraint",

 // CHECK-NEXT:        "detail": "(value: Value) -> Tuple<>",
 // CHECK-NEXT:        "kind": 8,
 // CHECK-NEXT:        "label": "ValueCst",
 // CHECK-NEXT:        "sortText": "2_ValueCst"
 // CHECK-NEXT:      }
 // CHECK-NEXT:    ]
 // CHECK-NEXT:  }
 // -----
+{"jsonrpc":"2.0","id":1,"method":"textDocument/completion","params":{
+  "textDocument":{"uri":"test:///foo.pdll"},
+  "position":{"line":0,"character":18}
+}}
+//      CHECK:  "id": 1
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "isIncomplete": false,
+// CHECK-NEXT:    "items": [
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "kind": 17,
+// CHECK-NEXT:        "label": "included.td"
+// CHECK-NEXT:      },
+// CHECK-NEXT:      {
+// CHECK-NEXT:        "kind": 17,
+// CHECK-NEXT:        "label": "included.pdll"
+// CHECK-NEXT:      }
+// CHECK-NEXT:    ]
+// CHECK-NEXT:  }
+// -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}