From 1ffcb90d2d8fb00d4d8ede58abea5279a5ee5186 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Mon, 18 Dec 2017 19:01:37 -0800
Subject: [PATCH] Ensure class arguments have the expected type

---
 crates/wasm-bindgen-cli-support/src/js.rs | 26 +++++++++++++++++------
 crates/wasm-bindgen-macro/src/lib.rs      | 18 ++++++++++++++--
 tests/classes.rs                          | 23 ++++++++++++++------
 3 files changed, 52 insertions(+), 15 deletions(-)

diff --git a/crates/wasm-bindgen-cli-support/src/js.rs b/crates/wasm-bindgen-cli-support/src/js.rs
index cb20af32a..77c7b22a1 100644
--- a/crates/wasm-bindgen-cli-support/src/js.rs
+++ b/crates/wasm-bindgen-cli-support/src/js.rs
@@ -7,6 +7,7 @@ pub struct Js {
     pub expose_get_string_from_wasm: bool,
     pub expose_pass_string_to_wasm: bool,
     pub expose_assert_num: bool,
+    pub expose_assert_class: bool,
     pub expose_token: bool,
     pub exports: Vec<(String, String)>,
     pub classes: Vec<String>,
@@ -128,18 +129,20 @@ impl Js {
                         ", i = i));
                     }
                 }
-                shared::Type::ByRef(_) |
-                shared::Type::ByMutRef(_) => {
+                shared::Type::ByRef(ref s) |
+                shared::Type::ByMutRef(ref s) => {
+                    self.expose_assert_class = true;
                     arg_conversions.push_str(&format!("\
-                        const ptr{i} = {arg}.__wasmPtr;
-                    ", i = i, arg = name));
+                        const ptr{i} = _assertClass({arg}, {struct_});
+                    ", i = i, arg = name, struct_ = s));
                     pass(&format!("ptr{}", i));
                 }
-                shared::Type::ByValue(_) => {
+                shared::Type::ByValue(ref s) => {
+                    self.expose_assert_class = true;
                     arg_conversions.push_str(&format!("\
-                        const ptr{i} = {arg}.__wasmPtr;
+                        const ptr{i} = _assertClass({arg}, {struct_});
                         {arg}.__wasmPtr = 0;
-                    ", i = i, arg = name));
+                    ", i = i, arg = name, struct_ = s));
                     pass(&format!("ptr{}", i));
                 }
             }
@@ -270,6 +273,15 @@ impl Js {
                 ");
             }
         }
+        if self.expose_assert_class {
+            globals.push_str("
+                function _assertClass(instance, klass) {
+                    if (!(instance instanceof klass))
+                        throw new Error(`expected instance of ${klass.name}`);
+                    return instance.__wasmPtr;
+                }
+            ");
+        }
 
         let mut exports = String::new();
         for class in self.classes.iter() {
diff --git a/crates/wasm-bindgen-macro/src/lib.rs b/crates/wasm-bindgen-macro/src/lib.rs
index bab7db437..2bf2eb001 100644
--- a/crates/wasm-bindgen-macro/src/lib.rs
+++ b/crates/wasm-bindgen-macro/src/lib.rs
@@ -26,6 +26,7 @@ macro_rules! my_quote {
 
 #[proc_macro]
 pub fn wasm_bindgen(input: TokenStream) -> TokenStream {
+    // Parse the input as a list of Rust items, reusing the `syn::File` parser.
     let file = syn::parse::<syn::File>(input)
         .expect("expected a set of valid Rust items");
 
@@ -36,24 +37,33 @@ pub fn wasm_bindgen(input: TokenStream) -> TokenStream {
         free_functions: Vec::new(),
     };
 
+    // Translate all input items into our own internal representation (the `ast`
+    // module). We'll be panicking here on anything that we can't process
+
     for item in file.items.iter() {
-        item.to_tokens(&mut ret);
         match *item {
             syn::Item::Fn(ref f) => {
+                item.to_tokens(&mut ret);
                 program.free_functions.push(ast::Function::from(f));
             }
             syn::Item::Struct(ref s) => {
+                item.to_tokens(&mut ret);
                 let s = ast::Struct::from(s);
                 if program.structs.iter().any(|a| a.name == s.name) {
                     panic!("redefinition of struct: {}", s.name);
                 }
                 program.structs.push(s);
             }
-            syn::Item::Impl(ref s) => program.push_impl(s),
+            syn::Item::Impl(ref s) => {
+                item.to_tokens(&mut ret);
+                program.push_impl(s);
+            }
             _ => panic!("unexpected item in bindgen macro"),
         }
     }
 
+    // Generate wrappers for all the items that we've found
+
     for function in program.free_functions.iter() {
         bindgen_fn(function, &mut ret);
     }
@@ -61,6 +71,10 @@ pub fn wasm_bindgen(input: TokenStream) -> TokenStream {
         bindgen_struct(s, &mut ret);
     }
 
+    // Finally generate a static which will eventually be what lives in a custom
+    // section of the wasm executable. For now it's just a plain old static, but
+    // we'll eventually have it actually in its own section.
+
     static CNT: AtomicUsize = ATOMIC_USIZE_INIT;
     let generated_static_name = format!("__WASM_BINDGEN_GENERATED{}",
                                         CNT.fetch_add(1, Ordering::SeqCst));
diff --git a/tests/classes.rs b/tests/classes.rs
index b52359c51..7443ce45a 100644
--- a/tests/classes.rs
+++ b/tests/classes.rs
@@ -131,6 +131,15 @@ fn exceptions() {
                     pub fn bar(&mut self, _: &mut A) {
                     }
                 }
+
+                pub struct B {
+                }
+
+                impl B {
+                    pub fn new() -> B {
+                        B {}
+                    }
+                }
             }
         "#)
         .file("test.js", r#"
@@ -144,12 +153,14 @@ fn exceptions() {
                 assert.throws(() => a.free(), /RuntimeError: unreachable/);
 
                 let b = wasm.A.new();
-                try {
-                    b.foo(b);
-                    assert.throws(() => b.bar(b), /RuntimeError: unreachable/);
-                } finally {
-                    b.free();
-                }
+                b.foo(b);
+                assert.throws(() => b.bar(b), /RuntimeError: unreachable/);
+
+                let c = wasm.A.new();
+                let d = wasm.B.new();
+                assert.throws(() => c.foo(d), /expected instance of A/);
+                d.free();
+                c.free();
             }
         "#)
         .test();