//===- AggressiveInstCombine.cpp ------------------------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements the aggressive expression pattern combiner classes.
// Currently, it handles expression patterns for:
//  * Truncate instruction
//
//===----------------------------------------------------------------------===//

#include "llvm/Transforms/AggressiveInstCombine/AggressiveInstCombine.h"
#include "AggressiveInstCombineInternal.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AliasAnalysis.h"
#include "llvm/Analysis/AssumptionCache.h"
#include "llvm/Analysis/BasicAliasAnalysis.h"
#include "llvm/Analysis/ConstantFolding.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/TargetLibraryInfo.h"
#include "llvm/Analysis/TargetTransformInfo.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/DataLayout.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/PatternMatch.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/ProfileData/InstrProf.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"
#include "llvm/Transforms/Utils/BuildLibCalls.h"
#include "llvm/Transforms/Utils/Local.h"

using namespace llvm;
using namespace PatternMatch;

#define DEBUG_TYPE "aggressive-instcombine"

STATISTIC(NumAnyOrAllBitsSet, "Number of any/all-bits-set patterns folded");
STATISTIC(NumGuardedRotates,
          "Number of guarded rotates transformed into funnel shifts");
STATISTIC(NumGuardedFunnelShifts,
          "Number of guarded funnel shifts transformed into funnel shifts");
STATISTIC(NumPopCountRecognized, "Number of popcount idioms recognized");
STATISTIC(NumSelectCTTZFolded,
          "Number of select-based split cttz patterns folded");
STATISTIC(NumSelectCTLZFolded,
          "Number of select-based split ctlz patterns folded");
STATISTIC(NumMemSetsGuarded, "Number of memsets guarded for a zero length");
STATISTIC(NumTableBasedLowBitsMask,
          "Number of low-bits mask table loads folded to (1 << i) - 1");

static cl::opt<unsigned> MaxInstrsToScan(
    "aggressive-instcombine-max-scan-instrs", cl::init(64), cl::Hidden,
    cl::desc("Max number of instructions to scan for aggressive instcombine."));

static cl::opt<unsigned> StrNCmpInlineThreshold(
    "strncmp-inline-threshold", cl::init(3), cl::Hidden,
    cl::desc("The maximum length of a constant string for a builtin string cmp "
             "call eligible for inlining. The default value is 3."));

static cl::opt<unsigned>
    MemChrInlineThreshold("memchr-inline-threshold", cl::init(3), cl::Hidden,
                          cl::desc("The maximum length of a constant string to "
                                   "inline a memchr call."));

namespace llvm {
extern cl::opt<bool> ProfcheckDisableMetadataFixes;
} // namespace llvm

/// Try to fold a select-based split cttz pattern into a single full-width cttz.
///
///   %lo = trunc iN %val to i(N/2)
///   %cmp = icmp eq i(N/2) %lo, 0
///   %shr = lshr iN %val, N/2
///   %hi = trunc iN %shr to i(N/2)
///   %cttz_hi = call i(N/2) @llvm.cttz.i(N/2)(i(N/2) %hi, ...)
///   %hi_plus = add/or_disjoint i(N/2) %cttz_hi, N/2
///   %cttz_lo = call i(N/2) @llvm.cttz.i(N/2)(i(N/2) %lo, ...)
///   %result = select i1 %cmp, i(N/2) %hi_plus, i(N/2) %cttz_lo
/// -->
///   %cttz_wide = call iN @llvm.cttz.iN(iN %val, i1 false)
///   %result = trunc iN %cttz_wide to i(N/2)
/// Alive proof (for i64/i32):  https://alive2.llvm.org/ce/z/-s14-s
// TrueVal/FalseVal are pre-normalized by the caller to the EQ/NE cases.
static bool foldSelectSplitCTTZ(Instruction &I, Value *LoTrunc, Value *HiResult,
                                Value *LoResult, Type *HalfTy) {
  unsigned HalfWidth = HalfTy->getIntegerBitWidth();
  unsigned FullWidth = HalfWidth * 2;

  // LoTrunc: trunc iN SrcVal to i(N/2)
  Value *SrcVal;
  if (!match(LoTrunc, m_Trunc(m_Value(SrcVal))))
    return false;
  if (!SrcVal->getType()->isIntegerTy(FullWidth))
    return false;

  // LoResult: cttz(trunc(SrcVal), _),  must use same truncated value
  if (!match(LoResult, m_OneUse(m_Cttz(m_Specific(LoTrunc), m_Value()))))
    return false;

  // HiResult: add/or_disjoint(cttz(trunc(lshr(SrcVal, N/2)), _), N/2)
  Value *CttzHiCall;
  if (!match(HiResult, m_OneUse(m_AddLike(m_Value(CttzHiCall),
                                          m_SpecificInt(HalfWidth)))))
    return false;

  Value *HiCttzArg;
  if (!match(CttzHiCall, m_OneUse(m_Cttz(m_Value(HiCttzArg), m_Value()))))
    return false;

  if (!match(HiCttzArg,
             m_Trunc(m_LShr(m_Specific(SrcVal), m_SpecificInt(HalfWidth)))))
    return false;

  // Match successful.
  IRBuilder<> Builder(&I);
  Value *CttzWide = Builder.CreateIntrinsic(
      Intrinsic::cttz, {SrcVal->getType()}, {SrcVal, Builder.getFalse()});
  Value *Trunc = Builder.CreateTrunc(CttzWide, HalfTy);

  I.replaceAllUsesWith(Trunc);
  ++NumSelectCTTZFolded;
  return true;
}

/// Same as foldSelectSplitCTTZ but for leading zeros (ctlz).
///
///   %shr = lshr iN %val, N/2
///   %hi = trunc iN %shr to i(N/2)
///   %cmp = icmp eq i(N/2) %hi, 0   (or icmp eq iN %shr, 0)
///   %lo = trunc iN %val to i(N/2)
///   %ctlz_lo = call i(N/2) @llvm.ctlz.i(N/2)(i(N/2) %lo, ...)
///   %lo_plus = add/or_disjoint i(N/2) %ctlz_lo, N/2
///   %ctlz_hi = call i(N/2) @llvm.ctlz.i(N/2)(i(N/2) %hi, ...)
///   %result = select i1 %cmp, i(N/2) %lo_plus, i(N/2) %ctlz_hi
/// -->
///   %ctlz_wide = call iN @llvm.ctlz.iN(iN %val, i1 false)
///   %result = trunc iN %ctlz_wide to i(N/2)
///
/// Alive proof (for i64/i32): https://alive2.llvm.org/ce/z/WfQepH
// TrueVal/FalseVal are pre-normalized by the caller to the EQ/NE cases.
static bool foldSelectSplitCTLZ(Instruction &I, Value *HiPart, Value *LoResult,
                                Value *HiResult, Type *HalfTy) {
  unsigned HalfWidth = HalfTy->getIntegerBitWidth();
  unsigned FullWidth = HalfWidth * 2;

  // Extract SrcVal from HiPart: either trunc(lshr(SrcVal, N/2)) or
  // lshr(SrcVal, N/2)
  Value *SrcVal;
  if (match(HiPart, m_Trunc(m_Value(SrcVal))))
    HiPart = SrcVal;

  if (!match(HiPart, m_LShr(m_Value(SrcVal), m_SpecificInt(HalfWidth))))
    return false;
  if (!SrcVal->getType()->isIntegerTy(FullWidth))
    return false;

  // HiResult: ctlz(trunc(lshr(SrcVal, N/2)), _)
  Value *HiCtlzArg;
  if (!match(HiResult, m_OneUse(m_Ctlz(m_Value(HiCtlzArg), m_Value()))))
    return false;

  if (!match(HiCtlzArg,
             m_Trunc(m_LShr(m_Specific(SrcVal), m_SpecificInt(HalfWidth)))))
    return false;

  // LoResult: add/or_disjoint(ctlz(trunc(SrcVal), _), N/2)
  Value *CtlzLoCall;
  if (!match(LoResult, m_OneUse(m_AddLike(m_Value(CtlzLoCall),
                                          m_SpecificInt(HalfWidth)))))
    return false;

  Value *LoCtlzArg;
  if (!match(CtlzLoCall, m_OneUse(m_Ctlz(m_Value(LoCtlzArg), m_Value()))))
    return false;

  if (!match(LoCtlzArg, m_Trunc(m_Specific(SrcVal))))
    return false;

  // Match successful.
  IRBuilder<> Builder(&I);
  Value *CtlzWide = Builder.CreateIntrinsic(
      Intrinsic::ctlz, {SrcVal->getType()}, {SrcVal, Builder.getFalse()});
  Value *Trunc = Builder.CreateTrunc(CtlzWide, HalfTy);

  I.replaceAllUsesWith(Trunc);
  ++NumSelectCTLZFolded;
  return true;
}

/// Common entry point for folding select-based split cttz/ctlz patterns.
/// Performs the initial select and type matching shared by both transforms,
/// then delegates to foldSelectSplitCTTZ and foldSelectSplitCTLZ.
static bool foldSelectSplitCTLZCTTZ(Instruction &I) {
  Value *Cond, *TrueVal, *FalseVal;
  if (!match(&I, m_Select(m_Value(Cond), m_Value(TrueVal), m_Value(FalseVal))))
    return false;

  Type *Ty = I.getType();
  if (!Ty->isIntegerTy())
    return false;

  // Bail out on very small types (i1, i2): the full-width cttz/ctlz can return
  // values not representable in the half type (e.g., cttz.i4 can return 4,
  // which doesn't fit in i2).
  if (Ty->getIntegerBitWidth() <= 2)
    return false;

  CmpPredicate Pred;
  Value *CmpOp;
  if (!match(Cond, m_ICmp(Pred, m_Value(CmpOp), m_ZeroInt())) ||
      !ICmpInst::isEquality(Pred))
    return false;

  // Canonicalize select operands.
  if (Pred == CmpInst::ICMP_NE)
    std::swap(TrueVal, FalseVal);

  return foldSelectSplitCTTZ(I, CmpOp, TrueVal, FalseVal, Ty) ||
         foldSelectSplitCTLZ(I, CmpOp, TrueVal, FalseVal, Ty);
}

/// Match a pattern for a bitwise funnel/rotate operation that partially guards
/// against undefined behavior by branching around the funnel-shift/rotation
/// when the shift amount is 0.
static bool foldGuardedFunnelShift(Instruction &I, const DominatorTree &DT) {
  if (I.getOpcode() != Instruction::PHI || I.getNumOperands() != 2)
    return false;

  // As with the one-use checks below, this is not strictly necessary, but we
  // are being cautious to avoid potential perf regressions on targets that
  // do not actually have a funnel/rotate instruction (where the funnel shift
  // would be expanded back into math/shift/logic ops).
  if (!isPowerOf2_32(I.getType()->getScalarSizeInBits()))
    return false;

  // Match V to funnel shift left/right and capture the source operands and
  // shift amount.
  auto matchFunnelShift = [](Value *V, Value *&ShVal0, Value *&ShVal1,
                             Value *&ShAmt) {
    unsigned Width = V->getType()->getScalarSizeInBits();

    // fshl(ShVal0, ShVal1, ShAmt)
    //  == (ShVal0 << ShAmt) | (ShVal1 >> (Width -ShAmt))
    if (match(V, m_OneUse(m_c_Or(
                     m_Shl(m_Value(ShVal0), m_Value(ShAmt)),
                     m_LShr(m_Value(ShVal1), m_Sub(m_SpecificInt(Width),
                                                   m_Deferred(ShAmt))))))) {
      return Intrinsic::fshl;
    }

    // fshr(ShVal0, ShVal1, ShAmt)
    //  == (ShVal0 >> ShAmt) | (ShVal1 << (Width - ShAmt))
    if (match(V,
              m_OneUse(m_c_Or(m_Shl(m_Value(ShVal0), m_Sub(m_SpecificInt(Width),
                                                           m_Value(ShAmt))),
                              m_LShr(m_Value(ShVal1), m_Deferred(ShAmt)))))) {
      return Intrinsic::fshr;
    }

    return Intrinsic::not_intrinsic;
  };

  // One phi operand must be a funnel/rotate operation, and the other phi
  // operand must be the source value of that funnel/rotate operation:
  // phi [ rotate(RotSrc, ShAmt), FunnelBB ], [ RotSrc, GuardBB ]
  // phi [ fshl(ShVal0, ShVal1, ShAmt), FunnelBB ], [ ShVal0, GuardBB ]
  // phi [ fshr(ShVal0, ShVal1, ShAmt), FunnelBB ], [ ShVal1, GuardBB ]
  PHINode &Phi = cast<PHINode>(I);
  unsigned FunnelOp = 0, GuardOp = 1;
  Value *P0 = Phi.getOperand(0), *P1 = Phi.getOperand(1);
  Value *ShVal0, *ShVal1, *ShAmt;
  Intrinsic::ID IID = matchFunnelShift(P0, ShVal0, ShVal1, ShAmt);
  if (IID == Intrinsic::not_intrinsic ||
      (IID == Intrinsic::fshl && ShVal0 != P1) ||
      (IID == Intrinsic::fshr && ShVal1 != P1)) {
    IID = matchFunnelShift(P1, ShVal0, ShVal1, ShAmt);
    if (IID == Intrinsic::not_intrinsic ||
        (IID == Intrinsic::fshl && ShVal0 != P0) ||
        (IID == Intrinsic::fshr && ShVal1 != P0))
      return false;
    assert((IID == Intrinsic::fshl || IID == Intrinsic::fshr) &&
           "Pattern must match funnel shift left or right");
    std::swap(FunnelOp, GuardOp);
  }

  // The incoming block with our source operand must be the "guard" block.
  // That must contain a cmp+branch to avoid the funnel/rotate when the shift
  // amount is equal to 0. The other incoming block is the block with the
  // funnel/rotate.
  BasicBlock *GuardBB = Phi.getIncomingBlock(GuardOp);
  BasicBlock *FunnelBB = Phi.getIncomingBlock(FunnelOp);
  Instruction *TermI = GuardBB->getTerminator();

  // Ensure that the shift values dominate each block.
  if (!DT.dominates(ShVal0, TermI) || !DT.dominates(ShVal1, TermI))
    return false;

  BasicBlock *PhiBB = Phi.getParent();
  if (!match(TermI, m_Br(m_SpecificICmp(CmpInst::ICMP_EQ, m_Specific(ShAmt),
                                        m_ZeroInt()),
                         m_SpecificBB(PhiBB), m_SpecificBB(FunnelBB))))
    return false;

  IRBuilder<> Builder(PhiBB, PhiBB->getFirstInsertionPt());

  if (ShVal0 == ShVal1)
    ++NumGuardedRotates;
  else
    ++NumGuardedFunnelShifts;

  // If this is not a rotate then the select was blocking poison from the
  // 'shift-by-zero' non-TVal, but a funnel shift won't - so freeze it.
  bool IsFshl = IID == Intrinsic::fshl;
  if (ShVal0 != ShVal1) {
    if (IsFshl && !llvm::isGuaranteedNotToBePoison(ShVal1))
      ShVal1 = Builder.CreateFreeze(ShVal1);
    else if (!IsFshl && !llvm::isGuaranteedNotToBePoison(ShVal0))
      ShVal0 = Builder.CreateFreeze(ShVal0);
  }

  // We matched a variation of this IR pattern:
  // GuardBB:
  //   %cmp = icmp eq i32 %ShAmt, 0
  //   br i1 %cmp, label %PhiBB, label %FunnelBB
  // FunnelBB:
  //   %sub = sub i32 32, %ShAmt
  //   %shr = lshr i32 %ShVal1, %sub
  //   %shl = shl i32 %ShVal0, %ShAmt
  //   %fsh = or i32 %shr, %shl
  //   br label %PhiBB
  // PhiBB:
  //   %cond = phi i32 [ %fsh, %FunnelBB ], [ %ShVal0, %GuardBB ]
  // -->
  // llvm.fshl.i32(i32 %ShVal0, i32 %ShVal1, i32 %ShAmt)
  Phi.replaceAllUsesWith(
      Builder.CreateIntrinsic(IID, Phi.getType(), {ShVal0, ShVal1, ShAmt}));
  return true;
}

