/*
 * Copyright 2022 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.
 */

//
// Remove params from signature/function types where possible.
//
// This differs from DeadArgumentElimination in that DAE will look at each
// function by itself, and cannot handle indirectly-called functions. This pass
// looks at each heap type at a time, and if all functions with a heap type do
// not use a particular param, will remove the param.
//
// Like in DAE, as part of pruning parameters this will find parameters that are
// always sent the same constant value. We can then apply that value in the
// function, making the parameter's value unused, which means we can prune it.
//

#include "ir/find_all.h"
#include "ir/intrinsics.h"
#include "ir/module-utils.h"
#include "ir/subtypes.h"
#include "ir/type-updating.h"
#include "param-utils.h"
#include "pass.h"
#include "support/insert_ordered.h"
#include "support/sorted_vector.h"
#include "wasm-type.h"
#include "wasm.h"

#ifndef DAE_STATS
#define DAE_STATS 0
#else
#endif

#if DAE_STATS
#include <iostream>
#endif

namespace wasm {

namespace {

struct SignaturePruning : public Pass {
  void run(Module* module) override {
#if DAE_STATS
    Index startParams = 0;
    for (auto& func : module->functions) {
      startParams += func->getNumParams();
    }
#endif // DAE_STATS

    if (!module->features.hasGC()) {
      return;
    }

    if (getPassOptions().worldMode == WorldMode::Open) {
      Fatal() << "SignaturePruning requires --closed-world";
    }

    if (!module->tables.empty()) {
      // When there are tables we must also take their types into account, which
      // would require us to take call_indirect, element segments, etc. into
      // account. For now, do nothing if there are tables.
      // TODO
      return;
    }

    // The first iteration may suggest additional work is possible. If so, run
    // another cycle. (Even more cycles may help, but limit ourselves to 2 for
    // now.)
    if (iteration(module)) {
      iteration(module);
    }

#if DAE_STATS
    Index endParams = 0;
    for (auto& func : module->functions) {
      endParams += func->getNumParams();
    }
    std::cout << "Removed parameters: " << (startParams - endParams) << "\n";
#endif // DAE_STATS
  }

