[embedded] Fix a memory leak caused by incorrect refcounting logic around doNotFreeBit

lib/IRGen/GenConstant.cpp

@@ -458,7 +458,8 @@ llvm::Constant *irgen::emitConstantObject(IRGenModule &IGM, ObjectInst *OI,
   if (IGM.canMakeStaticObjectReadOnly(OI->getType())) {
     if (!IGM.swiftImmortalRefCount) {
       if (IGM.Context.LangOpts.hasFeature(Feature::Embedded)) {
-        // = HeapObject.immortalRefCount
+        // = HeapObject.immortalRefCount | HeapObject.doNotFreeBit
+        // 0xffff_ffff on 32-bit, 0xffff_ffff_ffff_ffff on 64-bit
         IGM.swiftImmortalRefCount = llvm::ConstantInt::get(IGM.IntPtrTy, -1);
       } else {
         IGM.swiftImmortalRefCount = llvm::ConstantExpr::getPtrToInt(

stdlib/public/core/EmbeddedRuntime.swift

@@ -38,20 +38,21 @@ public struct HeapObject {
   // to think about supporting (or banning) weak and/or unowned references.
   var refcount: Int
 
+  // Note: The immortalRefCount value is also hard-coded in IRGen in `irgen::emitConstantObject`.
 #if _pointerBitWidth(_64)
-  static let doNotFreeBit = Int(bitPattern: 0x8000_0000_0000_0000)
-  static let refcountMask = Int(bitPattern: 0x7fff_ffff_ffff_ffff)
+  static let doNotFreeBit     = Int(bitPattern: 0x8000_0000_0000_0000)
+  static let refcountMask     = Int(bitPattern: 0x7fff_ffff_ffff_ffff)
+  static let immortalRefCount = Int(bitPattern: 0x7fff_ffff_ffff_ffff) // Make sure we don't have doNotFreeBit set
 #elseif _pointerBitWidth(_32)
-  static let doNotFreeBit = Int(bitPattern: 0x8000_0000)
-  static let refcountMask = Int(bitPattern: 0x7fff_ffff)
+  static let doNotFreeBit     = Int(bitPattern: 0x8000_0000)
+  static let refcountMask     = Int(bitPattern: 0x7fff_ffff)
+  static let immortalRefCount = Int(bitPattern: 0x7fff_ffff) // Make sure we don't have doNotFreeBit set
 #elseif _pointerBitWidth(_16)
-  static let doNotFreeBit = Int(bitPattern: 0x8000)
-  static let refcountMask = Int(bitPattern: 0x7fff)
+  static let doNotFreeBit     = Int(bitPattern: 0x8000)
+  static let refcountMask     = Int(bitPattern: 0x7fff)
+  static let immortalRefCount = Int(bitPattern: 0x7fff) // Make sure we don't have doNotFreeBit set
 #endif
 
-  // Note: The immortalRefCount value of -1 is also hard-coded in IRGen in `irgen::emitConstantObject`.
-  static let immortalRefCount = -1
-
 #if _pointerBitWidth(_64)
   static let immortalObjectPointerBit = UInt(0x8000_0000_0000_0000)
 #endif
@@ -158,7 +159,7 @@ public func swift_initStaticObject(metadata: Builtin.RawPointer, object: Builtin
 
 func swift_initStaticObject(metadata: UnsafeMutablePointer<ClassMetadata>, object: UnsafeMutablePointer<HeapObject>) -> UnsafeMutablePointer<HeapObject> {
   _swift_embedded_set_heap_object_metadata_pointer(object, metadata)
-  object.pointee.refcount = HeapObject.immortalRefCount
+  object.pointee.refcount = HeapObject.immortalRefCount | HeapObject.doNotFreeBit
   return object
 }
 
@@ -251,7 +252,7 @@ public func swift_retain_n(object: Builtin.RawPointer, n: UInt32) -> Builtin.Raw
 
 func swift_retain_n_(object: UnsafeMutablePointer<HeapObject>, n: UInt32) -> UnsafeMutablePointer<HeapObject> {
   let refcount = refcountPointer(for: object)
-  if loadRelaxed(refcount) == HeapObject.immortalRefCount {
+  if loadRelaxed(refcount) & HeapObject.refcountMask == HeapObject.immortalRefCount {
     return object
   }
 
@@ -294,7 +295,8 @@ func swift_release_n_(object: UnsafeMutablePointer<HeapObject>?, n: UInt32) {
   }
 
   let refcount = refcountPointer(for: object)
-  if loadRelaxed(refcount) == HeapObject.immortalRefCount {
+  let loadedRefcount = loadRelaxed(refcount)
+  if loadedRefcount & HeapObject.refcountMask == HeapObject.immortalRefCount {
     return
   }
 
@@ -309,7 +311,8 @@ func swift_release_n_(object: UnsafeMutablePointer<HeapObject>?, n: UInt32) {
     // There can only be one thread with a reference at this point (because
     // we're releasing the last existing reference), so a relaxed store is
     // enough.
-    storeRelaxed(refcount, newValue: HeapObject.immortalRefCount)
+    let doNotFree = (loadedRefcount & HeapObject.doNotFreeBit) != 0
+    storeRelaxed(refcount, newValue: HeapObject.immortalRefCount | (doNotFree ? HeapObject.doNotFreeBit : 0))
 
     _swift_embedded_invoke_heap_object_destroy(object)
   } else if resultingRefcount < 0 {

test/embedded/Inputs/debug-malloc.c

@@ -0,0 +1,38 @@
+#include <stdlib.h>
+#include <stdio.h>
+
+#define HEAP_SIZE (8 * 1024)
+
+__attribute__((aligned(16)))
+char heap[HEAP_SIZE] = {};
+
+size_t next_heap_index = 0;
+
+void *calloc(size_t count, size_t size) {
+  printf("malloc(%ld)", count);
+
+  if (next_heap_index + count * size > HEAP_SIZE) {
+    puts("HEAP EXHAUSTED\n");
+    __builtin_trap();
+  }
+  void *p = &heap[next_heap_index];
+  next_heap_index += count * size;
+  printf("-> %p\n", p);
+  return p;
+}
+
+void *malloc(size_t count) {
+  return calloc(count, 1);
+}
+
+int posix_memalign(void **memptr, size_t alignment, size_t size) {
+  *memptr = calloc(size + alignment, 1);
+  if (((uintptr_t)*memptr) % alignment == 0) return 0;
+  *(uintptr_t *)memptr += alignment - ((uintptr_t)*memptr % alignment);
+  return 0;
+}
+
+void free(void *ptr) {
+  // don't actually free
+  printf("free(%p)\n", ptr);
+}

test/embedded/refcounting-leak.swift

@@ -0,0 +1,35 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-swift-frontend -enable-experimental-feature Embedded -parse-as-library %s -c -o %t/a.o
+// RUN: %target-clang -x c -std=c11 -c %S/Inputs/debug-malloc.c -o %t/debug-malloc.o
+// RUN: %target-clang %t/a.o %t/debug-malloc.o -o %t/a.out
+// RUN: %target-run %t/a.out | %FileCheck %s
+
+// REQUIRES: executable_test
+// REQUIRES: swift_in_compiler
+// REQUIRES: optimized_stdlib
+// REQUIRES: OS=macosx
+
+var x: UInt = 1
+func randomNumber() -> UInt {
+    x = ((x >> 16) ^ x) &* 0x45d9f3b
+    x = ((x >> 16) ^ x) &* 0x45d9f3b
+    x = (x >> 16) ^ x
+    return x
+}
+
+@main
+struct Main {
+    static func main() {
+        for _ in 0 ..< 3 {
+            _ = randomNumber().description
+        }
+        print("OK!")
+        // CHECK: malloc
+        // CHECK: free
+        // CHECK: malloc
+        // CHECK: free
+        // CHECK: malloc
+        // CHECK: free
+        // CHECK: OK!
+    }
+}