From 292e25a23a5d841c9715a392110e9e1620fe5951 Mon Sep 17 00:00:00 2001
From: Chinedu Francis Nwafili <frankie.nwafili@gmail.com>
Date: Sun, 28 Apr 2024 23:07:32 -0400
Subject: [PATCH] Fix memory leak for `Option<SwiftType>`

This commit fixes a memory leak when passing `Option<SwiftType>` to
`extern "Rust"` functions.

For example, the following bridge module would lead to memory leaks:

```rust
#[swift_bridge::bridge]
mod ffi {
    extern "Swift" {
        type SomeSwiftType;
    }

    extern "Rust" {
        fn rust_fn_takes_swift_type(arg: Option<SomeSwiftType>);
    }
}
```

When the above `rust_fn_takes_swift_type` function was called the
`SomeSwiftType` would be retained twice.

When the Rust side later freed `SomeSwiftType` the retain count would
drop from 2 to 1.

Since there was still 1 retain, `SomeSwiftType` would never get freed.

## Impact

Support for passing `Option<SomeSwiftType>` from Swift to Rust was
introduced earlier today.

This functionality was released in `swift-bridge = "0.1.54"`.

W will be deploying `swift-bridge = "0.1.55" today. "0.1.55" will
contain a fix for the leak.

Because support for `Option<SomeSwiftType>` was released today, and the
fix will be released today, it should be unlikely that anyone has
experienced this leak.

We will yank "0.1.54" from "crates.io".

All users (probably zero) that are bridging `Option<SomeSwiftType>` should
upgrade from "0.1.54" to "0.1.55" to.

## Solution

This commit does the following:

- When passing ownership of a Swift type from Swift to Rust we now
  increment the retain count by one. Before this commit we were
  incrementing it by two.

- When passing ownership of a Swift type from Rust to Swift we avoid
  calling the the Rust handle's `Drop` implementation. We avoid this by
  using `std::mem::forget`.
  This ensures that the retain count does not get decremented.

- When passing ownership of a Swift type from Rust to Swift the Swift
  side does not increment the retain count.

With all of the above we should now be able to:

- Create an `Option<SwiftType>` Swift type on the Swift side

- Pass ownership of the `Option<SwiftType> to Rust. The retain count is now 1.

- Pass ownership back to Swift. The retain count is still 1.

Before this commit, when passing an `Option<SwiftType>` it was:

- Create a Swift type on the Swift side

- Pass `Option<SwiftType>` to Rust. Retain count is now 2

- Rust passes ownership back to Swift. Retain count is now 1

- Retain count never hits 0. Memory has been leaked.
---
 .../src/bridged_type/bridged_opaque_type.rs          | 12 +++++++++++-
 .../codegen/codegen_tests/option_codegen_tests.rs    |  3 ++-
 2 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/crates/swift-bridge-ir/src/bridged_type/bridged_opaque_type.rs b/crates/swift-bridge-ir/src/bridged_type/bridged_opaque_type.rs
index 5a0649cf..ea979fa5 100644
--- a/crates/swift-bridge-ir/src/bridged_type/bridged_opaque_type.rs
+++ b/crates/swift-bridge-ir/src/bridged_type/bridged_opaque_type.rs
@@ -324,8 +324,18 @@ impl BridgeableType for OpaqueForeignType {
                 HostLang::Swift => {
                     let ty = &self.ty;
 
+                    // Here we are converting a Swift type from its Rust representation to its FFI
+                    // representation.
+                    // When we drop the Rust representation we do not want to free the backing Swift
+                    // type since we are passing ownership to Swift.
+                    // So, we are transitioning this Swift type from its
+                    // `RustRepr -> FfiRepr -> SwiftRepr`.
+                    // This means that the Swift side will be responsible for freeing the Swift type
+                    // whenever it is done with it.
+                    // Here we use `ManuallyDrop` to avoid freeing the Swift type.
                     quote! {
                         if let Some(val) = #expression {
+                            let val = std::mem::ManuallyDrop::new(val);
                             val.0 as *mut super::#ty
                         } else {
                             std::ptr::null_mut()
@@ -434,7 +444,7 @@ impl BridgeableType for OpaqueForeignType {
                     format!("{{ if let val = {expression} {{ val.isOwned = false; return val.ptr }} else {{ return nil }} }}()", expression = expression,)
                 }
                 HostLang::Swift => {
-                    format!("{{ if let val = {expression} {{ return Unmanaged.passRetained(val).retain().toOpaque() }} else {{ return nil }} }}()")
+                    format!("{{ if let val = {expression} {{ return Unmanaged.passRetained(val).toOpaque() }} else {{ return nil }} }}()")
                 }
             }
         }
diff --git a/crates/swift-bridge-ir/src/codegen/codegen_tests/option_codegen_tests.rs b/crates/swift-bridge-ir/src/codegen/codegen_tests/option_codegen_tests.rs
index a23611af..9e1e20ed 100644
--- a/crates/swift-bridge-ir/src/codegen/codegen_tests/option_codegen_tests.rs
+++ b/crates/swift-bridge-ir/src/codegen/codegen_tests/option_codegen_tests.rs
@@ -676,6 +676,7 @@ mod extern_rust_fn_return_option_opaque_swift_type {
                 #[export_name = "__swift_bridge__$some_function"]
                 pub extern "C" fn __swift_bridge__some_function() -> *mut super::SomeSwiftType {
                     if let Some(val) = super::some_function() {
+                        let val = std::mem::ManuallyDrop::new(val);
                         val.0 as *mut super::SomeSwiftType
                     } else {
                         std::ptr::null_mut()
@@ -945,7 +946,7 @@ mod extern_rust_fn_with_option_opaque_swift_type_arg {
         ExpectedSwiftCode::ContainsAfterTrim(
             r#"
 func some_function(_ arg: Optional<SomeSwiftType>) {
-    __swift_bridge__$some_function({ if let val = arg { return Unmanaged.passRetained(val).retain().toOpaque() } else { return nil } }())
+    __swift_bridge__$some_function({ if let val = arg { return Unmanaged.passRetained(val).toOpaque() } else { return nil } }())
 }
 "#,
         )