/*

* Copyright 2018 WebAssembly Community Group participants

*

* Licensed under the Apache License, Version 2.0 (the "License");

* you may not use this file except in compliance with the License.

* You may obtain a copy of the License at

*

* http://www.apache.org/licenses/LICENSE-2.0

*

* Unless required by applicable law or agreed to in writing, software

* distributed under the License is distributed on an "AS IS" BASIS,

* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

* See the License for the specific language governing permissions and

* limitations under the License.

*/

//

// Optimizes call arguments in a whole-program manner. In particular, this

// removes ones that are not used (dead), but it also does more things:

//

// * Find functions for whom an argument is always passed the same

// constant. If so, we can just set that local to that constant

// in the function.

// * Find functions that don't use the value passed to an argument.

// If so, we can avoid even sending and receiving it. (Note how if

// the previous point was true for an argument, then the second

// must as well.)

// * Find return values ("return arguments" ;) that are never used.

// * Refine the types of arguments, that is make the argument type more

// specific if all the passed values allow that.

//

// This pass does not depend on flattening, but it may be more effective,

// as then call arguments never have side effects (which we need to

// watch for here).

//

#include <algorithm>

#include <unordered_map>

#include <unordered_set>

#include "ir/lubs.h"

#include "ir/return-utils.h"

#include "ir/type-updating.h"

#include "ir/utils.h"

#include "param-utils.h"

#include "pass.h"

#include "passes/opt-utils.h"

#include "support/sorted_vector.h"

#include "wasm-builder.h"

#include "wasm.h"

#ifndef DAE_DEBUG

#define DAE_DEBUG 0

#endif

#ifndef DAE_STATS

#define DAE_STATS 0

#endif

#if DAE_STATS

#include <iostream>

#endif // DAE_STATS

namespace wasm {

// Information for a function

struct DAEFunctionInfo {

// Whether this needs to be recomputed. This begins as true for the first

// computation, and we reset it every time we touch the function.

bool stale = true;

// The unused parameters, if any.

SortedVector unusedParams;

// Maps a function name to the calls going to it.

std::unordered_map<Name, std::vector<Call*>> calls;

// Map of all calls that are dropped, to their drops' locations (so that

// if we can optimize out the drop, we can replace the drop there).

std::unordered_map<Call*, Expression**> droppedCalls;

// Whether this function contains any tail calls (including indirect tail

// calls) and the set of functions this function tail calls. Tail-callers and

// tail-callees cannot have their dropped returns removed because of the

// constraint that tail-callees must have the same return type as

// tail-callers. Indirectly tail called functions are already not optimized

// because being in a table inhibits DAE. TODO: Allow the removal of dropped

// returns from tail-callers if their tail-callees can have their returns

// removed as well.

bool hasTailCalls = false;

std::unordered_set<Name> tailCallees;

// The set of functions that have calls from places that limit what we can do.

// For now, any call we don't see inhibits our optimizations, but TODO: an

// export could be worked around by exporting a thunk that adds the parameter.

//

// This is built up in parallel in each function, and combined at the end.

std::unordered_set<Name> hasUnseenCalls;

// Clears all data, which marks us as stale and in need of recomputation.

void clear() { *this = DAEFunctionInfo(); }

void markStale() { stale = true; }

};

using DAEFunctionInfoMap = std::unordered_map<Name, DAEFunctionInfo>;

struct DAEScanner

: public WalkerPass<PostWalker<DAEScanner, Visitor<DAEScanner>>> {

bool isFunctionParallel() override { return true; }

bool modifiesBinaryenIR() override { return false; }

std::unique_ptr<Pass> create() override {

return std::make_unique<DAEScanner>(infoMap);

}

DAEScanner(DAEFunctionInfoMap* infoMap) : infoMap(infoMap) {}

// The map of all infos for all functions.

DAEFunctionInfoMap* infoMap;

// The info for the function this instance operates on. We stash this as an

// optimization.

DAEFunctionInfo* info = nullptr;

void visitCall(Call* curr) {

if (!getModule()->getFunction(curr->target)->imported()) {

info->calls[curr->target].push_back(curr);

}

if (curr->isReturn) {

info->hasTailCalls = true;

info->tailCallees.insert(curr->target);

}

}

void visitCallIndirect(CallIndirect* curr) {

if (curr->isReturn) {

info->hasTailCalls = true;

}

}

void visitCallRef(CallRef* curr) {

if (curr->isReturn) {

info->hasTailCalls = true;

}

}

void visitDrop(Drop* curr) {

if (auto* call = curr->value->dynCast<Call>()) {

info->droppedCalls[call] = getCurrentPointer();

}

}

void visitRefFunc(RefFunc* curr) {

// RefFunc may be visited from either a function, in which case |info| was

// set, or module-level code (in which case we use the null function name in

// the infoMap).

auto* currInfo = info ? info : &(*infoMap)[Name()];

// Treat a ref.func as an unseen call, preventing us from changing the

// function's type. If we did change it, it could be an observable

// difference from the outside, if the reference escapes, for example.

// TODO: look for actual escaping?

// TODO: create a thunk for external uses that allow internal optimizations

currInfo->hasUnseenCalls.insert(curr->func);

}

// main entry point

void doWalkFunction(Function* func) {

// Set the info for this function.

info = &((*infoMap)[func->name]);

if (!info->stale) {

// Nothing changed since last time.

return;

}

// Clear the data, mark us as no longer stale, and recompute everything.

info->clear();

info->stale = false;

auto numParams = func->getNumParams();

PostWalker<DAEScanner, Visitor<DAEScanner>>::doWalkFunction(func);

// If there are params, check if they are used.

// TODO: This work could be avoided if we cannot optimize for other reasons.

// That would require deferring this to later and checking that.

if (numParams > 0) {

auto usedParams = ParamUtils::getUsedParams(func, getModule());

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

if (!usedParams.contains(i)) {

info->unusedParams.insert(i);

}

}

}

}

};

