fix(wasm): get next func index

xtuc · xtuc · commit 396ecae06c65 · 2018-03-14T18:11:13.000+01:00
diff --git a/lib/WebAssemblyGenerator.js b/lib/WebAssemblyGenerator.js
@@ -23,6 +23,8 @@ function compose(...fns) {
 
 // Utility functions
 const isGlobalImport = moduleImport => moduleImport.descr.type === "GlobalType";
+const isFuncImport = moduleImport =>
+	moduleImport.descr.type === "FuncImportDescr";
 const initFuncId = t.identifier("__init__");
 
 /**
@@ -41,6 +43,12 @@ const removeStartFunc = state => bin => {
 	});
 };
 
+/**
+ * Retrieve the start function
+ *
+ * @param {Object} ast - Module's AST
+ * @returns {t.Identifier | undefined} - node if any
+ */
 function getStartFuncIndex(ast) {
 	let startAtFuncIndex;
 
@@ -53,6 +61,12 @@ function getStartFuncIndex(ast) {
 	return startAtFuncIndex;
 }
 
+/**
+ * Get imported globals
+ *
+ * @param {Object} ast - Module's AST
+ * @returns {Array<t.ModuleImport>} - nodes
+ */
 function getImportedGlobals(ast) {
 	const importedGlobals = [];
 
@@ -67,6 +81,33 @@ function getImportedGlobals(ast) {
 	return importedGlobals;
 }
 
+/**
+ * Get next func index
+ *
+ * We need to respect the instantation order, first count the number of imported
+ * func and func declared in the module.
+ *
+ * @param {Object} ast - Module's AST
+ * @returns {t.indexLiteral} - index
+ */
+function getNextFuncIndex(ast) {
+	let count = 0;
+
+	t.traverse(ast, {
+		ModuleImport({ node }) {
+			if (isFuncImport(node) === true) {
+				count++;
+			}
+		},
+
+		Func({ node }) {
+			count++;
+		}
+	});
+
+	return t.indexLiteral(count);
+}
+
 /**
  * Rewrite the import globals:
  * - removes the ModuleImport instruction
@@ -119,12 +160,11 @@ const rewriteImportedGlobals = state => bin => {
 const addInitFunction = ({
 	startAtFuncIndex,
 	importedGlobals,
-	funcSectionMetadata
+	funcSectionMetadata,
+	nextFuncIndex
 }) => bin => {
 	debug("addInitFunction");
 
-	const nextTypeIndex = funcSectionMetadata.vectorOfSize;
-
 	const funcParams = importedGlobals.map(importedGlobal => {
 		// used for debugging
 		const id = t.identifier(`${importedGlobal.module}.${importedGlobal.name}`);
@@ -153,13 +193,9 @@ const addInitFunction = ({
 	const func = t.func(initFuncId, funcParams, funcResults, funcBody);
 
 	const functype = t.typeInstructionFunc(func.params, func.result);
-	const funcindex = t.indexInFuncSection(t.indexLiteral(nextTypeIndex));
+	const funcindex = t.indexInFuncSection(nextFuncIndex);
 
-	const moduleExport = t.moduleExport(
-		initFuncId.value,
-		"Func",
-		t.indexLiteral(nextTypeIndex)
-	);
+	const moduleExport = t.moduleExport(initFuncId.value, "Func", nextFuncIndex);
 
 	return add(bin, [func, moduleExport, funcindex, functype]);
 };
@@ -175,6 +211,7 @@ class WebAssemblyGenerator {
 		const importedGlobals = getImportedGlobals(ast);
 		const funcSectionMetadata = t.getSectionMetadata(ast, "func");
 		const startAtFuncIndex = getStartFuncIndex(ast);
+		const nextFuncIndex = getNextFuncIndex(ast);
 
 		const transform = compose(
 			removeStartFunc(),
@@ -184,7 +221,8 @@ class WebAssemblyGenerator {
 			addInitFunction({
 				importedGlobals,
 				funcSectionMetadata,
-				startAtFuncIndex
+				startAtFuncIndex,
+				nextFuncIndex
 			})
 		);
 
