From 249f7e60735d1e6331ce9b0d6395a2e1e4099d07 Mon Sep 17 00:00:00 2001
From: Kuba Mracek <mracek@apple.com>
Date: Tue, 3 Sep 2024 14:38:03 -0700
Subject: [PATCH] [embedded] Handle retain/retain ops inside deinit in Embedded
 Swift's swift_release

---
 stdlib/public/core/EmbeddedRuntime.swift |  7 +++++++
 test/embedded/deinit-release.swift       | 26 ++++++++++++++++++++++++
 2 files changed, 33 insertions(+)
 create mode 100644 test/embedded/deinit-release.swift

diff --git a/stdlib/public/core/EmbeddedRuntime.swift b/stdlib/public/core/EmbeddedRuntime.swift
index e16809211ba..8155c28f0f6 100644
--- a/stdlib/public/core/EmbeddedRuntime.swift
+++ b/stdlib/public/core/EmbeddedRuntime.swift
@@ -291,6 +291,10 @@ func swift_release_n_(object: UnsafeMutablePointer<HeapObject>?, n: UInt32) {
 
   let resultingRefcount = subFetchAcquireRelease(refcount, n: Int(n)) & HeapObject.refcountMask
   if resultingRefcount == 0 {
+    // Set the refcount to immortalRefCount before calling the object destroyer
+    // to prevent future retains/releases from having any effect.
+    storeRelaxed(refcount, newValue: HeapObject.immortalRefCount)
+
     _swift_embedded_invoke_heap_object_destroy(object)
   } else if resultingRefcount < 0 {
     fatalError("negative refcount")
@@ -350,6 +354,9 @@ fileprivate func storeRelease(_ atomic: UnsafeMutablePointer<Int>, newValue: Int
   Builtin.atomicstore_release_Word(atomic._rawValue, newValue._builtinWordValue)
 }
 
+fileprivate func storeRelaxed(_ atomic: UnsafeMutablePointer<Int>, newValue: Int) {
+  Builtin.atomicstore_monotonic_Word(atomic._rawValue, newValue._builtinWordValue)
+}
 
 /// Exclusivity checking
 
diff --git a/test/embedded/deinit-release.swift b/test/embedded/deinit-release.swift
new file mode 100644
index 00000000000..bce697e5525
--- /dev/null
+++ b/test/embedded/deinit-release.swift
@@ -0,0 +1,26 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend %s -enable-experimental-feature Embedded -O -c -o %t/main.o
+// RUN: %target-clang %t/main.o -o %t/a.out -dead_strip
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: swift_in_compiler
+// REQUIRES: executable_test
+// REQUIRES: OS=macosx || OS=linux-gnu
+
+class C {}
+
+struct Foo {
+    let foo: C = C()
+    let bar: C = C()
+}
+
+class Bar {}
+class SubBar: Bar {
+    var qwe = Foo()
+}
+
+var bar: SubBar? = SubBar()
+bar = nil
+print("OK!")
+
+// CHECK: OK!