add identity support and tests

swiftwinrt/Resources/Support/Aggregation.swift

@@ -11,7 +11,7 @@ import Foundation
 // to the swift object, to ensure it doesn't get cleaned up. The Swift object in turn holds a strong
 // reference to this object so that it stays alive.
 @_spi(WinRTInternal)
-public final class WinRTClassWeakReference<Class: WinRTClass> {
+public final class WinRTClassWeakReference<Class: WinRTObject> {
     fileprivate weak var instance: Class?
     public init(_ instance: Class){
         self.instance = instance
@@ -41,7 +41,7 @@ extension WinRTClassWeakReference: CustomAddRef {
 }
 
 extension WinRTClassWeakReference: AnyObjectWrapper {
-    var obj: AnyObject? { instance }
+    var obj: WinRTObject? { instance }
 }
 
 @_spi(WinRTInternal)

swiftwinrt/Resources/Support/IInspectable.swift

@@ -40,7 +40,7 @@ open class IInspectable: IUnknown {
 // }
 // internal typealias Composable = IBaseNoOverrides
 protocol AnyObjectWrapper {
-    var obj: AnyObject? { get }
+    var obj: WinRTObject? { get }
 }
 
 public enum __ABI_ {
@@ -53,7 +53,13 @@ public enum __ABI_ {
         } else {
           let vtblPtr = withUnsafeMutablePointer(to: &IInspectableVTable) { $0 }
           let cAbi = C_IInspectable(lpVtbl: vtblPtr)
-          super.init(cAbi, swift as AnyObject)
+          // For WinRTObject types, hold a weak reference to the swift obj since
+          // it will hold a strong reference to ourselves
+          if let winrtObj = swift as? WinRTObject {
+            super.init(cAbi, WinRTClassWeakReference(winrtObj))
+          } else {
+            super.init(cAbi, swift as AnyObject)
+          }
         }
       }
 
@@ -62,12 +68,14 @@ public enum __ABI_ {
         if let swiftAbi = swiftObj as? IInspectable {
            let abi: UnsafeMutablePointer<C_IInspectable> = RawPointer(swiftAbi)
            return try body(abi)
-        } else if let winrtObj = swiftObj as? WinRTObject {
+        } else if let winrtObjWrapper = swiftObj as? AnyObjectWrapper,
+          let winrtObj = winrtObjWrapper.obj {
           if let identity = winrtObj.identity {
-            return try body(identity)
+            return try body(identity.get())
           }
           return try super.toABI{
-            winrtObj.identity = $0
+            winrtObj.identity = .init($0)
+            (winrtObjWrapper as? CustomAddRef)?.release()
             return try body($0)
           }
         } else {
@@ -124,8 +132,9 @@ public enum __ABI_ {
 
         GetRuntimeClassName: {
             guard let instance = AnyWrapper.tryUnwrapFrom(raw: $0) else { return E_INVALIDARG }
-            guard let winrtClass = instance as? WinRTClass else {
-              let string = String(reflecting: type(of: instance))
+            let winrtObj = (instance as? AnyObjectWrapper)?.obj ?? instance
+            guard let winrtClass = winrtObj as? WinRTClass else {
+              let string = String(reflecting: type(of: winrtObj))
               $1!.pointee = try! HString(string).detach()
               return S_OK
             }

swiftwinrt/Resources/Support/WinRTProtocols.swift

@@ -17,16 +17,23 @@ open class WinRTObject: Equatable {
   public init() {}
 
   public static func == (lhs: WinRTObject, rhs: WinRTObject) -> Bool {
-    return lhs.identity == rhs.identity
+    return lhs.identity?.get() == rhs.identity?.get()
   }
 
-  // Weakref to the underlying COM object since it holds a strong ref back
-  // to ourselves
-  internal var identity: UnsafeMutablePointer<C_IInspectable>?
+  internal var identity: ComPtr<C_IInspectable>?
+
+  @_spi(WinRTImplements)
+  open func queryInterface(_ iid: SUPPORT_MODULE.IID) -> IUnknownRef? {
+    // the vtables for IInspectable will properly respond to all default QueryInterface
+    // calls for IInspectable and IUnknown. The queryInterface method exists for subclasses
+    // to provide custom handling of QueryInterface calls for other interfaces that the code
+    // gen is unaware of
+    return nil
+  }
 }
 
-open class WinRTClass : CustomQueryInterface, Equatable {
-    public init() {}
+open class WinRTClass : WinRTObject, CustomQueryInterface {
+    override public init() { super.init() }
 
     @_spi(WinRTInternal)
     public init(_ ptr: SUPPORT_MODULE.IInspectable) {
@@ -47,10 +54,8 @@ open class WinRTClass : CustomQueryInterface, Equatable {
     @_spi(WinRTInternal)
     public internal(set) var _inner: SUPPORT_MODULE.IInspectable!
 
-    var identity: ComPtr<C_IInspectable>?
-
     @_spi(WinRTImplements)
-    open func queryInterface(_ iid: SUPPORT_MODULE.IID) -> IUnknownRef? {
+    override open func queryInterface(_ iid: SUPPORT_MODULE.IID) -> IUnknownRef? {
         SUPPORT_MODULE.queryInterface(self, iid)
     }
 

tests/test_app/IdentityTests.swift

@@ -29,10 +29,50 @@ class IdentityTests: XCTestCase {
         XCTAssertEqual(abi1, abi2)
     }
   }
+
+  public func testIdentityCopyTo() throws {
+    let obj = MyObj()
+    let wrapper = try XCTUnwrap(__ABI_.AnyWrapper(obj))
+    let wrapper2 = try XCTUnwrap(__ABI_.AnyWrapper(obj))
+    var copy: UnsafeMutablePointer<C_IInspectable>?
+    var copy2: UnsafeMutablePointer<C_IInspectable>?
+    defer {
+      _ = copy?.pointee.lpVtbl.pointee.Release(copy)
+      _ = copy2?.pointee.lpVtbl.pointee.Release(copy2)
+    }
+
+    wrapper.copyTo(&copy)
+    wrapper2.copyTo(&copy2)
+    XCTAssertEqual(copy, copy2)
+  }
+
+  // This test verifies that even after releasing the last reference
+  // from WinRT, that the original identity still persists and is valid.
+  // This is common in collection accessors, where the object is retrieved
+  // and then immediately released.
+  public func testIdentityAfterRelease() throws {
+    let obj = MyObj()
+    var copy2: UnsafeMutablePointer<C_IInspectable>?
+
+    var wrapper:__ABI_.AnyWrapper? = try XCTUnwrap(__ABI_.AnyWrapper(obj))
+    var copy: UnsafeMutablePointer<C_IInspectable>?
+
+    wrapper?.copyTo(&copy)
+    _ = copy?.pointee.lpVtbl.pointee.Release(copy)
+    wrapper = nil
+
+
+    let wrapper2 = try XCTUnwrap(__ABI_.AnyWrapper(obj))
+    wrapper2.copyTo(&copy2)
+    XCTAssertEqual(copy, copy2)
+    _ = copy2?.pointee.lpVtbl.pointee.Release(copy2)
+  }
 }
 
 var identityTests: [XCTestCaseEntry] = [
   testCase([
-    ("testIdentity", IdentityTests.testIdentity)
+    ("testIdentity", IdentityTests.testIdentity),
+    ("testIdentityCopyTo", IdentityTests.testIdentityCopyTo),
+    ("testIdentityAfterRelease", IdentityTests.testIdentityAfterRelease)
   ])
 ]

tests/test_app/MemoryManagementTests.swift

@@ -51,6 +51,24 @@ class MemoryManagementTests : XCTestCase {
         }()
         XCTAssertNil(weakDerived)
     }
+
+    func testWinRTObject() throws {
+        weak var weakDerived: MyObj? = nil
+        try {
+            let obj = MyObj()
+            weakDerived = obj
+
+            // wrapping MyObj will create the identity pointer. make sure
+            // this doesn't cause a leak
+            let wrapper = try XCTUnwrap(__ABI_.AnyWrapper(obj))
+            var copy: UnsafeMutablePointer<C_IInspectable>?
+            defer {
+              _ = copy?.pointee.lpVtbl.pointee.Release(copy)
+            }
+            wrapper.copyTo(&copy)
+        }()
+        XCTAssertNil(weakDerived)
+    }
 }
 
 var memoryManagementTests: [XCTestCaseEntry] = [
@@ -59,5 +77,6 @@ var memoryManagementTests: [XCTestCaseEntry] = [
     ("testNonAggregatedObject", MemoryManagementTests.testNonAggregatedObject),
     ("testReturningAggregatedObject", MemoryManagementTests.testReturningAggregatedObject),
     ("testReturningNonAggregatedObject", MemoryManagementTests.testReturningNonAggregatedObject),
+    ("testWinRTObject", MemoryManagementTests.testWinRTObject)
   ])
 ]