[TableGen:LSP] Resolve "go-to-def" on a let field to the base definition

[rriddle]

Summary

This allows for go-to-def on the a let field to resolve to the definition
of the base class. This is kind of like how C++ works with go-to-def
from use->def->decl, with the decl in this case being the base definition
of the field.

Depends on D134259

[jpienaar]

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.h

37

SMRange is defined as half-open so contains here feels off. closureContains would be more accurate but may not be too specialized. containsOrLimit would be more direct.

I'm worries that this could be hiding off-by-one issues.

(I also appreciate that this is a pure move and we could look at follow on that is more mechanical)

[rriddle]

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.h

37

Agreed. Fixed the comparison here, given that I don't think anything actually relies on the inclusive aspect of it.

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.h

 /// supports identifier-like tokens, strings, etc.
 SMRange convertTokenLocToRange(SMLoc loc);
 
 /// Extract a documentation comment for the given location within the source
 /// manager. Returns None if no comment could be computed.
 Optional<std::string> extractSourceDocComment(llvm::SourceMgr &sourceMgr,
                                               SMLoc loc);
 
+/// Returns true if the given range contains the given source location. Note
+/// that this has different behavior than SMRange because it is inclusive of the
+/// end location.
+bool contains(SMRange range, SMLoc loc);
+
 //===----------------------------------------------------------------------===//
 // SourceMgrInclude
 //===----------------------------------------------------------------------===//
 
 /// This class represents a single include within a root file.
 struct SourceMgrInclude {
   SourceMgrInclude(const lsp::URIForFile &uri, const lsp::Range &range)
       : uri(uri), range(range) {}

mlir/lib/Tools/lsp-server-support/SourceMgrUtils.cpp

while (Optional<StringRef> line = popLastLine()) {
     commentLines.push_back(line->drop_while([](char c) { return c == '/'; }));
   }
 
   if (commentLines.empty())
     return llvm::None;
   return llvm::join(llvm::reverse(commentLines), "\n");
 }
 
+bool lsp::contains(SMRange range, SMLoc loc) {
+  return range.Start.getPointer() <= loc.getPointer() &&
+         loc.getPointer() < range.End.getPointer();
+}
+
 //===----------------------------------------------------------------------===//
 // SourceMgrInclude
 //===----------------------------------------------------------------------===//
 
 Hover SourceMgrInclude::buildHover() const {
   Hover hover(range);
   {
     llvm::raw_string_ostream hoverOS(hover.contents.value);

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

 }
 
 /// Returns a language server location from the given source range.
 static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMRange range,
                                         const lsp::URIForFile &uri) {
   return lsp::Location(getURIFromLoc(mgr, range, uri), lsp::Range(mgr, range));
 }
 
-/// Returns true if the given range contains the given source location. Note
-/// that this has different behavior than SMRange because it is inclusive of the
-/// end location.
-static bool contains(SMRange range, SMLoc loc) {
-  return range.Start.getPointer() <= loc.getPointer() &&
-         loc.getPointer() <= range.End.getPointer();
-}
-
 /// Convert the given MLIR diagnostic to the LSP form.
 static Optional<lsp::Diagnostic>
 getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr, const ast::Diagnostic &diag,
                         const lsp::URIForFile &uri) {
   lsp::Diagnostic lspDiag;
   lspDiag.source = "pdll";
 
   // FIXME: Right now all of the diagnostics are treated as parser issues, but
if (failed(astModule))
     return;
   SMRange rangeLoc = range.getAsSMRange(sourceMgr);
   if (!rangeLoc.isValid())
     return;
   (*astModule)->walk([&](const ast::Node *node) {
     SMRange loc = node->getLoc();
 
     // Check that the location of this node is within the input range.
-    if (!contains(rangeLoc, loc.Start) && !contains(rangeLoc, loc.End))
+    if (!lsp::contains(rangeLoc, loc.Start) &&
+        !lsp::contains(rangeLoc, loc.End))
       return;
 
     // Handle hints for various types of nodes.
     llvm::TypeSwitch<const ast::Node *>(node)
         .Case<ast::VariableDecl, ast::CallExpr, ast::OperationExpr>(
             [&](const auto *node) {
               this->getInlayHintsFor(node, uri, inlayHints);
             });

mlir/lib/Tools/tblgen-lsp-server/TableGenServer.cpp

 void TableGenTextFile::getLocationsOf(const lsp::URIForFile &uri,
                                       const lsp::Position &defPos,
                                       std::vector<lsp::Location> &locations) {
   SMLoc posLoc = defPos.getAsSMLoc(sourceMgr);
   const TableGenIndexSymbol *symbol = index.lookup(posLoc);
   if (!symbol)
     return;
 
+  // If this symbol is a record value and the def position is already the def of
+  // the symbol, check to see if the value has a base definition. This allows
+  // for a "go-to-def" on a "let" to resolve the definition in the base class.
+  auto *valSym = dyn_cast<TableGenRecordValSymbol>(symbol);
+  if (valSym && lsp::contains(valSym->defLoc, posLoc)) {
+    if (auto *val = getBaseValue(valSym->record, valSym->getValue()).second) {
+      locations.push_back(getLocationFromLoc(sourceMgr, val->getLoc(), uri));
+      return;
+    }
+  }
+
   locations.push_back(getLocationFromLoc(sourceMgr, symbol->defLoc, uri));
 }
 
 void TableGenTextFile::findReferencesOf(
     const lsp::URIForFile &uri, const lsp::Position &pos,
     std::vector<lsp::Location> &references) {
   SMLoc posLoc = pos.getAsSMLoc(sourceMgr);
   const TableGenIndexSymbol *symbol = index.lookup(posLoc);

mlir/test/tblgen-lsp-server/definition.test

 }}
 //      CHECK:  "id": 2
 // CHECK-NEXT:  "jsonrpc": "2.0",
 // CHECK-NEXT:  "result": [
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 12,
-// CHECK-NEXT:          "line": 4
+// CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 6,
-// CHECK-NEXT:          "line": 4
+// CHECK-NEXT:          "line": 1
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.td"
 // CHECK-NEXT:    }
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}