[LLVM][IR] Add location tracking to LLVM IR parser

llvm/include/llvm/AsmParser/AsmParserContext.h

@@ -0,0 +1,69 @@
+//===-- AsmParserContext.h --------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_ASMPARSER_ASMPARSER_STATE_H
+#define LLVM_ASMPARSER_ASMPARSER_STATE_H
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/IR/Value.h"
+#include <optional>
+
+namespace llvm {
+
+/// Registry of file location information for LLVM IR constructs
+///
+/// This class provides access to the file location information
+/// for various LLVM IR constructs. Currently, it supports Function,
+/// BasicBlock and Instruction locations.
+///
+/// When available, it can answer queries about what is at a given
+/// file location, as well as where in a file a given IR construct
+/// is.
+///
+/// This information is optionally emitted by the LLParser while
+/// it reads LLVM textual IR.
+class AsmParserContext {
+  DenseMap<Function *, FileLocRange> Functions;
+  DenseMap<BasicBlock *, FileLocRange> Blocks;
+  DenseMap<Instruction *, FileLocRange> Instructions;
+
+public:
+  std::optional<FileLocRange> getFunctionLocation(const Function *) const;
+  std::optional<FileLocRange> getBlockLocation(const BasicBlock *) const;
+  std::optional<FileLocRange> getInstructionLocation(const Instruction *) 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.
+  Function *getFunctionAtLocation(const FileLocRange &) const;
+  /// Get the function at the requested location.
+  /// If no function occupies the queried location, or the record is missing, a
+  /// nullptr is returned.
+  Function *getFunctionAtLocation(const FileLoc &) const;
+  /// Get the block at the requested location range.
+  /// If no single block occupies the queried range, or the record is missing, a
+  /// nullptr is returned.
+  BasicBlock *getBlockAtLocation(const FileLocRange &) const;
+  /// Get the block at the requested location.
+  /// If no block occupies the queried location, or the record is missing, a
+  /// nullptr is returned.
+  BasicBlock *getBlockAtLocation(const FileLoc &) const;
+  /// Get the instruction at the requested location range.
+  /// If no single instruction occupies the queried range, or the record is
+  /// missing, a nullptr is returned.
+  Instruction *getInstructionAtLocation(const FileLocRange &) const;
+  /// Get the instruction at the requested location.
+  /// If no instruction occupies the queried location, or the record is missing,
+  /// a nullptr is returned.
+  Instruction *getInstructionAtLocation(const FileLoc &) const;
+  bool addFunctionLocation(Function *, const FileLocRange &);
+  bool addBlockLocation(BasicBlock *, const FileLocRange &);
+  bool addInstructionLocation(Instruction *, const FileLocRange &);
+};
+} // namespace llvm
+
+#endif

llvm/include/llvm/AsmParser/LLLexer.h

@@ -17,18 +17,21 @@
 #include "llvm/ADT/APFloat.h"
 #include "llvm/ADT/APSInt.h"
 #include "llvm/Support/SMLoc.h"
