Add orion:supr_tramp

Author: kabiroberai

Sources/Orion/ClassHook+Deinit.swift

@@ -68,4 +68,8 @@ extension _GlueClassHookTrampoline {
     public func deinitSuprError(file: StaticString = #file, line: UInt = #line) -> Never {
         orionError("Do not call `supr.deinitializer()`.", file: file, line: line)
     }
+
+    public func trampOrigError(file: StaticString = #file, line: UInt = #line) -> Never {
+        orionError("Attempted to call orig on a supr_tramp method.", file: file, line: line)
+    }
 }

Sources/OrionProcessor/OrionData.swift

@@ -17,6 +17,12 @@ public struct OrionData {
             }
             return text
         }
+
+        fileprivate func fetchDirective<T: OrionDirective>(ofType type: T.Type) -> T? {
+            guard let dir = directives.compactMap({ $0 as? T }).last else { return nil }
+            dir.setUsed()
+            return dir
+        }
     }
 
     struct ClassHook {
@@ -35,16 +41,15 @@ public struct OrionData {
             var superClosure: Syntax // (UnsafeRawPointer, Selector, Blah) -> Blah
             var superClosureUnmanaged: Syntax // (UnsafeRawPointer, Selector, Blah) -> Unmanaged<Blah>
 
-            var isAddition: Bool {
-                if let directive = function.directives.compactMap({ $0 as? OrionDirectives.New }).last {
-                    directive.setUsed()
-                    return true
-                } else {
-                    return false
-                }
+            func isSuprTramp() -> Bool {
+                function.fetchDirective(ofType: OrionDirectives.SuprTramp.self) != nil
+            }
+
+            func isAddition() -> Bool {
+                function.fetchDirective(ofType: OrionDirectives.New.self) != nil
             }
 
-            var returnsRetained: Bool {
+            func returnsRetained() -> Bool {
                 // https://developer.apple.com/library/archive/documentation/Cocoa/Conceptual/MemoryMgmt/Articles/mmRules.html
                 // explains when a method should return a retained value based on the selector. Swift computes
                 // method selectors here:
@@ -53,8 +58,7 @@ public struct OrionData {
                 // the method identifier (Swift could add a With/And but that's not relevant to any of our
                 // prefixes, and additional selector parts have a colon before them.)
 
-                if let directive = function.directives.compactMap({ $0 as? OrionDirectives.ReturnsRetained }).last {
-                    directive.setUsed()
+                if let directive = function.fetchDirective(ofType: OrionDirectives.ReturnsRetained.self) {
                     return directive.mode == .retained
                 }
 

Sources/OrionProcessor/OrionDirective.swift

@@ -146,7 +146,8 @@ enum OrionDirectives {
     static let all: [OrionDirective.Type] = [
         ReturnsRetained.self,
         New.self,
-        Disable.self
+        Disable.self,
+        SuprTramp.self
     ]
 
     struct ReturnsRetained: OrionDirective {
@@ -185,7 +186,7 @@ enum OrionDirectives {
     }
 
     struct Disable: OrionDirective {
-        static var matchRule: OrionDirectiveMatchRule = .exact("disable")
+        static let matchRule: OrionDirectiveMatchRule = .exact("disable")
 
         let base: OrionDirectiveBase
         init(base: OrionDirectiveBase) throws {
@@ -195,4 +196,16 @@ enum OrionDirectives {
             }
         }
     }
+
+    struct SuprTramp: OrionDirective {
+        static let matchRule: OrionDirectiveMatchRule = .exact("supr_tramp")
+
+        let base: OrionDirectiveBase
+        init(base: OrionDirectiveBase) throws {
+            self.base = base
+            guard base.arguments.isEmpty else {
+                throw OrionDirectiveDiagnostic("supr_tramp directive expected zero arguments, got \(base.arguments.count)")
+            }
+        }
+    }
 }

Sources/OrionProcessor/OrionGenerator.swift

@@ -87,10 +87,10 @@ public final class OrionGenerator {
         from method: OrionData.ClassHook.Method,
         className: String,
         index: Int
-    ) -> (orig: String?, supr: String?, main: String, activation: String) {
+    ) -> (orig: String?, supr: String?, main: String, activation: String?) {
         let orig: String?
         let supr: String?
-        let register: String
+        let register: String?
 
         let args = arguments(for: method.function)
         let argsList = args.joined(separator: ", ")
@@ -99,11 +99,11 @@ public final class OrionGenerator {
         // order of operands matters here; we don't want to evaluate isAddition
         // if it's a deinitializer, so that Orion can notify the user if they
         // have a pointless `// orion:new`
-        let isAddition = !method.isDeinitializer && method.isAddition
+        let isAddition = !method.isDeinitializer && method.isAddition()
         let origIdent = "orion_\(isAddition ? "imp" : "orig")\(index)"
         let selIdent = "orion_sel\(index)"
 
-        let returnsRetained = !method.isDeinitializer && method.returnsRetained
+        let returnsRetained = !method.isDeinitializer && method.returnsRetained()
         let takeRetained = returnsRetained ? ".takeRetainedValue()" : ""
         let passRetained = returnsRetained ? "Unmanaged.passRetained" : ""
         let methodClosure = returnsRetained ? method.methodClosureUnmanaged : method.methodClosure
@@ -137,14 +137,16 @@ public final class OrionGenerator {
             addMethod(\(selIdent), \(origIdent), isClassMethod: \(method.isClassMethod))
             """
         } else {
+            let isTramp = method.isSuprTramp()
+
             // While we don't need to add @objc due to the class being @objcMembers (and the #selector
             // failing if the function can't be represented in objc), this results in better diagnostics
             // than merely having an error on the #selector line
             let funcOverride = "\(method.objcAttribute == nil ? "@objc " : "")\(method.function.function)"
 
             orig = """
             \(funcOverride) {
-                _Glue.\(origIdent)(target, _Glue.\(selIdent)\(commaArgs))\(takeRetained)
+                \(isTramp ? "trampOrigError()" : "_Glue.\(origIdent)(target, _Glue.\(selIdent)\(commaArgs))\(takeRetained)")
             }
             """
 
@@ -157,7 +159,7 @@ public final class OrionGenerator {
             }
             """
 
-            register = """
+            register = isTramp ? nil : """
             builder.addHook(\(selIdent), \(origIdent), isClassMethod: \(method.isClassMethod)) { \(origIdent) = $0 }
             """
         }
@@ -209,7 +211,7 @@ public final class OrionGenerator {
                 static let storage = initializeStorage()
         \(indentedMains)
                 static func activate(withClassHookBuilder builder: inout _GlueClassHookBuilder) {
-                    \(registers.joined(separator: "\n            "))
+                    \(registers.compactMap { $0 }.joined(separator: "\n            "))
                 }
             }
         }

Tests/OrionTests/ClassHookTests.swift

@@ -79,8 +79,8 @@ class InheritedHook: ClassHook<InheritedClass> {
 }
 
 class InitHook: ClassHook<InitClass> {
-    // just a placeholder to allow forwarding
-    func `init`() -> Target { orig.`init`() }
+    // orion:supr_tramp
+    func `init`() -> Target { fatalError() }
 
     func `init`(withX x: Int32) -> Target {
         let this = supr.`init`()

Tests/OrionTests/Generated.xc.swift

@@ -221,7 +221,7 @@ extension InitHook {
 
         final class OrigType: InitHook, _GlueClassHookTrampoline {
             @objc override func `init`() -> Target {
-                _Glue.orion_orig1(target, _Glue.orion_sel1)
+                trampOrigError()
             }
 
             @objc override func `init`(withX arg1: Int32) -> Target {
@@ -256,7 +256,6 @@ extension InitHook {
         }
 
         static func activate(withClassHookBuilder builder: inout _GlueClassHookBuilder) {
-            builder.addHook(orion_sel1, orion_orig1, isClassMethod: false) { orion_orig1 = $0 }
             builder.addHook(orion_sel2, orion_orig2, isClassMethod: false) { orion_orig2 = $0 }
         }
     }