[Swiftify] Support __sized_by on byte-sized pointee types

Previously we would emit a macro that would error on expansion when
trying to add a safe wrapper to a function with __sized_by on a type
that mapped to UnsafePointer<T> instead of UnsafeRawPointer or
OpaquePointer. __sized_by is acceptable when used on byte-sized pointee
types, so this adds machinery in the macro expansion to support that.
Meanwhile on the ClangImporter side, we add a check so that __sized_by
on pointee types with a size is ignored if that size is larger than 1
byte.

When _SwiftifyImport applies .sizedBy to a pointer of type
UnsafePointer<T> it will still map it to a
RawSpan/UnsafeRawBufferPointer in the safe overload. The assumption is
that any API using __sized_by is dealing with raw bytes, so raw pointers
are a better Swift abstraction than UnsafePointer<CChar> etc. It also
lets the user avoid doing a scary pointer cast from some potentially
larger-than-byte-sized pointer to a byte-sized pointer. Casts to
RawPointers are generally safer and more ergonomic.

rdar://150966684
rdar://150966021

Author: hnrklssn

lib/ClangImporter/ImportDecl.cpp

@@ -9145,16 +9145,38 @@ struct CountedByExpressionValidator
 };
 } // namespace
 
-static bool SwiftifiableCAT(const clang::CountAttributedType *CAT) {
-  return CAT && CountedByExpressionValidator().Visit(CAT->getCountExpr());
+// don't try to transform any Swift types that _SwiftifyImport doesn't know how to handle
+static bool SwiftifiableCountedByPointerType(Type swiftType) {
+  Type nonnullType = swiftType->lookThroughSingleOptionalType();
+  PointerTypeKind PTK;
+  return nonnullType->getAnyPointerElementType(PTK) &&
+    (PTK == PTK_UnsafePointer || PTK == PTK_UnsafeMutablePointer);
 }
-
-static bool SwiftifiablePointerType(Type swiftType) {
-  // don't try to transform any Swift types that _SwiftifyImport doesn't know how to handle
+static bool SwiftifiableSizedByPointerType(const clang::ASTContext &ctx,
+                                           Type swiftType,
+                                           const clang::CountAttributedType *CAT) {
   Type nonnullType = swiftType->lookThroughSingleOptionalType();
+  if (nonnullType->isOpaquePointer())
+    return true;
   PointerTypeKind PTK;
-  return nonnullType->isOpaquePointer() ||
-    (nonnullType->getAnyPointerElementType(PTK) && PTK != PTK_AutoreleasingUnsafeMutablePointer);
+  if (!nonnullType->getAnyPointerElementType(PTK))
+    return false;
+  if (PTK == PTK_UnsafeRawPointer || PTK == PTK_UnsafeMutableRawPointer)
+    return true;
+  if (PTK != PTK_UnsafePointer && PTK != PTK_UnsafeMutablePointer)
+    return false;
+  // We have a pointer to a type with a size. Verify that it is char-sized.
+  auto PtrT = CAT->getAs<clang::PointerType>();
+  auto PointeeT = PtrT->getPointeeType();
+  return ctx.getTypeSizeInChars(PointeeT).isOne();
+}
+static bool SwiftifiableCAT(const clang::ASTContext &ctx,
+                            const clang::CountAttributedType *CAT,
+                            Type swiftType) {
+  return CAT && CountedByExpressionValidator().Visit(CAT->getCountExpr()) &&
+    (CAT->isCountInBytes() ?
+       SwiftifiableSizedByPointerType(ctx, swiftType, CAT)
+     : SwiftifiableCountedByPointerType(swiftType));
 }
 
 void ClangImporter::Implementation::swiftify(FuncDecl *MappedDecl) {
@@ -9194,7 +9216,7 @@ void ClangImporter::Implementation::swiftify(FuncDecl *MappedDecl) {
     bool returnIsStdSpan = registerStdSpanTypeMapping(
         swiftReturnTy, ClangDecl->getReturnType());
     auto *CAT = ClangDecl->getReturnType()->getAs<clang::CountAttributedType>();
-    if (SwiftifiableCAT(CAT) && SwiftifiablePointerType(swiftReturnTy)) {
+    if (SwiftifiableCAT(getClangASTContext(), CAT, swiftReturnTy)) {
       printer.printCountedBy(CAT, SwiftifyInfoPrinter::RETURN_VALUE_INDEX);
       attachMacro = true;
     }
@@ -9211,7 +9233,7 @@ void ClangImporter::Implementation::swiftify(FuncDecl *MappedDecl) {
       Type swiftParamTy = swiftParam->getInterfaceType();
       bool paramHasBoundsInfo = false;
       auto *CAT = clangParamTy->getAs<clang::CountAttributedType>();
-      if (SwiftifiableCAT(CAT) && SwiftifiablePointerType(swiftParamTy)) {
+      if (SwiftifiableCAT(getClangASTContext(), CAT, swiftParamTy)) {
         printer.printCountedBy(CAT, index);
         attachMacro = paramHasBoundsInfo = true;
       }

lib/Macros/Sources/SwiftMacros/SwiftifyImportMacro.swift

@@ -329,10 +329,7 @@ func transformType(
   let text = name.text
   let isRaw = isRawPointerType(text: text)
   if isRaw && !isSizedBy {
-    throw DiagnosticError("raw pointers only supported for SizedBy", node: name)
-  }
-  if !isRaw && isSizedBy {
-    throw DiagnosticError("SizedBy only supported for raw pointers", node: name)
+    throw DiagnosticError("void pointers not supported for countedBy", node: name)
   }
 
   guard let kind: Mutability = getPointerMutability(text: text) else {
@@ -375,6 +372,33 @@ func isMutablePointerType(_ type: TypeSyntax) -> Bool {
   }
 }
 
+func getPointeeType(_ type: TypeSyntax) -> TypeSyntax? {
+  if let optType = type.as(OptionalTypeSyntax.self) {
+    return getPointeeType(optType.wrappedType)
+  }
+  if let impOptType = type.as(ImplicitlyUnwrappedOptionalTypeSyntax.self) {
+    return getPointeeType(impOptType.wrappedType)
+  }
+  if let attrType = type.as(AttributedTypeSyntax.self) {
+    return getPointeeType(attrType.baseType)
+  }
+
+  guard let idType = type.as(IdentifierTypeSyntax.self) else {
+    return nil
+  }
+  let text = idType.name.text
+  if text != "UnsafePointer" && text != "UnsafeMutablePointer" {
+    return nil
+  }
+  guard let x = idType.genericArgumentClause else {
+    return nil
+  }
+  guard let y = x.arguments.first else {
+    return nil
+  }
+  return y.argument.as(TypeSyntax.self)
+}
+
 protocol BoundsCheckedThunkBuilder {
   func buildFunctionCall(_ pointerArgs: [Int: ExprSyntax]) throws -> ExprSyntax
   // buildBasicBoundsChecks creates a variable with the same name as the parameter it replaced,
@@ -648,6 +672,7 @@ extension PointerBoundsThunkBuilder {
       return try transformType(oldType, generateSpan, isSizedBy, isParameter)
     }
   }
+
   var countLabel: String {
     return isSizedBy && generateSpan ? "byteCount" : "count"
   }
@@ -826,7 +851,7 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
     var args = argOverrides
     let argExpr = ExprSyntax("\(unwrappedName).baseAddress")
     assert(args[index] == nil)
-    args[index] = try castPointerToOpaquePointer(unwrapIfNonnullable(argExpr))
+    args[index] = try castPointerToTargetType(unwrapIfNonnullable(argExpr))
     let call = try base.buildFunctionCall(args)
     let ptrRef = unwrapIfNullable("\(name)")
 
@@ -871,11 +896,16 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
     return type
   }
 
-  func castPointerToOpaquePointer(_ baseAddress: ExprSyntax) throws -> ExprSyntax {
+  func castPointerToTargetType(_ baseAddress: ExprSyntax) throws -> ExprSyntax {
     let type = peelOptionalType(getParam(signature, index).type)
     if type.canRepresentBasicType(type: OpaquePointer.self) {
       return ExprSyntax("OpaquePointer(\(baseAddress))")
     }
+    if isSizedBy {
+      if let pointeeType = getPointeeType(type) {
+        return "\(baseAddress).assumingMemoryBound(to: \(pointeeType).self)"
+      }
+    }
     return baseAddress
   }
 
@@ -907,7 +937,7 @@ struct CountedOrSizedPointerThunkBuilder: ParamBoundsThunkBuilder, PointerBounds
       return unwrappedCall
     }
 
-    args[index] = try castPointerToOpaquePointer(getPointerArg())
+    args[index] = try castPointerToTargetType(getPointerArg())
     return try base.buildFunctionCall(args)
   }
 }

test/Interop/C/swiftify-import/Inputs/sized-by-lifetimebound.h

@@ -1,6 +1,10 @@
 #pragma once
 
+#include <stdint.h>
+
+#ifndef __sized_by
 #define __sized_by(x) __attribute__((__sized_by__(x)))
+#endif
 #define __lifetimebound __attribute__((lifetimebound))
 
 const void * __sized_by(len) simple(int len, int len2, const void * __sized_by(len2) __lifetimebound p);
@@ -18,4 +22,10 @@ const void * __sized_by(len) _Nullable nullable(int len, int len2, const void *
 typedef struct foo opaque_t;
 opaque_t * __sized_by(len) opaque(int len, int len2, opaque_t * __sized_by(len2) __lifetimebound p);
 
-const void * __sized_by(len) nonsizedLifetime(int len, const void * __lifetimebound p);
+const void * __sized_by(len) nonsizedLifetime(int len, const void * __lifetimebound p);
+
+uint8_t *__sized_by(size)  bytesized(int size, const uint8_t *__sized_by(size) __lifetimebound);
+
+char *__sized_by(size) charsized(char *__sized_by(size) __lifetimebound, int size);
+
+const uint16_t *__sized_by(size)  doublebytesized(uint16_t *__sized_by(size) __lifetimebound, int size);

test/Interop/C/swiftify-import/Inputs/sized-by-noescape.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <stdint.h>
+
 #define __sized_by(x) __attribute__((__sized_by__(x)))
 #define __noescape __attribute__((noescape))
 
@@ -23,3 +25,8 @@ const void * __sized_by(len) __noescape _Nonnull returnPointer(int len);
 typedef struct foo opaque_t;
 void opaque(int len, opaque_t * __sized_by(len) __noescape p);
 
+void bytesized(int size, const uint8_t *__sized_by(size) __noescape);
+
+void charsized(char *__sized_by(size) __noescape, int size);
+
+void doublebytesized(uint16_t *__sized_by(size) __noescape, int size);

test/Interop/C/swiftify-import/Inputs/sized-by.h

@@ -1,5 +1,7 @@
 #pragma once
 
+#include <stdint.h>
+
 #define __sized_by(x) __attribute__((__sized_by__(x)))
 
 void simple(int len, void * __sized_by(len) p);
@@ -21,3 +23,9 @@ void * __sized_by(len) returnPointer(int len);
 
 typedef struct foo opaque_t;
 void opaque(int len, opaque_t * __sized_by(len) p);
+
+void charsized(char *__sized_by(size), int size);
+
+uint8_t *__sized_by(size) bytesized(int size);
+
+void doublebytesized(uint16_t *__sized_by(size), int size);

test/Interop/C/swiftify-import/sized-by-lifetimebound.swift

@@ -13,6 +13,17 @@ import SizedByLifetimeboundClang
 
 // CHECK:      /// This is an auto-generated wrapper for safer interop
 // CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+// CHECK-NEXT: @lifetime(copy _bytesized_param1)
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func bytesized(_ _bytesized_param1: RawSpan) -> MutableRawSpan
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+// CHECK-NEXT: @lifetime(copy _charsized_param0)
+// CHECK-NEXT: @lifetime(_charsized_param0: copy _charsized_param0)
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func charsized(_ _charsized_param0: inout MutableRawSpan) -> MutableRawSpan
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
 // CHECK-NEXT: @lifetime(copy p)
 // CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func complexExpr(_ len: Int32, _ offset: Int32, _ p: RawSpan) -> RawSpan
 
@@ -93,3 +104,16 @@ public func callSimple(_ p: RawSpan) {
 public func callNonsizedLifetime(_ p: UnsafeRawPointer) {
   let _: RawSpan = nonsizedLifetime(73, p)
 }
+
+@available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+@inlinable
+public func callBytesized(_ p: RawSpan) {
+  let _: MutableRawSpan = bytesized(p)
+}
+
+@available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+@inlinable
+@lifetime(p: copy p)
+public func callCharsized(_ p: inout MutableRawSpan) {
+  let _: MutableRawSpan = charsized(&p)
+}

test/Interop/C/swiftify-import/sized-by-noescape.swift

@@ -12,6 +12,15 @@ import SizedByNoEscapeClang
 
 // CHECK:      /// This is an auto-generated wrapper for safer interop
 // CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func bytesized(_ _bytesized_param1: RawSpan)
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+// CHECK-NEXT: @lifetime(_charsized_param0: copy _charsized_param0)
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func charsized(_  _charsized_param0: inout MutableRawSpan)
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
 // CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func complexExpr(_ len: Int{{.*}}, _ offset: Int{{.*}}, _ p: RawSpan)
 
 // CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
@@ -93,3 +102,16 @@ public func callSimple(_ p: RawSpan) {
 public func callSwiftAttr(_ p: RawSpan) {
   swiftAttr(p)
 }
+
+@available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+@inlinable
+public func callBytesized(_ p: RawSpan) {
+  bytesized(p)
+}
+
+@available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
+@inlinable
+@lifetime(p: copy p)
+public func callCharsized(_ p: inout MutableRawSpan) {
+  charsized(&p)
+}

test/Interop/C/swiftify-import/sized-by.swift

@@ -3,14 +3,19 @@
 // RUN: %target-swift-ide-test -print-module -module-to-print=SizedByClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers | %FileCheck %s
 
 // swift-ide-test doesn't currently typecheck the macro expansions, so run the compiler as well
-// RUN: %empty-directory(%t)
-// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/SizedBy.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers %s
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -I %S/Inputs -enable-experimental-feature SafeInteropWrappers %s
 
 // Check that ClangImporter correctly infers and expands @_SwiftifyImport macros for functions with __sized_by parameters.
 import SizedByClang
 
 
 // CHECK:      /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func bytesized(_  size: Int{{.*}}) -> UnsafeMutableRawBufferPointer
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
+// CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func charsized(_  _charsized_param0: UnsafeMutableRawBufferPointer)
+
+// CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
 // CHECK-NEXT: @_alwaysEmitIntoClient @_disfavoredOverload public func complexExpr(_ len: Int{{.*}}, _ offset: Int{{.*}}, _ p: UnsafeMutableRawBufferPointer)
 
 // CHECK-NEXT: /// This is an auto-generated wrapper for safer interop
@@ -82,3 +87,13 @@ public func callSimple(_ p: UnsafeMutableRawBufferPointer) {
 public func callSwiftAttr(_ p: UnsafeMutableRawBufferPointer) {
   swiftAttr(p)
 }
+
+@inlinable
+public func callCharsized(_ p: UnsafeMutableRawBufferPointer) {
+  charsized(p)
+}
+
+@inlinable
+public func callBytesized() {
+  let _: UnsafeMutableRawBufferPointer = bytesized(37)
+}