Weak reference doesn't work as expected when passing it to a function's argument

I am already aware of the strong/weak reference concept in swift.
yet after running the next code, and tapping on the Button (and dismissing the screen), the TestViewModel stayed in memory!
I was expecting that using [weak viewmodel] will be enough to prevent it.
in the second example I managed to fix it - but I don't understand why it worked

import SwiftUI
import Resolver

struct TestScreen: View {
    
    @StateObject var viewmodel = TestViewModel()
    @Injected var testStruct: TestStruct
    @Environment(\.presentationMode) var presentationMode
    
    var body: some View {
        NavigationView {
            
            VStack(spacing: 0) {
                
                Button("go back") { [weak viewmodel] in
                        testStruct.saveActionGlobaly(onAsyncAction:  viewmodel?.someAsyncAction )
                        presentationMode.wrappedValue.dismiss()
                    }
            }
        }
    }
}


import Foundation
import Resolver
import SwiftUI

public class TestStruct {
   var onAsyncAction: (() async throws -> Void)?
    
    public func saveActionGlobaly(onAsyncAction: (() async throws -> Void)?) {
        self.onAsyncAction = onAsyncAction
    } 
}

EXAMPLE 2:
I managed to prevent the leak by changing the code this way:
(notice the changes in the callback passed to onAsyncAction)

import Resolver

struct TestScreen: View {
    
    @StateObject var viewmodel = TestViewModel()
    @Injected var testStruct: TestStruct
    @Environment(\.presentationMode) var presentationMode
    
    var body: some View {
        NavigationView {
            
            VStack(spacing: 0) {
                
                Button("go back") { [weak viewmodel] in
                        testStruct.saveActionGlobaly(onAsyncAction:  { await viewmodel?.someAsyncAction() } )
                        presentationMode.wrappedValue.dismiss()
                    }
            }
        }
    }
}

I dont understand why the second TestScreen managed to apply the weak reference and the first one didn't,
thanks (:

environment:
swift 5
xcode 14.2

This is a good question about how weak references interact with bound method references! The first thing to note is that when you form a bound reference to some instance method, you are implicitly creating a closure which captures the self param for the method. So in the following code:

class C {
  var x: Int = 0
  func f() { print(x) }
}

let c = C()
c.x = 3
let printX = c.f
printX() // 3

the line let printX = c.f is doing basically the same thing as if you had written:

let printX = { [c] in
  c.f()
}

If you think about what would happen after c goes out of scope, you can see why this has to be the case: printX is able to be passed around independently of c, and if it didn't have its own (strong) reference to the underlying instance, it would be possible for c to be released and deallocated out from underneath printX, which would suddenly present a memory safety issue (because printX purports to call the C.f method on an instance that no longer exists.

But why does this still happen in the TestViewModel case when we have a weak reference? The thing to notice here is that while we may have a weak reference externally, within TestViewModel all methods work with a strong reference to self. That is, for the duration of the execution of the method, it is not possible for self to disappear out from under us—it will only be possible for self to be deallocated after the method completes execution and the self parameter is released.

So in the case of a bound method reference, if we're able to form a reference to TestViewModel.someAsyncAction that can be invoked later, we must have a strong reference to TestViewModel to pass as the self parameter. By forming that method reference through our weak instance reference, we are attempting to form the strong reference at the time we form the onAsyncAction closure. The code you've written in the first example is roughly semantically equivalent to the following:

let onAsyncAction = if let viewModel = viewModel {
  viewModel.someAsyncAction
} else {
  nil
}
testStruct.saveActionGlobally(onAsyncAction: onAsyncAction)

By doing things this way, we have formed a function which captures a strong reference to the TestViewModel instance for the duration of the lifetime of the underlying function, so we cannot deallocate viewModel until after onAsyncAction is deallocated.

OTOH, in the second example, we are forming a closure which captures only captures a weak reference to TestViewModel. Yes, we still have to form a strong reference in order to invoke someAsyncAction, but the key here is that this strong reference is formed only at the time of execution of onAsyncAction, and only for the duration of execution of someAsyncAction. The more explicit semantic equivalent looks something like this:

testStruct.saveActionGlobaly(onAsyncAction:  {
  if let viewModel = viewModel {
    viewModel.someAsyncAction()
  }
})

Hopefully this version makes it a bit clearer that we are only forming the strong reference to TestViewModel within the closure passed to saveActionGlobally, and that this reference gets released once we exit the scope, allowing the TestViewModel instance to be released if there are only weak references remaining.

6 Likes

Thank you so much for a very descriptive answer!
I think your example, where c can be deallocated out from underneath printX - was very good.

before your answered I actually tried to explained it to my self by adding the compiler into the equation.
I am curious to know if you find the next explanation legit as well?

first of all, the weak-reference concept is a compiler game - meaning, the complier runs over the first example and translate it to:

let onAsyncAction = if let viewModel = viewModel {
  viewModel.someAsyncAction
} else {
  nil
}
testStruct.saveActionGlobally(onAsyncAction: onAsyncAction)

so, there is another step where the compiler translate onAsyncAction to be a closure like so (very abstract):

let onAsyncAction = { [viewmodel] in
  viewmodel.someAsyncAction()
}

and hands it over to onAsyncAction argument.
in other words, the returned closure from evaluating onAsyncAction is holding another implicit viewmodel reference. the compiler can't conclude the explicit and the implicit viewmodels are related, thus the weakness is not applied to the later

to conclude:
my explanation here implies that it is more a compiler issue rather than a pre-meditated design - or at least it could be solved somehow at this stage if swift developers wanted too

I think this explanation mostly makes sense, but there's a couple things that it misses:

What you've illustrated here is not at all specific to weak references—this 'compiler game' is simply the fully-explicit version of the optional chaining operator ?. and applies to any optional. If x has type Optional<T> for some type T then let z = x?.y evaluates

let z = if let x = x { x.y } else { nil }

The only thing here specific to weak references is that the conditions where a weak reference will evaluate to nil are different than for a non-weak variable of optional type.

while this transformation is basically correct, I think it's more accurate to think about this transformation being applied at the point where you form the bound method reference. That is, rather than thinking about it being something applied to onAsyncAction, you should think about the 'full' version of the transformation applied to let onAsyncAction = viewModel?.someAsyncAction as looking like:

let onAsyncAction = if let viewModel = viewModel {
  { [viewModel] in viewModel.someAsyncAction() }
} else {
  nil
}

I wouldn't think about this as so much that the compiler can't figure this out, but that it is at odds with how the compiler chooses to model weak references. Specifically, weak is something that applies to a particular declaration, and copying from a weak variable to a non-weak variable involves attempting to form a strong reference to the underlying instance. That is, in the transformed code above, when you have if let viewModel = viewModel, that expresses the action of "check if the instance referred to by viewModel still exists and if so form a new (strong) reference to it called viewModel."

When you perform this action outside the closure, you're saying "check if viewModel exists and if so form a strong function reference to the bound someAsyncAction method (which captures viewModel strongly) and store it in onAsyncAction" but when you perform this inside the closure you are saying "when onAsyncAction is executed, check if viewModel exists and if so form a strong reference to call the someAsyncAction method, then release the strong reference afterwards."

1 Like

I agree - thank you so much!