+#include "llvm/Support/SourceMgr.h"
 #include <string>
 
 namespace llvm {
   class Type;
   class SMDiagnostic;
-  class SourceMgr;
   class LLVMContext;
 
   class LLLexer {
     const char *CurPtr;
     StringRef CurBuf;
 
+    // The the end (exclusive) of the current token
+    const char *PrevTokEnd = nullptr;
+
     enum class ErrorPriority {
       None,   // No error message present.
       Parser, // Errors issued by parser.
@@ -62,9 +65,7 @@ namespace llvm {
     explicit LLLexer(StringRef StartBuf, SourceMgr &SM, SMDiagnostic &,
                      LLVMContext &C);
 
-    lltok::Kind Lex() {
-      return CurKind = LexToken();
-    }
+    lltok::Kind Lex() { return CurKind = LexToken(); }
 
     typedef SMLoc LocTy;
     LocTy getLoc() const { return SMLoc::getFromPointer(TokStart); }
@@ -79,6 +80,23 @@ namespace llvm {
       IgnoreColonInIdentifiers = val;
     }
 
+    // Get the line, column position of the start of the current token,
+    // zero-indexed
+    std::pair<unsigned, unsigned> getTokLineColumnPos() {
+      auto LC = SM.getLineAndColumn(SMLoc::getFromPointer(TokStart));
+      --LC.first;
+      --LC.second;
+      return LC;
+    }
+    // Get the line, column position of the end of the previous token,
+    // zero-indexed exclusive
+    std::pair<unsigned, unsigned> getPrevTokEndLineColumnPos() {
+      auto LC = SM.getLineAndColumn(SMLoc::getFromPointer(PrevTokEnd));
+      --LC.first;
+      --LC.second;
+      return LC;
+    }
+
     // This returns true as a convenience for the parser functions that return
     // true on error.
     bool ParseError(LocTy ErrorLoc, const Twine &Msg) {
@@ -93,7 +111,12 @@ namespace llvm {
   private:
     lltok::Kind LexToken();
 
+    // Return closest pointer after `Ptr` that is an end of a label.
+    // Returns nullptr if `Ptr` doesn't point into a label.
+    const char *getLabelTail(const char *Ptr);
     int getNextChar();
+    const char *skipNChars(unsigned N);
+    void advancePositionTo(const char *Ptr);
     void SkipLineComment();
     bool SkipCComment();
     lltok::Kind ReadString(lltok::Kind kind);

llvm/include/llvm/AsmParser/LLParser.h

@@ -13,6 +13,7 @@
 #ifndef LLVM_ASMPARSER_LLPARSER_H
 #define LLVM_ASMPARSER_LLPARSER_H
 
+#include "AsmParserContext.h"
 #include "LLLexer.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/AsmParser/NumberedValues.h"
@@ -177,6 +178,9 @@ namespace llvm {
     // Map of module ID to path.
     std::map<unsigned, StringRef> ModuleIdMap;
 
+    /// Keeps track of source locations for Values, BasicBlocks, and Functions
+    AsmParserContext *ParserContext;
+
     /// Only the llvm-as tool may set this to false to bypass
     /// UpgradeDebuginfo so it can generate broken bitcode.
     bool UpgradeDebugInfo;
@@ -189,10 +193,11 @@ namespace llvm {
   public:
     LLParser(StringRef F, SourceMgr &SM, SMDiagnostic &Err, Module *M,
              ModuleSummaryIndex *Index, LLVMContext &Context,
-             SlotMapping *Slots = nullptr)
+             SlotMapping *Slots = nullptr,
+             AsmParserContext *ParserContext = nullptr)
         : Context(Context), OPLex(F, SM, Err, Context),
           Lex(F, SM, Err, Context), M(M), Index(Index), Slots(Slots),
-          BlockAddressPFS(nullptr) {}
+          BlockAddressPFS(nullptr), ParserContext(ParserContext) {}
     bool Run(
         bool UpgradeDebugInfo,
         DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {

llvm/include/llvm/AsmParser/Parser.h

@@ -15,6 +15,7 @@
 
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/AsmParser/AsmParserContext.h"
 #include "llvm/Support/Compiler.h"
 #include <memory>
 #include <optional>
@@ -62,7 +63,8 @@ parseAssemblyFile(StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
 ///              parsing.
 LLVM_ABI std::unique_ptr<Module>
 parseAssemblyString(StringRef AsmString, SMDiagnostic &Err,
-                    LLVMContext &Context, SlotMapping *Slots = nullptr);
+                    LLVMContext &Context, SlotMapping *Slots = nullptr,
+                    AsmParserContext *ParserContext = nullptr);
 
 /// Holds the Module and ModuleSummaryIndex returned by the interfaces
 /// that parse both.
@@ -128,9 +130,9 @@ parseSummaryIndexAssemblyString(StringRef AsmString, SMDiagnostic &Err);
 LLVM_ABI std::unique_ptr<Module> parseAssembly(
     MemoryBufferRef F, SMDiagnostic &Err, LLVMContext &Context,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
-      return std::nullopt;
-    });
+    DataLayoutCallbackTy DataLayoutCallback =
+        [](StringRef, StringRef) { return std::nullopt; },
+    AsmParserContext *ParserContext = nullptr);
 
 /// Parse LLVM Assembly including the summary index from a MemoryBuffer.
 ///
@@ -169,9 +171,9 @@ parseSummaryIndexAssembly(MemoryBufferRef F, SMDiagnostic &Err);
 LLVM_ABI bool parseAssemblyInto(
     MemoryBufferRef F, Module *M, ModuleSummaryIndex *Index, SMDiagnostic &Err,
     SlotMapping *Slots = nullptr,
-    DataLayoutCallbackTy DataLayoutCallback = [](StringRef, StringRef) {
-      return std::nullopt;
-    });
+    DataLayoutCallbackTy DataLayoutCallback =
+        [](StringRef, StringRef) { return std::nullopt; },
+    AsmParserContext *ParserContext = nullptr);
 
 /// Parse a type and a constant value in the given string.
 ///

llvm/include/llvm/IR/Value.h

@@ -55,6 +55,39 @@ class User;
 
 using ValueName = StringMapEntry<Value *>;
 
+struct FileLoc {
+  unsigned Line;
+  unsigned Col;
+
+  bool operator<=(const FileLoc &RHS) const {
+    return Line < RHS.Line || (Line == RHS.Line && Col <= RHS.Col);
+  }
+
+  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) {}
+  FileLoc(std::pair<unsigned, unsigned> LC) : Line(LC.first), Col(LC.second) {}
+};
+
+struct FileLocRange {
+  FileLoc Start;
+  FileLoc End;
+
+  FileLocRange() : Start(0, 0), End(0, 0) {}
+
+  FileLocRange(FileLoc S, FileLoc E) : Start(S), End(E) {
+    assert(Start <= End);
+  }
+
+  bool contains(FileLoc L) const { return Start <= L && L <= End; }
+
+  bool contains(FileLocRange LR) const {
+    return contains(LR.Start) && contains(LR.End);
+  }
+};
+
 //===----------------------------------------------------------------------===//
 //                                 Value Class
 //===----------------------------------------------------------------------===//

llvm/include/llvm/IRReader/IRReader.h

@@ -15,6 +15,7 @@
 #define LLVM_IRREADER_IRREADER_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/AsmParser/AsmParserContext.h"
 #include "llvm/Bitcode/BitcodeReader.h"
 #include "llvm/Support/Compiler.h"
 #include <memory>
@@ -50,19 +51,19 @@ getLazyIRFileModule(StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
 /// for it.  Otherwise, attempt to parse it as LLVM Assembly and return
 /// a Module for it.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
-LLVM_ABI std::unique_ptr<Module> parseIR(MemoryBufferRef Buffer,
-                                         SMDiagnostic &Err,
-                                         LLVMContext &Context,
-                                         ParserCallbacks Callbacks = {});
+LLVM_ABI std::unique_ptr<Module>
+parseIR(MemoryBufferRef Buffer, SMDiagnostic &Err, LLVMContext &Context,
+        ParserCallbacks Callbacks = {},
+        AsmParserContext *ParserContext = nullptr);
 
 /// If the given file holds a bitcode image, return a Module for it.
 /// Otherwise, attempt to parse it as LLVM Assembly and return a Module
 /// for it.
 /// \param DataLayoutCallback Override datalayout in the llvm assembly.
-LLVM_ABI std::unique_ptr<Module> parseIRFile(StringRef Filename,
-                                             SMDiagnostic &Err,
-                                             LLVMContext &Context,
-                                             ParserCallbacks Callbacks = {});
+LLVM_ABI std::unique_ptr<Module>
+parseIRFile(StringRef Filename, SMDiagnostic &Err, LLVMContext &Context,
+            ParserCallbacks Callbacks = {},
+            AsmParserContext *ParserContext = nullptr);
 }
 
 #endif

llvm/lib/AsmParser/AsmParserContext.cpp

@@ -0,0 +1,89 @@
+//===-- AsmParserContext.cpp ------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/AsmParser/AsmParserContext.h"
+
+namespace llvm {
+
+std::optional<FileLocRange>
+AsmParserContext::getFunctionLocation(const Function *F) const {
+  if (!Functions.contains(F))
+    return std::nullopt;
+  return Functions.at(F);
+}
+
+std::optional<FileLocRange>
+AsmParserContext::getBlockLocation(const BasicBlock *BB) const {
+  if (!Blocks.contains(BB))
+    return std::nullopt;
+  return Blocks.at(BB);
+}
+
+std::optional<FileLocRange>
+AsmParserContext::getInstructionLocation(const Instruction *I) const {
+  if (!Instructions.contains(I))
+    return std::nullopt;
+  return Instructions.at(I);
+}
+
+Function *
+AsmParserContext::getFunctionAtLocation(const FileLocRange &Query) const {
+  for (auto &[F, Loc] : Functions) {
+    if (Loc.contains(Query))
+      return F;
+  }
+  return nullptr;
+}
+
+Function *AsmParserContext::getFunctionAtLocation(const FileLoc &Query) const {
+  return getFunctionAtLocation(FileLocRange(Query, Query));
+}
+
+BasicBlock *
+AsmParserContext::getBlockAtLocation(const FileLocRange &Query) const {
+  for (auto &[BB, Loc] : Blocks) {
+    if (Loc.contains(Query))
+      return BB;
+  }
+  return nullptr;
+}
+
+BasicBlock *AsmParserContext::getBlockAtLocation(const FileLoc &Query) const {
+  return getBlockAtLocation(FileLocRange(Query, Query));
+}
+
+Instruction *
+AsmParserContext::getInstructionAtLocation(const FileLocRange &Query) const {
+  for (auto &[I, Loc] : Instructions) {
+    if (Loc.contains(Query))
+      return I;
+  }
+  return nullptr;
+}
+
+Instruction *
+AsmParserContext::getInstructionAtLocation(const FileLoc &Query) const {
+  return getInstructionAtLocation(FileLocRange(Query, Query));
+}
+
+bool AsmParserContext::addFunctionLocation(Function *F,
+                                           const FileLocRange &Loc) {
+  return Functions.insert({F, Loc}).second;
+}
+
+bool AsmParserContext::addBlockLocation(BasicBlock *BB,
+                                        const FileLocRange &Loc) {
+  return Blocks.insert({BB, Loc}).second;
+}
+
+bool AsmParserContext::addInstructionLocation(Instruction *I,
+                                              const FileLocRange &Loc) {
+  return Instructions.insert({I, Loc}).second;
+}
+
+} // namespace llvm

llvm/lib/AsmParser/CMakeLists.txt

@@ -1,5 +1,6 @@
 # AsmParser
 add_llvm_component_library(LLVMAsmParser
+  AsmParserContext.cpp
   LLLexer.cpp
   LLParser.cpp
   Parser.cpp