/// This is used by foldAnyOrAllBitsSet() to capture a source value (Root) and
/// the bit indexes (Mask) needed by a masked compare. If we're matching a chain
/// of 'and' ops, then we also need to capture the fact that we saw an
/// "and X, 1", so that's an extra return value for that case.
namespace {
struct MaskOps {
  Value *Root = nullptr;
  APInt Mask;
  bool MatchAndChain;
  bool FoundAnd1 = false;

  MaskOps(unsigned BitWidth, bool MatchAnds)
      : Mask(APInt::getZero(BitWidth)), MatchAndChain(MatchAnds) {}
};
} // namespace

/// This is a recursive helper for foldAnyOrAllBitsSet() that walks through a
/// chain of 'and' or 'or' instructions looking for shift ops of a common source
/// value. Examples:
///   or (or (or X, (X >> 3)), (X >> 5)), (X >> 8)
/// returns { X, 0x129 }
///   and (and (X >> 1), 1), (X >> 4)
/// returns { X, 0x12 }
static bool matchAndOrChain(Value *V, MaskOps &MOps) {
  Value *Op0, *Op1;
  if (MOps.MatchAndChain) {
    // Recurse through a chain of 'and' operands. This requires an extra check
    // vs. the 'or' matcher: we must find an "and X, 1" instruction somewhere
    // in the chain to know that all of the high bits are cleared.
    if (match(V, m_And(m_Value(Op0), m_One()))) {
      MOps.FoundAnd1 = true;
      return matchAndOrChain(Op0, MOps);
    }
    if (match(V, m_And(m_Value(Op0), m_Value(Op1))))
      return matchAndOrChain(Op0, MOps) && matchAndOrChain(Op1, MOps);
  } else {
    // Recurse through a chain of 'or' operands.
    if (match(V, m_Or(m_Value(Op0), m_Value(Op1))))
      return matchAndOrChain(Op0, MOps) && matchAndOrChain(Op1, MOps);
  }

  // We need a shift-right or a bare value representing a compare of bit 0 of
  // the original source operand.
  Value *Candidate;
  const APInt *BitIndex = nullptr;
  if (!match(V, m_LShr(m_Value(Candidate), m_APInt(BitIndex))))
    Candidate = V;

  // Initialize result source operand.
  if (!MOps.Root)
    MOps.Root = Candidate;

  // The shift constant is out-of-range? This code hasn't been simplified.
  if (BitIndex && BitIndex->uge(MOps.Mask.getBitWidth()))
    return false;

  // Fill in the mask bit derived from the shift constant.
  MOps.Mask.setBit(BitIndex ? BitIndex->getZExtValue() : 0);
  return MOps.Root == Candidate;
}

/// Match patterns that correspond to "any-bits-set" and "all-bits-set".
/// These will include a chain of 'or' or 'and'-shifted bits from a
/// common source value:
/// and (or  (lshr X, C), ...), 1 --> (X & CMask) != 0
/// and (and (lshr X, C), ...), 1 --> (X & CMask) == CMask
/// Note: "any-bits-clear" and "all-bits-clear" are variations of these patterns
/// that differ only with a final 'not' of the result. We expect that final
/// 'not' to be folded with the compare that we create here (invert predicate).
static bool foldAnyOrAllBitsSet(Instruction &I) {
  // The 'any-bits-set' ('or' chain) pattern is simpler to match because the
  // final "and X, 1" instruction must be the final op in the sequence.
  bool MatchAllBitsSet;
  bool MatchTrunc;
  Value *X;
  if (I.getType()->isIntOrIntVectorTy(1)) {
    if (match(&I, m_Trunc(m_OneUse(m_And(m_Value(), m_Value())))))
      MatchAllBitsSet = true;
    else if (match(&I, m_Trunc(m_OneUse(m_Or(m_Value(), m_Value())))))
      MatchAllBitsSet = false;
    else
      return false;
    MatchTrunc = true;
    X = I.getOperand(0);
  } else {
    if (match(&I, m_c_And(m_OneUse(m_And(m_Value(), m_Value())), m_Value()))) {
      X = &I;
      MatchAllBitsSet = true;
    } else if (match(&I,
                     m_And(m_OneUse(m_Or(m_Value(), m_Value())), m_One()))) {
      X = I.getOperand(0);
      MatchAllBitsSet = false;
    } else
      return false;
    MatchTrunc = false;
  }
  Type *Ty = X->getType();

  MaskOps MOps(Ty->getScalarSizeInBits(), MatchAllBitsSet);
  if (!matchAndOrChain(X, MOps) ||
      (MatchAllBitsSet && !MatchTrunc && !MOps.FoundAnd1))
    return false;

  // The pattern was found. Create a masked compare that replaces all of the
  // shift and logic ops.
  IRBuilder<> Builder(&I);
  Constant *Mask = ConstantInt::get(Ty, MOps.Mask);
  Value *And = Builder.CreateAnd(MOps.Root, Mask);
  Value *Cmp = MatchAllBitsSet ? Builder.CreateICmpEQ(And, Mask)
                               : Builder.CreateIsNotNull(And);
  Value *Zext = MatchTrunc ? Cmp : Builder.CreateZExt(Cmp, Ty);
  I.replaceAllUsesWith(Zext);
  ++NumAnyOrAllBitsSet;
  return true;
}

/// Helper function to replace an instruction with a popcount intrinsic.
/// This creates the ctpop intrinsic with an optional truncation appended at the
/// end, and replaces all uses of the instruction.
static void replaceWithPopCount(Instruction &I, Value *Root) {
  LLVM_DEBUG(dbgs() << "Recognized popcount intrinsic\n");
  Type *RootTy = Root->getType();
  Type *OrigTy = I.getType();

  IRBuilder<> Builder(&I);
  Value *NewVal = Builder.CreateIntrinsic(Intrinsic::ctpop, RootTy, {Root});
  if (OrigTy != RootTy) {
    assert(RootTy->getScalarSizeInBits() > OrigTy->getScalarSizeInBits() &&
           "Only truncation is supported for now");
    NewVal = Builder.CreateTrunc(NewVal, OrigTy);
  }
  I.replaceAllUsesWith(NewVal);
  ++NumPopCountRecognized;
}

// Matches the common innermost steps of the Hacker's Delight popcount idiom:
//   V = ((x + (x >> 4)) & 0x0F...)
//   x = (y & 0x33...) + ((y >> 2) & 0x33...)  [or y - 3*((y>>2)&0x33...)]
//   y = Root - ((Root >> 1) & 0x55...)
// This computes the popcount for each byte.
// Returns Root on success, nullptr on failure.
static Value *matchPopCountBytes(Value *V, unsigned Len, const DataLayout &DL) {
  APInt Mask55 = APInt::getSplat(Len, APInt(8, 0x55));
  APInt Mask33 = APInt::getSplat(Len, APInt(8, 0x33));
  APInt Mask0F = APInt::getSplat(Len, APInt(8, 0x0F));

  Value *Add2;
  // Matching "((x + (x >> 4)) & 0x0F...)".
  if (!match(V, m_And(m_c_Add(m_LShr(m_Value(Add2), m_SpecificInt(4)),
                              m_Deferred(Add2)),
                      m_SpecificInt(Mask0F))))
    return nullptr;

  Value *Sub1;
  APInt NegThree(Len, -3, /*isSigned=*/true);
  // Match
  //   x = (x & 0x33333333) + ((x >> 2) & 0x33333333)"
  // Or
  //   x = x - 3*((x >> 2) & 0x33333333)
  if (!match(Add2, m_c_Add(m_And(m_LShr(m_Value(Sub1), m_SpecificInt(2)),
                                 m_SpecificInt(Mask33)),
                           m_And(m_Deferred(Sub1), m_SpecificInt(Mask33)))) &&
      !match(Add2, m_Add(m_Mul(m_And(m_LShr(m_Value(Sub1), m_SpecificInt(2)),
                                     m_SpecificInt(Mask33)),
                               m_SpecificInt(NegThree)),
                         m_Deferred(Sub1))))
    return nullptr;

  Value *Root, *LShr;
  const APInt *AndMask;
  // Matching "x - ((x >> 1) & 0x55...)".
  if (!match(Sub1,
             m_Sub(m_Value(Root), m_And(m_Value(LShr, m_LShr(m_Deferred(Root),
                                                             m_SpecificInt(1))),
                                        m_APInt(AndMask)))))
    return nullptr;

  if (*AndMask != Mask55) {
    // Accept a narrowed mask if missing bits are known zero in Root>>1.
    if (!AndMask->isSubsetOf(Mask55))
      return nullptr;
    APInt NeededMask = Mask55 & ~*AndMask;
    if (!MaskedValueIsZero(LShr, NeededMask, SimplifyQuery(DL)))
      return nullptr;
  }

  return Root;
}

// Try to recognize below function as popcount intrinsic.
// This is the "best" algorithm from
// http://graphics.stanford.edu/~seander/bithacks.html#CountBitsSetParallel
// Also used in TargetLowering::expandCTPOP().
//
// int popcount(unsigned int i) {
//   i = i - ((i >> 1) & 0x55555555);
//   i = (i & 0x33333333) + ((i >> 2) & 0x33333333);
//   i = ((i + (i >> 4)) & 0x0F0F0F0F);
//   return (i * 0x01010101) >> 24;
// }
static bool tryToRecognizePopCount(Instruction &I) {
  if (I.getOpcode() != Instruction::LShr)
    return false;

  Type *Ty = I.getType();
  if (!Ty->isIntOrIntVectorTy())
    return false;

  unsigned Len = Ty->getScalarSizeInBits();
  // Len==8 is handled by tryToRecognizePopCount2n3.
  // FIXME: other irregular type lengths.
  if (Len > 128 || Len <= 8 || Len % 8 != 0)
    return false;

  APInt Mask01 = APInt::getSplat(Len, APInt(8, 0x01));

  Value *Op0 = I.getOperand(0);
  Value *Op1 = I.getOperand(1);
  Value *MulOp0;
  // Matching "(i * 0x01010101...) >> 24".
  if (!match(Op0, m_Mul(m_Value(MulOp0), m_SpecificInt(Mask01))) ||
      !match(Op1, m_SpecificInt(Len - 8)))
    return false;

  Value *Root = matchPopCountBytes(MulOp0, Len, I.getDataLayout());
  if (!Root)
    return false;

  replaceWithPopCount(I, Root);
  return true;
}

// Try to recognize below function as popcount intrinsic.
// Ref. Hacker Delights
// int popcount32(unsigned int i) {
// uWord = (uWord & 0x55555555) + ((uWord>>1) & 0x55555555);
// uWord = (uWord & 0x33333333) + ((uWord>>2) & 0x33333333);
// uWord = (uWord & 0x0F0F0F0F) + ((uWord>>4) & 0x0F0F0F0F);
// uWord = (uWord & 0x00FF00FF) + ((uWord>>8) & 0x00FF00FF);
// return  (uWord & 0x0000FFFF) + (uWord>>16);
// }
// int popcount64(unsigned long i) {
// uWord = (uWord & 0x5555555555555555) + ((uWord>>1) & 0x5555555555555555);
// uWord = (uWord & 0x3333333333333333) + ((uWord>>2) & 0x3333333333333333);
// uWord = (uWord & 0x0F0F0F0F0F0F0F0F) + ((uWord>>4) & 0x0F0F0F0F0F0F0F0F);
// uWord = (uWord & 0x00FF00FF00FF00FF) + ((uWord>>8) & 0x00FF00FF00FF00FF);
// uWord = (uWord & 0x0000FFFF0000FFFF) + ((uWord>>16) & 0x0000FFFF0000FFFF);
// return  (uWord & 0x00000000FFFFFFFF) + (uWord>>32) & 0x00000000FFFFFFFF;
// }
//
// InstCombine may narrow AND masks when it can prove the removed bits are
// known zero (e.g. 0x0F0F0F0F -> 0x07070707). We accept such narrowed masks
// by checking they are subsets of the expected masks and verifying the missing
// bits are known zero via MaskedValueIsZero.
static bool tryToRecognizePopCount1(Instruction &I) {
  if (I.getOpcode() != Instruction::Add)
    return false;

  Type *Ty = I.getType();
  if (!Ty->isIntOrIntVectorTy())
    return false;

  unsigned Len = Ty->getScalarSizeInBits();
  if (Len > 64 || Len <= 8 || Len % 8 != 0)
    return false;

  // Len should be a power of 2 for the loop to work correctly
  if (!isPowerOf2_32(Len))
    return false;

  APInt Mask55 = APInt::getSplat(Len, APInt(8, 0x55));
  APInt Mask33 = APInt::getSplat(Len, APInt(8, 0x33));

  SimplifyQuery SQ(I.getDataLayout());

  // Check if CapturedMask is a valid (possibly narrowed) version of
  // ExpectedMask for the given Operand. Returns true if the masks match
  // exactly, or if CapturedMask is a subset and the missing bits are
  // known zero in the Operand.
  auto isValidNarrowedMask = [&](const APInt &CapturedMask,
                                 const APInt &ExpectedMask,
                                 Value *Operand) -> bool {
    if (CapturedMask == ExpectedMask)
      return true;
    if (!CapturedMask.isSubsetOf(ExpectedMask))
      return false;
    APInt NeededMask = ExpectedMask & ~CapturedMask;
    return MaskedValueIsZero(Operand, NeededMask, SQ);
  };

  // For "(x & M) + ((x >> S) & M)" patterns, both AND masks may be narrowed.
  // Require subsets of BaseMask and prove any implied missing bits are zero.
  auto narrowAddPairMasksOk = [&](const APInt &BaseMask, unsigned ShiftAmt,
                                  Value *Val, const APInt &AndMask1,
                                  const APInt &AndMask2) -> bool {
    if (!AndMask1.isSubsetOf(BaseMask) || !AndMask2.isSubsetOf(BaseMask))
      return false;
    APInt NeededShifted = (BaseMask & ~AndMask1).shl(ShiftAmt);
    APInt NeededUnshifted = BaseMask & ~AndMask2;
    APInt AllNeeded = NeededShifted | NeededUnshifted;
    return AllNeeded.isZero() || MaskedValueIsZero(Val, AllNeeded, SQ);
  };

  Value *ShiftOp;
  Value *Start = &I;
  for (unsigned I = Len; I >= 8; I = I / 2) {
    APInt Mask = APInt::getSplat(Len, APInt::getLowBitsSet(I, I / 2));
    const APInt *AndMask1 = nullptr, *AndMask2 = nullptr;

    // Matching "(uWord & Mask) + ((uWord>>I/2) & Mask)".
    // Both masks might have been narrowed by InstCombine.
    if (match(Start,
              m_c_Add(m_And(m_LShr(m_Value(ShiftOp), m_SpecificInt(I / 2)),
                            m_APInt(AndMask1)),
                      m_And(m_Deferred(ShiftOp), m_APInt(AndMask2))))) {
      if (!narrowAddPairMasksOk(Mask, I / 2, ShiftOp, *AndMask1, *AndMask2))
        return false;
    }
    // Matching "(uWord & Mask) + (uWord>>I/2)".
    // The mask might have been narrowed by InstCombine.
    else if (match(Start,
                   m_c_Add(m_LShr(m_Value(ShiftOp), m_SpecificInt(I / 2)),
                           m_And(m_Deferred(ShiftOp), m_APInt(AndMask1))))) {
      if (!isValidNarrowedMask(*AndMask1, Mask, ShiftOp))
        return false;
    } else
      return false;
    Start = ShiftOp;
  }

  // Matching "uWord = (uWord & Mask33) + ((uWord>>2) & Mask33)".
  const APInt *AndMask1 = nullptr, *AndMask2 = nullptr;
  if (!match(Start, m_c_Add(m_And(m_LShr(m_Value(ShiftOp), m_SpecificInt(2)),
                                  m_APInt(AndMask1)),
                            m_And(m_Deferred(ShiftOp), m_APInt(AndMask2)))))
    return false;
  if (!narrowAddPairMasksOk(Mask33, 2, ShiftOp, *AndMask1, *AndMask2))
    return false;

  Start = ShiftOp;
  Value *Root;
  // Matching "uWord = (uWord & Mask55) + ((uWord>>1) & Mask55)".
  AndMask1 = nullptr;
  AndMask2 = nullptr;
  if (!match(Start, m_c_Add(m_And(m_LShr(m_Value(Root), m_SpecificInt(1)),
                                  m_APInt(AndMask1)),
                            m_And(m_Deferred(Root), m_APInt(AndMask2)))))
    return false;
  if (!narrowAddPairMasksOk(Mask55, 1, Root, *AndMask1, *AndMask2))
    return false;

  replaceWithPopCount(I, Root);
  return true;
}

