[Swiftify] Support __sized_by on byte-sized pointee types

lib/ClangImporter/ImportDecl.cpp

@@ -9145,16 +9145,39 @@ 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 +9217,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 +9234,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,6 +1,10 @@
 #pragma once
 
+#include <stdint.h>
+
+#ifndef __sized_by
 #define __sized_by(x) __attribute__((__sized_by__(x)))
+#endif
 
 void simple(int len, void * __sized_by(len) p);
 
@@ -21,3 +25,15 @@ void * __sized_by(len) returnPointer(int len);
 
 typedef struct foo opaque_t;
 void opaque(int len, opaque_t * __sized_by(len) p);
+
+typedef opaque_t *opaqueptr_t;
+void opaqueptr(int len, opaqueptr_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);
+
+typedef uint8_t * bytesizedptr_t;
+void aliasedBytesized(bytesizedptr_t __sized_by(size) p, int size);

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

@@ -1,11 +1,11 @@
 // REQUIRES: swift_feature_SafeInteropWrappers
 // REQUIRES: swift_feature_LifetimeDependence
 
-// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByLifetimeboundClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByLifetimeboundClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -Xcc -Wno-nullability-completeness | %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/CountedByLifetimebound.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence %s
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/CountedByLifetimebound.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence -strict-memory-safety -warnings-as-errors -Xcc -Werror -Xcc -Wno-nullability-completeness %s
 
 // Check that ClangImporter correctly infers and expands @_SwiftifyImport macros for functions with __sized_by __lifetimebound parameters and return values.
 
@@ -98,7 +98,7 @@ public func callSimple(_ p: inout MutableSpan<CInt>) {
 @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
 @inlinable
 public func callNoncountedLifetime(_ p: UnsafeMutablePointer<CInt>) {
-  let _: MutableSpan<CInt> = noncountedLifetime(73, p)
+  let _: MutableSpan<CInt> = unsafe noncountedLifetime(73, p)
 }
 
 @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)

test/Interop/C/swiftify-import/counted-by-no-swiftify.swift

@@ -1,4 +1,4 @@
-// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -Xcc -Wno-nullability-completeness | %FileCheck %s
 
 // REQUIRES: swift_feature_SafeInteropWrappers
 

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

@@ -1,11 +1,11 @@
 // REQUIRES: swift_feature_SafeInteropWrappers
 // REQUIRES: swift_feature_LifetimeDependence
 
-// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByNoEscapeClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByNoEscapeClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence -Xcc -Wno-nullability-completeness | %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/CountedByNoEscape.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence %s
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/CountedByNoEscape.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers -enable-experimental-feature LifetimeDependence -strict-memory-safety -warnings-as-errors -Xcc -Werror -Xcc -Wno-nullability-completeness %s
 
 // Check that ClangImporter correctly infers and expands @_SwiftifyImport macros for functions with __counted_by __noescape parameters.
 
@@ -150,14 +150,14 @@ public func callNullable(_ p: inout MutableSpan<CInt>?) {
 @lifetime(p: copy p)
 @inlinable
 public func callReturnLifetimeBound(_ p: inout MutableSpan<CInt>) {
-  let a: MutableSpan<CInt> = returnLifetimeBound(2, &p)
+  let _: MutableSpan<CInt> = returnLifetimeBound(2, &p)
 }
 
 @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
 @inlinable
 public func callReturnPointer() {
-  let a: UnsafeMutableBufferPointer<CInt>? = returnPointer(4) // call wrapper
-  let b: UnsafeMutablePointer<CInt>? = returnPointer(4) // call unsafe interop
+  let _: UnsafeMutableBufferPointer<CInt>? = returnPointer(4) // call wrapper
+  let _: UnsafeMutablePointer<CInt>? = returnPointer(4) // call unsafe interop
 }
 
 @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)
