path: root/backend/lib/CodeGen/GlobalISel/GISelValueTracking.cpp

// SPDX-License-Identifier: GPL-3.0-or-later

#include "rcc/CodeGen/GlobalISel/GISelValueTracking.h"

#include "rcc/CodeGen/GenericOpcodes.h"
#include "rcc/CodeGen/GlobalISel/Utils.h"
#include "rcc/CodeGen/LowLevelType.h"
#include "rcc/CodeGen/MachineFunction.h"
#include "rcc/CodeGen/MachineInstr.h"
#include "rcc/CodeGen/MachineOperand.h"
#include "rcc/CodeGen/MachineRegisterInfo.h"
#include "rcc/CodeGen/TargetDescription.h"

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <optional>
#include <vector>

namespace rcc {
namespace {

[[nodiscard]] std::optional<Register> registerOperand(const MachineInstr &Instruction,
                                                      std::size_t Index) noexcept {
  const MachineOperand *Operand = Instruction.operandAt(Index);
  return Operand == nullptr ? std::nullopt : Operand->registerValue();
}

[[nodiscard]] Expected<APInt> zero(std::uint32_t BitWidth) {
  return APInt::fromUnsigned(BitWidth, 0);
}

[[nodiscard]] Expected<APInt> resizeMask(const APInt &Value, std::uint32_t BitWidth) {
  if (Value.bitWidth() == BitWidth)
    return Value;
  if (Value.bitWidth() == 0)
    return zero(BitWidth);
  if (Value.bitWidth() < BitWidth)
    return Value.zext(BitWidth);
  return Value.trunc(BitWidth);
}

[[nodiscard]] Expected<APInt> highBitsMask(std::uint32_t BitWidth, std::uint32_t LowBitCount) {
  auto LowBits = APInt::getLowBitsSet(BitWidth, LowBitCount);
  if (!LowBits)
    return std::unexpected(LowBits.error());
  return LowBits->bitwiseNot();
}

[[nodiscard]] Expected<APInt> mergeMasks(const APInt &Left, const APInt &Right) {
  return Left.bitwiseOr(Right);
}

[[nodiscard]] Expected<APInt> intersectMasks(const APInt &Left, const APInt &Right) {
  return Left.bitwiseAnd(Right);
}

[[nodiscard]] Expected<APInt> addMasks(const APInt &Left, const APInt &Right) {
  if (Left.bitWidth() == 0 || Left.bitWidth() != Right.bitWidth())
    return makeUnexpected(ErrorCode::InvalidArgument,
                          "known-bits add operands must have the same nonzero width");
  std::vector<std::uint64_t> Words(Left.words().size(), 0);
  unsigned __int128 Carry = 0;
  for (std::size_t Index = 0; Index != Words.size(); ++Index) {
    const unsigned __int128 Sum = static_cast<unsigned __int128>(Left.word(Index)) +
                                  static_cast<unsigned __int128>(Right.word(Index)) + Carry;
    Words[Index] = static_cast<std::uint64_t>(Sum);
    Carry = Sum >> 64U;
  }
  return APInt::create(Left.bitWidth(), Words);
}

struct MaskProduct final {
  APInt Value;
  bool Overflow = false;
};

[[nodiscard]] Expected<MaskProduct> multiplyMasks(const APInt &Left, const APInt &Right) {
  if (Left.bitWidth() == 0 || Left.bitWidth() != Right.bitWidth())
    return makeUnexpected(ErrorCode::InvalidArgument,
                          "known-bits multiply operands must have the same nonzero width");

  const std::size_t WordCount = Left.words().size();
  std::vector<std::uint64_t> Product(WordCount * 2U, 0);
  for (std::size_t LeftIndex = 0; LeftIndex != WordCount; ++LeftIndex) {
    for (std::size_t RightIndex = 0; RightIndex != WordCount; ++RightIndex) {
      const std::size_t ProductIndex = LeftIndex + RightIndex;
      const unsigned __int128 Partial = static_cast<unsigned __int128>(Left.word(LeftIndex)) *
                                        static_cast<unsigned __int128>(Right.word(RightIndex));
      unsigned __int128 Sum = static_cast<unsigned __int128>(Product[ProductIndex]) +
                              static_cast<std::uint64_t>(Partial);
      Product[ProductIndex] = static_cast<std::uint64_t>(Sum);
      unsigned __int128 Carry = (Partial >> 64U) + (Sum >> 64U);
      for (std::size_t Index = ProductIndex + 1U; Carry != 0 && Index != Product.size(); ++Index) {
        Sum = static_cast<unsigned __int128>(Product[Index]) + static_cast<std::uint64_t>(Carry);
        Product[Index] = static_cast<std::uint64_t>(Sum);
        Carry = (Carry >> 64U) + (Sum >> 64U);
      }
    }
  }

  bool Overflow = false;
  for (std::size_t Index = WordCount; Index != Product.size(); ++Index)
    Overflow |= Product[Index] != 0;
  const std::uint32_t HighBits = Left.bitWidth() % 64U;
  if (HighBits != 0) {
    const std::uint64_t HighMask = (std::uint64_t{1} << HighBits) - 1U;
    Overflow |= (Product[WordCount - 1U] & ~HighMask) != 0;
    Product[WordCount - 1U] &= HighMask;
  }
  Product.resize(WordCount);
  auto Value = APInt::create(Left.bitWidth(), Product);
  if (!Value)
    return std::unexpected(Value.error());
  return MaskProduct{.Value = std::move(*Value), .Overflow = Overflow};
}

[[nodiscard]] std::uint32_t countTrailingOnes(const APInt &Value) noexcept {
  std::uint32_t Count = 0;
  while (Count != Value.bitWidth() && Value.testBit(Count))
    ++Count;
  return Count;
}

[[nodiscard]] Expected<APInt> constantAtWidth(const APInt &Value, std::uint32_t BitWidth) {
  return resizeMask(Value, BitWidth);
}

[[nodiscard]] std::optional<std::uint32_t>
constantShiftAmount(Register Reg, const MachineRegisterInfo &MRI, std::uint32_t BitWidth) {
  const auto Constant = getIConstantVRegValWithLookThrough(Reg, MRI);
  if (!Constant || Constant->Value.activeBitWidth() > 32)
    return std::nullopt;
  const std::uint64_t Amount = Constant->Value.word(0);
  if (Amount >= BitWidth)
    return std::nullopt;
  return static_cast<std::uint32_t>(Amount);
}

[[nodiscard]] Expected<APInt> shiftMaskLeft(const APInt &Value, std::uint32_t Amount) {
  std::vector<std::uint64_t> Words(Value.words().size(), 0);
  for (std::uint32_t Bit = 0; Bit + Amount < Value.bitWidth(); ++Bit) {
    if (Value.testBit(Bit)) {
      const std::uint32_t DestinationBit = Bit + Amount;
      Words[DestinationBit / 64U] |= std::uint64_t{1} << (DestinationBit % 64U);
    }
  }
  return APInt::create(Value.bitWidth(), Words);
}

[[nodiscard]] Expected<APInt> shiftMaskRight(const APInt &Value, std::uint32_t Amount) {
  std::vector<std::uint64_t> Words(Value.words().size(), 0);
  for (std::uint32_t Bit = Amount; Bit < Value.bitWidth(); ++Bit) {
    if (Value.testBit(Bit)) {
      const std::uint32_t DestinationBit = Bit - Amount;
      Words[DestinationBit / 64U] |= std::uint64_t{1} << (DestinationBit % 64U);
    }
  }
  return APInt::create(Value.bitWidth(), Words);
}

[[nodiscard]] std::uint32_t constantNumSignBits(const APInt &Value) noexcept {
  if (Value.bitWidth() == 0)
    return 1;
  if (Value.isNegative())
    return static_cast<std::uint32_t>(Value.countlOne());
  auto Inverted = Value.bitwiseNot();
  return Inverted ? static_cast<std::uint32_t>(Inverted->countlOne()) : 1U;
}

} // namespace

GISelValueTracking::GISelValueTracking(const MachineFunction &Function,
                                       const TargetDescription &Target, unsigned MaxDepth) noexcept
    : Function(Function), MRI(Function.registerInfo()), Target(Target), MaxDepth(MaxDepth) {}

Expected<GISelValueTracking::KnownBits>
GISelValueTracking::makeUnknown(std::uint32_t BitWidth) const {
  auto Zero = zero(BitWidth);
  if (!Zero)
    return std::unexpected(Zero.error());
  auto One = zero(BitWidth);
  if (!One)
    return std::unexpected(One.error());
  return KnownBits{.Zero = std::move(*Zero), .One = std::move(*One)};
}

Expected<GISelValueTracking::KnownBits> GISelValueTracking::getKnownBits(Register Reg,
                                                                         unsigned Depth) const {
  auto Type = MRI.getType(Reg);
  if (!Type)
    return std::unexpected(Type.error());
  const std::uint32_t BitWidth = Type->scalarBitWidth();
  auto Unknown = makeUnknown(BitWidth);
  if (!Unknown)
    return std::unexpected(Unknown.error());
  if (Depth >= MaxDepth)
    return Unknown;

  const MachineInstr *Definition = MRI.getVRegDef(Reg);
  if (Definition == nullptr)
    return Unknown;

  const auto Opcode = static_cast<GenericOpcode>(Definition->opcode().value());
  if (Opcode == GenericOpcode::G_CONSTANT) {
    const MachineOperand *ValueOperand = Definition->operandAt(1);
    const auto *Immediate =
        ValueOperand == nullptr ? nullptr : ValueOperand->getIf<IntegerImmediateOperand>();
    if (Immediate == nullptr)
      return Unknown;
    auto Value = constantAtWidth(Immediate->Value, BitWidth);
    if (!Value)
      return std::unexpected(Value.error());
    auto KnownZero = Value->bitwiseNot();
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*Value)};
  }

  if (Opcode == GenericOpcode::COPY) {
    const auto Source = registerOperand(*Definition, 1);
    if (!Source || !Source->isVirtual())
      return Unknown;
    auto SourceBits = getKnownBits(*Source, Depth);
    if (!SourceBits)
      return std::unexpected(SourceBits.error());
    auto KnownZero = resizeMask(SourceBits->Zero, BitWidth);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto KnownOne = resizeMask(SourceBits->One, BitWidth);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  if (Opcode == GenericOpcode::G_ICMP || Opcode == GenericOpcode::G_FCMP) {
    if (!Type->isVector() && BitWidth > 1 &&
        Target.getBooleanContents(false, Opcode == GenericOpcode::G_FCMP) ==
            BooleanContent::ZeroOrOne) {
      auto KnownZero = highBitsMask(BitWidth, 1);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      Unknown->Zero = std::move(*KnownZero);
    }
    return Unknown;
  }

  switch (Opcode) {
  case GenericOpcode::G_UADDO:
  case GenericOpcode::G_UADDE:
  case GenericOpcode::G_USUBO:
  case GenericOpcode::G_USUBE:
  case GenericOpcode::G_SADDO:
  case GenericOpcode::G_SADDE:
  case GenericOpcode::G_SSUBO:
  case GenericOpcode::G_SSUBE:
  case GenericOpcode::G_UMULO:
  case GenericOpcode::G_SMULO: {
    const auto BooleanResult = registerOperand(*Definition, 1);
    if (BooleanResult && *BooleanResult == Reg && !Type->isVector() && BitWidth > 1 &&
        Target.getBooleanContents(false, false) == BooleanContent::ZeroOrOne) {
      auto KnownZero = highBitsMask(BitWidth, 1);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      Unknown->Zero = std::move(*KnownZero);
    }
    return Unknown;
  }
  default:
    break;
  }

  const auto Source = registerOperand(*Definition, 1);
  if (!Source)
    return Unknown;

  if (Opcode == GenericOpcode::G_TRUNC || Opcode == GenericOpcode::G_ANYEXT ||
      Opcode == GenericOpcode::G_ZEXT || Opcode == GenericOpcode::G_SEXT ||
      Opcode == GenericOpcode::G_ASSERT_ZEXT || Opcode == GenericOpcode::G_ASSERT_SEXT ||
      Opcode == GenericOpcode::G_SEXT_INREG) {
    auto SourceBits = getKnownBits(*Source, Depth + 1U);
    if (!SourceBits)
      return std::unexpected(SourceBits.error());
    auto SourceType = MRI.getType(*Source);
    if (!SourceType)
      return std::unexpected(SourceType.error());
    const std::uint32_t SourceWidth = SourceType->scalarBitWidth();
    auto KnownZero = resizeMask(SourceBits->Zero, BitWidth);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto KnownOne = resizeMask(SourceBits->One, BitWidth);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());

    if (Opcode == GenericOpcode::G_ZEXT && SourceWidth < BitWidth) {
      auto HighZero = highBitsMask(BitWidth, SourceWidth);
      if (!HighZero)
        return std::unexpected(HighZero.error());
      KnownZero = mergeMasks(*KnownZero, *HighZero);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
    }

    std::uint32_t ExtendedFrom = SourceWidth;
    if (Opcode == GenericOpcode::G_ASSERT_ZEXT || Opcode == GenericOpcode::G_ASSERT_SEXT ||
        Opcode == GenericOpcode::G_SEXT_INREG) {
      const MachineOperand *WidthOperand = Definition->operandAt(2);
      const auto *Width = WidthOperand == nullptr ? nullptr : WidthOperand->getIf<std::int64_t>();
      if (Width == nullptr || *Width <= 0 || static_cast<std::uint64_t>(*Width) > BitWidth)
        return Unknown;
      ExtendedFrom = static_cast<std::uint32_t>(*Width);
    }

    if (Opcode == GenericOpcode::G_ASSERT_ZEXT) {
      auto HighZero = highBitsMask(BitWidth, ExtendedFrom);
      if (!HighZero)
        return std::unexpected(HighZero.error());
      KnownZero = mergeMasks(*KnownZero, *HighZero);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
    }

    if ((Opcode == GenericOpcode::G_SEXT || Opcode == GenericOpcode::G_ASSERT_SEXT ||
         Opcode == GenericOpcode::G_SEXT_INREG) &&
        ExtendedFrom < BitWidth) {
      auto HighBits = highBitsMask(BitWidth, ExtendedFrom);
      if (!HighBits)
        return std::unexpected(HighBits.error());
      if (KnownZero->testBit(ExtendedFrom - 1U)) {
        KnownZero = mergeMasks(*KnownZero, *HighBits);
        if (!KnownZero)
          return std::unexpected(KnownZero.error());
      } else if (KnownOne->testBit(ExtendedFrom - 1U)) {
        KnownOne = mergeMasks(*KnownOne, *HighBits);
        if (!KnownOne)
          return std::unexpected(KnownOne.error());
      }
    }

    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  if (Opcode == GenericOpcode::G_AND || Opcode == GenericOpcode::G_OR ||
      Opcode == GenericOpcode::G_XOR) {
    const auto Right = registerOperand(*Definition, 2);
    if (!Right)
      return Unknown;
    auto LeftBits = getKnownBits(*Source, Depth + 1U);
    if (!LeftBits)
      return std::unexpected(LeftBits.error());
    auto RightBits = getKnownBits(*Right, Depth + 1U);
    if (!RightBits)
      return std::unexpected(RightBits.error());

    if (Opcode == GenericOpcode::G_AND) {
      auto KnownZero = mergeMasks(LeftBits->Zero, RightBits->Zero);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      auto KnownOne = intersectMasks(LeftBits->One, RightBits->One);
      if (!KnownOne)
        return std::unexpected(KnownOne.error());
      return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
    }
    if (Opcode == GenericOpcode::G_OR) {
      auto KnownZero = intersectMasks(LeftBits->Zero, RightBits->Zero);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      auto KnownOne = mergeMasks(LeftBits->One, RightBits->One);
      if (!KnownOne)
        return std::unexpected(KnownOne.error());
      return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
    }

    auto BothZero = intersectMasks(LeftBits->Zero, RightBits->Zero);
    if (!BothZero)
      return std::unexpected(BothZero.error());
    auto BothOne = intersectMasks(LeftBits->One, RightBits->One);
    if (!BothOne)
      return std::unexpected(BothOne.error());
    auto KnownZero = mergeMasks(*BothZero, *BothOne);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto LeftZeroRightOne = intersectMasks(LeftBits->Zero, RightBits->One);
    if (!LeftZeroRightOne)
      return std::unexpected(LeftZeroRightOne.error());
    auto LeftOneRightZero = intersectMasks(LeftBits->One, RightBits->Zero);
    if (!LeftOneRightZero)
      return std::unexpected(LeftOneRightZero.error());
    auto KnownOne = mergeMasks(*LeftZeroRightOne, *LeftOneRightZero);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  if (Opcode == GenericOpcode::G_ADD) {
    const auto Right = registerOperand(*Definition, 2);
    if (!Right)
      return Unknown;
    auto LeftBits = getKnownBits(*Source, Depth + 1U);
    if (!LeftBits)
      return std::unexpected(LeftBits.error());
    auto RightBits = getKnownBits(*Right, Depth + 1U);
    if (!RightBits)
      return std::unexpected(RightBits.error());

    auto LeftMaximum = LeftBits->Zero.bitwiseNot();
    if (!LeftMaximum)
      return std::unexpected(LeftMaximum.error());
    auto RightMaximum = RightBits->Zero.bitwiseNot();
    if (!RightMaximum)
      return std::unexpected(RightMaximum.error());
    auto PossibleSumZero = addMasks(*LeftMaximum, *RightMaximum);
    if (!PossibleSumZero)
      return std::unexpected(PossibleSumZero.error());
    auto PossibleSumOne = addMasks(LeftBits->One, RightBits->One);
    if (!PossibleSumOne)
      return std::unexpected(PossibleSumOne.error());

    auto CarryKnownZero = PossibleSumZero->bitwiseXor(LeftBits->Zero);
    if (!CarryKnownZero)
      return std::unexpected(CarryKnownZero.error());
    CarryKnownZero = CarryKnownZero->bitwiseXor(RightBits->Zero);
    if (!CarryKnownZero)
      return std::unexpected(CarryKnownZero.error());
    CarryKnownZero = CarryKnownZero->bitwiseNot();
    if (!CarryKnownZero)
      return std::unexpected(CarryKnownZero.error());
    auto CarryKnownOne = PossibleSumOne->bitwiseXor(LeftBits->One);
    if (!CarryKnownOne)
      return std::unexpected(CarryKnownOne.error());
    CarryKnownOne = CarryKnownOne->bitwiseXor(RightBits->One);
    if (!CarryKnownOne)
      return std::unexpected(CarryKnownOne.error());

    auto LeftKnown = mergeMasks(LeftBits->Zero, LeftBits->One);
    if (!LeftKnown)
      return std::unexpected(LeftKnown.error());
    auto RightKnown = mergeMasks(RightBits->Zero, RightBits->One);
    if (!RightKnown)
      return std::unexpected(RightKnown.error());
    auto CarryKnown = mergeMasks(*CarryKnownZero, *CarryKnownOne);
    if (!CarryKnown)
      return std::unexpected(CarryKnown.error());
    auto Known = intersectMasks(*LeftKnown, *RightKnown);
    if (!Known)
      return std::unexpected(Known.error());
    Known = intersectMasks(*Known, *CarryKnown);
    if (!Known)
      return std::unexpected(Known.error());

    auto NotPossibleSumZero = PossibleSumZero->bitwiseNot();
    if (!NotPossibleSumZero)
      return std::unexpected(NotPossibleSumZero.error());
    auto KnownZero = intersectMasks(*NotPossibleSumZero, *Known);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto KnownOne = intersectMasks(*PossibleSumOne, *Known);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  if (Opcode == GenericOpcode::G_MUL) {
    const auto Right = registerOperand(*Definition, 2);
    if (!Right)
      return Unknown;
    auto LeftBits = getKnownBits(*Source, Depth + 1U);
    if (!LeftBits)
      return std::unexpected(LeftBits.error());
    auto RightBits = getKnownBits(*Right, Depth + 1U);
    if (!RightBits)
      return std::unexpected(RightBits.error());

    auto LeftMaximum = LeftBits->Zero.bitwiseNot();
    if (!LeftMaximum)
      return std::unexpected(LeftMaximum.error());
    auto RightMaximum = RightBits->Zero.bitwiseNot();
    if (!RightMaximum)
      return std::unexpected(RightMaximum.error());
    auto MaximumProduct = multiplyMasks(*LeftMaximum, *RightMaximum);
    if (!MaximumProduct)
      return std::unexpected(MaximumProduct.error());
    const std::uint32_t LeadingZeroes =
        MaximumProduct->Overflow
            ? 0
            : BitWidth - static_cast<std::uint32_t>(MaximumProduct->Value.activeBitWidth());

    auto LeftKnown = mergeMasks(LeftBits->Zero, LeftBits->One);
    if (!LeftKnown)
      return std::unexpected(LeftKnown.error());
    auto RightKnown = mergeMasks(RightBits->Zero, RightBits->One);
    if (!RightKnown)
      return std::unexpected(RightKnown.error());
    const std::uint32_t LeftTrailingKnown = countTrailingOnes(*LeftKnown);
    const std::uint32_t RightTrailingKnown = countTrailingOnes(*RightKnown);
    const std::uint32_t LeftTrailingZeroes = countTrailingOnes(LeftBits->Zero);
    const std::uint32_t RightTrailingZeroes = countTrailingOnes(RightBits->Zero);
    const std::uint32_t TrailingZeroes = LeftTrailingZeroes + RightTrailingZeroes;
    const std::uint32_t SmallestOperand =
        std::min(LeftTrailingKnown - LeftTrailingZeroes, RightTrailingKnown - RightTrailingZeroes);
    const std::uint32_t ResultBitsKnown = std::min(SmallestOperand + TrailingZeroes, BitWidth);

    auto BottomKnown = multiplyMasks(LeftBits->One, RightBits->One);
    if (!BottomKnown)
      return std::unexpected(BottomKnown.error());
    auto LowMask = APInt::getLowBitsSet(BitWidth, ResultBitsKnown);
    if (!LowMask)
      return std::unexpected(LowMask.error());
    auto NotBottomKnown = BottomKnown->Value.bitwiseNot();
    if (!NotBottomKnown)
      return std::unexpected(NotBottomKnown.error());
    auto KnownZero = intersectMasks(*NotBottomKnown, *LowMask);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto HighZero = highBitsMask(BitWidth, BitWidth - LeadingZeroes);
    if (!HighZero)
      return std::unexpected(HighZero.error());
    KnownZero = mergeMasks(*KnownZero, *HighZero);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto KnownOne = intersectMasks(BottomKnown->Value, *LowMask);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  if (Opcode == GenericOpcode::G_SHL || Opcode == GenericOpcode::G_LSHR ||
      Opcode == GenericOpcode::G_ASHR) {
    const auto AmountReg = registerOperand(*Definition, 2);
    if (!AmountReg)
      return Unknown;
    const auto Amount = constantShiftAmount(*AmountReg, MRI, BitWidth);
    if (!Amount)
      return Unknown;
    auto SourceBits = getKnownBits(*Source, Depth + 1U);
    if (!SourceBits)
      return std::unexpected(SourceBits.error());

    if (Opcode == GenericOpcode::G_SHL) {
      auto KnownZero = shiftMaskLeft(SourceBits->Zero, *Amount);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      auto LowZero = APInt::getLowBitsSet(BitWidth, *Amount);
      if (!LowZero)
        return std::unexpected(LowZero.error());
      KnownZero = mergeMasks(*KnownZero, *LowZero);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
      auto KnownOne = shiftMaskLeft(SourceBits->One, *Amount);
      if (!KnownOne)
        return std::unexpected(KnownOne.error());
      return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
    }

    auto KnownZero = shiftMaskRight(SourceBits->Zero, *Amount);
    if (!KnownZero)
      return std::unexpected(KnownZero.error());
    auto KnownOne = shiftMaskRight(SourceBits->One, *Amount);
    if (!KnownOne)
      return std::unexpected(KnownOne.error());
    auto HighBits = highBitsMask(BitWidth, BitWidth - *Amount);
    if (!HighBits)
      return std::unexpected(HighBits.error());
    if (Opcode == GenericOpcode::G_LSHR || SourceBits->Zero.testBit(BitWidth - 1U)) {
      KnownZero = mergeMasks(*KnownZero, *HighBits);
      if (!KnownZero)
        return std::unexpected(KnownZero.error());
    } else if (SourceBits->One.testBit(BitWidth - 1U)) {
      KnownOne = mergeMasks(*KnownOne, *HighBits);
      if (!KnownOne)
        return std::unexpected(KnownOne.error());
    }
    return KnownBits{.Zero = std::move(*KnownZero), .One = std::move(*KnownOne)};
  }

  return Unknown;
}

Expected<APInt> GISelValueTracking::getKnownZeroes(Register Reg) const {
  auto Known = getKnownBits(Reg);
  if (!Known)
    return std::unexpected(Known.error());
  return std::move(Known->Zero);
}

Expected<std::uint32_t> GISelValueTracking::computeNumSignBits(Register Reg, unsigned Depth) const {
  auto Type = MRI.getType(Reg);
  if (!Type)
    return std::unexpected(Type.error());
  const std::uint32_t BitWidth = Type->scalarBitWidth();
  if (Depth >= MaxDepth)
    return 1U;

  const MachineInstr *Definition = MRI.getVRegDef(Reg);
  if (Definition == nullptr)
    return 1U;
  const auto Opcode = static_cast<GenericOpcode>(Definition->opcode().value());
  std::uint32_t FirstAnswer = 1U;

  if (Opcode == GenericOpcode::G_CONSTANT) {
    const MachineOperand *ValueOperand = Definition->operandAt(1);
    const auto *Immediate =
        ValueOperand == nullptr ? nullptr : ValueOperand->getIf<IntegerImmediateOperand>();
    if (Immediate == nullptr)
      return 1U;
    auto Value = constantAtWidth(Immediate->Value, BitWidth);
    if (!Value)
      return std::unexpected(Value.error());
    return constantNumSignBits(*Value);
  }

  const auto Source = registerOperand(*Definition, 1);
  if (Opcode == GenericOpcode::COPY && Source && Source->isVirtual())
    return computeNumSignBits(*Source, Depth);

  if (Opcode == GenericOpcode::G_SEXT && Source) {
    auto SourceType = MRI.getType(*Source);
    if (!SourceType)
      return std::unexpected(SourceType.error());
    auto SourceSignBits = computeNumSignBits(*Source, Depth + 1U);
    if (!SourceSignBits)
      return std::unexpected(SourceSignBits.error());
    return std::min(BitWidth, *SourceSignBits + BitWidth - SourceType->scalarBitWidth());
  }

  if ((Opcode == GenericOpcode::G_ASSERT_SEXT || Opcode == GenericOpcode::G_SEXT_INREG) && Source) {
    const MachineOperand *WidthOperand = Definition->operandAt(2);
    const auto *Width = WidthOperand == nullptr ? nullptr : WidthOperand->getIf<std::int64_t>();
    if (Width != nullptr && *Width > 0 && static_cast<std::uint64_t>(*Width) <= BitWidth) {
      auto SourceSignBits = computeNumSignBits(*Source, Depth + 1U);
      if (!SourceSignBits)
        return std::unexpected(SourceSignBits.error());
      return std::max(*SourceSignBits, BitWidth - static_cast<std::uint32_t>(*Width) + 1U);
    }
  }

  if ((Opcode == GenericOpcode::G_AND || Opcode == GenericOpcode::G_OR ||
       Opcode == GenericOpcode::G_XOR) &&
      Source) {
    const auto Right = registerOperand(*Definition, 2);
    if (Right) {
      auto LeftSignBits = computeNumSignBits(*Source, Depth + 1U);
      if (!LeftSignBits)
        return std::unexpected(LeftSignBits.error());
      if (*LeftSignBits != 1) {
        auto RightSignBits = computeNumSignBits(*Right, Depth + 1U);
        if (!RightSignBits)
          return std::unexpected(RightSignBits.error());
        if (*RightSignBits != 1)
          FirstAnswer = std::min(*LeftSignBits, *RightSignBits);
      }
    }
  }

  if (Opcode == GenericOpcode::G_ASHR && Source) {
    const auto Right = registerOperand(*Definition, 2);
    if (Right) {
      auto SourceSignBits = computeNumSignBits(*Source, Depth + 1U);
      if (!SourceSignBits)
        return std::unexpected(SourceSignBits.error());
      FirstAnswer = *SourceSignBits;
      if (const auto Amount = constantShiftAmount(*Right, MRI, BitWidth))
        FirstAnswer = std::min(BitWidth, FirstAnswer + *Amount);
    }
  }

  if (Opcode == GenericOpcode::G_TRUNC && Source) {
    auto SourceType = MRI.getType(*Source);
    if (!SourceType)
      return std::unexpected(SourceType.error());
    const std::uint32_t SourceWidth = SourceType->scalarBitWidth();
    auto SourceSignBits = computeNumSignBits(*Source, Depth + 1U);
    if (!SourceSignBits)
      return std::unexpected(SourceSignBits.error());
    if (SourceWidth > BitWidth && *SourceSignBits > SourceWidth - BitWidth)
      return *SourceSignBits - (SourceWidth - BitWidth);
  }

  if (Opcode == GenericOpcode::G_ICMP || Opcode == GenericOpcode::G_FCMP) {
    if (BitWidth > 1) {
      const BooleanContent Contents =
          Target.getBooleanContents(Type->isVector(), Opcode == GenericOpcode::G_FCMP);
      if (Contents == BooleanContent::ZeroOrNegativeOne)
        return BitWidth;
      if (Contents == BooleanContent::ZeroOrOne)
        return BitWidth - 1U;
    }
  }

  switch (Opcode) {
  case GenericOpcode::G_UADDO:
  case GenericOpcode::G_UADDE:
  case GenericOpcode::G_USUBO:
  case GenericOpcode::G_USUBE:
  case GenericOpcode::G_SADDO:
  case GenericOpcode::G_SADDE:
  case GenericOpcode::G_SSUBO:
  case GenericOpcode::G_SSUBE:
  case GenericOpcode::G_UMULO:
  case GenericOpcode::G_SMULO: {
    const auto BooleanResult = registerOperand(*Definition, 1);
    if (BooleanResult && *BooleanResult == Reg &&
        Target.getBooleanContents(Type->isVector(), false) == BooleanContent::ZeroOrNegativeOne)
      return BitWidth;
    break;
  }
  default:
    break;
  }

  auto Known = getKnownBits(Reg, Depth);
  if (!Known)
    return std::unexpected(Known.error());
  std::uint32_t KnownSignBits = 1;
  if (Known->Zero.testBit(BitWidth - 1U))
    KnownSignBits = static_cast<std::uint32_t>(Known->Zero.countlOne());
  else if (Known->One.testBit(BitWidth - 1U))
    KnownSignBits = static_cast<std::uint32_t>(Known->One.countlOne());
  return std::max(FirstAnswer, std::max(1U, KnownSignBits));
}

} // namespace rcc