// Try to recognize below function as popcount intrinsic.
// Ref. Hackers Delight
// int popcnt(unsigned x) {
// x = x - ((x >> 1) & 0x55555555);
// x = (x & 0x33333333) + ((x >> 2) & 0x33333333);
// x = (x + (x >> 4)) & 0x0F0F0F0F;
// x = x + (x >> 8);
// x = x + (x >> 16);
// return x & 0x0000003F;
// }

// int popcnt(unsigned x) {
// x = x - ((x >> 1) & 0x55555555);
// x = x - 3*((x >> 2) & 0x33333333);
// x = (x + (x >> 4)) & 0x0F0F0F0F;
// x = x + (x >> 8);
// x = x + (x >> 16);
// return x & 0x0000003F;
// }
static bool tryToRecognizePopCount2n3(Instruction &I) {
  if (I.getOpcode() != Instruction::And)
    return false;

  Type *Ty = I.getType();
  if (!Ty->isIntOrIntVectorTy())
    return false;

  unsigned Len = Ty->getScalarSizeInBits();
  Value *Add1;
  if (Len == 8) {
    // Special case for Len == 8, we only need to match the And at the end of
    // matchPopCountBytes.
    Add1 = &I;
  } else {
    const APInt *MaskRes;
    if (!match(&I, m_And(m_Value(Add1), m_APInt(MaskRes))))
      return false;

    // Since `(trunc (and x, C))` might be canonicalized into `(and (trunc x),
    // C)` we might loose the opportunity to recognize `(trunc (popcount y))`.
    // The following block tries to capture such truncation, update `Len`, and
    // append the truncation at the end of the emitting popcount, if there is
    // any.
    Value *TruncSrc;
    if (match(Add1, m_OneUse(m_Trunc(m_Value(TruncSrc))))) {
      Add1 = TruncSrc;
      Len = Add1->getType()->getScalarSizeInBits();
    }

    if (Len > 64 || Len <= 8 || Len % 8 != 0)
      return false;

    // Len should be a power of 2 for the loop to work correctly
    if (!isPowerOf2_32(Len))
      return false;

    // Number of bits needed to represent Len.
    unsigned NumLenBits = Log2_32(Len) + 1;
    // The "mask" here really only needs to fulfill two conditions:
    // (1) All ones for the lower NumLenBits-bits
    // (2) Zeros from bit 8 and onward.
    // Condition (1) is straightforward. The reason behind condition
    // (2) is that we don't care any 8-bit chunks but the first one
    // in the original divide-and-conquer algorithm.
    if (MaskRes->countTrailingOnes() < NumLenBits ||
        MaskRes->getActiveBits() > 8)
      return false;

    for (unsigned I = Len; I >= 16; I = I / 2) {
      Value *Add2;
      // Matching "x = x + (x >> I/2)" for I-bit.
      if (!match(Add1, m_c_Add(m_LShr(m_Value(Add2), m_SpecificInt(I / 2)),
                               m_Deferred(Add2))))
        return false;
      Add1 = Add2;
    }
  }

  Value *Root = matchPopCountBytes(Add1, Len, I.getDataLayout());
  if (!Root)
    return false;

  replaceWithPopCount(I, Root);
  return true;
}

/// Fold smin(smax(fptosi(x), C1), C2) to llvm.fptosi.sat(x), providing C1 and
/// C2 saturate the value of the fp conversion. The transform is not reversable
/// as the fptosi.sat is more defined than the input - all values produce a
/// valid value for the fptosi.sat, where as some produce poison for original
/// that were out of range of the integer conversion. The reversed pattern may
/// use fmax and fmin instead. As we cannot directly reverse the transform, and
/// it is not always profitable, we make it conditional on the cost being
/// reported as lower by TTI.
static bool tryToFPToSat(Instruction &I, TargetTransformInfo &TTI) {
  // Look for min(max(fptosi, converting to fptosi_sat.
  Value *In;
  const APInt *MinC, *MaxC;
  if (!match(&I, m_SMax(m_OneUse(m_SMin(m_OneUse(m_FPToSI(m_Value(In))),
                                        m_APInt(MinC))),
                        m_APInt(MaxC))) &&
      !match(&I, m_SMin(m_OneUse(m_SMax(m_OneUse(m_FPToSI(m_Value(In))),
                                        m_APInt(MaxC))),
                        m_APInt(MinC))))
    return false;

  // Check that the constants clamp a saturate.
  if (!(*MinC + 1).isPowerOf2() || -*MaxC != *MinC + 1)
    return false;

  Type *IntTy = I.getType();
  Type *FpTy = In->getType();
  Type *SatTy =
      IntegerType::get(IntTy->getContext(), (*MinC + 1).exactLogBase2() + 1);
  if (auto *VecTy = dyn_cast<VectorType>(IntTy))
    SatTy = VectorType::get(SatTy, VecTy->getElementCount());

  // Get the cost of the intrinsic, and check that against the cost of
  // fptosi+smin+smax
  InstructionCost SatCost = TTI.getIntrinsicInstrCost(
      IntrinsicCostAttributes(Intrinsic::fptosi_sat, SatTy, {In}, {FpTy}),
      TTI::TCK_RecipThroughput);
  SatCost += TTI.getCastInstrCost(Instruction::SExt, IntTy, SatTy,
                                  TTI::CastContextHint::None,
                                  TTI::TCK_RecipThroughput);

  InstructionCost MinMaxCost = TTI.getCastInstrCost(
      Instruction::FPToSI, IntTy, FpTy, TTI::CastContextHint::None,
      TTI::TCK_RecipThroughput);
  MinMaxCost += TTI.getIntrinsicInstrCost(
      IntrinsicCostAttributes(Intrinsic::smin, IntTy, {IntTy}),
      TTI::TCK_RecipThroughput);
  MinMaxCost += TTI.getIntrinsicInstrCost(
      IntrinsicCostAttributes(Intrinsic::smax, IntTy, {IntTy}),
      TTI::TCK_RecipThroughput);

  if (SatCost >= MinMaxCost)
    return false;

  IRBuilder<> Builder(&I);
  Value *Sat =
      Builder.CreateIntrinsic(Intrinsic::fptosi_sat, {SatTy, FpTy}, In);
  I.replaceAllUsesWith(Builder.CreateSExt(Sat, IntTy));
  return true;
}

/// Try to replace a mathlib call to sqrt with the LLVM intrinsic. This avoids
/// pessimistic codegen that has to account for setting errno and can enable
/// vectorization.
static bool foldSqrt(CallInst *Call, LibFunc Func, TargetTransformInfo &TTI,
                     TargetLibraryInfo &TLI, AssumptionCache &AC,
                     DominatorTree &DT) {
  // If (1) this is a sqrt libcall, (2) we can assume that NAN is not created
  // (because NNAN or the operand arg must not be less than -0.0) and (2) we
  // would not end up lowering to a libcall anyway (which could change the value
  // of errno), then:
  // (1) errno won't be set.
  // (2) it is safe to convert this to an intrinsic call.
  Type *Ty = Call->getType();
  Value *Arg = Call->getArgOperand(0);
  if (TTI.haveFastSqrt(Ty) &&
      (Call->hasNoNaNs() ||
       cannotBeOrderedLessThanZero(
           Arg, SimplifyQuery(Call->getDataLayout(), &TLI, &DT, &AC, Call)))) {
    IRBuilder<> Builder(Call);
    Value *NewSqrt =
        Builder.CreateIntrinsic(Intrinsic::sqrt, Ty, Arg, Call, "sqrt");
    Call->replaceAllUsesWith(NewSqrt);

    // Explicitly erase the old call because a call with side effects is not
    // trivially dead.
    Call->eraseFromParent();
    return true;
  }

  return false;
}

// Check if this array of constants represents a cttz table.
// Iterate over the elements from \p Table by trying to find/match all
// the numbers from 0 to \p InputBits that should represent cttz results.
static bool isCTTZTable(Constant *Table, const APInt &Mul, const APInt &Shift,
                        const APInt &AndMask, Type *AccessTy,
                        unsigned InputBits, const APInt &GEPIdxFactor,
                        const DataLayout &DL) {
  for (unsigned Idx = 0; Idx < InputBits; Idx++) {
    APInt Index =
        (APInt::getOneBitSet(InputBits, Idx) * Mul).lshr(Shift) & AndMask;
    ConstantInt *C = dyn_cast_or_null<ConstantInt>(
        ConstantFoldLoadFromConst(Table, AccessTy, Index * GEPIdxFactor, DL));
    if (!C || C->getValue() != Idx)
      return false;
  }

  return true;
}

// Try to recognize table-based ctz implementation.
// E.g., an example in C (for more cases please see the llvm/tests):
// int f(unsigned x) {
//    static const char table[32] =
//      {0, 1, 28, 2, 29, 14, 24, 3, 30,
//       22, 20, 15, 25, 17, 4, 8, 31, 27,
//       13, 23, 21, 19, 16, 7, 26, 12, 18, 6, 11, 5, 10, 9};
//    return table[((unsigned)((x & -x) * 0x077CB531U)) >> 27];
// }
// this can be lowered to `cttz` instruction.
// There is also a special case when the element is 0.
//
// The (x & -x) sets the lowest non-zero bit to 1. The multiply is a de-bruijn
// sequence that contains each pattern of bits in it. The shift extracts
// the top bits after the multiply, and that index into the table should
// represent the number of trailing zeros in the original number.
//
// Here are some examples or LLVM IR for a 64-bit target:
//
// CASE 1:
// %sub = sub i32 0, %x
// %and = and i32 %sub, %x
// %mul = mul i32 %and, 125613361
// %shr = lshr i32 %mul, 27
// %idxprom = zext i32 %shr to i64
// %arrayidx = getelementptr inbounds [32 x i8], [32 x i8]* @ctz1.table, i64 0,
//     i64 %idxprom
// %0 = load i8, i8* %arrayidx, align 1, !tbaa !8
//
// CASE 2:
// %sub = sub i32 0, %x
// %and = and i32 %sub, %x
// %mul = mul i32 %and, 72416175
// %shr = lshr i32 %mul, 26
// %idxprom = zext i32 %shr to i64
// %arrayidx = getelementptr inbounds [64 x i16], [64 x i16]* @ctz2.table,
//     i64 0, i64 %idxprom
// %0 = load i16, i16* %arrayidx, align 2, !tbaa !8
//
// CASE 3:
// %sub = sub i32 0, %x
// %and = and i32 %sub, %x
// %mul = mul i32 %and, 81224991
// %shr = lshr i32 %mul, 27
// %idxprom = zext i32 %shr to i64
// %arrayidx = getelementptr inbounds [32 x i32], [32 x i32]* @ctz3.table,
//     i64 0, i64 %idxprom
// %0 = load i32, i32* %arrayidx, align 4, !tbaa !8
//
// CASE 4:
// %sub = sub i64 0, %x
// %and = and i64 %sub, %x
// %mul = mul i64 %and, 283881067100198605
// %shr = lshr i64 %mul, 58
// %arrayidx = getelementptr inbounds [64 x i8], [64 x i8]* @table, i64 0,
//     i64 %shr
// %0 = load i8, i8* %arrayidx, align 1, !tbaa !8
//
// All these can be lowered to @llvm.cttz.i32/64 intrinsics.
//
// This shares its initial match (load from a GEP into a constant table with
// a single variable index) with tryToRecognizeTableBasedLog2() below; see
// tryToRecognizeTableBasedPatterns().
static bool tryToRecognizeTableBasedCttz(LoadInst *LI, Type *AccessType,
                                         GlobalVariable *GVTable, Value *GepIdx,
                                         const APInt &GEPScale,
                                         const DataLayout &DL) {
  Value *X1;
  const APInt *MulConst, *ShiftConst, *AndCst = nullptr;
  // Check that the gep variable index is ((x & -x) * MulConst) >> ShiftConst.
  // This might be extended to the pointer index type, and if the gep index type
  // has been replaced with an i8 then a new And (and different ShiftConst) will
  // be present.
  auto MatchInner = m_LShr(
      m_Mul(m_c_And(m_Neg(m_Value(X1)), m_Deferred(X1)), m_APInt(MulConst)),
      m_APInt(ShiftConst));
  if (!match(GepIdx, m_CastOrSelf(MatchInner)) &&
      !match(GepIdx, m_CastOrSelf(m_And(MatchInner, m_APInt(AndCst)))))
    return false;

  unsigned InputBits = X1->getType()->getScalarSizeInBits();
  if (InputBits != 16 && InputBits != 32 && InputBits != 64 && InputBits != 128)
    return false;

  if (!GEPScale.isIntN(InputBits) ||
      !isCTTZTable(GVTable->getInitializer(), *MulConst, *ShiftConst,
                   AndCst ? *AndCst : APInt::getAllOnes(InputBits), AccessType,
                   InputBits, GEPScale.zextOrTrunc(InputBits), DL))
    return false;

  ConstantInt *ZeroTableElem = cast<ConstantInt>(
      ConstantFoldLoadFromConst(GVTable->getInitializer(), AccessType, DL));
  bool DefinedForZero = ZeroTableElem->equalsInt(InputBits);

  IRBuilder<> B(LI);
  ConstantInt *BoolConst = B.getInt1(!DefinedForZero);
  Type *XType = X1->getType();
  auto Cttz = B.CreateIntrinsic(Intrinsic::cttz, {XType}, {X1, BoolConst});
  Value *Res = B.CreateZExtOrTrunc(Cttz, AccessType);

  if (!DefinedForZero) {
    // If the value in elem 0 isn't the same as InputBits, we still want to
    // produce the value from the table. Emit the select in AccessType with elem
    // 0 unchanged, as the table's element type may be wider than the input
    // type (and directly truncating ZeroTableElem into the input type could
    // incorrectly drop bits).
    auto Cmp = B.CreateICmpEQ(X1, ConstantInt::get(XType, 0));
    Res = B.CreateSelect(Cmp, ZeroTableElem, Res);

    // The true branch of select handles the cttz(0) case, which is rare.
    if (Instruction *SelectI = dyn_cast<Instruction>(Res))
      SelectI->setMetadata(
          LLVMContext::MD_prof,
          MDBuilder(SelectI->getContext()).createUnlikelyBranchWeights());

    // NOTE: If the table[0] is 0, but the cttz(0) is defined by the Target
    // it should be handled as: `cttz(x) & (typeSize - 1)`.
  }

  LI->replaceAllUsesWith(Res);

  return true;
}

