[clang][bytecode] Fix comparing the addresses of union members (#133852)

Union members get the same address, so we can't just use
`Pointer::getByteOffset()`.

4 files changed, 84 insertions, 2 deletions

diff --git a/clang/lib/AST/ByteCode/Interp.h b/clang/lib/AST/ByteCode/Interp.h
index 938077a..6fe1d4b 100644
--- a/clang/lib/AST/ByteCode/Interp.h
+++ b/clang/lib/AST/ByteCode/Interp.h
@@ -1070,9 +1070,18 @@ inline bool CmpHelperEQ<Pointer>(InterpState &S, CodePtr OpPC, CompareFn Fn) {
   }
 
   if (Pointer::hasSameBase(LHS, RHS)) {
+    if (LHS.inUnion() && RHS.inUnion()) {
+      // If the pointers point into a union, things are a little more
+      // complicated since the offset we save in interp::Pointer can't be used
+      // to compare the pointers directly.
+      size_t A = LHS.computeOffsetForComparison();
+      size_t B = RHS.computeOffsetForComparison();
+      S.Stk.push<BoolT>(BoolT::from(Fn(Compare(A, B))));
+      return true;
+    }
+
     unsigned VL = LHS.getByteOffset();
     unsigned VR = RHS.getByteOffset();
-
     // In our Pointer class, a pointer to an array and a pointer to the first
     // element in the same array are NOT equal. They have the same Base value,
     // but a different Offset. This is a pretty rare case, so we fix this here
diff --git a/clang/lib/AST/ByteCode/Pointer.cpp b/clang/lib/AST/ByteCode/Pointer.cpp
index 79b47c2..5b52261 100644
--- a/clang/lib/AST/ByteCode/Pointer.cpp
+++ b/clang/lib/AST/ByteCode/Pointer.cpp
@@ -339,6 +339,39 @@ void Pointer::print(llvm::raw_ostream &OS) const {
   }
 }
 
+/// Compute an integer that can be used to compare this pointer to
+/// another one.
+size_t Pointer::computeOffsetForComparison() const {
+  if (!isBlockPointer())
+    return Offset;
+
+  size_t Result = 0;
+  Pointer P = *this;
+  while (!P.isRoot()) {
+    if (P.isArrayRoot()) {
+      P = P.getBase();
+      continue;
+    }
+    if (P.isArrayElement()) {
+      P = P.expand();
+      Result += (P.getIndex() * P.elemSize());
+      P = P.getArray();
+      continue;
+    }
+
+    if (const Record *R = P.getBase().getRecord(); R && R->isUnion()) {
+      // Direct child of a union - all have offset 0.
+      P = P.getBase();
+      continue;
+    }
+
+    Result += P.getInlineDesc()->Offset;
+    P = P.getBase();
+  }
+
+  return Result;
+}
+
 std::string Pointer::toDiagnosticString(const ASTContext &Ctx) const {
   if (isZero())
     return "nullptr";
diff --git a/clang/lib/AST/ByteCode/Pointer.h b/clang/lib/AST/ByteCode/Pointer.h
index fd33ee9..988237d 100644
--- a/clang/lib/AST/ByteCode/Pointer.h
+++ b/clang/lib/AST/ByteCode/Pointer.h
@@ -417,7 +417,7 @@ public:
     return false;
   }
   bool inUnion() const {
-    if (isBlockPointer())
+    if (isBlockPointer() && asBlockPointer().Base >= sizeof(InlineDescriptor))
       return getInlineDesc()->InUnion;
     return false;
   };
@@ -727,6 +727,8 @@ public:
   /// Prints the pointer.
   void print(llvm::raw_ostream &OS) const;
 
+  size_t computeOffsetForComparison() const;
+
 private:
   friend class Block;
   friend class DeadBlock;
diff --git a/clang/test/AST/ByteCode/unions.cpp b/clang/test/AST/ByteCode/unions.cpp
index 66b8389..3911a2b 100644
--- a/clang/test/AST/ByteCode/unions.cpp
+++ b/clang/test/AST/ByteCode/unions.cpp
@@ -600,3 +600,41 @@ namespace MoveOrAssignOp {
   static_assert(foo());
 }
 #endif
+
+namespace AddressComparison {
+  union {
+    int a;
+    int c;
+  } U;
+  static_assert(__builtin_addressof(U.a) == (void*)__builtin_addressof(U.c));
+  static_assert(&U.a == &U.c);
+
+
+  struct {
+    union {
+      struct {
+        int a;
+        int b;
+      } a;
+      struct {
+        int b;
+        int a;
+      }b;
+    } u;
+    int b;
+  } S;
+
+  static_assert(&S.u.a.a == &S.u.b.b);
+  static_assert(&S.u.a.b != &S.u.b.b);
+  static_assert(&S.u.a.b == &S.u.b.b); // both-error {{failed}}
+
+
+  union {
+    int a[2];
+    int b[2];
+  } U2;
+
+  static_assert(&U2.a[0] == &U2.b[0]);
+  static_assert(&U2.a[0] != &U2.b[1]);
+  static_assert(&U2.a[0] == &U2.b[1]); // both-error {{failed}}
+}