From 61491eafbf6e958df97c3d296ae88a8e92667a50 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Wed, 15 Aug 2018 17:52:26 -0700
Subject: [PATCH] Add experimental support for `WeakRef`

This commit adds experimental support for `WeakRef` to be used to automatically
free wasm objects instead of having to always call the `free` function manually.
Note that when enabled the `free` function for all exported objects is still
generated, it's just optionally invoked by the application.

Support isn't exposed through a CLI flag right now due to the early stages of
the `WeakRef` proposal, but the env var `WASM_BINDGEN_WEAKREF` can be used to
enable this generation. Upon doing so the output can then be edited slightly as
well to work in the SpiderMonkey shell and it looks like this is working!

Closes #704
---
 crates/cli-support/src/js/mod.rs | 74 +++++++++++++++++++++++++++++---
 crates/cli-support/src/lib.rs    |  5 +++
 2 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/crates/cli-support/src/js/mod.rs b/crates/cli-support/src/js/mod.rs
index 220f6e70..4e4b5ca6 100644
--- a/crates/cli-support/src/js/mod.rs
+++ b/crates/cli-support/src/js/mod.rs
@@ -490,6 +490,39 @@ impl<'a> Context<'a> {
         let mut dst = format!("class {} {{\n", name);
         let mut ts_dst = format!("export {}", dst);
 
+        let (mkweakref, freeref) = if self.config.weak_refs {
+            // When weak refs are enabled we use them to automatically free the
+            // contents of an exported rust class when it's gc'd. Note that a
+            // manual `free` function still exists for deterministic
+            // destruction.
+            //
+            // This is implemented by using a `WeakRefGroup` to run finalizers
+            // for all `WeakRef` objects that it creates. Upon construction of
+            // a new wasm object we use `makeRef` with "holdings" of a thunk to
+            // free the wasm instance.  Once the `this` (the instance we're
+            // creating) is gc'd then the finalizer will run with the
+            // `WeakRef`, and we'll pull out the `holdings`, our pointer.
+            //
+            // Note, though, that if manual finalization happens we want to
+            // cancel the `WeakRef`-generated finalization, so we retain the
+            // `WeakRef` in a global map. This global map is then used to
+            // `drop()` the `WeakRef` (cancel finalization) whenever it is
+            // finalized.
+            self.expose_cleanup_groups();
+            let mk = format!("
+                const cleanup_ptr = this.ptr;
+                const ref = CLEANUPS.makeRef(this, () => free{}(cleanup_ptr));
+                CLEANUPS_MAP.set(this.ptr, ref);
+            ", name);
+            let free = "
+                CLEANUPS_MAP.get(ptr).drop();
+                CLEANUPS_MAP.delete(ptr);
+            ";
+            (mk, free)
+        } else {
+            (String::new(), "")
+        };
+
         if self.config.debug || class.constructor.is_some() {
             self.expose_constructor_token();
 
@@ -516,9 +549,11 @@ impl<'a> Context<'a> {
                     // This invocation of new will call this constructor with a ConstructorToken
                     let instance = {class}.{constructor}(...args);
                     this.ptr = instance.ptr;
+                    {mkweakref}
                     ",
                     class = name,
-                    constructor = constructor
+                    constructor = constructor,
+                    mkweakref = mkweakref,
                 ));
             } else {
                 dst.push_str(
@@ -537,9 +572,11 @@ impl<'a> Context<'a> {
 
                 constructor(ptr) {{
                     this.ptr = ptr;
+                    {}
                 }}
                 ",
-                name
+                name,
+                mkweakref,
             ));
         }
 
@@ -599,16 +636,29 @@ impl<'a> Context<'a> {
             }
         }
 
-        dst.push_str(&format!(
+        self.global(&format!(
             "
-            free() {{
-                const ptr = this.ptr;
-                this.ptr = 0;
+            function free{}(ptr) {{
+                {}
                 wasm.{}(ptr);
             }}
             ",
+            name,
+            freeref,
             shared::free_function(&name)
         ));
+        dst.push_str(&format!(
+            "
+            free() {{
+                if (this.ptr === 0)
+                    return;
+                const ptr = this.ptr;
+                this.ptr = 0;
+                free{}(this.ptr);
+            }}
+            ",
+            name,
+        ));
         ts_dst.push_str("free(): void;\n");
         dst.push_str(&class.contents);
         ts_dst.push_str(&class.typescript);
@@ -1594,6 +1644,18 @@ impl<'a> Context<'a> {
         ");
     }
 
+    fn expose_cleanup_groups(&mut self) {
+        if !self.exposed_globals.insert("cleanup_groups") {
+            return
+        }
+        self.global(
+            "
+                const CLEANUPS = new WeakRefGroup(x => x.holdings());
+                const CLEANUPS_MAP = new Map();
+            "
+        );
+    }
+
     fn gc(&mut self) -> Result<(), Error> {
         let module = mem::replace(self.module, Module::default());
         let module = module.parse_names().unwrap_or_else(|p| p.1);
diff --git a/crates/cli-support/src/lib.rs b/crates/cli-support/src/lib.rs
index 5933a691..59f4eb85 100644
--- a/crates/cli-support/src/lib.rs
+++ b/crates/cli-support/src/lib.rs
@@ -10,6 +10,7 @@ extern crate failure;
 
 use std::any::Any;
 use std::collections::BTreeSet;
+use std::env;
 use std::fmt;
 use std::fs;
 use std::mem;
@@ -33,6 +34,9 @@ pub struct Bindgen {
     typescript: bool,
     demangle: bool,
     keep_debug: bool,
+    // Experimental support for `WeakRefGroup`, an upcoming ECMAScript feature.
+    // Currently only enable-able through an env var.
+    weak_refs: bool,
 }
 
 enum Input {
@@ -55,6 +59,7 @@ impl Bindgen {
             typescript: false,
             demangle: true,
             keep_debug: false,
+            weak_refs: env::var("WASM_BINDGEN_WEAKREF").is_ok(),
         }
     }