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

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.

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

mlir/lib/Parser/AsmParserState.cpp

@@ -18,7 +18,8 @@ using namespace mlir;
 
 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) {
@@ -73,10 +74,12 @@ void AsmParserState::Impl::resolveSymbolUses(Operation *op,
                                           symbolOps)))
       continue;
 
-    for (const auto &symIt : llvm::zip(symbolOps, it.second)) {
-      auto opIt = operationToIdx.find(std::get<0>(symIt));
-      if (opIt != operationToIdx.end())
-        operations[opIt->second]->symbolUses.push_back(std::get<1>(symIt));
+    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));
+      }
     }
   }
 }
@@ -282,8 +285,8 @@ void AsmParserState::addUses(SymbolRefAttr refAttr,
 
   assert((refAttr.getNestedReferences().size() + 1) == locations.size() &&
          "expected the same number of references as provided locations");
-  (*impl->symbolUseScopes.back())[refAttr].append(locations.begin(),
-                                                  locations.end());
+  (*impl->symbolUseScopes.back())[refAttr].emplace_back(locations.begin(),
+                                                        locations.end());
 }
 
 void AsmParserState::refineDefinition(Value oldValue, Value newValue) {

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

@@ -5,7 +5,7 @@
   "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":{
@@ -33,11 +33,11 @@
 // 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"
@@ -64,6 +64,19 @@
 // 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:  ]
 // -----