// Check if this array of constants represents a log2 table.
// Iterate over the elements from \p Table by trying to find/match all
// the numbers from 0 to \p InputBits that should represent log2 results.
static bool isLog2Table(Constant *Table, const APInt &Mul, const APInt &Shift,
                        Type *AccessTy, unsigned InputBits,
                        const APInt &GEPIdxFactor, const DataLayout &DL) {
  for (unsigned Idx = 0; Idx < InputBits; Idx++) {
    APInt Index = (APInt::getLowBitsSet(InputBits, Idx + 1) * Mul).lshr(Shift);
    ConstantInt *C = dyn_cast_or_null<ConstantInt>(
        ConstantFoldLoadFromConst(Table, AccessTy, Index * GEPIdxFactor, DL));
    if (!C || C->getValue() != Idx)
      return false;
  }

  // Verify that an input of zero will select table index 0.
  APInt ZeroIndex = Mul.lshr(Shift);
  if (!ZeroIndex.isZero())
    return false;

  return true;
}

// Try to recognize table-based log2 implementation.
// E.g., an example in C (for more cases please the llvm/tests):
// int f(unsigned v) {
//    static const char table[32] =
//    {0, 9, 1, 10, 13, 21, 2, 29, 11, 14, 16, 18, 22, 25, 3, 30,
//     8, 12, 20, 28, 15, 17, 24, 7, 19, 27, 23, 6, 26, 5, 4, 31};
//
//    v |= v >> 1; // first round down to one less than a power of 2
//    v |= v >> 2;
//    v |= v >> 4;
//    v |= v >> 8;
//    v |= v >> 16;
//
//    return table[(unsigned)(v * 0x07C4ACDDU) >> 27];
// }
// this can be lowered to `ctlz` instruction.
// There is also a special case when the element is 0.
//
// The >> and |= sequence sets all bits below the most significant set bit. The
// multiply is a de-bruijn sequence that contains each pattern of bits in it.
// The shift extracts the top bits after the multiply, and that index into the
// table should represent the floor log base 2 of the original number.
//
// Here are some examples of LLVM IR for a 64-bit target.
//
// CASE 1:
// %shr = lshr i32 %v, 1
// %or = or i32 %shr, %v
// %shr1 = lshr i32 %or, 2
// %or2 = or i32 %shr1, %or
// %shr3 = lshr i32 %or2, 4
// %or4 = or i32 %shr3, %or2
// %shr5 = lshr i32 %or4, 8
// %or6 = or i32 %shr5, %or4
// %shr7 = lshr i32 %or6, 16
// %or8 = or i32 %shr7, %or6
// %mul = mul i32 %or8, 130329821
// %shr9 = lshr i32 %mul, 27
// %idxprom = zext nneg i32 %shr9 to i64
// %arrayidx = getelementptr inbounds i8, ptr @table, i64 %idxprom
// %0 = load i8, ptr %arrayidx, align 1
//
// CASE 2:
// %shr = lshr i64 %v, 1
// %or = or i64 %shr, %v
// %shr1 = lshr i64 %or, 2
// %or2 = or i64 %shr1, %or
// %shr3 = lshr i64 %or2, 4
// %or4 = or i64 %shr3, %or2
// %shr5 = lshr i64 %or4, 8
// %or6 = or i64 %shr5, %or4
// %shr7 = lshr i64 %or6, 16
// %or8 = or i64 %shr7, %or6
// %shr9 = lshr i64 %or8, 32
// %or10 = or i64 %shr9, %or8
// %mul = mul i64 %or10, 285870213051386505
// %shr11 = lshr i64 %mul, 58
// %arrayidx = getelementptr inbounds i8, ptr @table, i64 %shr11
// %0 = load i8, ptr %arrayidx, align 1
//
// CASE 3:
// A variant where the most-significant set bit of the OR-cascade result is
// isolated via subtraction before the multiply, i.e.
// table[((v - (v >> 1)) * MulConst) >> ShiftConst], analogous to how the
// cttz pattern isolates the least-significant set bit via `x & -x`:
//
// %shr = lshr i64 %v, 1
// %or = or i64 %shr, %v
// ... (rest of the OR-cascade, as above) ...
// %shr11 = lshr i64 %or10, 1
// %sub = sub i64 %or10, %shr11
// %mul = mul i64 %sub, 571347909858961602
// %shr12 = lshr i64 %mul, 58
// %arrayidx = getelementptr inbounds i8, ptr @table, i64 %shr12
// %0 = load i8, ptr %arrayidx, align 1
//
// All these can be lowered to @llvm.ctlz.i32/64 intrinsics and a subtract.
//
// This shares its initial match (load from a GEP into a constant table with
// a single variable index) with tryToRecognizeTableBasedCttz() above; see
// tryToRecognizeTableBasedPatterns().
static bool tryToRecognizeTableBasedLog2(LoadInst *LI, Type *AccessType,
                                         GlobalVariable *GVTable, Value *GepIdx,
                                         const APInt &GEPScale,
                                         const DataLayout &DL,
                                         TargetTransformInfo &TTI) {
  Value *X;
  const APInt *MulConst, *ShiftConst;
  // Check that the gep variable index is (x * MulConst) >> ShiftConst.
  auto MatchInner =
      m_LShr(m_Mul(m_Value(X), m_APInt(MulConst)), m_APInt(ShiftConst));
  if (!match(GepIdx, m_CastOrSelf(MatchInner)))
    return false;

  // The multiplied value may instead be the OR-cascade result with its
  // most-significant set bit isolated first via `v - (v >> 1)`: since every
  // bit below the MSB of an OR-cascade result is 1, this subtraction leaves
  // just the MSB, mirroring how tryToRecognizeTableBasedCttz() isolates the
  // least-significant set bit via `x & -x`.
  bool IsolatedMSB = false;
  Value *V;
  if (match(X, m_Sub(m_Value(V), m_LShr(m_Deferred(V), m_SpecificInt(1))))) {
    IsolatedMSB = true;
    X = V;
  }

  unsigned InputBits = X->getType()->getScalarSizeInBits();
  if (InputBits != 16 && InputBits != 32 && InputBits != 64 && InputBits != 128)
    return false;

  // Verify shift amount.
  // TODO: Allow other shift amounts when we have proper test coverage.
  if (*ShiftConst != InputBits - Log2_32(InputBits))
    return false;

  // Match the sequence of OR operations with right shifts by powers of 2.
  for (unsigned ShiftAmt = InputBits / 2; ShiftAmt != 0; ShiftAmt /= 2) {
    Value *Y;
    if (!match(X, m_c_Or(m_LShr(m_Value(Y), m_SpecificInt(ShiftAmt)),
                         m_Deferred(Y))))
      return false;
    X = Y;
  }

  if (!GEPScale.isIntN(InputBits))
    return false;

  if (IsolatedMSB) {
    // With the MSB isolated, the multiplicand for an input whose MSB is at bit
    // Idx is a single set bit rather than a run of low bits, which is exactly
    // what isCTTZTable() checks for (there is no additional masking here, so
    // pass an all-ones mask).
    if (!isCTTZTable(GVTable->getInitializer(), *MulConst, *ShiftConst,
                     APInt::getAllOnes(InputBits), AccessType, InputBits,
                     GEPScale.zextOrTrunc(InputBits), DL))
      return false;
  } else {
    if (!isLog2Table(GVTable->getInitializer(), *MulConst, *ShiftConst,
                     AccessType, InputBits, GEPScale.zextOrTrunc(InputBits),
                     DL))
      return false;
  }

  ConstantInt *ZeroTableElem = cast<ConstantInt>(
      ConstantFoldLoadFromConst(GVTable->getInitializer(), AccessType, DL));

  // Use InputBits - 1 - ctlz(X) to compute log2(X).
  IRBuilder<> B(LI);
  ConstantInt *BoolConst = B.getTrue();
  Type *XType = X->getType();

  // Check the the backend has an efficient ctlz instruction.
  // FIXME: Teach the backend to emit the original code when ctlz isn't
  // supported like we do for cttz.
  IntrinsicCostAttributes Attrs(
      Intrinsic::ctlz, XType,
      {PoisonValue::get(XType), /*is_zero_poison=*/BoolConst});
  InstructionCost Cost =
      TTI.getIntrinsicInstrCost(Attrs, TargetTransformInfo::TCK_SizeAndLatency);
  if (Cost > TargetTransformInfo::TCC_Basic)
    return false;

  Constant *InputBitsM1 = ConstantInt::get(XType, InputBits - 1);

  Value *Result;
  if (ZeroTableElem->getZExtValue() == InputBits - 1) {
    Value *Ctlz =
        B.CreateIntrinsic(Intrinsic::ctlz, {XType}, {X, B.getFalse()});
    Result = B.CreateAnd(B.CreateNot(Ctlz), InputBitsM1);
  } else {
    Value *Ctlz = B.CreateIntrinsic(Intrinsic::ctlz, {XType}, {X, BoolConst});
    Value *Sub = B.CreateSub(InputBitsM1, Ctlz);

    // The table won't produce a sensible result for 0.
    Value *Cmp = B.CreateICmpEQ(X, ConstantInt::get(XType, 0));
    Value *Select =
        B.CreateSelect(Cmp, B.CreateZExt(ZeroTableElem, XType), Sub);

    // The true branch of select handles the log2(0) case, which is rare.
    if (Instruction *SelectI = dyn_cast<Instruction>(Select))
      SelectI->setMetadata(
          LLVMContext::MD_prof,
          MDBuilder(SelectI->getContext()).createUnlikelyBranchWeights());

    Result = Select;
  }

  Value *ZExtOrTrunc = B.CreateZExtOrTrunc(Result, AccessType);

  LI->replaceAllUsesWith(ZExtOrTrunc);

  return true;
}

// Recognize a load from a "low bits mask" table, tbl[j] == (1 << j) - 1:
//
//   static const uintN_t tbl[K] = { 0, 1, 3, 7, ... };
//   ... x & tbl[i] ...
//
// The loaded value equals the arithmetic mask (1 << i) - 1, so replace the load
// with that expression and drop the table. On X86+BMI2 a surrounding `and` then
// selects to a single `bzhi`; on other targets it is a shift and a decrement.
//
// Shares the load/GEP/offset match with the cttz and log2 recognizers above;
// see tryToRecognizeTableBasedPatterns().
static bool tryToRecognizeTableBasedLowBitsMask(LoadInst *LI, Type *AccessType,
                                                GlobalVariable *GVTable,
                                                Value *GepIdx,
                                                const APInt &GEPScale,
                                                const DataLayout &DL) {
  // The value must be exactly the table element; refuse volatile/atomic loads.
  if (!LI->isSimple())
    return false;

  // The table must be [K x AccessType] of integers; a narrower load would read
  // only part of an element.
  auto *ArrTy = dyn_cast<ArrayType>(GVTable->getValueType());
  if (!ArrTy || ArrTy->getElementType() != AccessType ||
      !AccessType->isIntegerTy())
    return false;

  unsigned EltBits = AccessType->getIntegerBitWidth();
  uint64_t EltBytes = DL.getTypeAllocSize(AccessType).getFixedValue();
  uint64_t K = ArrTy->getNumElements();

  // Only fire when the element type fits in a legal integer, so the emitted
  // shift is a cheap (possibly promoted) instruction rather than a runtime
  // libcall (e.g. __ashlti3 for i128) -- which would be worse than the load
  // and is unavailable in freestanding environments that do not link a
  // runtime library.
  if (!DL.fitsInLegalInteger(EltBits))
    return false;

  // The rewrite is only valid for indices inside the table. The nusw the
  // shared matcher requires keeps the address computation from wrapping but
  // does not keep it inside the object: a nusw-only GEP with i >= K is a
  // well-defined pointer past the table, and a load through it (into
  // whatever follows) is not (1 << i) - 1. With inbounds such a GEP is poison
  // and the load is UB, so every executed load has 0 <= i < K.
  auto *GEP = cast<GetElementPtrInst>(LI->getPointerOperand());
  if (!GEP->isInBounds())
    return false;

  // K <= EltBits then guarantees i < EltBits, so the emitted `1 << i` never
  // shifts by >= bitwidth (no poison). A table with EltBits + 1 entries (whose
  // last element is the all-ones mask, needing a shift by EltBits) is
  // therefore rejected.
  if (K == 0 || K > EltBits)
    return false;

  // The index must step by exactly one element, so the runtime index value is
  // the shift amount; a different scale would load tbl[c * i].
  if (GEPScale != EltBytes)
    return false;

  // Every element must be the low-bits mask for its position.
  for (uint64_t J = 0; J < K; ++J) {
    Constant *Elt = ConstantFoldLoadFromConst(
        GVTable->getInitializer(), AccessType,
        APInt(GEPScale.getBitWidth(), J) * GEPScale, DL);
    auto *CI = dyn_cast_or_null<ConstantInt>(Elt);
    if (!CI || CI->getValue() != APInt::getLowBitsSet(EltBits, J))
      return false;
  }

  // Emit (1 << i) - 1 in the element type and replace the load. A later
  // InstCombine canonicalizes this to ~(-1 << i); X86 lowers both to bzhi.
  IRBuilder<> Builder(LI);
  Value *Idx = Builder.CreateZExtOrTrunc(GepIdx, AccessType);
  Value *Mask =
      Builder.CreateSub(Builder.CreateShl(ConstantInt::get(AccessType, 1), Idx),
                        ConstantInt::get(AccessType, 1));
  LI->replaceAllUsesWith(Mask);
  ++NumTableBasedLowBitsMask;
  return true;
}

