Merge pull request #28 from p-x9/feature/improve-key-implementation

Improve key implementation

Author: p-x9

AssociatedObject.podspec

@@ -21,7 +21,7 @@ Pod::Spec.new do |s|
 
   s.prepare_command = 'swift build -c release && cp -f .build/release/AssociatedObjectPlugin ./Binary'
 
-  s.source_files  = "Sources/AssociatedObject/AssociatedObject.swift"
+  s.source_files  = "Sources/AssociatedObject/AssociatedObject.swift", 'Sources/AssociatedObjectC/**/*.{c,h,m,swift}'
   s.swift_versions = "5.9"
 
   s.preserve_paths = ["Binary/AssociatedObjectPlugin"]

Package.swift

@@ -30,9 +30,13 @@ let package = Package(
         .target(
             name: "AssociatedObject",
             dependencies: [
+                "AssociatedObjectC",
                 "AssociatedObjectPlugin"
             ]
         ),
+        .target(
+            name: "AssociatedObjectC"
+        ),
         .macro(
             name: "AssociatedObjectPlugin",
             dependencies: [

Sources/AssociatedObject/AssociatedObject.swift

@@ -1,5 +1,9 @@
 @_exported import ObjectiveC
 
+#if canImport(AssociatedObjectC)
+@_exported import AssociatedObjectC
+#endif
+
 @attached(peer, names: arbitrary)
 @attached(accessor)
 public macro AssociatedObject(

Sources/AssociatedObjectC/associated_object_key.c

@@ -0,0 +1 @@
+

Sources/AssociatedObjectC/include/associated_object_key.h

@@ -0,0 +1,39 @@
+//
+//  associated_object_key.h
+//  
+//
+//  Created by p-x9 on 2024/01/26.
+//  
+//
+
+#ifndef associated_object_key_h
+#define associated_object_key_h
+
+#if defined(__wasm__)
+// Wasm can't access call frame for security purposes
+#define get_return_address() ((void*) 0)
+#elif __GNUC__
+#define get_return_address() __builtin_return_address(0)
+#elif _MSC_VER
+#include <intrin.h>
+#define get_return_address() _ReturnAddress()
+#else
+#error missing implementation for get_return_address
+#define get_return_address() ((void*) 0)
+#endif
+
+#ifdef __GNUC__
+#define NOINLINE __attribute__((noinline))
+#elif _MSC_VER
+#define NOINLINE __declspec(noinline)
+#else
+#error missing implementation for `NOINLINE`
+#define NOINLINE
+#endif
+
+NOINLINE
+const void *_associated_object_key() {
+    return get_return_address();
+}
+
+#endif /* associated_object_key_h */

Sources/AssociatedObjectPlugin/AssociatedObjectMacro.swift

@@ -49,13 +49,22 @@ extension AssociatedObjectMacro: PeerMacro {
             return []
         }
 
+        let keyAccessor = """
+        _associated_object_key()
+        """
+
         let keyDecl = VariableDeclSyntax(
+            attributes: [
+                .attribute("@inline(never)")
+            ],
             bindingSpecifier: .identifier("static var"),
             bindings: PatternBindingListSyntax {
                 PatternBindingSyntax(
                     pattern: IdentifierPatternSyntax(identifier: .identifier("__associated_\(identifier.trimmed)Key")),
-                    typeAnnotation: .init(type: IdentifierTypeSyntax(name: .identifier("UInt8"))),
-                    initializer: InitializerClauseSyntax(value: ExprSyntax(stringLiteral: "0"))
+                    typeAnnotation: .init(type: IdentifierTypeSyntax(name: .identifier("UnsafeRawPointer"))),
+                    accessorBlock: .init(
+                        accessors: .getter("\(raw: keyAccessor)")
+                    )
                 )
             }
         )
@@ -65,6 +74,7 @@ extension AssociatedObjectMacro: PeerMacro {
         ]
 
         if type.isOptional && defaultValue != nil {
+            let flagName = "__associated_\(identifier.trimmed)IsSet"
             let flagDecl = VariableDeclSyntax(
                 attributes: [
                     .attribute("@_AssociatedObject(.OBJC_ASSOCIATION_ASSIGN)")
@@ -73,22 +83,33 @@ extension AssociatedObjectMacro: PeerMacro {
                 bindings: PatternBindingListSyntax {
                     PatternBindingSyntax(
                         pattern: IdentifierPatternSyntax(
-                            identifier: .identifier("__associated_\(identifier.trimmed)IsSet")
+                            identifier: .identifier(flagName)
                         ),
                         typeAnnotation: .init(type: IdentifierTypeSyntax(name: .identifier("Bool"))),
                         initializer: InitializerClauseSyntax(value: BooleanLiteralExprSyntax(false))
                     )
                 }
             )
+
+            // nested peer macro will not expand
+            // https://github.com/apple/swift/issues/69073
+            let keyAccessor = """
+            _associated_object_key()
+            """
             let flagKeyDecl = VariableDeclSyntax(
+                attributes: [
+                    .attribute("@inline(never)")
+                ],
                 bindingSpecifier: .identifier("static var"),
                 bindings: PatternBindingListSyntax {
                     PatternBindingSyntax(
                         pattern: IdentifierPatternSyntax(
                             identifier: .identifier("__associated___associated_\(identifier.trimmed)IsSetKey")
                         ),
-                        typeAnnotation: .init(type: IdentifierTypeSyntax(name: .identifier("UInt8"))),
-                        initializer: InitializerClauseSyntax(value: ExprSyntax(stringLiteral: "0"))
+                        typeAnnotation: .init(type: IdentifierTypeSyntax(name: .identifier("UnsafeRawPointer"))),
+                        accessorBlock: .init(
+                            accessors: .getter("\(raw: keyAccessor)")
+                        )
                     )
                 }
             )
@@ -160,7 +181,7 @@ extension AssociatedObjectMacro: AccessorMacro {
             return []
         }
 
-        var associatedKey: ExprSyntax = "&Self.__associated_\(identifier.trimmed)Key"
+        var associatedKey: ExprSyntax = "Self.__associated_\(identifier.trimmed)Key"
         if case let .argumentList(arguments) = node.arguments,
            let element = arguments.first(where: { $0.label?.text == "key" }),
            let customKey = element.expression.as(ExprSyntax.self) {
@@ -236,7 +257,7 @@ extension AssociatedObjectMacro {
                         """
                         if let value = objc_getAssociatedObject(
                             self,
-                            &Self.__associated_\(identifier.trimmed)Key
+                            \(associatedKey)
                         ) as? \(varTypeWithoutOptional.trimmed) {
                             return value
                         } else {