/**

* @fileoverview A potential minimalistic shadow stack. Currently not used.

*

* Instruments a module's exports to track when the execution stack is fully

* unwound, and injects a call to `__autocollect` to be invoked when it is.

* Accounts for the currently in-flight managed return value from Wasm to the

* host by pushing it to a mini stack, essentially a stack of only one value,

* while `__autocollect` is executing.

*

* @license Apache-2.0

*/

import {

Pass

} from "./pass";

import {

BinaryOp,

expandType,

ExportRef,

ExpressionRef,

ExternalKind,

Index,

Module,

TypeRef,

UnaryOp

} from "../module";

import {

_BinaryenAddFunctionExport,

_BinaryenExportGetKind,

_BinaryenExportGetName,

_BinaryenExportGetValue,

_BinaryenFunctionGetName,

_BinaryenFunctionGetNumLocals,

_BinaryenFunctionGetParams,

_BinaryenFunctionGetResults,

_BinaryenGetExportByIndex,

_BinaryenGetFunction,

_BinaryenGetNumExports,

_BinaryenRemoveExport

} from "../glue/binaryen";

import {

Program

} from "../program";

const MINISTACK = "~lib/rt/__ministack";

const STACK_DEPTH = "~stack_depth";

const AUTOCOLLECT = "__autocollect";

/** Instruments a module with a minimalistic shadow stack for precise GC. */

export class MiniStack extends Pass {

/** Compiler reference. */

program: Program;

/** Exported functions returning managed values. */

managedReturns: Set<string> = new Set();

constructor(module: Module, program: Program) {

super(module);

this.program = program;

}

/** Notes the presence of an exported function with a managed return value. */

noteManagedReturn(exportName: string): void {

this.managedReturns.add(exportName);

}

/** Instruments a function export to also maintain stack depth. */

instrumentFunctionExport(ref: ExportRef): void {

assert(_BinaryenExportGetKind(ref) == ExternalKind.Function);

let module = this.module;

let internalNameRef = _BinaryenExportGetValue(ref);

let externalNameRef = _BinaryenExportGetName(ref);

let functionRef = _BinaryenGetFunction(module.ref, internalNameRef);

let originalName = module.readStringCached(_BinaryenFunctionGetName(functionRef))!;

let wrapperName = "export:" + originalName;

if (!module.hasFunction(wrapperName)) {

let params = _BinaryenFunctionGetParams(functionRef);

let results = _BinaryenFunctionGetResults(functionRef);

let numLocals = _BinaryenFunctionGetNumLocals(functionRef);

let vars = new Array<TypeRef>();

// Prepare a call to the original function

let paramTypes = expandType(params);

let numParams = paramTypes.length;

let operands = new Array<ExpressionRef>(numParams);

for (let i = 0; i < numParams; ++i) {

operands[i] = module.local_get(i, paramTypes[i]);

}

let call = module.call(originalName, operands, results);

// Create a wrapper function also maintaining stack depth

let stmts = new Array<ExpressionRef>();

if (numLocals) {

stmts.push(

module.global_set(STACK_DEPTH,

module.binary(BinaryOp.AddI32,

module.global_get(STACK_DEPTH, TypeRef.I32),

module.i32(1) // only need to know > 0

)

)

);

}

if (results == TypeRef.None) {

stmts.push(

call

);

} else {

vars.push(results);

stmts.push(

module.local_set(numParams, call, false) // no shadow stack here

);

}

if (numLocals) {

stmts.push(

module.global_set(STACK_DEPTH,

module.binary(BinaryOp.SubI32,

module.global_get(STACK_DEPTH, TypeRef.I32),

module.i32(1) // only need to know > 0

)

)

);

}

let exportName = module.readStringCached(externalNameRef)!;

stmts.push(

module.global_set(MINISTACK,

this.managedReturns.has(exportName)

? module.local_get(numParams, results)

: module.i32(0)

)

);

stmts.push(

module.if(

module.unary(UnaryOp.EqzI32,

module.global_get(STACK_DEPTH, TypeRef.I32)

),

module.call(AUTOCOLLECT, null, TypeRef.None)

)

);

if (results != TypeRef.None) {

stmts.push(

module.local_get(numParams, results)

);

}

module.addFunction(wrapperName, params, results, vars,

module.block(null, stmts, results)

);

}

// Replace the original export with the wrapped one

_BinaryenRemoveExport(module.ref, externalNameRef);

_BinaryenAddFunctionExport(module.ref, module.allocStringCached(wrapperName), externalNameRef);

}

/** Runs the pass. Returns `true` if the mini stack has been added. */

run(): bool {

let module = this.module;

let moduleRef = module.ref;

let numExports = _BinaryenGetNumExports(moduleRef);

if (numExports) {

let functionExportRefs = new Array<ExportRef>();

// We are going to modify the list of exports, so do this in two steps

for (let i: Index = 0; i < numExports; ++i) {

let exportRef = _BinaryenGetExportByIndex(moduleRef, i);

if (_BinaryenExportGetKind(exportRef) == ExternalKind.Function) {

functionExportRefs.push(exportRef);

}

}

let numFunctionExports = functionExportRefs.length;

if (numFunctionExports) {

for (let i = 0; i < numFunctionExports; ++i) {

this.instrumentFunctionExport(functionExportRefs[i]);

}

module.addGlobal(STACK_DEPTH, TypeRef.I32, true, module.i32(0));

return true;

}

}

return false;

}

}