// Match a table-based cttz, log2, or low-bits-mask implementation. These
// patterns share a load from a global table that we match first; then we
// try the specific matches.
static bool tryToRecognizeTableBasedPatterns(Instruction &I,
                                             const DataLayout &DL,
                                             TargetTransformInfo &TTI) {
  LoadInst *LI = dyn_cast<LoadInst>(&I);
  if (!LI)
    return false;

  Type *AccessType = LI->getType();
  if (!AccessType->isIntegerTy())
    return false;

  GetElementPtrInst *GEP = dyn_cast<GetElementPtrInst>(LI->getPointerOperand());
  if (!GEP || !GEP->hasNoUnsignedSignedWrap())
    return false;

  GlobalVariable *GVTable = dyn_cast<GlobalVariable>(GEP->getPointerOperand());
  if (!GVTable || !GVTable->isConstant() ||
      !GVTable->hasDefinitiveInitializer())
    return false;

  unsigned BW = DL.getIndexTypeSizeInBits(GEP->getType());
  APInt ModOffset(BW, 0);
  SmallMapVector<Value *, APInt, 4> VarOffsets;
  if (!GEP->collectOffset(DL, BW, VarOffsets, ModOffset) ||
      VarOffsets.size() != 1 || ModOffset != 0)
    return false;
  auto [GepIdx, GEPScale] = VarOffsets.front();

  if (tryToRecognizeTableBasedCttz(LI, AccessType, GVTable, GepIdx, GEPScale,
                                   DL))
    return true;

  if (tryToRecognizeTableBasedLog2(LI, AccessType, GVTable, GepIdx, GEPScale,
                                   DL, TTI))
    return true;

  return tryToRecognizeTableBasedLowBitsMask(LI, AccessType, GVTable, GepIdx,
                                             GEPScale, DL);
}

/// This is used by foldLoadsRecursive() to capture a Root Load node which is
/// of type or(load, load) and recursively build the wide load. Also capture the
/// shift amount, zero extend type and loadSize.
struct LoadOps {
  LoadInst *Root = nullptr;
  LoadInst *RootInsert = nullptr;
  bool FoundRoot = false;
  uint64_t LoadSize = 0;
  uint64_t Shift = 0;
  Type *ZextType;
  AAMDNodes AATags;
};

// Identify and Merge consecutive loads recursively which is of the form
// (ZExt(L1) << shift1) | (ZExt(L2) << shift2) -> ZExt(L3) << shift1
// (ZExt(L1) << shift1) | ZExt(L2) -> ZExt(L3)
static bool foldLoadsRecursive(Value *V, LoadOps &LOps, const DataLayout &DL,
                               AliasAnalysis &AA, bool IsRoot = false) {
  uint64_t ShAmt2;
  Value *X;
  Instruction *L1, *L2;

  // For the root instruction, allow multiple uses since the final result
  // may legitimately be used in multiple places. For intermediate values,
  // require single use to avoid creating duplicate loads.
  if (!IsRoot && !V->hasOneUse())
    return false;

  if (!match(V, m_c_Or(m_Value(X),
                       m_OneUse(m_ShlOrSelf(m_OneUse(m_ZExt(m_Instruction(L2))),
                                            ShAmt2)))))
    return false;

  if (!foldLoadsRecursive(X, LOps, DL, AA, /*IsRoot=*/false) && LOps.FoundRoot)
    // Avoid Partial chain merge.
    return false;

  // Check if the pattern has loads
  LoadInst *LI1 = LOps.Root;
  uint64_t ShAmt1 = LOps.Shift;
  if (LOps.FoundRoot == false &&
      match(X, m_OneUse(
                   m_ShlOrSelf(m_OneUse(m_ZExt(m_Instruction(L1))), ShAmt1)))) {
    LI1 = dyn_cast<LoadInst>(L1);
  }
  LoadInst *LI2 = dyn_cast<LoadInst>(L2);

  // Check if loads are same, atomic, volatile and having same address space.
  if (LI1 == LI2 || !LI1 || !LI2 || !LI1->isSimple() || !LI2->isSimple() ||
      LI1->getPointerAddressSpace() != LI2->getPointerAddressSpace())
    return false;

  // Check if Loads come from same BB.
  if (LI1->getParent() != LI2->getParent())
    return false;

  // Find the data layout
  bool IsBigEndian = DL.isBigEndian();

  // Check if loads are consecutive and same size.
  Value *Load1Ptr = LI1->getPointerOperand();
  APInt Offset1(DL.getIndexTypeSizeInBits(Load1Ptr->getType()), 0);
  Load1Ptr =
      Load1Ptr->stripAndAccumulateConstantOffsets(DL, Offset1,
                                                  /* AllowNonInbounds */ true);

  Value *Load2Ptr = LI2->getPointerOperand();
  APInt Offset2(DL.getIndexTypeSizeInBits(Load2Ptr->getType()), 0);
  Load2Ptr =
      Load2Ptr->stripAndAccumulateConstantOffsets(DL, Offset2,
                                                  /* AllowNonInbounds */ true);

  // Verify if both loads have same base pointers
  uint64_t LoadSize1 = LI1->getType()->getPrimitiveSizeInBits();
  uint64_t LoadSize2 = LI2->getType()->getPrimitiveSizeInBits();
  if (Load1Ptr != Load2Ptr)
    return false;

  // Make sure that there are no padding bits.
  if (!DL.typeSizeEqualsStoreSize(LI1->getType()) ||
      !DL.typeSizeEqualsStoreSize(LI2->getType()))
    return false;

  // Alias Analysis to check for stores b/w the loads.
  LoadInst *Start = LOps.FoundRoot ? LOps.RootInsert : LI1, *End = LI2;
  MemoryLocation Loc;
  if (!Start->comesBefore(End)) {
    std::swap(Start, End);
    // If LOps.RootInsert comes after LI2, since we use LI2 as the new insert
    // point, we should make sure whether the memory region accessed by LOps
    // isn't modified.
    if (LOps.FoundRoot)
      Loc = MemoryLocation(
          LOps.Root->getPointerOperand(),
          LocationSize::precise(DL.getTypeStoreSize(
              IntegerType::get(LI1->getContext(), LOps.LoadSize))),
          LOps.AATags);
    else
      Loc = MemoryLocation::get(End);
  } else
    Loc = MemoryLocation::get(End);
  unsigned NumScanned = 0;
  for (Instruction &Inst :
       make_range(Start->getIterator(), End->getIterator())) {
    if (Inst.mayWriteToMemory() && isModSet(AA.getModRefInfo(&Inst, Loc)))
      return false;

    if (++NumScanned > MaxInstrsToScan)
      return false;
  }

  // Make sure Load with lower Offset is at LI1
  bool Reverse = false;
  if (Offset2.slt(Offset1)) {
    std::swap(LI1, LI2);
    std::swap(ShAmt1, ShAmt2);
    std::swap(Offset1, Offset2);
    std::swap(Load1Ptr, Load2Ptr);
    std::swap(LoadSize1, LoadSize2);
    Reverse = true;
  }

  // Big endian swap the shifts
  if (IsBigEndian)
    std::swap(ShAmt1, ShAmt2);

  // First load is always LI1. This is where we put the new load.
  // Use the merged load size available from LI1 for forward loads.
  if (LOps.FoundRoot) {
    if (!Reverse)
      LoadSize1 = LOps.LoadSize;
    else
      LoadSize2 = LOps.LoadSize;
  }

  // Verify if shift amount and load index aligns and verifies that loads
  // are consecutive.
  uint64_t ShiftDiff = IsBigEndian ? LoadSize2 : LoadSize1;
  uint64_t PrevSize =
      DL.getTypeStoreSize(IntegerType::get(LI1->getContext(), LoadSize1));
  if ((ShAmt2 - ShAmt1) != ShiftDiff || (Offset2 - Offset1) != PrevSize)
    return false;

  // Reject if the combined size of the loads exceeds the target type size.
  // This avoids attempting to emit an invalid ZExt (from wider to narrower
  // type) when out-of-bounds shifts lead to matching too many loads.
  if (LoadSize1 + LoadSize2 > X->getType()->getScalarSizeInBits())
    return false;

  // Update LOps
  AAMDNodes AATags1 = LOps.AATags;
  AAMDNodes AATags2 = LI2->getAAMetadata();
  if (LOps.FoundRoot == false) {
    LOps.FoundRoot = true;
    AATags1 = LI1->getAAMetadata();
  }
  LOps.LoadSize = LoadSize1 + LoadSize2;
  LOps.RootInsert = Start;

  // Concatenate the AATags of the Merged Loads.
  LOps.AATags = AATags1.concat(AATags2);

  LOps.Root = LI1;
  LOps.Shift = ShAmt1;
  LOps.ZextType = X->getType();
  return true;
}

// For a given BB instruction, evaluate all loads in the chain that form a
// pattern which suggests that the loads can be combined. The one and only use
// of the loads is to form a wider load.
static bool foldConsecutiveLoads(Instruction &I, const DataLayout &DL,
                                 TargetTransformInfo &TTI, AliasAnalysis &AA,
                                 const DominatorTree &DT) {
  // Only consider load chains of scalar values.
  if (isa<VectorType>(I.getType()))
    return false;

  LoadOps LOps;
  if (!foldLoadsRecursive(&I, LOps, DL, AA, /*IsRoot=*/true) || !LOps.FoundRoot)
    return false;

  IRBuilder<> Builder(&I);
  LoadInst *NewLoad = nullptr, *LI1 = LOps.Root;

  // Allow a power of 2 number of bytes that fit in a legal integer type.
  bool Allowed = LOps.LoadSize >= 16 && isPowerOf2_64(LOps.LoadSize) &&
                 DL.fitsInLegalInteger(LOps.LoadSize);
  if (!Allowed)
    return false;

  unsigned AS = LI1->getPointerAddressSpace();
  unsigned Fast = 0;
  Allowed = TTI.allowsMisalignedMemoryAccesses(I.getContext(), LOps.LoadSize,
                                               AS, LI1->getAlign(), &Fast);
  if (!Allowed || !Fast)
    return false;

  // Get the Index and Ptr for the new GEP.
  Value *Load1Ptr = LI1->getPointerOperand();
  Builder.SetInsertPoint(LOps.RootInsert);
  if (!DT.dominates(Load1Ptr, LOps.RootInsert)) {
    APInt Offset1(DL.getIndexTypeSizeInBits(Load1Ptr->getType()), 0);
    Load1Ptr = Load1Ptr->stripAndAccumulateConstantOffsets(
        DL, Offset1, /* AllowNonInbounds */ true);
    Load1Ptr = Builder.CreatePtrAdd(Load1Ptr, Builder.getInt(Offset1));
  }
  // Generate wider load.
  IntegerType *WiderType = IntegerType::get(I.getContext(), LOps.LoadSize);
  NewLoad = Builder.CreateAlignedLoad(WiderType, Load1Ptr, LI1->getAlign(),
                                      LI1->isVolatile(), "");
  NewLoad->takeName(LI1);
  // Set the New Load AATags Metadata.
  if (LOps.AATags)
    NewLoad->setAAMetadata(LOps.AATags);

  Value *NewOp = NewLoad;
  // Zero extend if needed.
  NewOp = Builder.CreateZExt(NewOp, LOps.ZextType);

  // Check if shift needed. We need to shift with the amount of load1
  // shift if not zero.
  if (LOps.Shift)
    NewOp = Builder.CreateShl(NewOp, LOps.Shift);
  I.replaceAllUsesWith(NewOp);

  return true;
}

/// ValWidth bits starting at ValOffset of Val stored at PtrBase+PtrOffset.
struct PartStore {
  Value *PtrBase;
  APInt PtrOffset;
  Value *Val;
  uint64_t ValOffset;
  uint64_t ValWidth;
  StoreInst *Store;

  bool isCompatibleWith(const PartStore &Other) const {
    // Offset stripping looks through addrspacecasts, so an equal PtrBase does
    // not imply an equal address space, and thus not an equal PtrOffset width.
    return PtrBase == Other.PtrBase && Val == Other.Val &&
           Store->getPointerAddressSpace() ==
               Other.Store->getPointerAddressSpace();
  }

  bool operator<(const PartStore &Other) const {
    return PtrOffset.slt(Other.PtrOffset);
  }
};

static std::optional<PartStore> matchPartStore(Instruction &I,
                                               const DataLayout &DL) {
  auto *Store = dyn_cast<StoreInst>(&I);
  if (!Store || !Store->isSimple())
    return std::nullopt;

  Value *StoredVal = Store->getValueOperand();
  Type *StoredTy = StoredVal->getType();
  if (!StoredTy->isIntegerTy() || !DL.typeSizeEqualsStoreSize(StoredTy))
    return std::nullopt;

  uint64_t ValWidth = StoredTy->getPrimitiveSizeInBits();
  uint64_t ValOffset;
  Value *Val;
  if (!match(StoredVal, m_Trunc(m_LShrOrSelf(m_Value(Val), ValOffset))))
    return std::nullopt;

  Value *Ptr = Store->getPointerOperand();
  APInt PtrOffset(DL.getIndexTypeSizeInBits(Ptr->getType()), 0);
  Value *PtrBase = Ptr->stripAndAccumulateConstantOffsets(
      DL, PtrOffset, /*AllowNonInbounds=*/true);
  return {{PtrBase, PtrOffset, Val, ValOffset, ValWidth, Store}};
}

