[AsmParser] Track value references and function arguments (#174566)

This PR is part of the LLVM IR LSP server project:
https://discourse.llvm.org/t/rfc-ir-visualization-with-vs-code-extension-using-an-lsp-server/87773

To be able to implement goto definition in LSP. One must first have to
know what values are referenced on what positions.
This PR implements a Location -> Value* map that allows looking up
values referenced on queried locations.

For example
```llvm
%inst = add i32 1, %arg
```
Querying on 0:20 (the location of %arg) would return the Value* of the
`%arg`.

`getInstuctionLocation(%arg)` would still return the definition of that
value

Tracking for function arguments was also added in the same fashion a
function/instruction/basic block location tracking.

The tests were also updated to test the new features a bit.
This commit is contained in:
Bertik23 2026-02-02 12:21:45 +01:00 committed by GitHub
parent 8e18cdcb15
commit fc6ad39044
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 112 additions and 38 deletions

View File

@ -39,6 +39,7 @@ class AsmParserContext {
DenseMap<Function *, FileLocRange> Functions;
FMap::Allocator FAllocator;
FMap FunctionsInverse = FMap(FAllocator);
DenseMap<BasicBlock *, FileLocRange> Blocks;
using BBMap =
IntervalMap<FileLoc, BasicBlock *,
@ -46,13 +47,15 @@ class AsmParserContext {
IntervalMapHalfOpenInfo<FileLoc>>;
BBMap::Allocator BBAllocator;
BBMap BlocksInverse = BBMap(BBAllocator);
DenseMap<Instruction *, FileLocRange> Instructions;
using IMap =
IntervalMap<FileLoc, Instruction *,
IntervalMapImpl::NodeSizer<FileLoc, Instruction *>::LeafSize,
DenseMap<Value *, FileLocRange> InstructionsAndArguments;
using VMap =
IntervalMap<FileLoc, Value *,
IntervalMapImpl::NodeSizer<FileLoc, Value *>::LeafSize,
IntervalMapHalfOpenInfo<FileLoc>>;
IMap::Allocator IAllocator;
IMap InstructionsInverse = IMap(IAllocator);
VMap::Allocator VAllocator;
VMap InstructionsAndArgumentsInverse = VMap(VAllocator);
VMap ReferencedValues = VMap(VAllocator);
public:
LLVM_ABI std::optional<FileLocRange>
@ -60,7 +63,7 @@ public:
LLVM_ABI std::optional<FileLocRange>
getBlockLocation(const BasicBlock *) const;
LLVM_ABI std::optional<FileLocRange>
getInstructionLocation(const Instruction *) const;
getInstructionOrArgumentLocation(const Value *) const;
/// Get the function at the requested location range.
/// If no single function occupies the queried range, or the record is
/// missing, a nullptr is returned.
@ -77,17 +80,27 @@ public:
/// If no block occupies the queried location, or the record is missing, a
/// nullptr is returned.
LLVM_ABI BasicBlock *getBlockAtLocation(const FileLoc &) const;
/// Get the instruction at the requested location range.
/// Get the instruction or function argument at the requested location range.
/// If no single instruction occupies the queried range, or the record is
/// missing, a nullptr is returned.
LLVM_ABI Instruction *getInstructionAtLocation(const FileLocRange &) const;
/// Get the instruction at the requested location.
LLVM_ABI Value *
getInstructionOrArgumentAtLocation(const FileLocRange &) const;
/// Get the instruction or function argument at the requested location.
/// If no instruction occupies the queried location, or the record is missing,
/// a nullptr is returned.
LLVM_ABI Instruction *getInstructionAtLocation(const FileLoc &) const;
LLVM_ABI Value *getInstructionOrArgumentAtLocation(const FileLoc &) const;
/// Get value referenced at the requested location.
/// If no value occupies the queried location, or the record is missing,
/// a nullptr is returned.
LLVM_ABI Value *getValueReferencedAtLocation(const FileLoc &) const;
/// Get value referenced at the requested location range.
/// If no value occupies the queried location, or the record is missing,
/// a nullptr is returned.
LLVM_ABI Value *getValueReferencedAtLocation(const FileLocRange &) const;
LLVM_ABI bool addFunctionLocation(Function *, const FileLocRange &);
LLVM_ABI bool addBlockLocation(BasicBlock *, const FileLocRange &);
LLVM_ABI bool addInstructionLocation(Instruction *, const FileLocRange &);
LLVM_ABI bool addInstructionOrArgumentLocation(Value *, const FileLocRange &);
LLVM_ABI bool addValueReferenceAtLocation(Value *, const FileLocRange &);
};
} // namespace llvm

View File

@ -614,10 +614,12 @@ namespace llvm {
struct ArgInfo {
LocTy Loc;
Type *Ty;
std::optional<FileLocRange> IdentLoc;
AttributeSet Attrs;
std::string Name;
ArgInfo(LocTy L, Type *ty, AttributeSet Attr, const std::string &N)
: Loc(L), Ty(ty), Attrs(Attr), Name(N) {}
ArgInfo(LocTy L, Type *ty, std::optional<FileLocRange> IdentLoc,
AttributeSet Attr, const std::string &N)
: Loc(L), Ty(ty), IdentLoc(IdentLoc), Attrs(Attr), Name(N) {}
};
bool parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
SmallVectorImpl<unsigned> &UnnamedArgNums,

View File

@ -25,8 +25,10 @@ AsmParserContext::getBlockLocation(const BasicBlock *BB) const {
}
std::optional<FileLocRange>
AsmParserContext::getInstructionLocation(const Instruction *I) const {
if (auto IIt = Instructions.find(I); IIt != Instructions.end())
AsmParserContext::getInstructionOrArgumentLocation(const Value *IA) const {
assert(isa<Instruction>(IA) || isa<Argument>(IA));
if (auto IIt = InstructionsAndArguments.find(IA);
IIt != InstructionsAndArguments.end())
return IIt->second;
return std::nullopt;
}
@ -55,17 +57,30 @@ BasicBlock *AsmParserContext::getBlockAtLocation(const FileLoc &Query) const {
return BlocksInverse.lookup(Query, nullptr);
}
Instruction *
AsmParserContext::getInstructionAtLocation(const FileLocRange &Query) const {
auto It = InstructionsInverse.find(Query.Start);
Value *AsmParserContext::getInstructionOrArgumentAtLocation(
const FileLocRange &Query) const {
auto It = InstructionsAndArgumentsInverse.find(Query.Start);
if (It.stop() <= Query.End)
return *It;
return nullptr;
}
Instruction *
AsmParserContext::getInstructionAtLocation(const FileLoc &Query) const {
return InstructionsInverse.lookup(Query, nullptr);
Value *AsmParserContext::getInstructionOrArgumentAtLocation(
const FileLoc &Query) const {
return InstructionsAndArgumentsInverse.lookup(Query, nullptr);
}
Value *AsmParserContext::getValueReferencedAtLocation(
const FileLocRange &Query) const {
auto It = ReferencedValues.find(Query.Start);
if (It.stop() <= Query.End)
return *It;
return nullptr;
}
Value *
AsmParserContext::getValueReferencedAtLocation(const FileLoc &Query) const {
return ReferencedValues.lookup(Query, nullptr);
}
bool AsmParserContext::addFunctionLocation(Function *F,
@ -84,12 +99,19 @@ bool AsmParserContext::addBlockLocation(BasicBlock *BB,
return Inserted;
}
bool AsmParserContext::addInstructionLocation(Instruction *I,
const FileLocRange &Loc) {
bool Inserted = Instructions.insert({I, Loc}).second;
bool AsmParserContext::addInstructionOrArgumentLocation(
Value *IA, const FileLocRange &Loc) {
assert(isa<Instruction>(IA) || isa<Argument>(IA));
bool Inserted = InstructionsAndArguments.insert({IA, Loc}).second;
if (Inserted)
InstructionsInverse.insert(Loc.Start, Loc.End, I);
InstructionsAndArgumentsInverse.insert(Loc.Start, Loc.End, IA);
return Inserted;
}
bool AsmParserContext::addValueReferenceAtLocation(Value *V,
const FileLocRange &Loc) {
ReferencedValues.insert(Loc.Start, Loc.End, V);
return true;
}
} // namespace llvm

View File

@ -3423,18 +3423,26 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
return error(TypeLoc, "argument can not have void type");
std::string Name;
FileLoc IdentStart;
FileLoc IdentEnd;
bool Unnamed = false;
if (Lex.getKind() == lltok::LocalVar) {
Name = Lex.getStrVal();
IdentStart = Lex.getTokLineColumnPos();
Lex.Lex();
IdentEnd = Lex.getPrevTokEndLineColumnPos();
} else {
unsigned ArgID;
if (Lex.getKind() == lltok::LocalVarID) {
ArgID = Lex.getUIntVal();
IdentStart = Lex.getTokLineColumnPos();
if (checkValueID(TypeLoc, "argument", "%", CurValID, ArgID))
return true;
Lex.Lex();
IdentEnd = Lex.getPrevTokEndLineColumnPos();
} else {
ArgID = CurValID;
Unnamed = true;
}
UnnamedArgNums.push_back(ArgID);
CurValID = ArgID + 1;
@ -3443,9 +3451,11 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
if (!FunctionType::isValidArgumentType(ArgTy))
return error(TypeLoc, "invalid type for function argument");
ArgList.emplace_back(TypeLoc, ArgTy,
AttributeSet::get(ArgTy->getContext(), Attrs),
std::move(Name));
ArgList.emplace_back(
TypeLoc, ArgTy,
Unnamed ? std::nullopt
: std::make_optional(FileLocRange(IdentStart, IdentEnd)),
AttributeSet::get(ArgTy->getContext(), Attrs), std::move(Name));
} while (EatIfPresent(lltok::comma));
}
@ -6707,8 +6717,13 @@ bool LLParser::parseConstantValue(Type *Ty, Constant *&C) {
bool LLParser::parseValue(Type *Ty, Value *&V, PerFunctionState *PFS) {
V = nullptr;
ValID ID;
return parseValID(ID, PFS, Ty) ||
convertValIDToValue(Ty, ID, V, PFS);
FileLoc Start = Lex.getTokLineColumnPos();
bool Ret = parseValID(ID, PFS, Ty) || convertValIDToValue(Ty, ID, V, PFS);
FileLoc End = Lex.getPrevTokEndLineColumnPos();
if (!Ret && ParserContext)
ParserContext->addValueReferenceAtLocation(V, FileLocRange(Start, End));
return Ret;
}
bool LLParser::parseTypeAndValue(Value *&V, PerFunctionState *PFS) {
@ -6955,6 +6970,9 @@ bool LLParser::parseFunctionHeader(Function *&Fn, bool IsDefine,
// Add all of the arguments we parsed to the function.
Function::arg_iterator ArgIt = Fn->arg_begin();
for (unsigned i = 0, e = ArgList.size(); i != e; ++i, ++ArgIt) {
if (ParserContext && ArgList[i].IdentLoc)
ParserContext->addInstructionOrArgumentLocation(
&*ArgIt, ArgList[i].IdentLoc.value());
// If the argument has a name, insert it into the argument symbol table.
if (ArgList[i].Name.empty()) continue;
@ -7164,7 +7182,7 @@ bool LLParser::parseBasicBlock(PerFunctionState &PFS) {
BB->insertDbgRecordBefore(DR.release(), Inst->getIterator());
TrailingDbgRecord.clear();
if (ParserContext) {
ParserContext->addInstructionLocation(
ParserContext->addInstructionOrArgumentLocation(
Inst, FileLocRange(InstStart, Lex.getPrevTokEndLineColumnPos()));
}
} while (!Inst->isTerminator());

View File

@ -495,9 +495,9 @@ TEST(AsmParserTest, DIExpressionBodyAtBeginningWithSlotMappingParsing) {
} while (false)
TEST(AsmParserTest, ParserObjectLocations) {
StringRef Source = "define i32 @main() {\n"
StringRef Source = "define i32 @main(i32 %arg, i64) {\n"
"entry:\n"
" %a = add i32 1, 2\n"
" %a = add i32 1, %arg\n"
" ret i32 %a\n"
"}\n";
LLVMContext Ctx;
@ -527,17 +527,36 @@ TEST(AsmParserTest, ParserObjectLocations) {
ParserContext.getBlockAtLocation(*MaybeEntryBBLoc));
SmallVector<FileLocRange> InstructionLocations = {
FileLocRange(FileLoc{2, 4}, FileLoc{2, 21}),
FileLocRange(FileLoc{2, 4}, FileLoc{2, 24}),
FileLocRange(FileLoc{3, 4}, FileLoc{3, 14})};
for (const auto &[Inst, ExpectedLoc] : zip(EntryBB, InstructionLocations)) {
auto MaybeInstLoc = ParserContext.getInstructionLocation(&Inst);
auto MaybeInstLoc = ParserContext.getInstructionOrArgumentLocation(&Inst);
ASSERT_TRUE(MaybeMainLoc.has_value());
auto InstLoc = MaybeInstLoc.value();
ASSERT_EQ_LOC(InstLoc, ExpectedLoc);
ASSERT_EQ(ParserContext.getInstructionAtLocation(MaybeInstLoc->Start),
ParserContext.getInstructionAtLocation(*MaybeInstLoc));
ASSERT_EQ(
ParserContext.getInstructionOrArgumentAtLocation(MaybeInstLoc->Start),
ParserContext.getInstructionOrArgumentAtLocation(*MaybeInstLoc));
}
SmallVector<std::optional<FileLocRange>> FunctionArgumentLocations = {
FileLocRange(FileLoc{0, 21}, FileLoc{0, 25}), std::nullopt};
for (const auto &[Arg, ExpectedLoc] :
zip(MainFn->args(), FunctionArgumentLocations)) {
auto MaybeArgLoc = ParserContext.getInstructionOrArgumentLocation(&Arg);
ASSERT_EQ(MaybeArgLoc.has_value(), ExpectedLoc.has_value());
if (!ExpectedLoc.has_value())
continue;
auto ArgLoc = MaybeArgLoc.value();
ASSERT_EQ_LOC(ArgLoc, ExpectedLoc.value());
ASSERT_EQ(ParserContext.getInstructionOrArgumentAtLocation(ArgLoc.Start),
ParserContext.getInstructionOrArgumentAtLocation(ArgLoc));
}
ASSERT_EQ(&*MainFn->arg_begin(),
ParserContext.getValueReferencedAtLocation(FileLoc(2, 22)));
ASSERT_EQ(&*EntryBB.begin(),
ParserContext.getValueReferencedAtLocation(FileLoc(3, 13)));
}
} // end anonymous namespace