feat(runtime): weak collections key native objects by lifetime in canary mode

NativeScript/runtime/Interop.mm

@@ -18,7 +18,6 @@
 #include "Reference.h"
 #include "Runtime.h"
 #include "RuntimeConfig.h"
-#include "SymbolIterator.h"
 #include "SymbolLoader.h"
 #include "UnmanagedType.h"
 #include "robin_hood.h"
@@ -1399,11 +1398,6 @@ static bool SnapshotArrayLike(Isolate* isolate, Local<Object> object, Local<v8::
       [result release];
     }
 
-    if ([result isKindOfClass:[NSArray class]]) {
-      // attach Symbol.iterator to the instance
-      SymbolIterator::Set(context, jsResult);
-    }
-
     tns::DeleteWrapperIfUnused(isolate, jsResult, wrapper);
 
     return jsResult;

NativeScript/runtime/MetadataBuilder.mm

@@ -14,6 +14,7 @@
 #include "SymbolLoader.h"
 #include "Tasks.h"
 #include "Worker.h"
+#include "SymbolIterator.h"
 #include "WrapperLifetimeCanary.h"
 
 using namespace v8;
@@ -603,6 +604,9 @@ throw NativeScriptException(
   ctorFuncTemplate->PrototypeTemplate()->Set(
       tns::ToV8String(isolate, "toString"),
       FunctionTemplate::New(isolate, MetadataBuilder::ToStringFunctionCallback));
+  if (meta->type() == MetaType::Interface && strcmp(meta->name(), "NSArray") == 0) {
+    SymbolIterator::Install(isolate, ctorFuncTemplate->PrototypeTemplate());
+  }
 
   if (meta->type() == MetaType::Interface) {
     const InterfaceMeta* interfaceMeta = static_cast<const InterfaceMeta*>(meta);

NativeScript/runtime/SymbolIterator.h

@@ -7,7 +7,8 @@ namespace tns {
 
 class SymbolIterator {
 public:
-    static void Set(v8::Local<v8::Context> context, v8::Local<v8::Value> object);
+    // On NSArray's prototype template, so every array wrapper inherits it.
+    static void Install(v8::Isolate* isolate, v8::Local<v8::ObjectTemplate> prototypeTemplate);
 private:
     static v8::Local<v8::Object> CreateIteratorObject(v8::Local<v8::Context> context, v8::Local<v8::Value> object);
     static void NextCallback(const v8::FunctionCallbackInfo<v8::Value>& args);

NativeScript/runtime/SymbolIterator.mm

@@ -7,26 +7,13 @@
 
 namespace tns {
 
-void SymbolIterator::Set(Local<Context> context, Local<Value> object) {
-  Isolate* isolate = v8::Isolate::GetCurrent();
-
-  // Iterates its receiver rather than capturing `object`: in the canary
-  // wrapper lifetime mode this property lives on the object's state, which a
-  // reference back to the object would keep alive until native death.
-  Local<v8::Function> iterator;
-  bool success = v8::Function::New(context,
-                                   [](const FunctionCallbackInfo<Value>& args) {
-                                     Isolate* isolate = args.GetIsolate();
-                                     Local<Context> context = isolate->GetCurrentContext();
-                                     Local<Object> result =
-                                         CreateIteratorObject(context, args.This());
-                                     args.GetReturnValue().Set(result);
-                                   })
-                     .ToLocal(&iterator);
-  tns::Assert(success, isolate);
-
-  Local<Value> symbolIteratorKey = Symbol::GetIterator(isolate);
-  success = object.As<Object>()->Set(context, symbolIteratorKey, iterator).FromMaybe(false);
+void SymbolIterator::Install(Isolate* isolate, Local<ObjectTemplate> prototypeTemplate) {
+  prototypeTemplate->Set(
+      Symbol::GetIterator(isolate),
+      FunctionTemplate::New(isolate, [](const FunctionCallbackInfo<Value>& args) {
+        Local<Context> context = args.GetIsolate()->GetCurrentContext();
+        args.GetReturnValue().Set(CreateIteratorObject(context, args.This()));
+      }));
 }
 
 Local<Object> SymbolIterator::CreateIteratorObject(Local<Context> context, Local<Value> object) {

NativeScript/runtime/WeakRef.cpp

@@ -2,6 +2,8 @@
 
 #include "BuiltinLoader.h"
 #include "Helpers.h"
+#include "WrapperLifetime.h"
+#include "WrapperLifetimeCanary.h"
 
 using namespace v8;
 
@@ -18,6 +20,18 @@ void WeakRef::Init(Local<Context> context) {
     tns::LogError(isolate, tc);
   }
   tns::Assert(success, isolate);
+
+  if (GetWrapperLifetimeMode() != WrapperLifetimeMode::kCanary) {
+    return;
+  }
+  Local<Object> binding = Object::New(isolate);
+  success = WrapperLifetimeCanary::RegisterWrapFunctions(context, binding) &&
+            BuiltinLoader::RunBuiltin(context, BuiltinId::kWeakCollections, binding)
+                .ToLocal(&result);
+  if (!success && tc.HasCaught()) {
+    tns::LogError(isolate, tc);
+  }
+  tns::Assert(success, isolate);
 }
 
 }  // namespace tns

NativeScript/runtime/WrapperLifetimeCanary.h

@@ -47,6 +47,11 @@ class WrapperLifetimeCanary {
   // the mode is canary.
   static void ConfigureInstanceTemplate(v8::Local<v8::ObjectTemplate> instanceTemplate);
 
+  // Adds `unwrap` (a canary's native object -> its weak-collection key) and
+  // `wrap` (key -> live or new canary, or undefined once the object is gone)
+  // to `binding`, for the weak-collections builtin. Canary mode only.
+  static bool RegisterWrapFunctions(v8::Local<v8::Context> context, v8::Local<v8::Object> binding);
+
   // The record of an object that already has state in this isolate, for a new
   // canary to share. Null otherwise.
   static std::shared_ptr<CanaryRecord> Find(v8::Isolate* isolate, id object);

NativeScript/runtime/WrapperLifetimeCanary.mm

@@ -7,6 +7,7 @@
 #include <mutex>
 #include <unordered_map>
 #include <vector>
+#include "ArgConverter.h"
 #include "Caches.h"
 #include "DataWrapper.h"
 #include "EventLoop.h"
@@ -32,6 +33,7 @@
 #endif
 
 extern "C" id objc_initWeakOrNil(id* location, id value);
+extern "C" id objc_loadWeakRetained(id* location);
 extern "C" void objc_destroyWeak(id* location);
 
 using namespace v8;
@@ -51,8 +53,10 @@
   // tells, when a sweep reads it.
   bool reportsDeath = false;
   // An ObjC weak reference to the object, zeroed by the runtime when it is
-  // deallocated, on any thread. Set for records nothing reports.
+  // deallocated, on any thread. Set for records nothing reports, and for any
+  // record whose object became a weak-collection key (Wrap reads it).
   id object = nil;
+  bool weakSet = false;
 
   ~CanaryRecord() { objc_destroyWeak(&object); }
 
@@ -90,6 +94,10 @@ bool ObjectAlive() const {
 
 struct CanaryState {
   std::shared_ptr<CanaryRegistry> registry = std::make_shared<CanaryRegistry>();
+  // The weak collections' key for a native object hangs off its state, and
+  // names the object it stands for.
+  Global<Private> keyOfState;
+  Global<Private> objectOfKey;
 };
 
 void DrainDead(const std::shared_ptr<CanaryRegistry>& registry) {
@@ -320,10 +328,13 @@ bool StateForWrite(Isolate* isolate, Local<Object> holder, Local<Object>* out) {
       if (slot == nullptr) {
         auto record = std::make_shared<CanaryRecord>();
         record->reportsDeath = extended || NS_CANARY_DEATH_SIGNAL == NS_CANARY_DEATH_ASSOCIATION;
-        if (!record->reportsDeath && objc_initWeakOrNil(&record->object, object) == nil) {
-          // The class refuses weak references.
-          registry->records.erase((__bridge void*)object);
-          return false;
+        if (!record->reportsDeath) {
+          if (objc_initWeakOrNil(&record->object, object) == nil) {
+            // The class refuses weak references.
+            registry->records.erase((__bridge void*)object);
+            return false;
+          }
+          record->weakSet = true;
         }
         record->state.Reset(isolate, Object::New(isolate, Null(isolate), nullptr, nullptr, 0));
         slot = std::move(record);
@@ -502,6 +513,98 @@ void PostSweepAfterGC(Isolate* isolate, GCType, GCCallbackFlags) {
 }
 #endif
 
+// The weak collections' key for a canary's native object: a frozen,
+// null-prototype object that lives exactly as long as the object's state, that
+// is, until the native object is deallocated. The holder itself for anything
+// that is not a canary.
+void UnwrapCallback(const FunctionCallbackInfo<Value>& info) {
+  Isolate* isolate = info.GetIsolate();
+  Local<Value> value = info[0];
+  info.GetReturnValue().Set(value);
+  CanaryState* canaryState = Caches::StateFor<CanaryState>(isolate);
+  Local<Object> state;
+  if (canaryState == nullptr || !value->IsObject() ||
+      !StateForWrite(isolate, value.As<Object>(), &state)) {
+    return;
+  }
+  Local<Context> context = isolate->GetCurrentContext();
+  Local<Private> keyOfState = canaryState->keyOfState.Get(isolate);
+  Local<Value> key;
+  if (state->GetPrivate(context, keyOfState).ToLocal(&key) && key->IsObject()) {
+    info.GetReturnValue().Set(key);
+    return;
+  }
+  ObjectWeakCallbackState* canary = CanaryOf(isolate, value.As<Object>());
+  CanaryRecord* record = canary->canaryRecord_.get();
+  if (!record->weakSet) {
+    if (objc_initWeakOrNil(&record->object, canary->lifetimeTarget_) == nil) {
+      return;
+    }
+    record->weakSet = true;
+  }
+  Local<Object> newKey = Object::New(isolate, Null(isolate), nullptr, nullptr, 0);
+  uint64_t address = reinterpret_cast<uintptr_t>((__bridge void*)canary->lifetimeTarget_);
+  if (!newKey
+           ->SetPrivate(context, canaryState->objectOfKey.Get(isolate),
+                        BigInt::NewFromUnsigned(isolate, address))
+           .FromMaybe(false) ||
+      !newKey->SetIntegrityLevel(context, IntegrityLevel::kFrozen).FromMaybe(false) ||
+      !state->SetPrivate(context, keyOfState, newKey).FromMaybe(false)) {
+    return;
+  }
+  info.GetReturnValue().Set(newKey);
+}
+
+// The object a key stands for: its live canary or a new one, or undefined once
+// the native object is gone. Anything that is not a key passes through.
+void WrapCallback(const FunctionCallbackInfo<Value>& info) {
+  Isolate* isolate = info.GetIsolate();
+  Local<Value> value = info[0];
+  info.GetReturnValue().Set(value);
+  CanaryState* canaryState = Caches::StateFor<CanaryState>(isolate);
+  if (canaryState == nullptr || !value->IsObject()) {
+    return;
+  }
+  Local<Context> context = isolate->GetCurrentContext();
+  Local<Object> key = value.As<Object>();
+  Local<Value> address;
+  if (!key->GetPrivate(context, canaryState->objectOfKey.Get(isolate)).ToLocal(&address) ||
+      !address->IsBigInt()) {
+    return;
+  }
+  info.GetReturnValue().SetUndefined();
+  std::shared_ptr<CanaryRecord> record;
+  {
+    std::lock_guard<std::mutex> lock(canaryState->registry->mutex);
+    auto it = canaryState->registry->records.find(
+        reinterpret_cast<void*>(static_cast<uintptr_t>(address.As<BigInt>()->Uint64Value())));
+    if (it == canaryState->registry->records.end()) {
+      return;
+    }
+    record = it->second;
+  }
+  if (!Usable(record) || !record->weakSet) {
+    return;
+  }
+  // A record at that address may belong to a later object: it must be the
+  // key's own.
+  Local<Value> recordKey;
+  Local<Object> state = Local<Object>::New(isolate, record->state);
+  if (!state->GetPrivate(context, canaryState->keyOfState.Get(isolate)).ToLocal(&recordKey) ||
+      !recordKey->StrictEquals(key)) {
+    return;
+  }
+  id object = objc_loadWeakRetained(&record->object);
+  if (object == nil) {
+    return;
+  }
+  ObjCDataWrapper* wrapper = new ObjCDataWrapper(object);
+  Local<Value> result = ArgConverter::ConvertArgument(context, wrapper);
+  tns::DeleteWrapperIfUnused(isolate, result, wrapper);
+  [object release];
+  info.GetReturnValue().Set(result);
+}
+
 }  // namespace
 
 void WrapperLifetimeCanary::Init(Isolate* isolate) {
@@ -513,6 +616,8 @@ void PostSweepAfterGC(Isolate* isolate, GCType, GCCallbackFlags) {
     return;
   }
   state->registry->isolate = isolate;
+  state->keyOfState.Reset(isolate, Private::New(isolate));
+  state->objectOfKey.Reset(isolate, Private::New(isolate));
   if (Runtime* runtime = Runtime::GetRuntime(isolate)) {
     state->registry->loop = runtime->GetEventLoop();
   }
@@ -526,6 +631,16 @@ void PostSweepAfterGC(Isolate* isolate, GCType, GCCallbackFlags) {
 #endif
 }
 
+bool WrapperLifetimeCanary::RegisterWrapFunctions(Local<Context> context, Local<Object> binding) {
+  Isolate* isolate = v8::Isolate::GetCurrent();
+  Local<v8::Function> unwrap;
+  Local<v8::Function> wrap;
+  return v8::Function::New(context, UnwrapCallback).ToLocal(&unwrap) &&
+         v8::Function::New(context, WrapCallback).ToLocal(&wrap) &&
+         binding->Set(context, tns::ToV8String(isolate, "unwrap"), unwrap).FromMaybe(false) &&
+         binding->Set(context, tns::ToV8String(isolate, "wrap"), wrap).FromMaybe(false);
+}
+
 void WrapperLifetimeCanary::InstallDeallocHook(Class extendedClass) {
   if (GetWrapperLifetimeMode() != WrapperLifetimeMode::kCanary) {
     return;