static bool mergeConsecutivePartStores(ArrayRef<PartStore> Parts,
                                       unsigned Width, const DataLayout &DL,
                                       TargetTransformInfo &TTI) {
  if (Parts.size() < 2)
    return false;

  // Check whether combining the stores is profitable.
  // FIXME: We could generate smaller stores if we can't produce a large one.
  const PartStore &First = Parts.front();
  LLVMContext &Ctx = First.Store->getContext();
  unsigned Fast = 0;
  bool Allowed =
      Width >= 16 && isPowerOf2_64(Width) && DL.fitsInLegalInteger(Width);
  if (!Allowed ||
      !TTI.allowsMisalignedMemoryAccesses(Ctx, Width,
                                          First.Store->getPointerAddressSpace(),
                                          First.Store->getAlign(), &Fast) ||
      !Fast)
    return false;

  // Generate the combined store.
  IRBuilder<> Builder(First.Store);
  Type *NewTy = Type::getIntNTy(Ctx, Width);
  Value *Val = First.Val;
  if (First.ValOffset != 0)
    Val = Builder.CreateLShr(Val, First.ValOffset);
  Val = Builder.CreateZExtOrTrunc(Val, NewTy);
  StoreInst *Store = Builder.CreateAlignedStore(
      Val, First.Store->getPointerOperand(), First.Store->getAlign());

  // Merge various metadata onto the new store.
  AAMDNodes AATags = First.Store->getAAMetadata();
  SmallVector<Instruction *> Stores = {First.Store};
  Stores.reserve(Parts.size());
  SmallVector<DebugLoc> DbgLocs = {First.Store->getDebugLoc()};
  DbgLocs.reserve(Parts.size());
  for (const PartStore &Part : drop_begin(Parts)) {
    AATags = AATags.concat(Part.Store->getAAMetadata());
    Stores.push_back(Part.Store);
    DbgLocs.push_back(Part.Store->getDebugLoc());
  }
  Store->setAAMetadata(AATags);
  Store->mergeDIAssignID(Stores);
  Store->setDebugLoc(DebugLoc::getMergedLocations(DbgLocs));

  // Remove the old stores.
  for (const PartStore &Part : Parts)
    Part.Store->eraseFromParent();

  return true;
}

static bool mergePartStores(SmallVectorImpl<PartStore> &Parts,
                            const DataLayout &DL, TargetTransformInfo &TTI) {
  if (Parts.size() < 2)
    return false;

  // We now have multiple parts of the same value stored to the same pointer.
  // Sort the parts by pointer offset, and make sure they are consistent with
  // the value offsets. Also check that the value is fully covered without
  // overlaps.
  bool Changed = false;
  llvm::sort(Parts);
  int64_t LastEndOffsetFromFirst = 0;
  const PartStore *First = &Parts[0];
  for (const PartStore &Part : Parts) {
    APInt PtrOffsetFromFirst = Part.PtrOffset - First->PtrOffset;
    int64_t ValOffsetFromFirst = Part.ValOffset - First->ValOffset;
    if (PtrOffsetFromFirst * 8 != ValOffsetFromFirst ||
        LastEndOffsetFromFirst != ValOffsetFromFirst) {
      Changed |= mergeConsecutivePartStores(ArrayRef(First, &Part),
                                            LastEndOffsetFromFirst, DL, TTI);
      First = &Part;
      LastEndOffsetFromFirst = Part.ValWidth;
      continue;
    }

    LastEndOffsetFromFirst = ValOffsetFromFirst + Part.ValWidth;
  }

  Changed |= mergeConsecutivePartStores(ArrayRef(First, Parts.end()),
                                        LastEndOffsetFromFirst, DL, TTI);
  return Changed;
}

static bool foldConsecutiveStores(BasicBlock &BB, const DataLayout &DL,
                                  TargetTransformInfo &TTI, AliasAnalysis &AA) {
  // FIXME: Add big endian support.
  if (DL.isBigEndian())
    return false;

  BatchAAResults BatchAA(AA);
  SmallVector<PartStore, 8> Parts;
  bool MadeChange = false;
  for (Instruction &I : make_early_inc_range(BB)) {
    if (std::optional<PartStore> Part = matchPartStore(I, DL)) {
      if (Parts.empty() || Part->isCompatibleWith(Parts[0])) {
        Parts.push_back(std::move(*Part));
        continue;
      }

      MadeChange |= mergePartStores(Parts, DL, TTI);
      Parts.clear();
      Parts.push_back(std::move(*Part));
      continue;
    }

    if (Parts.empty())
      continue;

    if (I.mayThrow() ||
        (I.mayReadOrWriteMemory() &&
         isModOrRefSet(BatchAA.getModRefInfo(
             &I, MemoryLocation::getBeforeOrAfter(Parts[0].PtrBase))))) {
      MadeChange |= mergePartStores(Parts, DL, TTI);
      Parts.clear();
      continue;
    }
  }

  MadeChange |= mergePartStores(Parts, DL, TTI);
  return MadeChange;
}

/// Combine away instructions providing they are still equivalent when compared
/// against 0. i.e do they have any bits set.
static Value *optimizeShiftInOrChain(Value *V, IRBuilder<> &Builder) {
  auto *I = dyn_cast<Instruction>(V);
  if (!I || I->getOpcode() != Instruction::Or || !I->hasOneUse())
    return nullptr;

  Value *A;

  // Look deeper into the chain of or's, combining away shl (so long as they are
  // nuw or nsw).
  Value *Op0 = I->getOperand(0);
  if (match(Op0, m_CombineOr(m_NSWShl(m_Value(A), m_Value()),
                             m_NUWShl(m_Value(A), m_Value()))))
    Op0 = A;
  else if (auto *NOp = optimizeShiftInOrChain(Op0, Builder))
    Op0 = NOp;

  Value *Op1 = I->getOperand(1);
  if (match(Op1, m_CombineOr(m_NSWShl(m_Value(A), m_Value()),
                             m_NUWShl(m_Value(A), m_Value()))))
    Op1 = A;
  else if (auto *NOp = optimizeShiftInOrChain(Op1, Builder))
    Op1 = NOp;

  if (Op0 != I->getOperand(0) || Op1 != I->getOperand(1))
    return Builder.CreateOr(Op0, Op1);
  return nullptr;
}

static bool foldICmpOrChain(Instruction &I, const DataLayout &DL,
                            TargetTransformInfo &TTI, AliasAnalysis &AA,
                            const DominatorTree &DT) {
  CmpPredicate Pred;
  Value *Op0;
  if (!match(&I, m_ICmp(Pred, m_Value(Op0), m_Zero())) ||
      !ICmpInst::isEquality(Pred))
    return false;

  // If the chain or or's matches a load, combine to that before attempting to
  // remove shifts.
  if (auto OpI = dyn_cast<Instruction>(Op0))
    if (OpI->getOpcode() == Instruction::Or)
      if (foldConsecutiveLoads(*OpI, DL, TTI, AA, DT))
        return true;

  IRBuilder<> Builder(&I);
  // icmp eq/ne or(shl(a), b), 0 -> icmp eq/ne or(a, b), 0
  if (auto *Res = optimizeShiftInOrChain(Op0, Builder)) {
    I.replaceAllUsesWith(Builder.CreateICmp(Pred, Res, I.getOperand(1)));
    return true;
  }

  return false;
}

// Calculate GEP Stride and accumulated const ModOffset. Return Stride and
// ModOffset
static std::pair<APInt, APInt>
getStrideAndModOffsetOfGEP(Value *PtrOp, const DataLayout &DL) {
  unsigned BW = DL.getIndexTypeSizeInBits(PtrOp->getType());
  std::optional<APInt> Stride;
  APInt ModOffset(BW, 0);
  // Return a minimum gep stride, greatest common divisor of consective gep
  // index scales(c.f. Bézout's identity).
  while (auto *GEP = dyn_cast<GEPOperator>(PtrOp)) {
    SmallMapVector<Value *, APInt, 4> VarOffsets;
    if (!GEP->collectOffset(DL, BW, VarOffsets, ModOffset))
      break;

    for (auto [V, Scale] : VarOffsets) {
      // Only keep a power of two factor for non-inbounds
      if (!GEP->hasNoUnsignedSignedWrap())
        Scale = APInt::getOneBitSet(Scale.getBitWidth(), Scale.countr_zero());

      if (!Stride)
        Stride = Scale;
      else
        Stride = APIntOps::GreatestCommonDivisor(*Stride, Scale);
    }

    PtrOp = GEP->getPointerOperand();
  }

  // Check whether pointer arrives back at Global Variable via at least one GEP.
  // Even if it doesn't, we can check by alignment.
  if (!isa<GlobalVariable>(PtrOp) || !Stride)
    return {APInt(BW, 1), APInt(BW, 0)};

  // In consideration of signed GEP indices, non-negligible offset become
  // remainder of division by minimum GEP stride.
  ModOffset = ModOffset.srem(*Stride);
  if (ModOffset.isNegative())
    ModOffset += *Stride;

  return {*Stride, ModOffset};
}

/// If C is a constant patterned array and all valid loaded results for given
/// alignment are same to a constant, return that constant.
static bool foldPatternedLoads(Instruction &I, const DataLayout &DL) {
  auto *LI = dyn_cast<LoadInst>(&I);
  if (!LI || LI->isVolatile())
    return false;

  // We can only fold the load if it is from a constant global with definitive
  // initializer. Skip expensive logic if this is not the case.
  auto *PtrOp = LI->getPointerOperand();
  auto *GV = dyn_cast<GlobalVariable>(getUnderlyingObject(PtrOp));
  if (!GV || !GV->isConstant() || !GV->hasDefinitiveInitializer())
    return false;

  // Bail for large initializers in excess of 4K to avoid too many scans.
  Constant *C = GV->getInitializer();
  uint64_t GVSize = DL.getTypeAllocSize(C->getType());
  if (!GVSize || 4096 < GVSize)
    return false;

  Type *LoadTy = LI->getType();
  unsigned BW = DL.getIndexTypeSizeInBits(PtrOp->getType());
  auto [Stride, ConstOffset] = getStrideAndModOffsetOfGEP(PtrOp, DL);

  // Any possible offset could be multiple of GEP stride. And any valid
  // offset is multiple of load alignment, so checking only multiples of bigger
  // one is sufficient to say results' equality.
  if (auto LA = LI->getAlign();
      LA <= GV->getAlign().valueOrOne() && Stride.getZExtValue() < LA.value()) {
    ConstOffset = APInt(BW, 0);
    Stride = APInt(BW, LA.value());
  }

  Constant *Ca = ConstantFoldLoadFromConst(C, LoadTy, ConstOffset, DL);
  if (!Ca)
    return false;

  unsigned E = GVSize - DL.getTypeStoreSize(LoadTy);
  for (; ConstOffset.getZExtValue() <= E; ConstOffset += Stride)
    if (Ca != ConstantFoldLoadFromConst(C, LoadTy, ConstOffset, DL))
      return false;

  I.replaceAllUsesWith(Ca);

  return true;
}

namespace {
class StrNCmpInliner {
public:
  StrNCmpInliner(CallInst *CI, LibFunc Func, DomTreeUpdater *DTU,
                 const DataLayout &DL)
      : CI(CI), Func(Func), DTU(DTU), DL(DL) {}

  bool optimizeStrNCmp();

private:
  void inlineCompare(Value *LHS, StringRef RHS, uint64_t N, bool Swapped);

  CallInst *CI;
  LibFunc Func;
  DomTreeUpdater *DTU;
  const DataLayout &DL;
};

} // namespace

/// First we normalize calls to strncmp/strcmp to the form of
/// compare(s1, s2, N), which means comparing first N bytes of s1 and s2
/// (without considering '\0').
///
/// Examples:
///
/// \code
///   strncmp(s, "a", 3) -> compare(s, "a", 2)
///   strncmp(s, "abc", 3) -> compare(s, "abc", 3)
///   strncmp(s, "a\0b", 3) -> compare(s, "a\0b", 2)
///   strcmp(s, "a") -> compare(s, "a", 2)
///
///   char s2[] = {'a'}
///   strncmp(s, s2, 3) -> compare(s, s2, 3)
///
///   char s2[] = {'a', 'b', 'c', 'd'}
///   strncmp(s, s2, 3) -> compare(s, s2, 3)
/// \endcode
///
/// We only handle cases where N and exactly one of s1 and s2 are constant.
/// Cases that s1 and s2 are both constant are already handled by the
/// instcombine pass.
///
/// We do not handle cases where N > StrNCmpInlineThreshold.
///
/// We also do not handles cases where N < 2, which are already
/// handled by the instcombine pass.
///
bool StrNCmpInliner::optimizeStrNCmp() {
  if (StrNCmpInlineThreshold < 2)
    return false;

  if (!isOnlyUsedInZeroComparison(CI))
    return false;

  Value *Str1P = CI->getArgOperand(0);
  Value *Str2P = CI->getArgOperand(1);
  // Should be handled elsewhere.
  if (Str1P == Str2P)
    return false;

  StringRef Str1, Str2;
  bool HasStr1 = getConstantStringInfo(Str1P, Str1, /*TrimAtNul=*/false);
  bool HasStr2 = getConstantStringInfo(Str2P, Str2, /*TrimAtNul=*/false);
  if (HasStr1 == HasStr2)
    return false;

  // Note that '\0' and characters after it are not trimmed.
  StringRef Str = HasStr1 ? Str1 : Str2;
  Value *StrP = HasStr1 ? Str2P : Str1P;

  size_t Idx = Str.find('\0');
  uint64_t N = Idx == StringRef::npos ? UINT64_MAX : Idx + 1;
  if (Func == LibFunc_strncmp) {
    if (auto *ConstInt = dyn_cast<ConstantInt>(CI->getArgOperand(2)))
      N = std::min(N, ConstInt->getZExtValue());
    else
      return false;
  }
  // Now N means how many bytes we need to compare at most.
  if (N > Str.size() || N < 2 || N > StrNCmpInlineThreshold)
    return false;

  // Cases where StrP has two or more dereferenceable bytes might be better
  // optimized elsewhere.
  bool CanBeNull = false;
  if (StrP->getPointerDereferenceableBytes(DL, CanBeNull,
                                           /*CanBeFreed=*/nullptr) > 1)
    return false;
  inlineCompare(StrP, Str, N, HasStr1);
  return true;
}

