Improve loader function table handling

Author: dcodeIO

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

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. */

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;
 }
 

lib/loader/tests/build/untouched.wasm

@@ -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);