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