[mlir-lsp-server] Fix bug in symbol use/def tracking

[rriddle]

Summary

We were accidentally only using the first found reference, instead of all of them. This revision fixes this by properly tracking all references to a symbol.

Depends On D103729

mlir/lib/Parser/AsmParserState.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // AsmParserState::Impl
 //===----------------------------------------------------------------------===//
 
 struct AsmParserState::Impl {
   /// A map from a SymbolRefAttr to a range of uses.
-  using SymbolUseMap = DenseMap<Attribute, SmallVector<llvm::SMRange>>;
+  using SymbolUseMap =
+      DenseMap<Attribute, SmallVector<SmallVector<llvm::SMRange>, 0>>;
 
   struct PartialOpDef {
     explicit PartialOpDef(const OperationName &opName) {
       const auto *abstractOp = opName.getAbstractOperation();
       if (abstractOp && abstractOp->hasTrait<OpTrait::SymbolTable>())
         symbolTable = std::make_unique<SymbolUseMap>();
     }
 
void AsmParserState::Impl::resolveSymbolUses(Operation *op,
 
   SmallVector<Operation *> symbolOps;
   for (auto &it : *opDef.symbolTable) {
     symbolOps.clear();
     if (failed(symbolTable.lookupSymbolIn(op, it.first.cast<SymbolRefAttr>(),
                                           symbolOps)))
       continue;
 
-    for (const auto &symIt : llvm::zip(symbolOps, it.second)) {
+    for (ArrayRef<llvm::SMRange> useRange : it.second) {
+      for (const auto &symIt : llvm::zip(symbolOps, useRange)) {
         auto opIt = operationToIdx.find(std::get<0>(symIt));
         if (opIt != operationToIdx.end())
           operations[opIt->second]->symbolUses.push_back(std::get<1>(symIt));
       }
     }
   }
+}
 
 //===----------------------------------------------------------------------===//
 // AsmParserState
 //===----------------------------------------------------------------------===//
 
 AsmParserState::AsmParserState() : impl(std::make_unique<Impl>()) {}
 AsmParserState::~AsmParserState() {}
 AsmParserState &AsmParserState::operator=(AsmParserState &&other) {

 void AsmParserState::addUses(SymbolRefAttr refAttr,
                              ArrayRef<llvm::SMRange> locations) {
   // Ignore this symbol if no scopes are active.
   if (impl->symbolUseScopes.empty())
     return;
 
   assert((refAttr.getNestedReferences().size() + 1) == locations.size() &&
          "expected the same number of references as provided locations");
-  (*impl->symbolUseScopes.back())[refAttr].append(locations.begin(),
+  (*impl->symbolUseScopes.back())[refAttr].emplace_back(locations.begin(),
                                                         locations.end());
 }
 
 void AsmParserState::refineDefinition(Value oldValue, Value newValue) {
   auto it = impl->placeholderValueUses.find(oldValue);
   assert(it != impl->placeholderValueUses.end() &&
          "expected `oldValue` to be a placeholder");
   addUses(newValue, it->second);
   impl->placeholderValueUses.erase(oldValue);
 }

mlir/test/mlir-lsp-server/references.test

 // RUN: mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
 {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
 // -----
 {"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
   "uri":"test:///foo.mlir",
   "languageId":"mlir",
   "version":1,
-  "text":"func @foo() -> i1 {\n%value = constant true\nreturn %value : i1\n}"
+  "text":"func @foo() -> i1 {\n%value = constant true\n%result = call @foo() : () -> i1\nreturn %value : i1\n}"
 }}}
 // -----
 {"jsonrpc":"2.0","id":1,"method":"textDocument/references","params":{
   "textDocument":{"uri":"test:///foo.mlir"},
   "position":{"line":1,"character":2},
   "context":{"includeDeclaration": false}
 }}
 //      CHECK:  "id": 1

 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    },
 // CHECK-NEXT:    {
 // CHECK-NEXT:      "range": {
 // CHECK-NEXT:        "end": {
 // CHECK-NEXT:          "character": 13,
-// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:          "line": 3
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 7,
-// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:          "line": 3
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  ]
 // -----
 {"jsonrpc":"2.0","id":2,"method":"textDocument/references","params":{
   "textDocument":{"uri":"test:///foo.mlir"},

 // CHECK-NEXT:          "line": 0
 // CHECK-NEXT:        },
 // CHECK-NEXT:        "start": {
 // CHECK-NEXT:          "character": 5,
 // CHECK-NEXT:          "line": 0
 // CHECK-NEXT:        }
 // CHECK-NEXT:      },
 // CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
+// CHECK-NEXT:    },
+// CHECK-NEXT:    {
+// CHECK-NEXT:      "range": {
+// CHECK-NEXT:        "end": {
+// CHECK-NEXT:          "character": 19,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        },
+// CHECK-NEXT:        "start": {
+// CHECK-NEXT:          "character": 15,
+// CHECK-NEXT:          "line": 2
+// CHECK-NEXT:        }
+// CHECK-NEXT:      },
+// CHECK-NEXT:      "uri": "{{.*}}/foo.mlir"
 // CHECK-NEXT:    }
 // CHECK-NEXT:  ]
 // -----
 {"jsonrpc":"2.0","id":3,"method":"shutdown"}
 // -----
 {"jsonrpc":"2.0","method":"exit"}