Implement goto definition

llvm/include/llvm/AsmParser/AsmParserContext.h

@@ -10,11 +10,45 @@
 #define LLVM_ASMPARSER_ASMPARSER_STATE_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/IntervalMap.h"
 #include "llvm/IR/Value.h"
+#include "llvm/Support/Allocator.h"
+#include "llvm/Support/AllocatorBase.h"
+#include <llvm/Support/FormatVariadic.h>
 #include <optional>
 
 namespace llvm {
 
+template <> struct DenseMapInfo<FileLocRange> {
+  static constexpr FileLocRange getEmptyKey() {
+    return FileLocRange(FileLoc(-1, -1), FileLoc(-1, -1));
+  }
+  static constexpr FileLocRange getTombstoneKey() {
+    return FileLocRange(FileLoc(-2, -2), FileLoc(-2, -2));
+  }
+  static unsigned getHashValue(const FileLocRange &Val) {
+    return (Val.Start.Line * 31) ^ (Val.Start.Col * 37) ^ (Val.End.Line * 41) ^
+           (Val.End.Col * 43);
+  }
+  static bool isEqual(const FileLocRange &LHS, const FileLocRange &RHS) {
+    return LHS.contains(RHS) && RHS.contains(LHS);
+  }
+};
+
+template <> struct format_provider<FileLoc> {
+  static void format(const FileLoc &Loc, raw_ostream &Stream) {
+    Stream << Loc.Line << ":" << Loc.Col;
+  }
+};
+
+template <> struct format_provider<FileLocRange> {
+  static void format(const FileLocRange &Range, raw_ostream &Stream) {
+    llvm::format_provider<FileLoc>::format(Range.Start, Stream);
+    Stream << "-";
+    llvm::format_provider<FileLoc>::format(Range.End, Stream);
+  }
+};
+
 /// Registry of file location information for LLVM IR constructs
 ///
 /// This class provides access to the file location information
@@ -32,19 +66,28 @@ class AsmParserContext {
   std::optional<FileLocRange> getFunctionLocation(const Function *) const;
   std::optional<FileLocRange> getBlockLocation(const BasicBlock *) const;
   std::optional<FileLocRange> getInstructionLocation(const Instruction *) const;
+  std::optional<FileLocRange>
+  getFunctionArgumentLocation(const Argument *) const;
+  std::optional<Value *> getValueAtLocation(const FileLocRange &) const;
+  std::optional<Value *> getValueAtLocation(const FileLoc &) const;
   std::optional<Function *> getFunctionAtLocation(const FileLocRange &) const;
   std::optional<Function *> getFunctionAtLocation(const FileLoc &) const;
   std::optional<BasicBlock *> getBlockAtLocation(const FileLocRange &) const;
   std::optional<BasicBlock *> getBlockAtLocation(const FileLoc &) const;
   std::optional<Instruction *>
   getInstructionAtLocation(const FileLocRange &) const;
   std::optional<Instruction *> getInstructionAtLocation(const FileLoc &) const;
+  Value *getValueReferencedAtLocation(const FileLoc &);
   bool addFunctionLocation(Function *, const FileLocRange &);
   bool addBlockLocation(BasicBlock *, const FileLocRange &);
   bool addInstructionLocation(Instruction *, const FileLocRange &);
+  bool addFunctionArgumentLocation(Argument *, const FileLocRange &);
+  bool addValueReferenceOnLocation(Value *, const FileLocRange &);
 
 private:
+  DenseMap<FileLocRange, Value *> LocRangeValueMap;
   DenseMap<Function *, FileLocRange> Functions;
+  DenseMap<Argument *, FileLocRange> FunctionArguments;
   DenseMap<BasicBlock *, FileLocRange> Blocks;
   DenseMap<Instruction *, FileLocRange> Instructions;
 };

llvm/include/llvm/AsmParser/LLParser.h

@@ -612,6 +612,7 @@ namespace llvm {
     struct ArgInfo {
       LocTy Loc;
       Type *Ty;
+      FileLocRange IdentLoc;
       AttributeSet Attrs;
       std::string Name;
       ArgInfo(LocTy L, Type *ty, AttributeSet Attr, const std::string &N)

llvm/include/llvm/IR/Value.h

@@ -59,24 +59,25 @@ struct FileLoc {
   unsigned Line;
   unsigned Col;
 
-  bool operator<=(const FileLoc &RHS) const {
+  constexpr bool operator<=(const FileLoc &RHS) const {
     return Line < RHS.Line || (Line == RHS.Line && Col <= RHS.Col);
   }
 
-  bool operator<(const FileLoc &RHS) const {
+  constexpr bool operator<(const FileLoc &RHS) const {
     return Line < RHS.Line || (Line == RHS.Line && Col < RHS.Col);
   }
 
-  FileLoc(unsigned L, unsigned C) : Line(L), Col(C) {}
+  constexpr FileLoc() : Line(0), Col(0) {}
+  constexpr FileLoc(unsigned L, unsigned C) : Line(L), Col(C) {}
 };
 
 struct FileLocRange {
   FileLoc Start;
   FileLoc End;
 
-  FileLocRange() : Start(0, 0), End(0, 0) {}
+  constexpr FileLocRange() : Start(0, 0), End(0, 0) {}
 
-  FileLocRange(FileLoc S, FileLoc E) : Start(S), End(E) {
+  constexpr FileLocRange(FileLoc S, FileLoc E) : Start(S), End(E) {
     assert(Start <= End);
   }
 
@@ -85,6 +86,10 @@ struct FileLocRange {
   bool contains(FileLocRange LR) const {
     return contains(LR.Start) && contains(LR.End);
   }
+
+  bool operator==(const FileLocRange &RHS) {
+    return contains(RHS) && RHS.contains(*this);
+  }
 };
 
 //===----------------------------------------------------------------------===//

llvm/include/llvm/Support/LSP/Protocol.h

@@ -545,6 +545,7 @@ llvm::json::Value toJSON(const MarkupContent &mc);
 //===----------------------------------------------------------------------===//
 
 struct Hover {
+  Hover() = default;
   /// Construct a default hover with the given range that uses Markdown content.
   Hover(Range range) : contents{MarkupKind::Markdown, ""}, range(range) {}
 

llvm/lib/AsmParser/AsmParserContext.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/AsmParser/AsmParserContext.h"
+#include "llvm/IR/Value.h"
 
 namespace llvm {
 
@@ -31,6 +32,27 @@ AsmParserContext::getInstructionLocation(const Instruction *I) const {
   return Instructions.at(I);
 }
 
+std::optional<FileLocRange>
+AsmParserContext::getFunctionArgumentLocation(const Argument *FA) const {
+  if (!FunctionArguments.contains(FA))
+    return std::nullopt;
+  return FunctionArguments.at(FA);
+}
+
+std::optional<Value *>
+AsmParserContext::getValueAtLocation(const FileLocRange &Query) const {
+  for (auto &[Loc, V] : LocRangeValueMap) {
+    if (Loc.contains(Query))
+      return V;
+  }
+  return std::nullopt;
+}
+
+std::optional<Value *>
+AsmParserContext::getValueAtLocation(const FileLoc &Query) const {
+  return getValueAtLocation(FileLocRange(Query, Query));
+}
+
 std::optional<Function *>
 AsmParserContext::getFunctionAtLocation(const FileLocRange &Query) const {
   for (auto &[F, Loc] : Functions) {
@@ -73,11 +95,24 @@ AsmParserContext::getInstructionAtLocation(const FileLoc &Query) const {
   return getInstructionAtLocation(FileLocRange(Query, Query));
 }
 
+Value *AsmParserContext::getValueReferencedAtLocation(const FileLoc &Query) {
+  for (const auto &[Loc, V] : LocRangeValueMap) {
+    if (Loc.contains(Query))
+      return V;
+  }
+  return nullptr;
+}
+
 bool AsmParserContext::addFunctionLocation(Function *F,
                                            const FileLocRange &Loc) {
   return Functions.insert({F, Loc}).second;
 }
 
+bool AsmParserContext::addFunctionArgumentLocation(Argument *FA,
+                                                   const FileLocRange &Loc) {
+  return FunctionArguments.insert({FA, Loc}).second;
+}
+
 bool AsmParserContext::addBlockLocation(BasicBlock *BB,
                                         const FileLocRange &Loc) {
   return Blocks.insert({BB, Loc}).second;
@@ -88,4 +123,9 @@ bool AsmParserContext::addInstructionLocation(Instruction *I,
   return Instructions.insert({I, Loc}).second;
 }
 
+bool AsmParserContext::addValueReferenceOnLocation(Value *V,
+                                                   const FileLocRange &Loc) {
+  return LocRangeValueMap.insert({Loc, V}).second;
+}
+
 } // namespace llvm

llvm/lib/AsmParser/LLParser.cpp

@@ -3372,16 +3372,28 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
         return error(TypeLoc, "argument can not have void type");
 
       std::string Name;
+      unsigned NameStartCol = 0;
+      unsigned NameStartLine = 0;
+      unsigned NameEndCol = 0;
+      unsigned NameEndLine = 0;
       if (Lex.getKind() == lltok::LocalVar) {
         Name = Lex.getStrVal();
+        NameStartCol = Lex.getTokColNum();
+        NameStartLine = Lex.getTokLineNum();
         Lex.Lex();
+        NameEndCol = Lex.getPrevTokEndColNum();
+        NameEndLine = Lex.getPrevTokEndLineNum();
       } else {
         unsigned ArgID;
         if (Lex.getKind() == lltok::LocalVarID) {
           ArgID = Lex.getUIntVal();
+          NameStartCol = Lex.getTokColNum();
+          NameStartLine = Lex.getTokLineNum();
           if (checkValueID(TypeLoc, "argument", "%", CurValID, ArgID))
             return true;
           Lex.Lex();
+          NameEndCol = Lex.getPrevTokEndColNum();
+          NameEndLine = Lex.getPrevTokEndLineNum();
         } else {
           ArgID = CurValID;
         }
@@ -3395,6 +3407,8 @@ bool LLParser::parseArgumentList(SmallVectorImpl<ArgInfo> &ArgList,
       ArgList.emplace_back(TypeLoc, ArgTy,
                            AttributeSet::get(ArgTy->getContext(), Attrs),
                            std::move(Name));
+      ArgList.back().IdentLoc = FileLocRange({NameStartLine, NameStartCol},
+                                             {NameEndLine, NameEndCol});
     } while (EatIfPresent(lltok::comma));
   }
 
@@ -6596,8 +6610,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);
+  auto Start = FileLoc(Lex.getTokLineNum(), Lex.getTokColNum());
+
+  auto Ret = parseValID(ID, PFS, Ty) || convertValIDToValue(Ty, ID, V, PFS);
+  auto End = FileLoc(Lex.getPrevTokEndLineNum(), Lex.getPrevTokEndColNum());
+  if (ParserContext)
+    ParserContext->addValueReferenceOnLocation(V, FileLocRange(Start, End));
+  return Ret;
 }
 
 bool LLParser::parseTypeAndValue(Value *&V, PerFunctionState *PFS) {
@@ -6844,6 +6863,10 @@ 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) {
+      ParserContext->addFunctionArgumentLocation(&*ArgIt, ArgList[i].IdentLoc);
+    }
+
     // If the argument has a name, insert it into the argument symbol table.
     if (ArgList[i].Name.empty()) continue;
 

llvm/tools/llvm-lsp/llvm-lsp-server.cpp

@@ -11,6 +11,7 @@
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
+#include "llvm/Support/LSP/Protocol.h"
 #include "llvm/Support/Program.h"
 
 #include "IRDocument.h"
@@ -37,6 +38,15 @@ static lsp::Range llvmFileLocRangeToLspRange(const FileLocRange &Range) {
                     llvmFileLocToLspPosition(Range.End));
 }
 
+static FileLoc lspPositionToLlvmFileLoc(const lsp::Position &Pos) {
+  return FileLoc(Pos.line, Pos.character);
+}
+
+static FileLocRange lspRangeToLlvmFileLocRange(const lsp::Range &Range) {
+  return FileLocRange(lspPositionToLlvmFileLoc(Range.start),
+                      lspPositionToLlvmFileLoc(Range.end));
+}
+
 llvm::Error LspServer::run() {
   registerMessageHandlers();
   return Transport.run(MessageHandler);
@@ -49,6 +59,16 @@ void LspServer::sendInfo(const std::string &Message) {
 void LspServer::sendError(const std::string &Message) {
   ShowMessageSender(lsp::ShowMessageParams(lsp::MessageType::Error, Message));
 }
+template <typename T>
+void fileNotOpenError(lsp::Callback<T> &Reply,
+                      llvm::lsp::TextDocumentIdentifier File) {
+  lsp::Logger::error(
+      "Document in textDocument/documentSymbol request not open: {}",
+      File.uri.file());
+  return Reply(make_error<lsp::LSPError>(
+      formatv("Did not open file previously {}", File.uri.file()),
+      lsp::ErrorCode::InvalidParams));
+}
 
 void LspServer::handleRequestInitialize(
     const lsp::InitializeParams &Params,
@@ -69,6 +89,8 @@ void LspServer::handleRequestInitialize(
         {"referencesProvider", true},
         {"codeActionProvider", true},
         {"documentSymbolProvider", true},
+        {"hoverProvider", true},
+        {"definitionProvider", true}
       }
     }
   };
@@ -183,6 +205,61 @@ void LspServer::handleRequestCodeAction(const lsp::CodeActionParams &Params,
       json::Object{{"title", "Open CFG Preview"}, {"command", "llvm.cfg"}}});
 }
 
+void LspServer::handleRequestTextDocumentHover(
+    const lsp::TextDocumentPositionParams &Params,
+    lsp::Callback<lsp::Hover> Reply) {
+  if (!OpenDocuments.contains(Params.textDocument.uri.file().str())) {
+    return fileNotOpenError(Reply, Params.textDocument);
+  }
+  return;
+  // sendInfo("Searching for values at this position");
+  // auto NumVals = 0u;
+  // for (const auto &[Loc, Val] :
+  //      OpenDocuments[Params.textDocument.uri.file().str()]
+  //          ->ParserContext.LocRangeValueMap) {
+  //   if (Loc.contains(lspPositionToLlvmFileLoc(Params.position))) {
+  //     sendInfo("Value on this position found");
+  //     NumVals++;
+  //   }
+  // }
+  // lsp::Hover Result;
+  // Result.contents = {lsp::MarkupKind::PlainText,
+  //                    formatv("Number of vals on this position: {}",
+  //                    NumVals)};
+  // Reply(Result);
+}
+
+void LspServer::handleRequestTextDocumentDefinition(
+    const lsp::TextDocumentPositionParams &Params,
+    lsp::Callback<std::optional<lsp::Location>> Reply) {
+  if (!OpenDocuments.contains(Params.textDocument.uri.file().str())) {
+    return fileNotOpenError(Reply, Params.textDocument);
+  }
+  sendInfo("Searching for definition at this position");
+  auto Query = lspPositionToLlvmFileLoc(Params.position);
+  auto MaybeVal = OpenDocuments[Params.textDocument.uri.file().str()]
+                      ->ParserContext.getValueAtLocation(Query);
+  if (!MaybeVal)
+    return Reply(std::nullopt);
+  auto *Val = MaybeVal.value();
+  sendInfo("Value on this position found");
+  if (isa<Instruction>(Val))
+    return Reply(lsp::Location(
+        Params.textDocument.uri,
+        llvmFileLocRangeToLspRange(
+            OpenDocuments[Params.textDocument.uri.file().str()]
+                ->ParserContext.getInstructionLocation(cast<Instruction>(Val))
+                .value())));
+  if (isa<Argument>(Val))
+    return Reply(lsp::Location(
+        Params.textDocument.uri,
+        llvmFileLocRangeToLspRange(
+            OpenDocuments[Params.textDocument.uri.file().str()]
+                ->ParserContext.getFunctionArgumentLocation(cast<Argument>(Val))
+                .value())));
+  Reply(std::nullopt);
+}
+
 void LspServer::handleRequestGetCFG(const lsp::GetCfgParams &Params,
                                     lsp::Callback<lsp::CFG> Reply) {
   // TODO: have a flag to force regenerating the artifacts
@@ -343,6 +420,13 @@ bool LspServer::registerMessageHandlers() {
                         &LspServer::handleRequestTextDocumentDocumentSymbol);
   MessageHandler.method("textDocument/codeAction", this,
                         &LspServer::handleRequestCodeAction);
+
+  MessageHandler.method("textDocument/hover", this,
+                        &LspServer::handleRequestTextDocumentHover);
+
+  MessageHandler.method("textDocument/definition", this,
+                        &LspServer::handleRequestTextDocumentDefinition);
+
   // Custom messages
   MessageHandler.method("llvm/getCfg", this, &LspServer::handleRequestGetCFG);
   MessageHandler.method("llvm/bbLocation", this,