/// Convert
///
/// \code
///   ret = compare(s1, s2, N)
/// \endcode
///
/// into
///
/// \code
///   ret = (int)s1[0] - (int)s2[0]
///   if (ret != 0)
///     goto NE
///   ...
///   ret = (int)s1[N-2] - (int)s2[N-2]
///   if (ret != 0)
///     goto NE
///   ret = (int)s1[N-1] - (int)s2[N-1]
///   NE:
/// \endcode
///
/// CFG before and after the transformation:
///
/// (before)
/// BBCI
///
/// (after)
/// BBCI -> BBSubs[0] (sub,icmp) --NE-> BBNE -> BBTail
///                 |                    ^
///                 E                    |
///                 |                    |
///        BBSubs[1] (sub,icmp) --NE-----+
///                ...                   |
///        BBSubs[N-1]    (sub) ---------+
///
void StrNCmpInliner::inlineCompare(Value *LHS, StringRef RHS, uint64_t N,
                                   bool Swapped) {
  auto &Ctx = CI->getContext();
  IRBuilder<> B(Ctx);
  // We want these instructions to be recognized as inlined instructions for the
  // compare call, but we don't have a source location for the definition of
  // that function, since we're generating that code now. Because the generated
  // code is a viable point for a memory access error, we make the pragmatic
  // choice here to directly use CI's location so that we have useful
  // attribution for the generated code.
  B.SetCurrentDebugLocation(CI->getDebugLoc());

  BasicBlock *BBCI = CI->getParent();
  BasicBlock *BBTail =
      SplitBlock(BBCI, CI, DTU, nullptr, nullptr, BBCI->getName() + ".tail");

  SmallVector<BasicBlock *> BBSubs;
  for (uint64_t I = 0; I < N; ++I)
    BBSubs.push_back(
        BasicBlock::Create(Ctx, "sub_" + Twine(I), BBCI->getParent(), BBTail));
  BasicBlock *BBNE = BasicBlock::Create(Ctx, "ne", BBCI->getParent(), BBTail);

  cast<UncondBrInst>(BBCI->getTerminator())->setSuccessor(BBSubs[0]);

  B.SetInsertPoint(BBNE);
  PHINode *Phi = B.CreatePHI(CI->getType(), N);
  B.CreateBr(BBTail);

  Value *Base = LHS;
  for (uint64_t i = 0; i < N; ++i) {
    B.SetInsertPoint(BBSubs[i]);
    Value *VL =
        B.CreateZExt(B.CreateLoad(B.getInt8Ty(),
                                  B.CreateInBoundsPtrAdd(Base, B.getInt64(i))),
                     CI->getType());
    Value *VR =
        ConstantInt::get(CI->getType(), static_cast<unsigned char>(RHS[i]));
    Value *Sub = Swapped ? B.CreateSub(VR, VL) : B.CreateSub(VL, VR);
    if (i < N - 1) {
      CondBrInst *CondBrInst = B.CreateCondBr(
          B.CreateICmpNE(Sub, ConstantInt::get(CI->getType(), 0)), BBNE,
          BBSubs[i + 1]);

      Function *F = CI->getFunction();
      assert(F && "Instruction does not belong to a function!");
      std::optional<uint64_t> EC = F->getEntryCount();
      if (EC && *EC > 0)
        setExplicitlyUnknownBranchWeights(*CondBrInst, DEBUG_TYPE);
    } else {
      B.CreateBr(BBNE);
    }

    Phi->addIncoming(Sub, BBSubs[i]);
  }

  CI->replaceAllUsesWith(Phi);
  CI->eraseFromParent();

  if (DTU) {
    SmallVector<DominatorTree::UpdateType, 8> Updates;
    Updates.push_back({DominatorTree::Insert, BBCI, BBSubs[0]});
    for (uint64_t i = 0; i < N; ++i) {
      if (i < N - 1)
        Updates.push_back({DominatorTree::Insert, BBSubs[i], BBSubs[i + 1]});
      Updates.push_back({DominatorTree::Insert, BBSubs[i], BBNE});
    }
    Updates.push_back({DominatorTree::Insert, BBNE, BBTail});
    Updates.push_back({DominatorTree::Delete, BBCI, BBTail});
    DTU->applyUpdates(Updates);
  }
}

/// Convert memchr with a small constant string into a switch
static bool foldMemChr(CallInst *Call, DomTreeUpdater *DTU,
                       const DataLayout &DL) {
  if (isa<Constant>(Call->getArgOperand(1)))
    return false;

  StringRef Str;
  Value *Base = Call->getArgOperand(0);
  if (!getConstantStringInfo(Base, Str, /*TrimAtNul=*/false))
    return false;

  uint64_t N = Str.size();
  if (auto *ConstInt = dyn_cast<ConstantInt>(Call->getArgOperand(2))) {
    uint64_t Val = ConstInt->getZExtValue();
    // Ignore the case that n is larger than the size of string.
    if (Val > N)
      return false;
    N = Val;
  } else
    return false;

  if (N > MemChrInlineThreshold)
    return false;

  BasicBlock *BB = Call->getParent();
  BasicBlock *BBNext = SplitBlock(BB, Call, DTU);
  IRBuilder<> IRB(BB);
  IRB.SetCurrentDebugLocation(Call->getDebugLoc());
  IntegerType *ByteTy = IRB.getInt8Ty();
  BB->getTerminator()->eraseFromParent();
  SwitchInst *SI = IRB.CreateSwitch(
      IRB.CreateTrunc(Call->getArgOperand(1), ByteTy), BBNext, N);
  // We can't know the precise weights here, as they would depend on the value
  // distribution of Call->getArgOperand(1). So we just mark it as "unknown".
  setExplicitlyUnknownBranchWeightsIfProfiled(*SI, DEBUG_TYPE);
  Type *IndexTy = DL.getIndexType(Call->getType());
  SmallVector<DominatorTree::UpdateType, 8> Updates;

  BasicBlock *BBSuccess = BasicBlock::Create(
      Call->getContext(), "memchr.success", BB->getParent(), BBNext);
  IRB.SetInsertPoint(BBSuccess);
  PHINode *IndexPHI = IRB.CreatePHI(IndexTy, N, "memchr.idx");
  Value *FirstOccursLocation = IRB.CreateInBoundsPtrAdd(Base, IndexPHI);
  IRB.CreateBr(BBNext);
  if (DTU)
    Updates.push_back({DominatorTree::Insert, BBSuccess, BBNext});

  SmallPtrSet<ConstantInt *, 4> Cases;
  for (uint64_t I = 0; I < N; ++I) {
    ConstantInt *CaseVal =
        ConstantInt::get(ByteTy, static_cast<unsigned char>(Str[I]));
    if (!Cases.insert(CaseVal).second)
      continue;

    BasicBlock *BBCase = BasicBlock::Create(Call->getContext(), "memchr.case",
                                            BB->getParent(), BBSuccess);
    SI->addCase(CaseVal, BBCase);
    IRB.SetInsertPoint(BBCase);
    IndexPHI->addIncoming(ConstantInt::get(IndexTy, I), BBCase);
    IRB.CreateBr(BBSuccess);
    if (DTU) {
      Updates.push_back({DominatorTree::Insert, BB, BBCase});
      Updates.push_back({DominatorTree::Insert, BBCase, BBSuccess});
    }
  }

  PHINode *PHI =
      PHINode::Create(Call->getType(), 2, Call->getName(), BBNext->begin());
  PHI->addIncoming(Constant::getNullValue(Call->getType()), BB);
  PHI->addIncoming(FirstOccursLocation, BBSuccess);

  Call->replaceAllUsesWith(PHI);
  Call->eraseFromParent();

  if (DTU)
    DTU->applyUpdates(Updates);

  return true;
}

static bool foldLibCalls(Instruction &I, TargetTransformInfo &TTI,
                         TargetLibraryInfo &TLI, AssumptionCache &AC,
                         DominatorTree &DT, const DataLayout &DL,
                         bool &MadeCFGChange) {

  auto *CI = dyn_cast<CallInst>(&I);
  if (!CI || CI->isNoBuiltin())
    return false;

  Function *CalledFunc = CI->getCalledFunction();
  if (!CalledFunc)
    return false;

  LibFunc LF = TLI.getLibFunc(*CalledFunc);
  if (!isLibFuncEmittable(CI->getModule(), &TLI, LF))
    return false;

  DomTreeUpdater DTU(&DT, DomTreeUpdater::UpdateStrategy::Lazy);

  switch (LF) {
  case LibFunc_sqrt:
  case LibFunc_sqrtf:
  case LibFunc_sqrtl:
    return foldSqrt(CI, LF, TTI, TLI, AC, DT);
  case LibFunc_strcmp:
  case LibFunc_strncmp:
    if (StrNCmpInliner(CI, LF, &DTU, DL).optimizeStrNCmp()) {
      MadeCFGChange = true;
      return true;
    }
    break;
  case LibFunc_memchr:
    if (foldMemChr(CI, &DTU, DL)) {
      MadeCFGChange = true;
      return true;
    }
    break;
  default:;
  }
  return false;
}