struct DAE : public Pass {

// This pass changes locals and parameters.

// FIXME DWARF updating does not handle local changes yet.

bool invalidatesDWARF() override { return true; }

bool optimize = false;

Index numFunctions;

// Map of function names to indexes. This lets us use indexes below for speed.

std::unordered_map<Name, Index> indexes;

void run(Module* module) override {

#if DAE_STATS

Index startParams = 0, startResults = 0;

for (auto& func : module->functions) {

startParams += func->getNumParams();

startResults += func->getResults().size();

}

#endif // DAE_STATS

DAEFunctionInfoMap infoMap;

// Ensure all entries exist so the parallel threads don't modify the data

// structure.

for (auto& func : module->functions) {

infoMap.try_emplace(func->name);

}

// The null name represents module-level code (not in a function).

infoMap.try_emplace(Name());

numFunctions = module->functions.size();

for (Index i = 0; i < numFunctions; i++) {

indexes[module->functions[i]->name] = i;

}

// Iterate to convergence.

while (1) {

if (!iteration(module, infoMap)) {

break;

}

}

#if DAE_STATS

Index endParams = 0, endResults = 0;

for (auto& func : module->functions) {

endParams += func->getNumParams();

endResults += func->getResults().size();

}

std::cout << "Removed parameters: " << (startParams - endParams) << "\n";

std::cout << "Removed results: " << (startResults - endResults) << "\n";

#endif // DAE_STATS

}

// For each function, the set of callers. This is used to propagate changes,

// e.g. if we remove a return value from a function, the calls might benefit

// from optimization. It is ok if this is an over-approximation, that is, if

// we think there are more callers than there are, as it would just lead to

// unneeded extra scanning of calling functions (in the example just given, if

// a caller did not actually call, they would not benefit from the extra

// optimization, but no harm is done, and no optimization missed). Such over-

// approximation can happen in later optimization iterations: We may manage to

// remove a call from a function to another (say, after applying a constant

// param, we see the call is not reached). This is somewhat rare, and the cost

// of computing this map is significant, so we compute it once at the start

// and then use that possibly-over-approximating data.

std::vector<std::vector<Name>> callers;

// A count of how many iterations we saw unprofitable removals of parameters.

// An unprofitable removal is one where we only manage to remove from a single

// call, that is, from one call target and it has a single call going to it.

// Such calls are not very interesting, as when there is a single call like

// that then inlining will handle it anyhow, in most cases, and inlining

// does so far more efficiently in situations of call chains:

//

// a -> b -> c -> d

//

// Imagine we remove a param from d, and so we remove it from the call in c.

// If c received that as a parameter, and only ever used it to call d, then

// now we can remove a param from c, and from the call in b, and so forth -

// requiring a full iteration each time to find the small amount of progress.

// (Inlining, otoh, will inline b into a, then c into a, and d into a,

// efficiently.)

Index unprofitableRemovalIters = 0;

bool iteration(Module* module, DAEFunctionInfoMap& infoMap) {

allDroppedCalls.clear();

#if DAE_DEBUG

// Enable this path to mark all contents as stale at the start of each

// iteration, which can be used to check for staleness bugs (that is, bugs

// where something should have been marked stale, but wasn't). Note, though,

// that staleness bugs can easily cause serious issues with validation (e.g.

// if data is stale we may miss that there is an additional caller, that

// prevents refining argument types etc.), so this may not be terribly

// helpful.

if (getenv("ALWAYS_MARK_STALE")) {

for (auto& [_, info] : infoMap) {

info.markStale();

}

}

#endif

DAEScanner scanner(&infoMap);

scanner.walkModuleCode(module);

// Scan all the functions.

scanner.run(getPassRunner(), module);

// Combine all the info from the scan.

std::vector<std::vector<Call*>> allCalls(numFunctions);

std::vector<bool> tailCallees(numFunctions);

std::vector<bool> hasUnseenCalls(numFunctions);

for (auto& [func, info] : infoMap) {

for (auto& [name, calls] : info.calls) {

auto& allCallsToName = allCalls[indexes[name]];

allCallsToName.insert(allCallsToName.end(), calls.begin(), calls.end());

}

for (auto& callee : info.tailCallees) {

tailCallees[indexes[callee]] = true;

}

for (auto& [call, dropp] : info.droppedCalls) {

allDroppedCalls[call] = dropp;

}

for (auto& name : info.hasUnseenCalls) {

hasUnseenCalls[indexes[name]] = true;

}

}

// Exports are considered unseen calls.

for (auto& curr : module->exports) {

if (curr->kind == ExternalKind::Function) {

hasUnseenCalls[indexes[*curr->getInternalName()]] = true;

}

}

// See comment above, we compute callers once and never again.

if (callers.empty()) {

// Compute first as sets, to deduplicate.

std::vector<std::unordered_set<Name>> callersSets(numFunctions);

for (auto& [func, info] : infoMap) {

for (auto& [name, calls] : info.calls) {

callersSets[indexes[name]].insert(func);

}

}

// Copy into efficient vectors.

callers.resize(numFunctions);

for (Index i = 0; i < numFunctions; ++i) {

auto& set = callersSets[i];

callers[i] = std::vector<Name>(set.begin(), set.end());

}

}

// Track which functions we changed that are worth re-optimizing at the end.

std::unordered_set<Function*> worthOptimizing;

// If we refine return types then we will need to do more type updating

// at the end.

bool refinedReturnTypes = false;

// If we find that localizing call arguments can help (by moving their

// effects outside, so ParamUtils::removeParameters can handle them), then

// we do that at the end and perform another cycle. It is simpler to just do

// another cycle than to track the locations of calls, which is tricky as

// localization might move a call (if a call happens to be another call's

// param). In practice it is rare to find call arguments we want to remove,

// and even more rare to find effects get in the way, so this should not

// cause much overhead.

//

// This set tracks the functions for whom calls to it should be modified.

std::unordered_set<Name> callTargetsToLocalize;

// As we optimize, we mark things as stale.

auto markStale = [&](Name func) {

// We only ever mark functions stale (not the global scope, which we never

// modify). An attempt to modify the global scope, identified by a null

// function name, is a logic bug.

assert(func.is());

infoMap[func].markStale();

};

auto markCallersStale = [&](Index index) {

for (auto caller : callers[index]) {

markStale(caller);

}

};

// We now have a mapping of all call sites for each function, and can look

// for optimization opportunities.

for (Index index = 0; index < numFunctions; index++) {

auto* func = module->functions[index].get();

if (func->imported()) {

continue;

}

// We can only optimize if we see all the calls and can modify them.

if (hasUnseenCalls[index]) {

continue;

}

auto& calls = allCalls[index];

if (calls.empty()) {

// Nothing calls this, so it is not worth optimizing.

continue;

}

// Refine argument types before doing anything else. This does not

// affect whether an argument is used or not, it just refines the type

// where possible.

auto name = func->name;

if (refineArgumentTypes(func, calls, module, infoMap[name])) {

worthOptimizing.insert(func);

markStale(func->name);

}

// Refine return types as well.

if (refineReturnTypes(func, calls, module)) {

refinedReturnTypes = true;

markStale(name);

markCallersStale(index);

}

auto optimizedIndexes =

ParamUtils::applyConstantValues({func}, calls, {}, module);

for (auto i : optimizedIndexes) {

// Mark it as unused, which we know it now is (no point to re-scan just

// for that).

infoMap[name].unusedParams.insert(i);

}

if (!optimizedIndexes.empty()) {

markStale(func->name);

}

}

if (refinedReturnTypes) {

// Changing a call expression's return type can propagate out to its

// parents, and so we must refinalize.

// TODO: We could track in which functions we actually make changes.

ReFinalize().run(getPassRunner(), module);

}

// We now know which parameters are unused, and can potentially remove them.

// Only do so if we didn't run into unprofitable removals - if so, leave

// any further removals for other invocations of this pass. (This avoids us

// getting stuck in long unprofitable call chains as mentioned in the

// comment earlier; note that we do process one unprofitable iteration

// before giving up here, so we do make progress at least.)

if (!unprofitableRemovalIters) {

Index removals = 0;

Index singleCallerRemovals = 0;

for (Index index = 0; index < numFunctions; index++) {

auto* func = module->functions[index].get();

if (func->imported()) {

continue;

}

if (hasUnseenCalls[index]) {

continue;

}

auto numParams = func->getNumParams();

if (numParams == 0) {

continue;

}

auto& calls = allCalls[index];

if (calls.empty()) {

continue;

}

auto name = func->name;

auto [removedIndexes, outcome] =

ParamUtils::removeParameters({func},

infoMap[name].unusedParams,

calls,

{},

module,

getPassRunner());

if (!removedIndexes.empty()) {

// Success!

worthOptimizing.insert(func);

markStale(name);

markCallersStale(index);

if (calls.size() == 1) {

singleCallerRemovals++;

}

removals++;

}

if (outcome == ParamUtils::RemovalOutcome::Failure) {

callTargetsToLocalize.insert(name);

}

}

if (removals == 1 && singleCallerRemovals == 1 &&

callTargetsToLocalize.empty()) {

// We only removed parameters from one function, and it had a single

// caller, and we don't have other pending actions (call targets we

// need to localize), so this was unprofitable as mentioned earlier.

unprofitableRemovalIters++;

}

}

// We can also tell which calls have all their return values dropped. Note

// that we can't do this if we changed anything so far, as we may have

// modified allCalls (we can't modify a call site twice in one iteration,

// once to remove a param, once to drop the return value).

if (worthOptimizing.empty()) {

for (Index index = 0; index < numFunctions; index++) {

auto& func = module->functions[index];

if (func->imported()) {

continue;

}

if (func->getResults() == Type::none) {

continue;

}

if (hasUnseenCalls[index]) {

continue;

}

auto name = func->name;

if (infoMap[name].hasTailCalls) {

continue;

}

if (tailCallees[index]) {

continue;

}

auto& calls = allCalls[index];

if (calls.empty()) {

continue;

}

bool allDropped =

std::all_of(calls.begin(), calls.end(), [&](Call* call) {

return allDroppedCalls.contains(call);

});

if (!allDropped) {

continue;

}

if (removeReturnValue(func.get(), calls, module)) {

// We should optimize the callers.

for (auto caller : callers[index]) {

worthOptimizing.insert(module->getFunction(caller));

}

}

// TODO Removing a drop may also open optimization opportunities in the

// callers.

worthOptimizing.insert(func.get());

markStale(name);

markCallersStale(index);

}

}

if (!callTargetsToLocalize.empty()) {

ParamUtils::localizeCallsTo(

callTargetsToLocalize, *module, getPassRunner(), [&](Function* func) {

markStale(func->name);

});

}

if (optimize && !worthOptimizing.empty()) {

OptUtils::optimizeAfterInlining(worthOptimizing, module, getPassRunner());

}

return !worthOptimizing.empty() || refinedReturnTypes ||

!callTargetsToLocalize.empty();

}

private:

std::unordered_map<Call*, Expression**> allDroppedCalls;

// Returns `true` if the caller should be optimized.

bool

removeReturnValue(Function* func, std::vector<Call*>& calls, Module* module) {

// If the result type is uninhabitable, then the caller knows the call will

// never return. That useful information would be lost if we did nothing

// else when removing the return value, but we will insert an `unreachable`

// after the call in the caller to preserve the optimization effect. TODO:

// Do this for more complicated uninhabitable types such as non-nullable

// references to structs with non-nullable reference cycles.

bool wasReturnUninhabitable =

func->getResults().isNull() && func->getResults().isNonNullable();

func->setResults(Type::none);

// Remove the drops on the calls. Note that we must do this before updating

// returns in ReturnUpdater, as there may be recursive calls of this

// function to itself. So we first use the information in allDroppedCalls

// before the ReturnUpdater potentially invalidates that information as it

// modifies the function.

for (auto* call : calls) {

auto iter = allDroppedCalls.find(call);

assert(iter != allDroppedCalls.end());

Expression** location = iter->second;

if (wasReturnUninhabitable) {

Builder builder(*module);

*location = builder.makeSequence(call, builder.makeUnreachable());

} else {

*location = call;

}

// Update the call's type.

if (call->type != Type::unreachable) {

call->type = Type::none;

}

}

// Remove any return values.

ReturnUtils::removeReturns(func, *module);

// It's definitely worth optimizing the caller after inserting the

// unreachable.

return wasReturnUninhabitable;

}

// Given a function and all the calls to it, see if we can refine the type of

// its arguments. If we only pass in a subtype, we may as well refine the type

// to that.

//

// This assumes that the function has no calls aside from |calls|, that is, it

// is not exported or called from the table or by reference.

bool refineArgumentTypes(Function* func,

const std::vector<Call*>& calls,

Module* module,

const DAEFunctionInfo& info) {

if (!module->features.hasGC()) {

return false;

}

auto numParams = func->getNumParams();

std::vector<Type> newParamTypes;

newParamTypes.reserve(numParams);

std::vector<LUBFinder> lubs(numParams);

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

auto originalType = func->getLocalType(i);

// If the parameter type is not a reference, there is nothing to refine.

// And if it is unused, also do nothing, as we can leave it to the other

// parts of this pass to optimize it properly, which avoids having to

// think about corner cases involving refining the type of an unused

// param (in particular, unused params are turned into locals, which means

// we'd need to think about defaultability etc.).

if (!originalType.isRef() || info.unusedParams.has(i)) {

newParamTypes.push_back(originalType);

continue;

}

auto& lub = lubs[i];

for (auto* call : calls) {

auto* operand = call->operands[i];

lub.note(operand->type);

if (lub.getLUB() == originalType) {

// We failed to refine this parameter to anything more specific.

break;

}

}

// Nothing is sent here at all; leave such optimizations to DCE.

if (!lub.noted()) {

return false;

}

newParamTypes.push_back(lub.getLUB());

}

// Check if we are able to optimize here before we do the work to scan the

// function body.

auto newParams = Type(newParamTypes);

if (newParams == func->getParams()) {

return false;

}

// We can do this!

TypeUpdating::updateParamTypes(func, newParamTypes, *module);

// Update the function's type.

func->setParams(newParams);

return true;

}

// See if the types returned from a function allow us to define a more refined

// return type for it. If so, we can update it and all calls going to it.

//

// This assumes that the function has no calls aside from |calls|, that is, it

// is not exported or called from the table or by reference. Exports should be

// fine, as should indirect calls in principle, but VMs will need to support

// function subtyping in indirect calls. TODO: relax this when possible

//

// Returns whether we optimized.

//

// TODO: We may be missing a global optimum here, as e.g. if a function calls

// itself and returns that value, then we would not do any change here,

// as one of the return values is exactly what it already is. Similar

// unoptimality can happen with multiple functions, more local code in

// the middle, etc.

bool refineReturnTypes(Function* func,

const std::vector<Call*>& calls,

Module* module) {

auto lub = LUB::getResultsLUB(func, *module);

if (!lub.noted()) {

return false;

}

auto newType = lub.getLUB();

if (newType != func->getResults()) {

func->setResults(newType);

for (auto* call : calls) {

if (call->type != Type::unreachable) {

call->type = newType;

}

}

return true;

}

return false;

}

};

Pass* createDAEPass() { return new DAE(); }

Pass* createDAEOptimizingPass() {

auto* ret = new DAE();

ret->optimize = true;

return ret;

}

} // namespace wasm