From 4683599a82acfe86297460c5f1e50af7a5898dfd Mon Sep 17 00:00:00 2001
From: dcodeIO <dcode@dcode.io>
Date: Sat, 20 Oct 2018 15:40:23 +0200
Subject: [PATCH] Improve loader function table handling

---
 lib/loader/README.md                  |   7 +++++
 lib/loader/index.d.ts                 |   5 ++++
 lib/loader/index.js                   |  39 +++++++++++++-------------
 lib/loader/tests/build/untouched.wasm | Bin 8470 -> 8542 bytes
 lib/loader/tests/index.js             |   2 +-
 5 files changed, 33 insertions(+), 20 deletions(-)

diff --git a/lib/loader/README.md b/lib/loader/README.md
index 88a27dbf..d19e91b0 100644
--- a/lib/loader/README.md
+++ b/lib/loader/README.md
@@ -83,6 +83,13 @@ Instances are automatically populated with useful utility:
 * **freeArray**(ptr: `number`): `void`<br />
   Frees a typed array in the module's memory. Must not be accessed anymore afterwards.
 
+* **getFunction**(ptr: `number`): `function`<br />
+  Gets a function by its pointer.
+
+* **newFunction**(fn: `function`): `number`<br />
+  Creates a new function in the module's table and returns its pointer. Note that only actual
+  WebAssembly functions, i.e. as exported by the module, are supported.
+
 <sup>1</sup> This feature has not yet landed in any VM as of this writing.
 
 Examples
diff --git a/lib/loader/index.d.ts b/lib/loader/index.d.ts
index c9476c08..693b015e 100644
--- a/lib/loader/index.d.ts
+++ b/lib/loader/index.d.ts
@@ -66,6 +66,11 @@ interface ASUtil {
   freeArray(ptr: number): void;
   /** Gets a function by its pointer. */
   getFunction<R = any>(ptr: number): (...args: any[]) => R;
+  /**
+   * Creates a new function in the module's table and returns its pointer. Note that only actual
+   * WebAssembly functions, i.e. as exported by the module, are supported.
+   */
+  newFunction(fn: (...args: any[]) => any): number;
 }
 
 /** Instantiates an AssemblyScript module using the specified imports. */
diff --git a/lib/loader/index.js b/lib/loader/index.js
index 470521f6..1eb22281 100644
--- a/lib/loader/index.js
+++ b/lib/loader/index.js
@@ -155,8 +155,12 @@ function postInstantiate(baseModule, instance) {
 
   baseModule.freeArray = freeArray;
 
-  /** Creates a new function in the module's table and returns its pointer. */
+  /**
+   * Creates a new function in the module's table and returns its pointer. Note that only actual
+   * WebAssembly functions, i.e. as exported by the module, are supported.
+   */
   function newFunction(fn) {
+    if (typeof fn.original === "function") fn = fn.original;
     var index = table.length;
     table.grow(1);
     table.set(index, fn);
@@ -167,11 +171,7 @@ function postInstantiate(baseModule, instance) {
 
   /** Gets a function by its pointer. */
   function getFunction(ptr) {
-    var fn = table.get(ptr);
-    return (...args) => {
-      setargc(args.length);
-      return fn(...args);
-    };
+    return wrapFunction(table.get(ptr), setargc);
   }
 
   baseModule.getFunction = getFunction;
@@ -191,6 +191,18 @@ function postInstantiate(baseModule, instance) {
   }));
 }
 
+/** Wraps a WebAssembly function while also taking care of variable arguments. */
+function wrapFunction(fn, setargc) {
+  var wrap = (...args) => {
+    setargc(args.length);
+    return fn(...args);
+  }
+  // adding a function to the table with `newFunction` is limited to actual WebAssembly functions,
+  // hence we can't use the wrapper and instead need to provide a reference to the original
+  wrap.original = fn;
+  return wrap;
+}
+
 /** Instantiates an AssemblyScript module using the specified imports. */
 function instantiate(module, imports) {
   return postInstantiate(
@@ -266,7 +278,7 @@ function demangle(exports, baseModule) {
           });
         }
       } else {
-        curr[name] = wrapFunction(elem);
+        curr[name] = wrapFunction(elem, setargc);
       }
     } else {
       if (/^(get|set):/.test(name)) {
@@ -278,24 +290,13 @@ function demangle(exports, baseModule) {
           });
         }
       } else if (typeof elem === "function") {
-        curr[name] = wrapFunction(elem);
+        curr[name] = wrapFunction(elem, setargc);
       } else {
         curr[name] = elem;
       }
     }
   }
 
-  function wrapFunction(fn) {
-    var ret = function(...args) {
-      setargc(args.length);
-      return fn(...args);
-    };
-    // adding a function to the table with `newFunction` is limited to actual exported WebAssembly
-    // functions, hence we can't use the wrapper for that and instead need to pass a workaround:
-    ret.constructor = fn;
-    return ret;
-  }
-
   return module;
 }
 
diff --git a/lib/loader/tests/build/untouched.wasm b/lib/loader/tests/build/untouched.wasm
index 4c19d9412bf7285d1725e1c5fbf54fe511a72b76..18ec0b8e80a9d7b7ff62668f9f0e89a3e14d55dd 100644
GIT binary patch
delta 281
zcmbQ{bkAu+p)}`fZkD{n+*HQLo2#U6FxN|$q!yPH>n9c$r{*T*RO)BurKDEq7nc;}
zq~<BeVwKEIEXmMIOi58x!7A;XSfresp9eIeG`S?dNJ#^ys`S*7$vfqF*`X>XAD35V
zSAwXX{90a?T^Y*eRFGp=fU-3d{Ou&Lnpj+#tBKt=Wr;;V=W3}zjaaOp$SwnAA5{nh
E0178(f&c&j

delta 186
zcmccTG|g#4p)}`WZkD{n+*HQdo2#U6FdK>{78j@HCgoJ>XXd4(R_GU(6y>DmDTt#9
z=O&h9=q0A4D9TNqDlgBdG<mnYGLU=<BxNS^E66IyqG@nWEK*L+$uCZI$<Hq;R+E|>
jq2RA0gr=ytG*=Vd0cD9rKzp<lCO=hBWE7jss~89X?eRci

diff --git a/lib/loader/tests/index.js b/lib/loader/tests/index.js
index 3b331b80..b85fafe1 100644
--- a/lib/loader/tests/index.js
+++ b/lib/loader/tests/index.js
@@ -57,5 +57,5 @@ assert.strictEqual(fn(2, 3), 5);
 assert.strictEqual(fn(2), 4);
 
 // should be able to create a new function and call it from WASM
-ptr = module.newFunction(module.varadd.constructor); // must be an actual exported wasm function
+ptr = module.newFunction(module.varadd);
 assert.strictEqual(module.calladd(ptr, 2, 3), 5);