llvm/lib/Target/WebAssembly/WebAssemblyOptimizeReturned.cpp

//===-- WebAssemblyOptimizeReturned.cpp - Optimize "returned" attributes --===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
///
/// \file
/// Optimize calls with "returned" attributes for WebAssembly.
///
//===----------------------------------------------------------------------===//

#include "WebAssembly.h"
#include "llvm/IR/Analysis.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/InstVisitor.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
using namespace llvm;

#define DEBUG_TYPE "wasm-optimize-returned"

namespace {
class WebAssemblyOptimizeReturnedImpl
    : public InstVisitor<WebAssemblyOptimizeReturnedImpl> {
  DominatorTree *DT = nullptr;

public:
  WebAssemblyOptimizeReturnedImpl(DominatorTree *DT) : DT(DT) {}
  bool runOnFunction(Function &F);
  void visitCallBase(CallBase &CB);
};

class WebAssemblyOptimizeReturnedLegacy final : public FunctionPass {
  StringRef getPassName() const override {
    return "WebAssembly Optimize Returned";
  }

  void getAnalysisUsage(AnalysisUsage &AU) const override {
    AU.setPreservesCFG();
    AU.addRequired<DominatorTreeWrapperPass>();
    AU.addPreserved<DominatorTreeWrapperPass>();
    FunctionPass::getAnalysisUsage(AU);
  }

  bool runOnFunction(Function &F) override;

public:
  static char ID;
  WebAssemblyOptimizeReturnedLegacy() : FunctionPass(ID) {}

  void visitCallBase(CallBase &CB);
};
} // End anonymous namespace

char WebAssemblyOptimizeReturnedLegacy::ID = 0;
INITIALIZE_PASS(WebAssemblyOptimizeReturnedLegacy, DEBUG_TYPE,
                "Optimize calls with \"returned\" attributes for WebAssembly",
                false, false)

FunctionPass *llvm::createWebAssemblyOptimizeReturnedLegacyPass() {
  return new WebAssemblyOptimizeReturnedLegacy();
}

void WebAssemblyOptimizeReturnedImpl::visitCallBase(CallBase &CB) {
  for (unsigned I = 0, E = CB.arg_size(); I < E; ++I)
    if (CB.paramHasAttr(I, Attribute::Returned)) {
      Value *Arg = CB.getArgOperand(I);
      // Ignore constants, globals, undef, etc.
      if (isa<Constant>(Arg))
        continue;
      // Like replaceDominatedUsesWith but using Instruction/Use dominance.
      Arg->replaceUsesWithIf(&CB, [&](Use &U) {
        auto *I = cast<Instruction>(U.getUser());
        return !I->isLifetimeStartOrEnd() && DT->dominates(&CB, U);
      });
    }
}

bool WebAssemblyOptimizeReturnedImpl::runOnFunction(Function &F) {
  LLVM_DEBUG(dbgs() << "********** Optimize returned Attributes **********\n"
                       "********** Function: "
                    << F.getName() << '\n');

  visit(F);
  return true;
}

bool WebAssemblyOptimizeReturnedLegacy::runOnFunction(Function &F) {
  DominatorTree *DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
  WebAssemblyOptimizeReturnedImpl Impl(DT);
  return Impl.runOnFunction(F);
}

PreservedAnalyses
WebAssemblyOptimizeReturnedPass::run(Function &F,
                                     FunctionAnalysisManager &FAM) {
  DominatorTree *DT = &FAM.getResult<DominatorTreeAnalysis>(F);
  WebAssemblyOptimizeReturnedImpl Impl(DT);
  return Impl.runOnFunction(F)
             ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
             : PreservedAnalyses::all();
}