[llvm] 5ccff5a - ADT: Fix pointer comparison UB in SmallVector

[Duncan P. N. Exon Smith via llvm-commits]

Author: Duncan P. N. Exon Smith
Date: 2021-01-11T15:31:04-08:00
New Revision: 5ccff5aaa68ab789834c4463ce05b05e57593b34

URL: https://github.com/llvm/llvm-project/commit/5ccff5aaa68ab789834c4463ce05b05e57593b34
DIFF: https://github.com/llvm/llvm-project/commit/5ccff5aaa68ab789834c4463ce05b05e57593b34.diff

LOG: ADT: Fix pointer comparison UB in SmallVector

The standard requires comparisons of pointers to unrelated storage to
use `std::less`. Split out some helpers that do that and update all the
code that was comparing using `<` and friends (mostly assertions).

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

Added: 
    

Modified: 
    llvm/include/llvm/ADT/SmallVector.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/SmallVector.h b/llvm/include/llvm/ADT/SmallVector.h
index 3bbde2d9c0fb..803588143d81 100644
--- a/llvm/include/llvm/ADT/SmallVector.h
+++ b/llvm/include/llvm/ADT/SmallVector.h
@@ -136,11 +136,32 @@ class SmallVectorTemplateCommon
     this->Size = this->Capacity = 0; // FIXME: Setting Capacity to 0 is suspect.
   }
 
+  /// Return true if V is an internal reference to the given range.
+  bool isReferenceToRange(const void *V, const void *First, const void *Last) const {
+    // Use std::less to avoid UB.
+    std::less<> LessThan;
+    return !LessThan(V, First) && LessThan(V, Last);
+  }
+
+  /// Return true if V is an internal reference to this vector.
+  bool isReferenceToStorage(const void *V) const {
+    return isReferenceToRange(V, this->begin(), this->end());
+  }
+
+  /// Return true if First and Last form a valid (possibly empty) range in this
+  /// vector's storage.
+  bool isRangeInStorage(const void *First, const void *Last) const {
+    // Use std::less to avoid UB.
+    std::less<> LessThan;
+    return !LessThan(First, this->begin()) && !LessThan(Last, First) &&
+           !LessThan(this->end(), Last);
+  }
+
   /// Return true unless Elt will be invalidated by resizing the vector to
   /// NewSize.
   bool isSafeToReferenceAfterResize(const void *Elt, size_t NewSize) {
     // Past the end.
-    if (LLVM_LIKELY(Elt < this->begin() || Elt >= this->end()))
+    if (LLVM_LIKELY(!isReferenceToStorage(Elt)))
       return true;
 
     // Return false if Elt will be destroyed by shrinking.
@@ -579,8 +600,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     // Just cast away constness because this is a non-const member function.
     iterator I = const_cast<iterator>(CI);
 
-    assert(I >= this->begin() && "Iterator to erase is out of bounds.");
-    assert(I < this->end() && "Erasing at past-the-end iterator.");
+    assert(this->isReferenceToStorage(CI) && "Iterator to erase is out of bounds.");
 
     iterator N = I;
     // Shift all elts down one.
@@ -595,9 +615,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     iterator S = const_cast<iterator>(CS);
     iterator E = const_cast<iterator>(CE);
 
-    assert(S >= this->begin() && "Range to erase is out of bounds.");
-    assert(S <= E && "Trying to erase invalid range.");
-    assert(E <= this->end() && "Trying to erase past the end.");
+    assert(this->isRangeInStorage(S, E) && "Range to erase is out of bounds.");
 
     iterator N = S;
     // Shift all elts down.
@@ -615,8 +633,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
       return this->end()-1;
     }
 
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
+    assert(this->isReferenceToStorage(I) && "Insertion iterator is out of bounds.");
 
     // Check that adding an element won't invalidate Elt.
     this->assertSafeToAdd(&Elt);
@@ -635,7 +652,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
     // If we just moved the element we're inserting, be sure to update
     // the reference.
     std::remove_reference_t<ArgType> *EltPtr = &Elt;
-    if (I <= EltPtr && EltPtr < this->end())
+    if (this->isReferenceToRange(EltPtr, I, this->end()))
       ++EltPtr;
 
     *I = ::std::forward<ArgType>(*EltPtr);
@@ -658,8 +675,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
       return this->begin()+InsertElt;
     }
 
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
+    assert(this->isReferenceToStorage(I) && "Insertion iterator is out of bounds.");
 
     // Check that adding NumToInsert elements won't invalidate Elt.
     this->assertSafeToAdd(&Elt, NumToInsert);
@@ -716,8 +732,7 @@ class SmallVectorImpl : public SmallVectorTemplateBase<T> {
       return this->begin()+InsertElt;
     }
 
-    assert(I >= this->begin() && "Insertion iterator is out of bounds.");
-    assert(I <= this->end() && "Inserting past the end of the vector.");
+    assert(this->isReferenceToStorage(I) && "Insertion iterator is out of bounds.");
 
     // Check that the reserve that follows doesn't invalidate the iterators.
     this->assertSafeToAddRange(From, To);