[Mlir-commits] [mlir] 2db4701 - [mlir-lsp-server] Fix bug in symbol use/def tracking

[River Riddle]

Author: River Riddle
Date: 2021-06-07T14:07:41-07:00
New Revision: 2db4701cafd42d20f32c6c9a5dbd15484523fcda

URL: https://github.com/llvm/llvm-project/commit/2db4701cafd42d20f32c6c9a5dbd15484523fcda
DIFF: https://github.com/llvm/llvm-project/commit/2db4701cafd42d20f32c6c9a5dbd15484523fcda.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Parser/AsmParserState.cpp
    mlir/test/mlir-lsp-server/references.test

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Parser/AsmParserState.cpp b/mlir/lib/Parser/AsmParserState.cpp
index d85c068d493f..31b868458a7d 100644
--- a/mlir/lib/Parser/AsmParserState.cpp
+++ b/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) {

diff  --git a/mlir/test/mlir-lsp-server/references.test b/mlir/test/mlir-lsp-server/references.test
index 8d3c1b6f25c5..ce34cd6c6d8f 100644
--- a/mlir/test/mlir-lsp-server/references.test
+++ b/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:  ]
 // -----