  // Returns true if more work is possible.
  bool iteration(Module* module) {
    // First, find all the information we need. Start by collecting inside each
    // function in parallel.

    struct Info {
      std::vector<Call*> calls;
      std::vector<CallRef*> callRefs;

      std::unordered_set<Index> usedParams;

      // If we set this to false, we may not attempt to perform any optimization
      // whatsoever on this data.
      bool optimizable = true;
    };

    ModuleUtils::ParallelFunctionAnalysis<Info> analysis(
      *module, [&](Function* func, Info& info) {
        if (func->imported()) {
          // Imports cannot be modified.
          info.optimizable = false;
          return;
        }

        info.calls = std::move(FindAll<Call>(func->body).list);
        info.callRefs = std::move(FindAll<CallRef>(func->body).list);
        info.usedParams = ParamUtils::getUsedParams(func, module);
      });

    // A map of types to all the information combined over all the functions
    // with that type.
    std::unordered_map<HeapType, Info> allInfo;

    // Map heap types to all functions with that type.
    InsertOrderedMap<HeapType, std::vector<Function*>> sigFuncs;

    // Heap types of call targets that we found we should localize calls to, in
    // order to fully handle them. (See similar code in DeadArgumentElimination
    // for individual functions; here we handle a HeapType at a time.) A slight
    // complication is that we cannot track heap types here: heap types are
    // rewritten using |GlobalTypeRewriter::updateSignatures| below, and even
    // types that we do not modify end up replaced (as the entire set of types
    // becomes one new big rec group). We therefore need something more stable
    // to track here, which we do using either a Call or a Call Ref.
    std::unordered_set<Expression*> callTargetsToLocalize;

    // Combine all the information we gathered into that map, iterating in a
    // deterministic order as we build up vectors where the order matters.
    for (auto& f : module->functions) {
      auto* func = f.get();
      auto& info = analysis.map[func];

      // For direct calls, add each call to the type of the function being
      // called.
      for (auto* call : info.calls) {
        allInfo[module->getFunction(call->target)->type.getHeapType()]
          .calls.push_back(call);

        // Intrinsics limit our ability to optimize in some cases. We will avoid
        // modifying any type that is used by call.without.effects, to avoid
        // the complexity of handling that. After intrinsics are lowered,
        // this optimization will be able to run at full power anyhow.
        if (Intrinsics(*module).isCallWithoutEffects(call)) {
          // The last operand is the actual call target.
          auto* target = call->operands.back();
          if (target->type != Type::unreachable) {
            allInfo[target->type.getHeapType()].optimizable = false;
          }
        }
      }

      // For indirect calls, add each call_ref to the type the call_ref uses.
      for (auto* callRef : info.callRefs) {
        auto calledType = callRef->target->type;
        if (calledType != Type::unreachable) {
          allInfo[calledType.getHeapType()].callRefs.push_back(callRef);
        }
      }

      // A parameter used in this function is used in the heap type - just one
      // function is enough to prevent the parameter from being removed.
      auto& allUsedParams = allInfo[func->type.getHeapType()].usedParams;
      for (auto index : info.usedParams) {
        allUsedParams.insert(index);
      }

      if (!info.optimizable) {
        allInfo[func->type.getHeapType()].optimizable = false;
      }

      sigFuncs[func->type.getHeapType()].push_back(func);
    }

    // Find the public types, which cannot be modified.
    for (auto type :
         ModuleUtils::getPublicHeapTypes(*module, getPassOptions().worldMode)) {
      if (type.isFunction()) {
        allInfo[type].optimizable = false;
      }
    }

    // Similarly, we cannot yet modify types used in exception handling or stack
    // switching tags. TODO.
    for (auto& tag : module->tags) {
      allInfo[tag->type].optimizable = false;
    }

    // Continuations must not have params refined, because we do not update
    // their users (e.g. cont.bind, resume) with new types.
    // TODO: support refining continuations
    if (module->features.hasStackSwitching()) {
      for (auto type : ModuleUtils::collectHeapTypes(*module)) {
        if (type.isContinuation()) {
          allInfo[type.getContinuation().type].optimizable = false;
        }
      }
    }

    // Signature-called functions must also not be modified.
    // TODO: Explore whether removing parameters from the end could be
    //       beneficial (check if it does not regress call performance with JS).
    for (auto func : Intrinsics(*module).getJSCalledFunctions()) {
      allInfo[module->getFunction(func)->type.getHeapType()].optimizable =
        false;
    }

    // A type must have the same number of parameters and results as its
    // supertypes and subtypes, so we only attempt to modify types without
    // supertypes or subtypes.
    // TODO We could handle "cycles" where we remove fields from a group of
    //      types with subtyping relations at once.
    SubTypes subTypes(*module);

    // Maps each heap type to the possible pruned signature. We will fill this
    // during analysis and then use it while doing an update of the types. If a
    // type has no improvement that we can find, it will not appear in this map.
    std::unordered_map<HeapType, Signature> newSignatures;

    // Find parameters to prune.
    //
    // TODO: The order matters here, and more than one cycle can find more work
    //       in some cases, as finding a parameter is a constant and removing it
    //       can lead to another call (that receives that parameter's value) to
    //       now have constant parameters as well, and so it becomes
    //       optimizable. We could do a topological sort or greatest fixed point
    //       analysis to be optimal (that could handle a recursive call with a
    //       constant).
    for (auto& [type, funcs] : sigFuncs) {
      auto sig = type.getSignature();
      auto& info = allInfo[type];
      auto& usedParams = info.usedParams;
      auto numParams = sig.params.size();

      if (!info.optimizable) {
        continue;
      }

      if (!subTypes.getImmediateSubTypes(type).empty()) {
        continue;
      }
      if (auto super = type.getDeclaredSuperType()) {
        if (super->isSignature()) {
          continue;
        }
      }

      // Apply constant indexes: find the parameters that are always sent a
      // constant value, and apply that value in the function. That then makes
      // the parameter unused (since the applied value makes us ignore the value
      // arriving in the parameter).
      auto optimizedIndexes = ParamUtils::applyConstantValues(
        funcs, info.calls, info.callRefs, module);
      for (auto i : optimizedIndexes) {
        usedParams.erase(i);
      }

      if (usedParams.size() == numParams) {
        // All parameters are used, give up on this one.
        continue;
      }

      // We found possible work! Find the specific params that are unused & try
      // to prune them.
      SortedVector<Index> unusedParams;
      for (Index i = 0; i < numParams; i++) {
        if (!usedParams.contains(i)) {
          unusedParams.insert(i);
        }
      }

      auto oldParams = sig.params;
      auto [removedIndexes, outcome] =
        ParamUtils::removeParameters(funcs,
                                     unusedParams,
                                     info.calls,
                                     info.callRefs,
                                     module,
                                     getPassRunner());
      if (outcome == ParamUtils::RemovalOutcome::Failure) {
        // Use either a Call or a CallRef that has this type (see explanation
        // above on |callTargetsToLocalize|.
        if (!info.calls.empty()) {
          callTargetsToLocalize.insert(info.calls[0]);
        } else {
          assert(!info.callRefs.empty());
          callTargetsToLocalize.insert(info.callRefs[0]);
        }
      }
      if (removedIndexes.empty()) {
        continue;
      }

      // Success! Update the types.
      std::vector<Type> newParams;
      for (Index i = 0; i < numParams; i++) {
        if (!removedIndexes.has(i)) {
          newParams.push_back(oldParams[i]);
        }
      }

      // Create a new signature. When the TypeRewriter operates below it will
      // modify the existing heap type in place to change its signature to this
      // one. TypeRewriter will also ensure that distinct types remain
      // distinct, even if they have the same signature after optimization.
      newSignatures[type] = Signature(Type(newParams), sig.results);

      // removeParameters() updates the type as it goes, but in this pass we
      // need the type to be updated in all locations at once. That is, we need
      // all the functions of a particular type to still have the same type
      // after this operation, and that must be the exact same type at the
      // relevant call_refs and so forth. The TypeRewriter below will do the
      // right thing as it rewrites everything all at once, so we do not want
      // the type to be modified by removeParameters(), and so we undo the type
      // it made.
      //
      // Note that we cannot just ask removeParameters() to not update the type,
      // as it adds a new local there, whose index depends on the type (which
      // contains the # of parameters, and that determine where non-parameter
      // local indexes begin). Rather than have it update the type and then undo
      // that, which would add more complexity in that method, undo the change
      // here.
      for (auto* func : funcs) {
        func->type = func->type.with(type);
      }
    }

    // Rewrite the types.
    GlobalTypeRewriter::updateSignatures(
      newSignatures, *module, getPassOptions().worldMode);

    if (callTargetsToLocalize.empty()) {
      return false;
    }

    // Localize after updating signatures, to not interfere with that
    // operation (localization adds locals, and the indexes of locals must be
    // taken into account in |GlobalTypeRewriter::updateSignatures| (as var
    // indexes change when params are pruned).
    std::unordered_set<HeapType> callTargetTypes;
    for (auto* call : callTargetsToLocalize) {
      HeapType type;
      if (auto* c = call->dynCast<Call>()) {
        type = module->getFunction(c->target)->type.getHeapType();
      } else if (auto* c = call->dynCast<CallRef>()) {
        type = c->target->type.getHeapType();
      } else {
        WASM_UNREACHABLE("bad call");
      }
      callTargetTypes.insert(type);
    }

    ParamUtils::localizeCallsTo(callTargetTypes, *module, getPassRunner());

    return true;
  }
};

} // anonymous namespace

Pass* createSignaturePruningPass() { return new SignaturePruning(); }

} // namespace wasm