[IR][AArch64] Add "ptrauth(...)" Constant to represent signed pointers.

This defines a new kind of IR Constant that represents a ptrauth
signed pointer, as used in AArch64 PAuth.

It allows representing most kinds of signed pointer constants
used thus far in the llvm ptrauth implementations, notably those
used in the Darwin and ELF ABIs being implemented for c/c++.
These signed pointer constants are then lowered to ELF/MachO
relocations.

These can be simply thought of as a constant `llvm.ptrauth.sign`,
with the interesting addition of discriminator computation:
the `ptrauth` constant can also represent a combined blend,
when both address and integer discriminator operands are used.

Co-Authored-by: Tim Northover <[email protected]>

Author: ahmedbougacha

llvm/docs/LangRef.rst

@@ -4748,6 +4748,33 @@ reference to the CFI jump table in the ``LowerTypeTests`` pass. These constants
 may be useful in low-level programs, such as operating system kernels, which
 need to refer to the actual function body.
 
+.. _ptrauth
+
+Authenticated Pointers
+----------------------
+
+``ptrauth (ptr CST, i32 KEY, ptr ADDRDISC, i16 DISC)
+
+A '``ptrauth``' constant represents a pointer with a cryptographic
+authentication signature embedded into some bits. Its type is the same as the
+first argument.
+
+
+If the address disciminator is ``null`` then the expression is equivalent to
+
+.. code-block:llvm
+    %tmp = call i64 @llvm.ptrauth.sign.i64(i64 ptrtoint (ptr CST to i64), i32 KEY, i64 DISC)
+    %val = inttoptr i64 %tmp to ptr
+
+If the address discriminator is present, then it is
+
+.. code-block:llvm
+    %tmp1 = call i64 @llvm.ptrauth.blend.i64(i64 ptrtoint (ptr ADDRDISC to i64), i64 DISC)
+    %tmp2 = call i64 @llvm.ptrauth.sign.i64(i64 ptrtoint (ptr CST to i64), i64  %tmp1)
+    %val = inttoptr i64 %tmp2 to ptr
+
+    %tmp = call i64 @llvm.ptrauth.blend.i64
+
 .. _constantexprs:
 
 Constant Expressions

llvm/include/llvm-c/Core.h

