diff --git a/mlir/lib/Parser/AsmParserState.cpp b/mlir/lib/Parser/AsmParserState.cpp --- a/mlir/lib/Parser/AsmParserState.cpp +++ b/mlir/lib/Parser/AsmParserState.cpp @@ -18,7 +18,8 @@ struct AsmParserState::Impl { /// A map from a SymbolRefAttr to a range of uses. - using SymbolUseMap = DenseMap>; + using SymbolUseMap = + DenseMap, 0>>; struct PartialOpDef { explicit PartialOpDef(const OperationName &opName) { @@ -73,10 +74,12 @@ 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 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)); + } } } } @@ -287,8 +290,8 @@ 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 --- 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: ] // -----