/// Match high part of long multiplication.
///
/// Considering a multiply made up of high and low parts, we can split the
/// multiply into:
///  x * y == (xh*T + xl) * (yh*T + yl)
/// where xh == x>>32 and xl == x & 0xffffffff. T = 2^32.
/// This expands to
///  xh*yh*T*T + xh*yl*T + xl*yh*T + xl*yl
/// which can be drawn as
/// [  xh*yh  ]
///      [  xh*yl  ]
///      [  xl*yh  ]
///           [  xl*yl  ]
/// We are looking for the "high" half, which is xh*yh + xh*yl>>32 + xl*yh>>32 +
/// some carrys. The carry makes this difficult and there are multiple ways of
/// representing it. The ones we attempt to support here are:
///  Carry:  xh*yh + carry + lowsum
///          carry = lowsum < xh*yl ? 0x1000000 : 0
///          lowsum = xh*yl + xl*yh + (xl*yl>>32)
///  Ladder: xh*yh + c2>>32 + c3>>32
///          c2 = xh*yl + (xl*yl>>32); c3 = c2&0xffffffff + xl*yh
///       or c2 = (xl*yh&0xffffffff) + xh*yl + (xl*yl>>32); c3 = xl*yh
///  Carry4: xh*yh + carry + crosssum>>32 + (xl*yl + crosssum&0xffffffff) >> 32
///          crosssum = xh*yl + xl*yh
///          carry = crosssum < xh*yl ? 0x1000000 : 0
///  Ladder4: xh*yh + (xl*yh)>>32 + (xh*yl)>>32 + low>>32;
///          low = (xl*yl)>>32 + (xl*yh)&0xffffffff + (xh*yl)&0xffffffff
///
/// They all start by matching xh*yh + 2 or 3 other operands. The bottom of the
/// tree is xh*yh, xh*yl, xl*yh and xl*yl.
static bool foldMulHigh(Instruction &I) {
  Type *Ty = I.getType();
  if (!Ty->isIntOrIntVectorTy())
    return false;

  unsigned BitWidth = Ty->getScalarSizeInBits();
  APInt LowMask = APInt::getLowBitsSet(BitWidth, BitWidth / 2);
  if (BitWidth % 2 != 0)
    return false;

  auto CreateMulHigh = [&](Value *X, Value *Y) {
    IRBuilder<> Builder(&I);
    Type *NTy = Ty->getWithNewBitWidth(BitWidth * 2);
    Value *XExt = Builder.CreateZExt(X, NTy);
    Value *YExt = Builder.CreateZExt(Y, NTy);
    Value *Mul = Builder.CreateMul(XExt, YExt, "", /*HasNUW=*/true);
    Value *High = Builder.CreateLShr(Mul, BitWidth);
    Value *Res = Builder.CreateTrunc(High, Ty, "", /*HasNUW=*/true);
    Res->takeName(&I);
    I.replaceAllUsesWith(Res);
    LLVM_DEBUG(dbgs() << "Created long multiply from parts of " << *X << " and "
                      << *Y << "\n");
    return true;
  };

  // Common check routines for X_lo*Y_lo and X_hi*Y_lo
  auto CheckLoLo = [&](Value *XlYl, Value *X, Value *Y) {
    return match(XlYl, m_c_Mul(m_And(m_Specific(X), m_SpecificInt(LowMask)),
                               m_And(m_Specific(Y), m_SpecificInt(LowMask))));
  };
  auto CheckHiLo = [&](Value *XhYl, Value *X, Value *Y) {
    return match(XhYl,
                 m_c_Mul(m_LShr(m_Specific(X), m_SpecificInt(BitWidth / 2)),
                         m_And(m_Specific(Y), m_SpecificInt(LowMask))));
  };

  auto FoldMulHighCarry = [&](Value *X, Value *Y, Instruction *Carry,
                              Instruction *B) {
    // Looking for LowSum >> 32 and carry (select)
    if (Carry->getOpcode() != Instruction::Select)
      std::swap(Carry, B);

    // Carry = LowSum < XhYl ? 0x100000000 : 0
    Value *LowSum, *XhYl;
    if (!match(Carry,
               m_OneUse(m_Select(
                   m_OneUse(m_SpecificICmp(ICmpInst::ICMP_ULT, m_Value(LowSum),
                                           m_Value(XhYl))),
                   m_SpecificInt(APInt::getOneBitSet(BitWidth, BitWidth / 2)),
                   m_Zero()))))
      return false;

    // XhYl can be Xh*Yl or Xl*Yh
    if (!CheckHiLo(XhYl, X, Y)) {
      if (CheckHiLo(XhYl, Y, X))
        std::swap(X, Y);
      else
        return false;
    }
    if (XhYl->hasNUsesOrMore(3))
      return false;

    // B = LowSum >> 32
    if (!match(B, m_OneUse(m_LShr(m_Specific(LowSum),
                                  m_SpecificInt(BitWidth / 2)))) ||
        LowSum->hasNUsesOrMore(3))
      return false;

    // LowSum = XhYl + XlYh + XlYl>>32
    Value *XlYh, *XlYl;
    auto XlYlHi = m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2));
    if (!match(LowSum,
               m_c_Add(m_Specific(XhYl),
                       m_OneUse(m_c_Add(m_OneUse(m_Value(XlYh)), XlYlHi)))) &&
        !match(LowSum, m_c_Add(m_OneUse(m_Value(XlYh)),
                               m_OneUse(m_c_Add(m_Specific(XhYl), XlYlHi)))) &&
        !match(LowSum,
               m_c_Add(XlYlHi, m_OneUse(m_c_Add(m_Specific(XhYl),
                                                m_OneUse(m_Value(XlYh)))))))
      return false;

    // Check XlYl and XlYh
    if (!CheckLoLo(XlYl, X, Y))
      return false;
    if (!CheckHiLo(XlYh, Y, X))
      return false;

    return CreateMulHigh(X, Y);
  };

  auto FoldMulHighLadder = [&](Value *X, Value *Y, Instruction *A,
                               Instruction *B) {
    //  xh*yh + c2>>32 + c3>>32
    //    c2 = xh*yl + (xl*yl>>32); c3 = c2&0xffffffff + xl*yh
    // or c2 = (xl*yh&0xffffffff) + xh*yl + (xl*yl>>32); c3 = xh*yl
    Value *XlYh, *XhYl, *XlYl, *C2, *C3;
    // Strip off the two expected shifts.
    if (!match(A, m_LShr(m_Value(C2), m_SpecificInt(BitWidth / 2))) ||
        !match(B, m_LShr(m_Value(C3), m_SpecificInt(BitWidth / 2))))
      return false;

    if (match(C3, m_c_Add(m_Add(m_Value(), m_Value()), m_Value())))
      std::swap(C2, C3);
    // Try to match c2 = (xl*yh&0xffffffff) + xh*yl + (xl*yl>>32)
    if (match(C2,
              m_c_Add(m_c_Add(m_And(m_Specific(C3), m_SpecificInt(LowMask)),
                              m_Value(XlYh)),
                      m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2)))) ||
        match(C2, m_c_Add(m_c_Add(m_And(m_Specific(C3), m_SpecificInt(LowMask)),
                                  m_LShr(m_Value(XlYl),
                                         m_SpecificInt(BitWidth / 2))),
                          m_Value(XlYh))) ||
        match(C2, m_c_Add(m_c_Add(m_LShr(m_Value(XlYl),
                                         m_SpecificInt(BitWidth / 2)),
                                  m_Value(XlYh)),
                          m_And(m_Specific(C3), m_SpecificInt(LowMask))))) {
      XhYl = C3;
    } else {
      // Match c3 = c2&0xffffffff + xl*yh
      if (!match(C3, m_c_Add(m_And(m_Specific(C2), m_SpecificInt(LowMask)),
                             m_Value(XlYh))))
        std::swap(C2, C3);
      if (!match(C3, m_c_Add(m_OneUse(
                                 m_And(m_Specific(C2), m_SpecificInt(LowMask))),
                             m_Value(XlYh))) ||
          !C3->hasOneUse() || C2->hasNUsesOrMore(3))
        return false;

      // Match c2 = xh*yl + (xl*yl >> 32)
      if (!match(C2, m_c_Add(m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2)),
                             m_Value(XhYl))))
        return false;
    }

    // Match XhYl and XlYh - they can appear either way around.
    if (!CheckHiLo(XlYh, Y, X))
      std::swap(XlYh, XhYl);
    if (!CheckHiLo(XlYh, Y, X))
      return false;
    if (!CheckHiLo(XhYl, X, Y))
      return false;
    if (!CheckLoLo(XlYl, X, Y))
      return false;

    return CreateMulHigh(X, Y);
  };

  auto FoldMulHighLadder4 = [&](Value *X, Value *Y, Instruction *A,
                                Instruction *B, Instruction *C) {
    ///  Ladder4: xh*yh + (xl*yh)>>32 + (xh+yl)>>32 + low>>32;
    ///           low = (xl*yl)>>32 + (xl*yh)&0xffffffff + (xh*yl)&0xffffffff

    // Find A = Low >> 32 and B/C = XhYl>>32, XlYh>>32.
    auto ShiftAdd =
        m_LShr(m_Add(m_Value(), m_Value()), m_SpecificInt(BitWidth / 2));
    if (!match(A, ShiftAdd))
      std::swap(A, B);
    if (!match(A, ShiftAdd))
      std::swap(A, C);
    Value *Low;
    if (!match(A, m_LShr(m_OneUse(m_Value(Low)), m_SpecificInt(BitWidth / 2))))
      return false;

    // Match B == XhYl>>32 and C == XlYh>>32
    Value *XhYl, *XlYh;
    if (!match(B, m_LShr(m_Value(XhYl), m_SpecificInt(BitWidth / 2))) ||
        !match(C, m_LShr(m_Value(XlYh), m_SpecificInt(BitWidth / 2))))
      return false;
    if (!CheckHiLo(XhYl, X, Y))
      std::swap(XhYl, XlYh);
    if (!CheckHiLo(XhYl, X, Y) || XhYl->hasNUsesOrMore(3))
      return false;
    if (!CheckHiLo(XlYh, Y, X) || XlYh->hasNUsesOrMore(3))
      return false;

    // Match Low as XlYl>>32 + XhYl&0xffffffff + XlYh&0xffffffff
    Value *XlYl;
    if (!match(
            Low,
            m_c_Add(
                m_OneUse(m_c_Add(
                    m_OneUse(m_And(m_Specific(XhYl), m_SpecificInt(LowMask))),
                    m_OneUse(m_And(m_Specific(XlYh), m_SpecificInt(LowMask))))),
                m_OneUse(
                    m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2))))) &&
        !match(
            Low,
            m_c_Add(
                m_OneUse(m_c_Add(
                    m_OneUse(m_And(m_Specific(XhYl), m_SpecificInt(LowMask))),
                    m_OneUse(
                        m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2))))),
                m_OneUse(m_And(m_Specific(XlYh), m_SpecificInt(LowMask))))) &&
        !match(
            Low,
            m_c_Add(
                m_OneUse(m_c_Add(
                    m_OneUse(m_And(m_Specific(XlYh), m_SpecificInt(LowMask))),
                    m_OneUse(
                        m_LShr(m_Value(XlYl), m_SpecificInt(BitWidth / 2))))),
                m_OneUse(m_And(m_Specific(XhYl), m_SpecificInt(LowMask))))))
      return false;
    if (!CheckLoLo(XlYl, X, Y))
      return false;

    return CreateMulHigh(X, Y);
  };

  auto FoldMulHighCarry4 = [&](Value *X, Value *Y, Instruction *Carry,
                               Instruction *B, Instruction *C) {
    //  xh*yh + carry + crosssum>>32 + (xl*yl + crosssum&0xffffffff) >> 32
    //  crosssum = xh*yl+xl*yh
    //  carry = crosssum < xh*yl ? 0x1000000 : 0
    if (Carry->getOpcode() != Instruction::Select)
      std::swap(Carry, B);
    if (Carry->getOpcode() != Instruction::Select)
      std::swap(Carry, C);

    // Carry = CrossSum < XhYl ? 0x100000000 : 0
    Value *CrossSum, *XhYl;
    if (!match(Carry,
               m_OneUse(m_Select(
                   m_OneUse(m_SpecificICmp(ICmpInst::ICMP_ULT,
                                           m_Value(CrossSum), m_Value(XhYl))),
                   m_SpecificInt(APInt::getOneBitSet(BitWidth, BitWidth / 2)),
                   m_Zero()))))
      return false;

    if (!match(B, m_LShr(m_Specific(CrossSum), m_SpecificInt(BitWidth / 2))))
      std::swap(B, C);
    if (!match(B, m_LShr(m_Specific(CrossSum), m_SpecificInt(BitWidth / 2))))
      return false;

    Value *XlYl, *LowAccum;
    if (!match(C, m_LShr(m_Value(LowAccum), m_SpecificInt(BitWidth / 2))) ||
        !match(LowAccum, m_c_Add(m_OneUse(m_LShr(m_Value(XlYl),
                                                 m_SpecificInt(BitWidth / 2))),
                                 m_OneUse(m_And(m_Specific(CrossSum),
                                                m_SpecificInt(LowMask))))) ||
        LowAccum->hasNUsesOrMore(3))
      return false;
    if (!CheckLoLo(XlYl, X, Y))
      return false;

    if (!CheckHiLo(XhYl, X, Y))
      std::swap(X, Y);
    if (!CheckHiLo(XhYl, X, Y))
      return false;
    Value *XlYh;
    if (!match(CrossSum, m_c_Add(m_Specific(XhYl), m_OneUse(m_Value(XlYh)))) ||
        !CheckHiLo(XlYh, Y, X) || CrossSum->hasNUsesOrMore(4) ||
        XhYl->hasNUsesOrMore(3))
      return false;

    return CreateMulHigh(X, Y);
  };

  // X and Y are the two inputs, A, B and C are other parts of the pattern
  // (crosssum>>32, carry, etc).
  Value *X, *Y;
  Instruction *A, *B, *C;
  auto HiHi = m_OneUse(m_Mul(m_LShr(m_Value(X), m_SpecificInt(BitWidth / 2)),
                             m_LShr(m_Value(Y), m_SpecificInt(BitWidth / 2))));
  if ((match(&I, m_c_Add(HiHi, m_OneUse(m_Add(m_Instruction(A),
                                              m_Instruction(B))))) ||
       match(&I, m_c_Add(m_Instruction(A),
                         m_OneUse(m_c_Add(HiHi, m_Instruction(B)))))) &&
      A->hasOneUse() && B->hasOneUse())
    if (FoldMulHighCarry(X, Y, A, B) || FoldMulHighLadder(X, Y, A, B))
      return true;

  if ((match(&I, m_c_Add(HiHi, m_OneUse(m_c_Add(
                                   m_Instruction(A),
                                   m_OneUse(m_Add(m_Instruction(B),
                                                  m_Instruction(C))))))) ||
       match(&I, m_c_Add(m_Instruction(A),
                         m_OneUse(m_c_Add(
                             HiHi, m_OneUse(m_Add(m_Instruction(B),
                                                  m_Instruction(C))))))) ||
       match(&I, m_c_Add(m_Instruction(A),
                         m_OneUse(m_c_Add(
                             m_Instruction(B),
                             m_OneUse(m_c_Add(HiHi, m_Instruction(C))))))) ||
       match(&I,
             m_c_Add(m_OneUse(m_c_Add(HiHi, m_Instruction(A))),
                     m_OneUse(m_Add(m_Instruction(B), m_Instruction(C)))))) &&
      A->hasOneUse() && B->hasOneUse() && C->hasOneUse())
    return FoldMulHighCarry4(X, Y, A, B, C) ||
           FoldMulHighLadder4(X, Y, A, B, C);

  return false;
}

/// Guard a memset whose nonconstant length is known to be in [0, 1].
/// Inserts a conditional branch around the memset and specialises the
/// executed path to a one-byte store.
static bool foldMemSetZeroOrOneLength(Instruction &I, const DataLayout &DL,
                                      TargetLibraryInfo &TLI,
                                      AssumptionCache &AC, DominatorTree &DT,
                                      bool &MadeCFGChange) {
  auto *MI = dyn_cast<MemSetInst>(&I);
  if (!MI || isa<ConstantInt>(MI->getLength()))
    return false;

  SimplifyQuery SQ(DL, &TLI, &DT, &AC, MI);
  KnownBits KnownLen = computeKnownBits(MI->getLength(), SQ);
  if (!KnownLen.getMaxValue().isOne())
    return false;

  uint64_t TotalCount;
  SmallVector<InstrProfValueData> MemsetVPMetadata = getValueProfDataFromInst(
      I, InstrProfValueKind::IPVK_MemOPSize, 2, TotalCount);
  std::optional<uint64_t> ZeroCount = std::nullopt;
  std::optional<uint64_t> OneCount = std::nullopt;
  for (const auto [MemOpSize, SizeFrequency] : MemsetVPMetadata) {
    if (MemOpSize == 0)
      ZeroCount = SizeFrequency;
    else if (MemOpSize == 1)
      OneCount = SizeFrequency;
  }
  // If we only have one value in the profile, we assume that the other is zero.
  if (MemsetVPMetadata.size() == 1) {
    if (ZeroCount.has_value())
      OneCount = 0;
    else if (OneCount.has_value())
      ZeroCount = 0;
  }

  BasicBlock *HeadBlock = MI->getIterator()->getParent();
  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
  IRBuilder<> B(MI);
  Value *IsNonZero = B.CreateIsNotNull(MI->getLength(), "memset.notzero");
  Instruction *ThenTerm = SplitBlockAndInsertIfThen(
      IsNonZero, MI->getIterator(), /*Unreachable=*/false,
      /*BranchWeights=*/nullptr, &DTU);

  Instruction &IsNonZeroBranch = *HeadBlock->getTerminator();
  if (!ProfcheckDisableMetadataFixes && OneCount.has_value() &&
      ZeroCount.has_value() && (*OneCount + *ZeroCount > 0))
    setFittedBranchWeights(IsNonZeroBranch, {*OneCount, *ZeroCount}, false);
  else
    setExplicitlyUnknownBranchWeightsIfProfiled(IsNonZeroBranch, DEBUG_TYPE);

  IRBuilder<> StoreBuilder(ThenTerm);
  StoreInst *Store = StoreBuilder.CreateAlignedStore(
      MI->getValue(), MI->getDest(), MI->getDestAlign(), MI->isVolatile());
  Store->copyMetadata(*MI, LLVMContext::MD_DIAssignID);
  MI->eraseFromParent();
  ++NumMemSetsGuarded;
  MadeCFGChange = true;
  return true;
}

/// This is the entry point for folds that could be implemented in regular
/// InstCombine, but they are separated because they are not expected to
/// occur frequently and/or have more than a constant-length pattern match.
static bool foldUnusualPatterns(Function &F, DominatorTree &DT,
                                TargetTransformInfo &TTI,
                                TargetLibraryInfo &TLI, AliasAnalysis &AA,
                                AssumptionCache &AC, bool &MadeCFGChange) {
  bool MadeChange = false;
  for (BasicBlock &BB : F) {
    // Ignore unreachable basic blocks.
    if (!DT.isReachableFromEntry(&BB))
      continue;

    const DataLayout &DL = F.getDataLayout();

    // Walk the block backwards for efficiency. We're matching a chain of
    // use->defs, so we're more likely to succeed by starting from the bottom.
    // Also, we want to avoid matching partial patterns.
    // TODO: It would be more efficient if we removed dead instructions
    // iteratively in this loop rather than waiting until the end.
    for (Instruction &I : make_early_inc_range(llvm::reverse(BB))) {
      MadeChange |= foldAnyOrAllBitsSet(I);
      MadeChange |= foldGuardedFunnelShift(I, DT);
      MadeChange |= foldSelectSplitCTLZCTTZ(I);
      MadeChange |= tryToRecognizePopCount(I);
      MadeChange |= tryToRecognizePopCount1(I);
      MadeChange |= tryToRecognizePopCount2n3(I);
      MadeChange |= tryToFPToSat(I, TTI);
      MadeChange |= tryToRecognizeTableBasedPatterns(I, DL, TTI);
      MadeChange |= foldConsecutiveLoads(I, DL, TTI, AA, DT);
      MadeChange |= foldPatternedLoads(I, DL);
      MadeChange |= foldICmpOrChain(I, DL, TTI, AA, DT);
      MadeChange |= foldMulHigh(I);
      // These folds can erase the instruction `I`, so they need to be called
      // at the end of this sequence.
      if (foldLibCalls(I, TTI, TLI, AC, DT, DL, MadeCFGChange)) {
        MadeChange = true;
        continue;
      }
      if (foldMemSetZeroOrOneLength(I, DL, TLI, AC, DT, MadeCFGChange))
        MadeChange = true;
    }

    // Do this separately to avoid redundantly scanning stores multiple times.
    MadeChange |= foldConsecutiveStores(BB, DL, TTI, AA);
  }

  // We're done with transforms, so remove dead instructions.
  if (MadeChange)
    for (BasicBlock &BB : F)
      SimplifyInstructionsInBlock(&BB);

  return MadeChange;
}

/// This is the entry point for all transforms. Pass manager differences are
/// handled in the callers of this function.
static bool runImpl(Function &F, AssumptionCache &AC, TargetTransformInfo &TTI,
                    TargetLibraryInfo &TLI, DominatorTree &DT,
                    AliasAnalysis &AA, bool &MadeCFGChange) {
  bool MadeChange = false;
  const DataLayout &DL = F.getDataLayout();
  TruncInstCombine TIC(AC, TLI, DL, DT);
  MadeChange |= TIC.run(F);
  MadeChange |= foldUnusualPatterns(F, DT, TTI, TLI, AA, AC, MadeCFGChange);
  return MadeChange;
}

PreservedAnalyses AggressiveInstCombinePass::run(Function &F,
                                                 FunctionAnalysisManager &AM) {
  auto &AC = AM.getResult<AssumptionAnalysis>(F);
  auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
  auto &TTI = AM.getResult<TargetIRAnalysis>(F);
  auto &AA = AM.getResult<AAManager>(F);
  bool MadeCFGChange = false;
  if (!runImpl(F, AC, TTI, TLI, DT, AA, MadeCFGChange)) {
    // No changes, all analyses are preserved.
    return PreservedAnalyses::all();
  }
  // Mark all the analyses that instcombine updates as preserved.
  PreservedAnalyses PA;
  if (MadeCFGChange)
    PA.preserve<DominatorTreeAnalysis>();
  else
    PA.preserveSet<CFGAnalyses>();
  return PA;
}