@@ -286,6 +286,7 @@ typedef enum {
   LLVMInstructionValueKind,
   LLVMPoisonValueValueKind,
   LLVMConstantTargetNoneValueKind,
+  LLVMConstantPtrAuthValueKind,
 } LLVMValueKind;
 
 typedef enum {

llvm/include/llvm/AsmParser/LLToken.h

@@ -343,6 +343,7 @@ enum Kind {
   kw_insertvalue,
   kw_blockaddress,
   kw_dso_local_equivalent,
+  kw_ptrauth,
   kw_no_cfi,
 
   kw_freeze,

llvm/include/llvm/Bitcode/LLVMBitCodes.h

@@ -411,6 +411,8 @@ enum ConstantsCodes {
                               //                 sideeffect|alignstack|
                               //                 asmdialect|unwind,
                               //                 asmstr,conststr]
+  CST_CODE_SIGNED_PTR = 31,   // CE_SIGNED_PTR: [ptrty, ptr, key,
+                              //                 addrdiscty, addrdisc, disc]
 };
 
 /// CastOpcodes - These are values used in the bitcode files to encode which

llvm/include/llvm/IR/Constants.h

@@ -1006,6 +1006,69 @@ struct OperandTraits<NoCFIValue> : public FixedNumOperandTraits<NoCFIValue, 1> {
 
 DEFINE_TRANSPARENT_OPERAND_ACCESSORS(NoCFIValue, Value)
 
+/// A signed pointer
+///
+class ConstantPtrAuth final : public Constant {
+  friend struct ConstantPtrAuthKeyType;
+  friend class Constant;
+
+  ConstantPtrAuth(Constant *Ptr, ConstantInt *Key, Constant *AddrDisc,
+                  ConstantInt *Disc);
+
+  void *operator new(size_t s) { return User::operator new(s, 4); }
+
+  void destroyConstantImpl();
+  Value *handleOperandChangeImpl(Value *From, Value *To);
+
+public:
+  /// Return a pointer authenticated with the specified parameters.
+  static ConstantPtrAuth *get(Constant *Ptr, ConstantInt *Key,
+                              Constant *AddrDisc, ConstantInt *Disc);
+
+  /// Produce a new ptrauth expression signing the given value using
+  /// the same schema as is stored in one.
+  ConstantPtrAuth *getWithSameSchema(Constant *Pointer) const;
+
+  /// Transparently provide more efficient getOperand methods.
+  DECLARE_TRANSPARENT_OPERAND_ACCESSORS(Constant);
+
+  /// The pointer that is authenticated in this authenticated global reference.
+  Constant *getPointer() const { return (Constant *)Op<0>().get(); }
+
+  /// The Key ID, an i32 constant.
+  ConstantInt *getKey() const { return (ConstantInt *)Op<1>().get(); }
+
+  /// The address discriminator if any, or the null constant.
+  /// If present, this must be a value equivalent to the storage location of
+  /// the only user of the authenticated ptrauth global.
+  Constant *getAddrDiscriminator() const { return (Constant *)Op<2>().get(); }
+
+  /// The discriminator.
+  ConstantInt *getDiscriminator() const { return (ConstantInt *)Op<3>().get(); }
+
+  /// Whether there is any non-null address discriminator.
+  bool hasAddressDiversity() const {
+    return !getAddrDiscriminator()->isNullValue();
+  }
+
+  /// Check whether an authentication operation with key \p KeyV and (possibly
+  /// blended) discriminator \p DiscriminatorV is compatible with this
+  /// authenticated global reference.
+  bool isCompatibleWith(const Value *Key, const Value *Discriminator,
+                        const DataLayout &DL) const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast:
+  static bool classof(const Value *V) {
+    return V->getValueID() == ConstantPtrAuthVal;
+  }
+};
+
+template <>
+struct OperandTraits<ConstantPtrAuth>
+    : public FixedNumOperandTraits<ConstantPtrAuth, 4> {};
+
+DEFINE_TRANSPARENT_OPERAND_ACCESSORS(ConstantPtrAuth, Constant)
+
 //===----------------------------------------------------------------------===//
 /// A constant value that is initialized with an expression using
 /// other constant values.

llvm/include/llvm/IR/Value.def

@@ -78,6 +78,7 @@ HANDLE_GLOBAL_VALUE(GlobalAlias)
 HANDLE_GLOBAL_VALUE(GlobalIFunc)
 HANDLE_GLOBAL_VALUE(GlobalVariable)
 HANDLE_CONSTANT(BlockAddress)
+HANDLE_CONSTANT(ConstantPtrAuth)
 HANDLE_CONSTANT(ConstantExpr)
 HANDLE_CONSTANT_EXCLUDE_LLVM_C_API(DSOLocalEquivalent)
 HANDLE_CONSTANT_EXCLUDE_LLVM_C_API(NoCFIValue)

llvm/lib/Analysis/ValueTracking.cpp

@@ -2900,6 +2900,10 @@ bool isKnownNonZero(const Value *V, const APInt &DemandedElts, unsigned Depth,
       return true;
     }
 
+    // Constant ptrauth can be null, iff the base pointer can be.
+    if (auto *CPA = dyn_cast<ConstantPtrAuth>(V))
+      return isKnownNonZero(CPA->getPointer(), DemandedElts, Depth, Q);
+
     // A global variable in address space 0 is non null unless extern weak
     // or an absolute symbol reference. Other address spaces may have null as a
     // valid address for a global, so we can't assume anything.
@@ -6993,6 +6997,9 @@ static bool isGuaranteedNotToBeUndefOrPoison(
         isa<ConstantPointerNull>(C) || isa<Function>(C))
       return true;
 
+    if (isa<ConstantPtrAuth>(C))
+      return true;
+
     if (C->getType()->isVectorTy() && !isa<ConstantExpr>(C))
       return (!includesUndef(Kind) ? !C->containsPoisonElement()
                                    : !C->containsUndefOrPoisonElement()) &&

llvm/lib/AsmParser/LLLexer.cpp

@@ -708,6 +708,7 @@ lltok::Kind LLLexer::LexIdentifier() {
   KEYWORD(blockaddress);
   KEYWORD(dso_local_equivalent);
   KEYWORD(no_cfi);
+  KEYWORD(ptrauth);
 
   // Metadata types.
   KEYWORD(distinct);

llvm/lib/AsmParser/LLParser.cpp

@@ -3998,6 +3998,48 @@ bool LLParser::parseValID(ValID &ID, PerFunctionState *PFS, Type *ExpectedTy) {
     ID.NoCFI = true;
     return false;
   }
+  case lltok::kw_ptrauth: {
+    // ValID ::= 'ptrauth' '(' ptr @foo ',' i32 <key> ','
+    //                         ptr addrdisc ',' i64 <disc> ')'
+    Lex.Lex();
+
+    Constant *Ptr, *Key, *AddrDisc, *Disc;
+
+    if (parseToken(lltok::lparen,
+                   "expected '(' in signed pointer expression") ||
+        parseGlobalTypeAndValue(Ptr) ||
+        parseToken(lltok::comma,
+                   "expected comma in signed pointer expression") ||
+        parseGlobalTypeAndValue(Key) ||
+        parseToken(lltok::comma,
+                   "expected comma in signed pointer expression") ||
+        parseGlobalTypeAndValue(AddrDisc) ||
+        parseToken(lltok::comma,
+                   "expected comma in signed pointer expression") ||
+        parseGlobalTypeAndValue(Disc) ||
+        parseToken(lltok::rparen, "expected ')' in signed pointer expression"))
+      return true;
+
+    if (!Ptr->getType()->isPointerTy())
+      return error(ID.Loc, "signed pointer must be a pointer");
+
+    auto KeyC = dyn_cast<ConstantInt>(Key);
+    if (!KeyC || KeyC->getBitWidth() != 32)
+      return error(ID.Loc, "signed pointer key must be i32 constant integer");
+
+    if (!AddrDisc->getType()->isPointerTy())
+      return error(ID.Loc,
+                   "signed pointer address discriminator must be a pointer");
+
+    auto DiscC = dyn_cast<ConstantInt>(Disc);
+    if (!DiscC || DiscC->getBitWidth() != 64)
+      return error(ID.Loc,
+                   "signed pointer discriminator must be i64 constant integer");
+
+    ID.ConstantVal = ConstantPtrAuth::get(Ptr, KeyC, AddrDisc, DiscC);
+    ID.Kind = ValID::t_Constant;
+    return false;
+  }
 
   case lltok::kw_trunc:
   case lltok::kw_bitcast:

llvm/lib/Bitcode/Reader/BitcodeReader.cpp

@@ -504,7 +504,8 @@ class BitcodeConstant final : public Value,
   static constexpr uint8_t NoCFIOpcode = 252;
   static constexpr uint8_t DSOLocalEquivalentOpcode = 251;
   static constexpr uint8_t BlockAddressOpcode = 250;
-  static constexpr uint8_t FirstSpecialOpcode = BlockAddressOpcode;
+  static constexpr uint8_t ConstantPtrAuthOpcode = 249;
+  static constexpr uint8_t FirstSpecialOpcode = ConstantPtrAuthOpcode;
 
   // Separate struct to make passing different number of parameters to
   // BitcodeConstant::create() more convenient.
@@ -1528,6 +1529,18 @@ Expected<Value *> BitcodeReader::materializeValue(unsigned StartValID,
         C = ConstantExpr::get(BC->Opcode, ConstOps[0], ConstOps[1], BC->Flags);
       } else {
         switch (BC->Opcode) {
+        case BitcodeConstant::ConstantPtrAuthOpcode: {
+          auto *Key = dyn_cast<ConstantInt>(ConstOps[1]);
+          if (!Key)
+            return error("ptrauth key operand must be ConstantInt");
+
+          auto *Disc = dyn_cast<ConstantInt>(ConstOps[3]);
+          if (!Disc)
+            return error("ptrauth disc operand must be ConstantInt");
+
+          C = ConstantPtrAuth::get(ConstOps[0], Key, ConstOps[2], Disc);
+          break;
+        }
         case BitcodeConstant::NoCFIOpcode: {
           auto *GV = dyn_cast<GlobalValue>(ConstOps[0]);
           if (!GV)
@@ -3596,6 +3609,21 @@ Error BitcodeReader::parseConstants() {
                                   Record[1]);
       break;
     }
+    case bitc::CST_CODE_SIGNED_PTR: {
+      if (Record.size() < 6)
+        return error("Invalid record");
+      Type *PtrTy = getTypeByID(Record[0]);
+      if (!PtrTy)
+        return error("Invalid record");
+
+      // PtrTy, Ptr, Key, AddrDiscTy, AddrDisc, Disc
+      V = BitcodeConstant::create(
+        Alloc, CurTy, BitcodeConstant::ConstantPtrAuthOpcode,
+        {(unsigned)Record[1], (unsigned)Record[2], (unsigned)Record[4],
+         (unsigned)Record[5]});
+
+      break;
+    }
     }
 
     assert(V->getType() == getTypeByID(CurTyID) && "Incorrect result type ID");

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -2800,6 +2800,14 @@ void ModuleBitcodeWriter::writeConstants(unsigned FirstVal, unsigned LastVal,
       Record.push_back(VE.getTypeID(BA->getFunction()->getType()));
       Record.push_back(VE.getValueID(BA->getFunction()));
       Record.push_back(VE.getGlobalBasicBlockID(BA->getBasicBlock()));
+    } else if (const ConstantPtrAuth *SP = dyn_cast<ConstantPtrAuth>(C)) {
+      Code = bitc::CST_CODE_SIGNED_PTR;
+      Record.push_back(VE.getTypeID(SP->getPointer()->getType()));
+      Record.push_back(VE.getValueID(SP->getPointer()));
+      Record.push_back(VE.getValueID(SP->getKey()));
+      Record.push_back(VE.getTypeID(SP->getAddrDiscriminator()->getType()));
+      Record.push_back(VE.getValueID(SP->getAddrDiscriminator()));
+      Record.push_back(VE.getValueID(SP->getDiscriminator()));
     } else if (const auto *Equiv = dyn_cast<DSOLocalEquivalent>(C)) {
       Code = bitc::CST_CODE_DSO_LOCAL_EQUIVALENT;
       Record.push_back(VE.getTypeID(Equiv->getGlobalValue()->getType()));

llvm/lib/IR/AsmWriter.cpp

@@ -1578,6 +1578,21 @@ static void WriteConstantInternal(raw_ostream &Out, const Constant *CV,
     return;
   }
 
+  if (const ConstantPtrAuth *SP = dyn_cast<ConstantPtrAuth>(CV)) {
+    Out << "ptrauth (";
+
+    for (unsigned i = 0; i < SP->getNumOperands(); ++i) {
+      WriterCtx.TypePrinter->print(SP->getOperand(i)->getType(), Out);
+      Out << ' ';
+      WriteAsOperandInternal(Out, SP->getOperand(i), WriterCtx);
+      if (i != SP->getNumOperands() - 1)
+        Out << ", ";
+    }
+
+    Out << ')';
+    return;
+  }
+
   if (const ConstantArray *CA = dyn_cast<ConstantArray>(CV)) {
     Type *ETy = CA->getType()->getElementType();
     Out << '[';

llvm/lib/IR/ConstantFold.cpp

@@ -1154,6 +1154,9 @@ static ICmpInst::Predicate evaluateICmpRelation(Constant *V1, Constant *V2) {
                                 GV->getType()->getAddressSpace()))
         return ICmpInst::ICMP_UGT;
     }
+  } else if (const ConstantPtrAuth *SP = dyn_cast<ConstantPtrAuth>(V1)) {
+    // FIXME: ahmedbougacha: implement ptrauth cst comparison
+    return ICmpInst::BAD_ICMP_PREDICATE;
   } else {
     // Ok, the LHS is known to be a constantexpr.  The RHS can be any of a
     // constantexpr, a global, block address, or a simple constant.