[CodeGen][arm64e] Add methods and data members to Address, which are needed to authenticate signed pointers (#67454)

To authenticate pointers, CodeGen needs access to the key and
discriminators that were used to sign the pointer. That information is
sometimes known from the context, but not always, which is why `Address`
needs to hold that information.

This patch adds methods and data members to `Address`, which will be
needed in subsequent patches to authenticate signed pointers, and uses
the newly added methods throughout CodeGen. Although this patch isn't
strictly NFC as it causes CodeGen to use different code paths in some
cases (e.g., `mergeAddressesInConditionalExpr`), it doesn't cause any
changes in functionality as it doesn't add any information needed for
authentication.

In addition to the changes mentioned above, this patch introduces class
`RawAddress`, which contains a pointer that we know is unsigned, and
adds several new functions for creating `Address` and `LValue` objects.

Author: ahatanak

clang/lib/CodeGen/ABIInfoImpl.cpp

@@ -187,7 +187,7 @@ CodeGen::emitVoidPtrDirectVAArg(CodeGenFunction &CGF, Address VAListAddr,
   CharUnits FullDirectSize = DirectSize.alignTo(SlotSize);
   Address NextPtr =
       CGF.Builder.CreateConstInBoundsByteGEP(Addr, FullDirectSize, "argp.next");
-  CGF.Builder.CreateStore(NextPtr.getPointer(), VAListAddr);
+  CGF.Builder.CreateStore(NextPtr.emitRawPointer(CGF), VAListAddr);
 
   // If the argument is smaller than a slot, and this is a big-endian
   // target, the argument will be right-adjusted in its slot.
@@ -239,8 +239,8 @@ Address CodeGen::emitMergePHI(CodeGenFunction &CGF, Address Addr1,
                               const llvm::Twine &Name) {
   assert(Addr1.getType() == Addr2.getType());
   llvm::PHINode *PHI = CGF.Builder.CreatePHI(Addr1.getType(), 2, Name);
-  PHI->addIncoming(Addr1.getPointer(), Block1);
-  PHI->addIncoming(Addr2.getPointer(), Block2);
+  PHI->addIncoming(Addr1.emitRawPointer(CGF), Block1);
+  PHI->addIncoming(Addr2.emitRawPointer(CGF), Block2);
   CharUnits Align = std::min(Addr1.getAlignment(), Addr2.getAlignment());
   return Address(PHI, Addr1.getElementType(), Align);
 }
@@ -400,7 +400,7 @@ Address CodeGen::EmitVAArgInstr(CodeGenFunction &CGF, Address VAListAddr,
     llvm::Type *ElementTy = CGF.ConvertTypeForMem(Ty);
     llvm::Type *BaseTy = llvm::PointerType::getUnqual(ElementTy);
     llvm::Value *Addr =
-        CGF.Builder.CreateVAArg(VAListAddr.getPointer(), BaseTy);
+        CGF.Builder.CreateVAArg(VAListAddr.emitRawPointer(CGF), BaseTy);
     return Address(Addr, ElementTy, TyAlignForABI);
   } else {
     assert((AI.isDirect() || AI.isExtend()) &&
@@ -416,7 +416,7 @@ Address CodeGen::EmitVAArgInstr(CodeGenFunction &CGF, Address VAListAddr,
            "Unexpected CoerceToType seen in arginfo in generic VAArg emitter!");
 
     Address Temp = CGF.CreateMemTemp(Ty, "varet");
-    Val = CGF.Builder.CreateVAArg(VAListAddr.getPointer(),
+    Val = CGF.Builder.CreateVAArg(VAListAddr.emitRawPointer(CGF),
                                   CGF.ConvertTypeForMem(Ty));
     CGF.Builder.CreateStore(Val, Temp);
     return Temp;

clang/lib/CodeGen/Address.h

@@ -15,35 +15,49 @@
 #define LLVM_CLANG_LIB_CODEGEN_ADDRESS_H
 
 #include "clang/AST/CharUnits.h"
+#include "clang/AST/Type.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/Support/MathExtras.h"
 
 namespace clang {
 namespace CodeGen {
 
+class Address;
+class CGBuilderTy;
+class CodeGenFunction;
+class CodeGenModule;
+
 // Indicates whether a pointer is known not to be null.
 enum KnownNonNull_t { NotKnownNonNull, KnownNonNull };
 
-/// An aligned address.
-class Address {
+/// An abstract representation of an aligned address. This is designed to be an
+/// IR-level abstraction, carrying just the information necessary to perform IR
+/// operations on an address like loads and stores.  In particular, it doesn't
+/// carry C type information or allow the representation of things like
+/// bit-fields; clients working at that level should generally be using
+/// `LValue`.
+/// The pointer contained in this class is known to be unsigned.
+class RawAddress {
   llvm::PointerIntPair<llvm::Value *, 1, bool> PointerAndKnownNonNull;
   llvm::Type *ElementType;
   CharUnits Alignment;
 
 protected:
-  Address(std::nullptr_t) : ElementType(nullptr) {}
+  RawAddress(std::nullptr_t) : ElementType(nullptr) {}
 
 public:
-  Address(llvm::Value *Pointer, llvm::Type *ElementType, CharUnits Alignment,
-          KnownNonNull_t IsKnownNonNull = NotKnownNonNull)
+  RawAddress(llvm::Value *Pointer, llvm::Type *ElementType, CharUnits Alignment,
+             KnownNonNull_t IsKnownNonNull = NotKnownNonNull)
       : PointerAndKnownNonNull(Pointer, IsKnownNonNull),
         ElementType(ElementType), Alignment(Alignment) {
     assert(Pointer != nullptr && "Pointer cannot be null");
     assert(ElementType != nullptr && "Element type cannot be null");
   }
 
-  static Address invalid() { return Address(nullptr); }
+  inline RawAddress(Address Addr);
+
+  static RawAddress invalid() { return RawAddress(nullptr); }
   bool isValid() const {
     return PointerAndKnownNonNull.getPointer() != nullptr;
   }
@@ -80,6 +94,133 @@ class Address {
     return Alignment;
   }
 
+  /// Return address with different element type, but same pointer and
+  /// alignment.
+  RawAddress withElementType(llvm::Type *ElemTy) const {
+    return RawAddress(getPointer(), ElemTy, getAlignment(), isKnownNonNull());
+  }
+
+  KnownNonNull_t isKnownNonNull() const {
+    assert(isValid());
+    return (KnownNonNull_t)PointerAndKnownNonNull.getInt();
+  }
+};
+
+/// Like RawAddress, an abstract representation of an aligned address, but the
+/// pointer contained in this class is possibly signed.
+class Address {
+  friend class CGBuilderTy;
+
+  // The boolean flag indicates whether the pointer is known to be non-null.
+  llvm::PointerIntPair<llvm::Value *, 1, bool> Pointer;
+
+  /// The expected IR type of the pointer. Carrying accurate element type
+  /// information in Address makes it more convenient to work with Address
+  /// values and allows frontend assertions to catch simple mistakes.
+  llvm::Type *ElementType = nullptr;
+
+  CharUnits Alignment;
+
+  /// Offset from the base pointer.
+  llvm::Value *Offset = nullptr;
+
+  llvm::Value *emitRawPointerSlow(CodeGenFunction &CGF) const;
+
+protected:
+  Address(std::nullptr_t) : ElementType(nullptr) {}
+
+public:
+  Address(llvm::Value *pointer, llvm::Type *elementType, CharUnits alignment,
+          KnownNonNull_t IsKnownNonNull = NotKnownNonNull)
+      : Pointer(pointer, IsKnownNonNull), ElementType(elementType),
+        Alignment(alignment) {
+    assert(pointer != nullptr && "Pointer cannot be null");
+    assert(elementType != nullptr && "Element type cannot be null");
+    assert(!alignment.isZero() && "Alignment cannot be zero");
+  }
+
+  Address(llvm::Value *BasePtr, llvm::Type *ElementType, CharUnits Alignment,
+          llvm::Value *Offset, KnownNonNull_t IsKnownNonNull = NotKnownNonNull)
+      : Pointer(BasePtr, IsKnownNonNull), ElementType(ElementType),
+        Alignment(Alignment), Offset(Offset) {}
+
+  Address(RawAddress RawAddr)
+      : Pointer(RawAddr.isValid() ? RawAddr.getPointer() : nullptr),
+        ElementType(RawAddr.isValid() ? RawAddr.getElementType() : nullptr),
+        Alignment(RawAddr.isValid() ? RawAddr.getAlignment()
+                                    : CharUnits::Zero()) {}
+
+  static Address invalid() { return Address(nullptr); }
+  bool isValid() const { return Pointer.getPointer() != nullptr; }
+
+  /// This function is used in situations where the caller is doing some sort of
+  /// opaque "laundering" of the pointer.
+  void replaceBasePointer(llvm::Value *P) {
+    assert(isValid() && "pointer isn't valid");
+    assert(P->getType() == Pointer.getPointer()->getType() &&
+           "Pointer's type changed");
+    Pointer.setPointer(P);
+    assert(isValid() && "pointer is invalid after replacement");
+  }
+
+  CharUnits getAlignment() const { return Alignment; }
+
+  void setAlignment(CharUnits Value) { Alignment = Value; }
+
+  llvm::Value *getBasePointer() const {
+    assert(isValid() && "pointer isn't valid");
+    return Pointer.getPointer();
+  }
+
+  /// Return the type of the pointer value.
+  llvm::PointerType *getType() const {
+    return llvm::PointerType::get(
+        ElementType,
+        llvm::cast<llvm::PointerType>(Pointer.getPointer()->getType())
+            ->getAddressSpace());
+  }
+
+  /// Return the type of the values stored in this address.
+  llvm::Type *getElementType() const {
+    assert(isValid());
+    return ElementType;
+  }
+
+  /// Return the address space that this address resides in.
+  unsigned getAddressSpace() const { return getType()->getAddressSpace(); }
+
+  /// Return the IR name of the pointer value.
+  llvm::StringRef getName() const { return Pointer.getPointer()->getName(); }
+
+  // This function is called only in CGBuilderBaseTy::CreateElementBitCast.
+  void setElementType(llvm::Type *Ty) {
+    assert(hasOffset() &&
+           "this funcion shouldn't be called when there is no offset");
+    ElementType = Ty;
+  }
+
+  /// Whether the pointer is known not to be null.
+  KnownNonNull_t isKnownNonNull() const {
+    assert(isValid());
+    return (KnownNonNull_t)Pointer.getInt();
+  }
+
+  Address setKnownNonNull() {
+    assert(isValid());
+    Pointer.setInt(KnownNonNull);
+    return *this;
+  }
+
+  bool hasOffset() const { return Offset; }
+
+  llvm::Value *getOffset() const { return Offset; }
+
+  /// Return the pointer contained in this class after authenticating it and
+  /// adding offset to it if necessary.
+  llvm::Value *emitRawPointer(CodeGenFunction &CGF) const {
+    return getBasePointer();
+  }
+
   /// Return address with different pointer, but same element type and
   /// alignment.
   Address withPointer(llvm::Value *NewPointer,
@@ -91,61 +232,59 @@ class Address {
   /// Return address with different alignment, but same pointer and element
   /// type.
   Address withAlignment(CharUnits NewAlignment) const {
-    return Address(getPointer(), getElementType(), NewAlignment,
+    return Address(Pointer.getPointer(), getElementType(), NewAlignment,
                    isKnownNonNull());
   }
 
   /// Return address with different element type, but same pointer and
   /// alignment.
   Address withElementType(llvm::Type *ElemTy) const {
-    return Address(getPointer(), ElemTy, getAlignment(), isKnownNonNull());
-  }
-
-  /// Whether the pointer is known not to be null.
-  KnownNonNull_t isKnownNonNull() const {
-    assert(isValid());
-    return (KnownNonNull_t)PointerAndKnownNonNull.getInt();
-  }
-
-  /// Set the non-null bit.
-  Address setKnownNonNull() {
-    assert(isValid());
-    PointerAndKnownNonNull.setInt(true);
-    return *this;
+    if (!hasOffset())
+      return Address(getBasePointer(), ElemTy, getAlignment(), nullptr,
+                     isKnownNonNull());
+    Address A(*this);
+    A.ElementType = ElemTy;
+    return A;
   }
 };
 
+inline RawAddress::RawAddress(Address Addr)
+    : PointerAndKnownNonNull(Addr.isValid() ? Addr.getBasePointer() : nullptr,
+                             Addr.isValid() ? Addr.isKnownNonNull()
+                                            : NotKnownNonNull),
+      ElementType(Addr.isValid() ? Addr.getElementType() : nullptr),
+      Alignment(Addr.isValid() ? Addr.getAlignment() : CharUnits::Zero()) {}
+
 /// A specialization of Address that requires the address to be an
 /// LLVM Constant.
-class ConstantAddress : public Address {
-  ConstantAddress(std::nullptr_t) : Address(nullptr) {}
+class ConstantAddress : public RawAddress {
+  ConstantAddress(std::nullptr_t) : RawAddress(nullptr) {}
 
 public:
   ConstantAddress(llvm::Constant *pointer, llvm::Type *elementType,
                   CharUnits alignment)
-      : Address(pointer, elementType, alignment) {}
+      : RawAddress(pointer, elementType, alignment) {}
 
   static ConstantAddress invalid() {
     return ConstantAddress(nullptr);
   }
 
   llvm::Constant *getPointer() const {
-    return llvm::cast<llvm::Constant>(Address::getPointer());
+    return llvm::cast<llvm::Constant>(RawAddress::getPointer());
   }
 
   ConstantAddress withElementType(llvm::Type *ElemTy) const {
     return ConstantAddress(getPointer(), ElemTy, getAlignment());
   }
 
-  static bool isaImpl(Address addr) {
+  static bool isaImpl(RawAddress addr) {
     return llvm::isa<llvm::Constant>(addr.getPointer());
   }
-  static ConstantAddress castImpl(Address addr) {
+  static ConstantAddress castImpl(RawAddress addr) {
     return ConstantAddress(llvm::cast<llvm::Constant>(addr.getPointer()),
                            addr.getElementType(), addr.getAlignment());
   }
 };
-
 }
 
 // Present a minimal LLVM-like casting interface.