[clang][modules] Don't allocate when searching for FileID

Author: jansvoboda11

clang/lib/Serialization/ASTReader.cpp

@@ -1491,19 +1491,16 @@ int ASTReader::getSLocEntryID(SourceLocation::UIntTy SLocOffset) {
          "Corrupted global sloc offset map");
   ModuleFile *F = SLocMapI->second;
 
-  std::vector<unsigned> Indices(F->LocalNumSLocEntries);
-  for (unsigned I = 0; I != F->LocalNumSLocEntries; ++I)
-    Indices[I] = I;
-
-  auto It = llvm::upper_bound(Indices, SLocOffset,
-                    [&](SourceLocation::UIntTy Offset, unsigned Index) {
-                      if (F->SLocEntryOffsetLoaded[Index] == -1U) {
-                        auto MaybeEntryOffset = readSLocOffset(F, Index);
-                        assert(MaybeEntryOffset && "Corrupted AST file");
-                        F->SLocEntryOffsetLoaded[Index] = *MaybeEntryOffset;
-                      }
-                      return Offset < F->SLocEntryOffsetLoaded[Index];
-                    });
+  auto It = llvm::upper_bound(
+      llvm::index_range(0, F->LocalNumSLocEntries), SLocOffset,
+      [&](SourceLocation::UIntTy Offset, std::size_t Index) {
+        if (F->SLocEntryOffsetLoaded[Index] == -1U) {
+          auto MaybeEntryOffset = readSLocOffset(F, Index);
+          assert(MaybeEntryOffset && "Corrupted AST file");
+          F->SLocEntryOffsetLoaded[Index] = *MaybeEntryOffset;
+        }
+        return Offset < F->SLocEntryOffsetLoaded[Index];
+      });
   // The iterator points to the first entry with start offset greater than the
   // offset of interest. The previous entry must contain the offset of interest.
   It = std::prev(It);

llvm/include/llvm/ADT/STLExtras.h

@@ -2261,43 +2261,68 @@ template <typename... Refs> struct enumerator_result<std::size_t, Refs...> {
   mutable range_reference_tuple Storage;
 };
 
-/// Infinite stream of increasing 0-based `size_t` indices.
-struct index_stream {
-  struct iterator : iterator_facade_base<iterator, std::forward_iterator_tag,
-                                         const iterator> {
-    iterator &operator++() {
-      assert(Index != std::numeric_limits<std::size_t>::max() &&
-             "Attempting to increment end iterator");
-      ++Index;
-      return *this;
-    }
+struct index_iterator
+    : llvm::iterator_facade_base<index_iterator,
+                                 std::random_access_iterator_tag, std::size_t,
+                                 std::ptrdiff_t> {
+  index_iterator(std::size_t Index) : Index(Index) {}
+
+  index_iterator &operator+=(std::ptrdiff_t N) {
+    Index += N;
+    return *this;
+  }
 
-    // Note: This dereference operator returns a value instead of a reference
-    // and does not strictly conform to the C++17's definition of forward
-    // iterator. However, it satisfies all the forward_iterator requirements
-    // that the `zip_common` depends on and fully conforms to the C++20
-    // definition of forward iterator.
-    std::size_t operator*() const { return Index; }
+  index_iterator &operator-=(std::ptrdiff_t N) {
+    Index -= N;
+    return *this;
+  }
 
-    friend bool operator==(const iterator &Lhs, const iterator &Rhs) {
-      return Lhs.Index == Rhs.Index;
-    }
+  std::ptrdiff_t operator-(const index_iterator &R) const {
+    return Index - R.Index;
+  }
 
-    std::size_t Index = 0;
-  };
+  // Note: This dereference operator returns a value instead of a reference
+  // and does not strictly conform to the C++17's definition of forward
+  // iterator. However, it satisfies all the forward_iterator requirements
+  // that the `zip_common` depends on and fully conforms to the C++20
+  // definition of forward iterator.
+  std::size_t operator*() const { return Index; }
 
-  iterator begin() const { return {}; }
-  iterator end() const {
+  friend bool operator==(const index_iterator &Lhs, const index_iterator &Rhs) {
+    return Lhs.Index == Rhs.Index;
+  }
+
+  friend bool operator<(const index_iterator &Lhs, const index_iterator &Rhs) {
+    return Lhs.Index < Rhs.Index;
+  }
+
+private:
+  std::size_t Index;
+};
+
+/// Infinite stream of increasing 0-based `size_t` indices.
+struct index_stream {
+  index_iterator begin() const { return {0}; }
+  index_iterator end() const {
     // We approximate 'infinity' with the max size_t value, which should be good
     // enough to index over any container.
-    iterator It;
-    It.Index = std::numeric_limits<std::size_t>::max();
-    return It;
+    return index_iterator{std::numeric_limits<std::size_t>::max()};
   }
 };
 
 } // end namespace detail
 
+/// Increasing range of `size_t` indices.
+class index_range {
+  std::size_t Begin;
+  std::size_t End;
+
+public:
+  index_range(std::size_t Begin, std::size_t End) : Begin(Begin), End(End) {}
+  detail::index_iterator begin() const { return {Begin}; }
+  detail::index_iterator end() const { return {End}; }
+};
+
 /// Given two or more input ranges, returns a new range whose values are are
 /// tuples (A, B, C, ...), such that A is the 0-based index of the item in the
 /// sequence, and B, C, ..., are the values from the original input ranges. All