[AST] Add Unsafe[...]Buffer types as known decls

Author: hamishknight

include/swift/AST/Decl.h

@@ -3220,6 +3220,14 @@ static inline bool isRawPointerKind(PointerTypeKind PTK) {
   llvm_unreachable("Unhandled PointerTypeKind in switch.");
 }
 
+// Kinds of buffer pointer types.
+enum BufferPointerTypeKind : unsigned {
+  BPTK_UnsafeMutableRawBufferPointer,
+  BPTK_UnsafeRawBufferPointer,
+  BPTK_UnsafeMutableBufferPointer,
+  BPTK_UnsafeBufferPointer,
+};
+
 enum KeyPathTypeKind : unsigned char {
   KPTK_AnyKeyPath,
   KPTK_PartialKeyPath,

include/swift/AST/KnownStdlibTypes.def

@@ -76,6 +76,11 @@ KNOWN_STDLIB_TYPE_DECL(UnsafePointer, NominalTypeDecl, 1)
 KNOWN_STDLIB_TYPE_DECL(OpaquePointer, NominalTypeDecl, 0)
 KNOWN_STDLIB_TYPE_DECL(AutoreleasingUnsafeMutablePointer, NominalTypeDecl, 1)
 
+KNOWN_STDLIB_TYPE_DECL(UnsafeBufferPointer, NominalTypeDecl, 1)
+KNOWN_STDLIB_TYPE_DECL(UnsafeMutableBufferPointer, NominalTypeDecl, 1)
+KNOWN_STDLIB_TYPE_DECL(UnsafeRawBufferPointer, NominalTypeDecl, 0)
+KNOWN_STDLIB_TYPE_DECL(UnsafeMutableRawBufferPointer, NominalTypeDecl, 0)
+
 KNOWN_STDLIB_TYPE_DECL(Unmanaged, NominalTypeDecl, 1)
 
 KNOWN_STDLIB_TYPE_DECL(Never, NominalTypeDecl, 0)

include/swift/AST/Types.h

@@ -49,6 +49,7 @@ enum class AllocationArena;
 class ArchetypeType;
 class AssociatedTypeDecl;
 class ASTContext;
+enum BufferPointerTypeKind : unsigned;
 class ClassDecl;
 class DependentMemberType;
 class GenericTypeParamDecl;
@@ -713,6 +714,14 @@ class alignas(1 << TypeAlignInBits) TypeBase {
   /// current type.
   Type wrapInPointer(PointerTypeKind kind);
 
+  /// Determines the element type of a known Unsafe[Mutable][Raw]BufferPointer
+  /// variant, or returns null if the type is not a buffer pointer.
+  Type getAnyBufferPointerElementType(BufferPointerTypeKind &BPTK);
+  Type getAnyBufferPointerElementType() {
+    BufferPointerTypeKind Ignore;
+    return getAnyBufferPointerElementType(Ignore);
+  }
+
   /// Determine whether the given type is "specialized", meaning that
   /// it involves generic types for which generic arguments have been provided.
   /// For example, the types Vector<Int> and Vector<Int>.Element are both

lib/AST/ASTContext.cpp

@@ -1384,6 +1384,10 @@ bool ASTContext::hasPointerArgumentIntrinsics() const {
     && getUnsafeMutablePointerDecl()
     && getUnsafePointerDecl()
     && (!LangOpts.EnableObjCInterop || getAutoreleasingUnsafeMutablePointerDecl())
+    && getUnsafeBufferPointerDecl()
+    && getUnsafeMutableBufferPointerDecl()
+    && getUnsafeRawBufferPointerDecl()
+    && getUnsafeMutableRawBufferPointerDecl()
     && getConvertPointerToPointerArgument()
     && getConvertMutableArrayToPointerArgument()
     && getConvertConstArrayToPointerArgument()

lib/AST/Type.cpp

@@ -645,6 +645,30 @@ Type TypeBase::wrapInPointer(PointerTypeKind kind) {
   return BoundGenericType::get(pointerDecl, /*parent*/nullptr, Type(this));
 }
 
+Type TypeBase::getAnyBufferPointerElementType(BufferPointerTypeKind &BPTK) {
+  auto &C = getASTContext();
+  if (auto nominalTy = getAs<NominalType>()) {
+    if (nominalTy->getDecl() == C.getUnsafeMutableRawBufferPointerDecl()) {
+      BPTK = BPTK_UnsafeMutableRawBufferPointer;
+    } else if (nominalTy->getDecl() == C.getUnsafeRawBufferPointerDecl()) {
+      BPTK = BPTK_UnsafeRawBufferPointer;
+    } else {
+      return Type();
+    }
+    return C.TheEmptyTupleType;
+  }
+  if (auto boundTy = getAs<BoundGenericType>()) {
+    if (boundTy->getDecl() == C.getUnsafeMutableBufferPointerDecl()) {
+      BPTK = BPTK_UnsafeMutableBufferPointer;
+    } else if (boundTy->getDecl() == C.getUnsafeBufferPointerDecl()) {
+      BPTK = BPTK_UnsafeBufferPointer;
+    } else {
+      return Type();
+    }
+    return boundTy->getGenericArgs()[0];
+  }
+  return Type();
+}
 
 Type TypeBase::lookThroughAllOptionalTypes() {
   Type type(this);