@@ -207,7 +207,7 @@ public func callFunc(_ p: inout MutableSpan<CInt>?) {
 @lifetime(p: copy p)
 @inlinable
 public func callFuncRenameKeyword(_ p: inout MutableSpan<CInt>?) {
-  funcRenamed(func: &p, extension: 1, init: 2, open: 3, var: 4, is: 5, as: 6, in: 7, guard: 8, where: 9)
+  let _ = funcRenamed(func: &p, extension: 1, init: 2, open: 3, var: 4, is: 5, as: 6, in: 7, guard: 8, where: 9)
 }
 
 @available(visionOS 1.1, tvOS 12.2, watchOS 5.2, iOS 12.2, macOS 10.14.4, *)

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

@@ -1,10 +1,10 @@
 // REQUIRES: swift_feature_SafeInteropWrappers
 
-// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers | %FileCheck %s
+// RUN: %target-swift-ide-test -print-module -module-to-print=CountedByClang -plugin-path %swift-plugin-dir -I %S/Inputs -source-filename=x -enable-experimental-feature SafeInteropWrappers -Xcc -Wno-nullability-completeness | %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/CountedBy.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers %s
+// RUN: %target-swift-frontend -emit-module -plugin-path %swift-plugin-dir -o %t/CountedBy.swiftmodule -I %S/Inputs -enable-experimental-feature SafeInteropWrappers -strict-memory-safety -warnings-as-errors -Xcc -Werror -Xcc -Wno-nullability-completeness %s
 
 // Check that ClangImporter correctly infers and expands @_SwiftifyImport macros for functions with __counted_by parameters.
 
@@ -67,73 +67,73 @@ import CountedByClang
 
 @inlinable
 public func callComplexExpr(_ p: UnsafeMutableBufferPointer<CInt>) {
-  complexExpr(CInt(p.count), 1, p)
+  unsafe complexExpr(CInt(p.count), 1, p)
 }
 
 @inlinable
 public func callConstInt(_ p: UnsafeMutableBufferPointer<CInt>) {
-  constInt(p)
+  unsafe constInt(p)
 }
 
 @inlinable
 public func callNonnull(_ p: UnsafeMutableBufferPointer<CInt>) {
-  nonnull(p)
+  unsafe nonnull(p)
 }
 
 @inlinable
 public func callNullUnspecified(_ p: UnsafeMutableBufferPointer<CInt>) {
-  nullUnspecified(p)
+  unsafe nullUnspecified(p)
 }
 
 @inlinable
 public func callNullable(_ p: UnsafeMutableBufferPointer<CInt>?) {
-  nullable(p)
+  unsafe nullable(p)
 }
 
 @inlinable
 public func callOffByOne(_ p: UnsafeMutableBufferPointer<CInt>) {
-  offByOne(0, p)
+  unsafe offByOne(0, p)
 }
 
 @inlinable
 public func callReturnPointer() {
-  let a: UnsafeMutableBufferPointer<CInt>? = returnPointer(4) // call wrapper
-  let b: UnsafeMutablePointer<CInt>? = returnPointer(4) // call unsafe interop
+  let _: UnsafeMutableBufferPointer<CInt>? = returnPointer(4) // call wrapper
+  let _: UnsafeMutablePointer<CInt>? = returnPointer(4) // call unsafe interop
 }
 
 @inlinable
 public func callScalar(_ p: UnsafeMutableBufferPointer<CInt>) {
-  scalar(4, 2, p)
+  unsafe scalar(4, 2, p)
 }
 
 @inlinable
 public func callShared(_ p: UnsafeMutableBufferPointer<CInt>, _ p2: UnsafeMutableBufferPointer<CInt>) {
-  shared(p, p2)
+  unsafe shared(p, p2)
 }
 
 @inlinable
 public func callSimple(_ p: UnsafeMutableBufferPointer<CInt>) {
-  simple(p)
+  unsafe simple(p)
 }
 
 @inlinable
 public func callSimpleIndirectOriginal(_ p: UnsafeMutablePointer<CInt>) {
-  let f = simple
-  f(13, p)
+  let f = unsafe simple
+  unsafe f(13, p)
 }
 
 @inlinable
 public func callSimpleIndirectOverload(_ p: UnsafeMutableBufferPointer<CInt>) {
-  let f: (UnsafeMutableBufferPointer<CInt>) -> Void = simple
-  f(p)
+  let f: (UnsafeMutableBufferPointer<CInt>) -> Void = unsafe simple
+  unsafe f(p)
 }
 
 @inlinable
 public func callSimpleFlipped(_ p: UnsafeMutableBufferPointer<CInt>) {
-  simpleFlipped(p)
+  unsafe simpleFlipped(p)
 }
 
 @inlinable
 public func callSwiftAttr(_ p: UnsafeMutableBufferPointer<CInt>) {
-  swiftAttr(p)
+  unsafe swiftAttr(p)
 }