[SourceMgr/MLIR diagnostics] Introduce a new method to speed things up

Summary:
This introduces a new SourceMgr::FindLocForLineAndColumn method that
uses the OffsetCache in SourceMgr::SrcBuffer to do do a constant time
lookup for the line number (once the cache is populated).

Use this method in MLIR's SourceMgrDiagnosticHandler::convertLocToSMLoc,
replacing the O(n) scanning logic.  This resolves a long standing TODO
in MLIR, and makes one of my usecases go dramatically faster (which is
currently producing many diagnostics in a 40MB SourceBuffer).

NFC, this is just a performance speedup and cleanup.

Reviewers: rriddle!, ftynse!

Subscribers: hiraditya, mehdi_amini, rriddle, jpienaar, shauheen, antiagainst, nicolasvasilache, arpith-jacob, mgester, lucyrfox, liufengdb, Joonsoo, grosul1, frgossen, Kayjukh, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D78868

Author: lattner

llvm/include/llvm/Support/SourceMgr.h

@@ -61,10 +61,9 @@ class SourceMgr {
     /// into relatively small files (often smaller than 2^8 or 2^16 bytes),
     /// we select the offset vector element type dynamically based on the
     /// size of Buffer.
-    using VariableSizeOffsets = PointerUnion<std::vector<uint8_t> *,
-                                             std::vector<uint16_t> *,
-                                             std::vector<uint32_t> *,
-                                             std::vector<uint64_t> *>;
+    using VariableSizeOffsets =
+        PointerUnion<std::vector<uint8_t> *, std::vector<uint16_t> *,
+                     std::vector<uint32_t> *, std::vector<uint64_t> *>;
 
     /// Vector of offsets into Buffer at which there are line-endings
     /// (lazily populated). Once populated, the '\n' that marks the end of
@@ -74,12 +73,17 @@ class SourceMgr {
     /// offset corresponding to a particular SMLoc).
     mutable VariableSizeOffsets OffsetCache;
 
-    /// Populate \c OffsetCache and look up a given \p Ptr in it, assuming
-    /// it points somewhere into \c Buffer. The static type parameter \p T
-    /// must be an unsigned integer type from uint{8,16,32,64}_t large
-    /// enough to store offsets inside \c Buffer.
-    template<typename T>
+    /// Look up a given \p Ptr in in the buffer, determining which line it came
+    /// from.
     unsigned getLineNumber(const char *Ptr) const;
+    template <typename T>
+    unsigned getLineNumberSpecialized(const char *Ptr) const;
+
+    /// Return a pointer to the first character of the specified line number or
+    /// null if the line number is invalid.
+    const char *getPointerForLineNumber(unsigned LineNo) const;
+    template <typename T>
+    const char *getPointerForLineNumberSpecialized(unsigned LineNo) const;
 
     /// This is the location of the parent include, or null if at the top level.
     SMLoc IncludeLoc;
@@ -134,9 +138,7 @@ class SourceMgr {
     return Buffers[i - 1].Buffer.get();
   }
 
-  unsigned getNumBuffers() const {
-    return Buffers.size();
-  }
+  unsigned getNumBuffers() const { return Buffers.size(); }
 
   unsigned getMainFileID() const {
     assert(getNumBuffers());
@@ -184,12 +186,16 @@ class SourceMgr {
   std::pair<unsigned, unsigned> getLineAndColumn(SMLoc Loc,
                                                  unsigned BufferID = 0) const;
 
+  /// Given a line and column number in a mapped buffer, turn it into an SMLoc.
+  /// This will return a null SMLoc if the line/column location is invalid.
+  SMLoc FindLocForLineAndColumn(unsigned BufferID, unsigned LineNo,
+                                unsigned ColNo);
+
   /// Emit a message about the specified location with the specified string.
   ///
   /// \param ShowColors Display colored messages if output is a terminal and
   /// the default error handler is used.
-  void PrintMessage(raw_ostream &OS, SMLoc Loc, DiagKind Kind,
-                    const Twine &Msg,
+  void PrintMessage(raw_ostream &OS, SMLoc Loc, DiagKind Kind, const Twine &Msg,
                     ArrayRef<SMRange> Ranges = None,
                     ArrayRef<SMFixIt> FixIts = None,
                     bool ShowColors = true) const;
@@ -234,13 +240,13 @@ class SMFixIt {
 public:
   // FIXME: Twine.str() is not very efficient.
   SMFixIt(SMLoc Loc, const Twine &Insertion)
-    : Range(Loc, Loc), Text(Insertion.str()) {
+      : Range(Loc, Loc), Text(Insertion.str()) {
     assert(Loc.isValid());
   }
 
   // FIXME: Twine.str() is not very efficient.
   SMFixIt(SMRange R, const Twine &Replacement)
-    : Range(R), Text(Replacement.str()) {
+      : Range(R), Text(Replacement.str()) {
     assert(R.isValid());
   }
 
@@ -274,13 +280,12 @@ class SMDiagnostic {
   SMDiagnostic() = default;
   // Diagnostic with no location (e.g. file not found, command line arg error).
   SMDiagnostic(StringRef filename, SourceMgr::DiagKind Knd, StringRef Msg)
-    : Filename(filename), LineNo(-1), ColumnNo(-1), Kind(Knd), Message(Msg) {}
+      : Filename(filename), LineNo(-1), ColumnNo(-1), Kind(Knd), Message(Msg) {}
 
   // Diagnostic with a location.
-  SMDiagnostic(const SourceMgr &sm, SMLoc L, StringRef FN,
-               int Line, int Col, SourceMgr::DiagKind Kind,
-               StringRef Msg, StringRef LineStr,
-               ArrayRef<std::pair<unsigned,unsigned>> Ranges,
+  SMDiagnostic(const SourceMgr &sm, SMLoc L, StringRef FN, int Line, int Col,
+               SourceMgr::DiagKind Kind, StringRef Msg, StringRef LineStr,
+               ArrayRef<std::pair<unsigned, unsigned>> Ranges,
                ArrayRef<SMFixIt> FixIts = None);
 
   const SourceMgr *getSourceMgr() const { return SM; }
@@ -293,13 +298,9 @@ class SMDiagnostic {
   StringRef getLineContents() const { return LineContents; }
   ArrayRef<std::pair<unsigned, unsigned>> getRanges() const { return Ranges; }
 
-  void addFixIt(const SMFixIt &Hint) {
-    FixIts.push_back(Hint);
-  }
+  void addFixIt(const SMFixIt &Hint) { FixIts.push_back(Hint); }
 
-  ArrayRef<SMFixIt> getFixIts() const {
-    return FixIts;
-  }
+  ArrayRef<SMFixIt> getFixIts() const { return FixIts; }
 
   void print(const char *ProgName, raw_ostream &S, bool ShowColors = true,
              bool ShowKindLabel = true) const;