Skip to content

Instantly share code, notes, and snippets.

Show Gist options
  • Save ColinEberhardt/0ab739f18576e6e90ce784a765c26b75 to your computer and use it in GitHub Desktop.
Save ColinEberhardt/0ab739f18576e6e90ce784a765c26b75 to your computer and use it in GitHub Desktop.
WebAssembly transform giving support for multi return values
const fs = require("fs");
const { parse } = require("@webassemblyjs/wast-parser");
const { print } = require("@webassemblyjs/wast-printer");
const identifiedToIndex = require("@webassemblyjs/ast/lib/transform/wast-identifier-to-index/index");
const t = require("@webassemblyjs/ast");
const traverse = t.traverse;
const insert = (a1, a2, index) => [
...a1.slice(0, index),
...a2,
...a1.slice(index)
];
// load and parse the wat file
const sourceWat = fs.readFileSync("in.wat", "utf8");
const ast = parse(sourceWat);
const wasmModule = ast.body[0];
// transform all call-by-identifier instrs into call-by-index
identifiedToIndex.transform(ast);
// find the maximum number of returns
let maxReturns = 0;
traverse(ast, {
Func(path) {
const results = path.node.signature.results.length;
if (results > maxReturns) {
maxReturns = results;
}
}
});
// add globals
const globalDefs = [];
for (var i = 0; i < maxReturns; i++) {
const global = t.global(t.globalType("i32", "var"), [
t.objectInstruction("const", "i32", [t.numberLiteralFromRaw(0)])
]);
globalDefs.push(global);
}
// locate the last index of any imports
let lastImportIndex = -1;
wasmModule.fields.forEach((field, index) => {
if (t.isModuleImport(field)) {
lastImportIndex = index;
}
});
// insert the globals after the last import
wasmModule.fields = insert(wasmModule.fields, globalDefs, lastImportIndex + 1);
const updateCallSites = (funcIndex, resultsCount) => {
traverse(ast, {
CallInstruction(path) {
if (path.node.index.value === funcIndex) {
// find the parent property that holds the instructions
const parentNode = path.parentPath.node;
const instructions = parentNode.instrArgs;
// find the index of this node within the parent
const index = instructions.findIndex(i => i === path.node);
// create the get global statements
const globalGetters = [];
for (var i = 0; i < resultsCount; i++) {
globalGetters.push(t.instruction("get_global", [t.indexLiteral(i)]));
}
globalGetters.reverse();
// insert the get global instruction
parentNode.instrArgs = insert(instructions, globalGetters, index + 1);
}
}
});
};
// transform multi-values functions
let funcIndex = 0;
traverse(ast, {
FuncImportDescr() {
funcIndex++;
},
Func(path) {
let results = path.node.signature.results.length;
if (results > 1) {
// remove all return values
path.node.signature.results = [];
// add global setters to the end of the function body
for (var i = 0; i < results; i++) {
path.node.body.push(t.instruction("set_global", [t.indexLiteral(i)]));
}
updateCallSites(funcIndex, results);
}
funcIndex++;
}
});
const generatedWat = print(ast);
fs.writeFileSync("out.wat", generatedWat);
@phreda4
Copy link

phreda4 commented Nov 13, 2018

Thank's for you work.
I need this feature to compile my language (forth-like), although I'm not sure if v8 will optimize the data stack that I use, and I need, explicitly. The other alternative that I have is to simulate the data stack with an array, although it loses possibility of optimization for the v8, I guess.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment