mlir/test/lib/Dialect/Test/TestOpDefs.cpp

//===- TestOpDefs.cpp - MLIR Test Dialect Operation Hooks -----------------===//
//
// 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
//
//===----------------------------------------------------------------------===//

#include "TestDialect.h"
#include "TestOps.h"
#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/IR/Verifier.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/MemorySlotInterfaces.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/TypeSwitch.h"

using namespace mlir;
using namespace test;

//===----------------------------------------------------------------------===//
// OverridenSymbolVisibilityOp
//===----------------------------------------------------------------------===//

SymbolTable::Visibility OverriddenSymbolVisibilityOp::getVisibility() {
  return SymbolTable::Visibility::Private;
}

StringAttr OverriddenSymbolVisibilityOp::getNameAttr() {
  return getSymNameAttr();
}

static StringLiteral getVisibilityString(SymbolTable::Visibility visibility) {
  switch (visibility) {
  case SymbolTable::Visibility::Private:
    return "private";
  case SymbolTable::Visibility::Nested:
    return "nested";
  case SymbolTable::Visibility::Public:
    return "public";
  }
  llvm_unreachable("Unknown SymbolTable::Visibility");
}

void OverriddenSymbolVisibilityOp::setVisibility(
    SymbolTable::Visibility visibility) {

  emitOpError("cannot change visibility of symbol to ")
      << getVisibilityString(visibility);
}

//===----------------------------------------------------------------------===//
// TestBranchOp
//===----------------------------------------------------------------------===//

SuccessorOperands TestBranchOp::getSuccessorOperands(unsigned index) {
  assert(index == 0 && "invalid successor index");
  return SuccessorOperands(getTargetOperandsMutable());
}

//===----------------------------------------------------------------------===//
// TestProducingBranchOp
//===----------------------------------------------------------------------===//

SuccessorOperands TestProducingBranchOp::getSuccessorOperands(unsigned index) {
  assert(index <= 1 && "invalid successor index");
  if (index == 1)
    return SuccessorOperands(getFirstOperandsMutable());
  return SuccessorOperands(getSecondOperandsMutable());
}

//===----------------------------------------------------------------------===//
// TestInternalBranchOp
//===----------------------------------------------------------------------===//

SuccessorOperands TestInternalBranchOp::getSuccessorOperands(unsigned index) {
  assert(index <= 1 && "invalid successor index");
  if (index == 0)
    return SuccessorOperands(0, getSuccessOperandsMutable());
  return SuccessorOperands(1, getErrorOperandsMutable());
}

//===----------------------------------------------------------------------===//
// TestCallOp
//===----------------------------------------------------------------------===//

LogicalResult TestCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  // Check that the callee attribute was specified.
  auto fnAttr = getCalleeAttr();
  if (!fnAttr)
    return emitOpError("requires a 'callee' symbol reference attribute");
  if (!symbolTable.lookupNearestSymbolFrom<FunctionOpInterface>(*this, fnAttr))
    return emitOpError() << "'" << fnAttr.getValue()
                         << "' does not reference a valid function";
  return success();
}

//===----------------------------------------------------------------------===//
// FoldToCallOp
//===----------------------------------------------------------------------===//

namespace {
struct FoldToCallOpPattern : public OpRewritePattern<FoldToCallOp> {
  using OpRewritePattern<FoldToCallOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(FoldToCallOp op,
                                PatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<func::CallOp>(op, TypeRange(),
                                              op.getCalleeAttr(), ValueRange());
    return success();
  }
};
} // namespace

void FoldToCallOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.add<FoldToCallOpPattern>(context);
}

//===----------------------------------------------------------------------===//
// IsolatedRegionOp - test parsing passthrough operands
//===----------------------------------------------------------------------===//

ParseResult IsolatedRegionOp::parse(OpAsmParser &parser,
                                    OperationState &result) {
  // Parse the input operand.
  OpAsmParser::Argument argInfo;
  argInfo.type = parser.getBuilder().getIndexType();
  if (parser.parseOperand(argInfo.ssaName) ||
      parser.resolveOperand(argInfo.ssaName, argInfo.type, result.operands))
    return failure();

  // Parse the body region, and reuse the operand info as the argument info.
  Region *body = result.addRegion();
  return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/true);
}

void IsolatedRegionOp::print(OpAsmPrinter &p) {
  p << ' ';
  p.printOperand(getOperand());
  p.shadowRegionArgs(getRegion(), getOperand());
  p << ' ';
  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
}

//===----------------------------------------------------------------------===//
// SSACFGRegionOp
//===----------------------------------------------------------------------===//

RegionKind SSACFGRegionOp::getRegionKind(unsigned index) {
  return RegionKind::SSACFG;
}

//===----------------------------------------------------------------------===//
// GraphRegionOp
//===----------------------------------------------------------------------===//

RegionKind GraphRegionOp::getRegionKind(unsigned index) {
  return RegionKind::Graph;
}

//===----------------------------------------------------------------------===//
// IsolatedGraphRegionOp
//===----------------------------------------------------------------------===//

RegionKind IsolatedGraphRegionOp::getRegionKind(unsigned index) {
  return RegionKind::Graph;
}

//===----------------------------------------------------------------------===//
// AffineScopeOp
//===----------------------------------------------------------------------===//

ParseResult AffineScopeOp::parse(OpAsmParser &parser, OperationState &result) {
  // Parse the body region, and reuse the operand info as the argument info.
  Region *body = result.addRegion();
  return parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{});
}

void AffineScopeOp::print(OpAsmPrinter &p) {
  p << " ";
  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
}

//===----------------------------------------------------------------------===//
// TestRemoveOpWithInnerOps
//===----------------------------------------------------------------------===//

namespace {
struct TestRemoveOpWithInnerOps
    : public OpRewritePattern<TestOpWithRegionPattern> {
  using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;

  void initialize() { setDebugName("TestRemoveOpWithInnerOps"); }

  LogicalResult matchAndRewrite(TestOpWithRegionPattern op,
                                PatternRewriter &rewriter) const override {
    rewriter.eraseOp(op);
    return success();
  }
};
} // namespace

//===----------------------------------------------------------------------===//
// TestOpWithRegionPattern
//===----------------------------------------------------------------------===//

void TestOpWithRegionPattern::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<TestRemoveOpWithInnerOps>(context);
}

//===----------------------------------------------------------------------===//
// TestOpWithRegionFold
//===----------------------------------------------------------------------===//

OpFoldResult TestOpWithRegionFold::fold(FoldAdaptor adaptor) {
  return getOperand();
}

//===----------------------------------------------------------------------===//
// TestOpConstant
//===----------------------------------------------------------------------===//

OpFoldResult TestOpConstant::fold(FoldAdaptor adaptor) { return getValue(); }

//===----------------------------------------------------------------------===//
// TestOpWithVariadicResultsAndFolder
//===----------------------------------------------------------------------===//

LogicalResult TestOpWithVariadicResultsAndFolder::fold(
    FoldAdaptor adaptor, SmallVectorImpl<OpFoldResult> &results) {
  for (Value input : this->getOperands()) {
    results.push_back(input);
  }
  return success();
}

//===----------------------------------------------------------------------===//
// TestOpInPlaceFold
//===----------------------------------------------------------------------===//

OpFoldResult TestOpInPlaceFold::fold(FoldAdaptor adaptor) {
  // Exercise the fact that an operation created with createOrFold should be
  // allowed to access its parent block.
  assert(getOperation()->getBlock() &&
         "expected that operation is not unlinked");

  if (adaptor.getOp() && !getProperties().attr) {
    // The folder adds "attr" if not present.
    getProperties().attr = dyn_cast_or_null<IntegerAttr>(adaptor.getOp());
    return getResult();
  }
  return {};
}

//===----------------------------------------------------------------------===//
// OpWithInferTypeInterfaceOp
//===----------------------------------------------------------------------===//

LogicalResult OpWithInferTypeInterfaceOp::inferReturnTypes(
    MLIRContext *, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  if (operands[0].getType() != operands[1].getType()) {
    return emitOptionalError(location, "operand type mismatch ",
                             operands[0].getType(), " vs ",
                             operands[1].getType());
  }
  inferredReturnTypes.assign({operands[0].getType()});
  return success();
}

//===----------------------------------------------------------------------===//
// OpWithShapedTypeInferTypeInterfaceOp
//===----------------------------------------------------------------------===//

LogicalResult OpWithShapedTypeInferTypeInterfaceOp::inferReturnTypeComponents(
    MLIRContext *context, std::optional<Location> location,
    ValueShapeRange operands, DictionaryAttr attributes, PropertyRef properties,
    RegionRange regions,
    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
  // Create return type consisting of the last element of the first operand.
  auto operandType = operands.front().getType();
  auto sval = dyn_cast<ShapedType>(operandType);
  if (!sval)
    return emitOptionalError(location, "only shaped type operands allowed");
  int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamic;
  auto type = IntegerType::get(context, 17);

  Attribute encoding;
  if (auto rankedTy = dyn_cast<RankedTensorType>(sval))
    encoding = rankedTy.getEncoding();
  inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));
  return success();
}

LogicalResult OpWithShapedTypeInferTypeInterfaceOp::reifyReturnTypeShapes(
    OpBuilder &builder, ValueRange operands,
    llvm::SmallVectorImpl<Value> &shapes) {
  shapes = SmallVector<Value, 1>{
      builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
  return success();
}

//===----------------------------------------------------------------------===//
// OpWithResultShapeInterfaceOp
//===----------------------------------------------------------------------===//

LogicalResult OpWithResultShapeInterfaceOp::reifyReturnTypeShapes(
    OpBuilder &builder, ValueRange operands,
    llvm::SmallVectorImpl<Value> &shapes) {
  Location loc = getLoc();
  shapes.reserve(operands.size());
  for (Value operand : llvm::reverse(operands)) {
    auto rank = cast<RankedTensorType>(operand.getType()).getRank();
    auto currShape = llvm::map_to_vector<4>(
        llvm::seq<int64_t>(0, rank), [&](int64_t dim) -> Value {
          return builder.createOrFold<tensor::DimOp>(loc, operand, dim);
        });
    shapes.push_back(tensor::FromElementsOp::create(
        builder, getLoc(),
        RankedTensorType::get({rank}, builder.getIndexType()), currShape));
  }
  return success();
}

//===----------------------------------------------------------------------===//
// ReifyShapedTypeUsingReifyResultShapesOp
//===----------------------------------------------------------------------===//

LogicalResult ReifyShapedTypeUsingReifyResultShapesOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  Location loc = getLoc();
  shapes.reserve(getNumOperands());
  for (Value operand : llvm::reverse(getOperands())) {
    auto tensorType = cast<RankedTensorType>(operand.getType());
    auto currShape = llvm::map_to_vector<4>(
        llvm::seq<int64_t>(0, tensorType.getRank()),
        [&](int64_t dim) -> OpFoldResult {
          return tensorType.isDynamicDim(dim)
                     ? static_cast<OpFoldResult>(
                           builder.createOrFold<tensor::DimOp>(loc, operand,
                                                               dim))
                     : static_cast<OpFoldResult>(
                           builder.getIndexAttr(tensorType.getDimSize(dim)));
        });
    shapes.emplace_back(std::move(currShape));
  }
  return success();
}

//===----------------------------------------------------------------------===//
// ReifyShapedTypeUsingReifyShapeOfResultOp
//===----------------------------------------------------------------------===//

LogicalResult ReifyShapedTypeUsingReifyShapeOfResultOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  return failure();
}

FailureOr<SmallVector<OpFoldResult>>
ReifyShapedTypeUsingReifyShapeOfResultOp::reifyShapeOfResult(OpBuilder &builder,
                                                             int resultIndex) {
  Location loc = getLoc();
  Value sourceOperand = getOperand(getNumOperands() - 1 - resultIndex);
  SmallVector<OpFoldResult> shape =
      tensor::getMixedSizes(builder, loc, sourceOperand);
  return shape;
}

//===----------------------------------------------------------------------===//
// ReifyShapedTypeUsingReifyDimOfResultOp
//===----------------------------------------------------------------------===//

LogicalResult ReifyShapedTypeUsingReifyDimOfResultOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  return failure();
}

FailureOr<SmallVector<OpFoldResult>>
ReifyShapedTypeUsingReifyDimOfResultOp::reifyShapeOfResult(OpBuilder &builder,
                                                           int resultIndex) {
  return failure();
}

FailureOr<OpFoldResult>
ReifyShapedTypeUsingReifyDimOfResultOp::reifyDimOfResult(OpBuilder &builder,
                                                         int resultIndex,
                                                         int dim) {
  Location loc = getLoc();
  Value sourceOperand = getOperand(getNumOperands() - 1 - resultIndex);
  OpFoldResult shape = tensor::getMixedSize(builder, loc, sourceOperand, dim);
  return shape;
}

//===----------------------------------------------------------------------===//
// UnreifableResultShapesOp
//===----------------------------------------------------------------------===//

LogicalResult UnreifiableResultShapesOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  Location loc = getLoc();
  shapes.resize(1);
  shapes[0] = {tensor::getMixedSize(builder, loc, getOperand(), 0),
               OpFoldResult()};
  return success();
}

//===----------------------------------------------------------------------===//
// UnreifableResultShapeOp
//===----------------------------------------------------------------------===//

LogicalResult UnreifiableResultShapeOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  return failure();
}

FailureOr<SmallVector<OpFoldResult>>
UnreifiableResultShapeOp::reifyShapeOfResult(OpBuilder &builder,
                                             int resultIndex) {
  SmallVector<OpFoldResult> shape = {
      tensor::getMixedSize(builder, getLoc(), getOperand(), 0), OpFoldResult()};
  return shape;
}

//===----------------------------------------------------------------------===//
// UnreifableResultShapeOp
//===----------------------------------------------------------------------===//

LogicalResult UnreifiableDimOfResultShapeOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &shapes) {
  return failure();
}

FailureOr<SmallVector<OpFoldResult>>
UnreifiableDimOfResultShapeOp::reifyShapeOfResult(OpBuilder &builder,
                                                  int resultIndex) {
  return failure();
}

FailureOr<OpFoldResult>
UnreifiableDimOfResultShapeOp::reifyDimOfResult(OpBuilder &builder,
                                                int resultIndex, int dim) {
  if (dim == 0)
    return tensor::getMixedSize(builder, getLoc(), getOperand(), 0);
  return failure();
}

//===----------------------------------------------------------------------===//
// SideEffectOp
//===----------------------------------------------------------------------===//

namespace {
/// A test resource for side effects.
struct TestResource : public SideEffects::Resource::Base<TestResource> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestResource)

  StringRef getName() const final { return "<Test>"; }
};
} // namespace

void SideEffectOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  // Check for an effects attribute on the op instance.
  ArrayAttr effectsAttr =
      (*this)->getDiscardableAttrOfType<ArrayAttr>("effects");
  if (!effectsAttr)
    return;

  for (Attribute element : effectsAttr) {
    DictionaryAttr effectElement = cast<DictionaryAttr>(element);

    // Get the specific memory effect.
    MemoryEffects::Effect *effect =
        StringSwitch<MemoryEffects::Effect *>(
            cast<StringAttr>(effectElement.get("effect")).getValue())
            .Case("allocate", MemoryEffects::Allocate::get())
            .Case("free", MemoryEffects::Free::get())
            .Case("read", MemoryEffects::Read::get())
            .Case("write", MemoryEffects::Write::get());

    // Check for a non-default resource to use.
    SideEffects::Resource *resource = SideEffects::DefaultResource::get();
    if (effectElement.get("test_resource"))
      resource = TestResource::get();
    else if (effectElement.get("test_nonaddressable_resource"))
      resource = TestNonAddressableResource::get();
    else if (effectElement.get("test_nonaddressable_resource_a"))
      resource = TestNonAddressableResourceA::get();
    else if (effectElement.get("test_nonaddressable_resource_b"))
      resource = TestNonAddressableResourceB::get();

    // Check for a result to affect.
    if (effectElement.get("on_result"))
      effects.emplace_back(effect, getOperation()->getOpResults()[0], resource);
    else if (Attribute ref = effectElement.get("on_reference"))
      effects.emplace_back(effect, cast<SymbolRefAttr>(ref), resource);
    else
      effects.emplace_back(effect, resource);
  }
}

void SideEffectOp::getEffects(
    SmallVectorImpl<TestEffects::EffectInstance> &effects) {
  testSideEffectOpGetEffect(getOperation(), effects);
}

//===----------------------------------------------------------------------===//
// ConditionalSideEffectOp
//===----------------------------------------------------------------------===//

void ConditionalSideEffectOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  if (!getHasEffects())
    return;

  SideEffects::Resource *resource = SideEffects::DefaultResource::get();
  effects.emplace_back(MemoryEffects::Read::get(), resource);
  effects.emplace_back(MemoryEffects::Write::get(), resource);
  effects.emplace_back(MemoryEffects::Free::get(), resource);
  effects.emplace_back(MemoryEffects::Allocate::get(), resource);
}

void SideEffectWithRegionOp::getEffects(
    SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
  // Check for an effects attribute on the op instance.
  ArrayAttr effectsAttr =
      (*this)->getDiscardableAttrOfType<ArrayAttr>("effects");
  if (!effectsAttr)
    return;

  for (Attribute element : effectsAttr) {
    DictionaryAttr effectElement = cast<DictionaryAttr>(element);

    // Get the specific memory effect.
    MemoryEffects::Effect *effect =
        StringSwitch<MemoryEffects::Effect *>(
            cast<StringAttr>(effectElement.get("effect")).getValue())
            .Case("allocate", MemoryEffects::Allocate::get())
            .Case("free", MemoryEffects::Free::get())
            .Case("read", MemoryEffects::Read::get())
            .Case("write", MemoryEffects::Write::get());

    // Check for a non-default resource to use.
    SideEffects::Resource *resource = SideEffects::DefaultResource::get();
    if (effectElement.get("test_resource"))
      resource = TestResource::get();
    else if (effectElement.get("test_nonaddressable_resource"))
      resource = TestNonAddressableResource::get();
    else if (effectElement.get("test_nonaddressable_resource_a"))
      resource = TestNonAddressableResourceA::get();
    else if (effectElement.get("test_nonaddressable_resource_b"))
      resource = TestNonAddressableResourceB::get();

    // Check for a result to affect.
    if (effectElement.get("on_result"))
      effects.emplace_back(effect, getOperation()->getOpResults()[0], resource);
    else if (effectElement.get("on_operand"))
      effects.emplace_back(effect, &getOperation()->getOpOperands()[0],
                           resource);
    else if (effectElement.get("on_argument"))
      effects.emplace_back(effect, getOperation()->getRegion(0).getArgument(0),
                           resource);
    else if (Attribute ref = effectElement.get("on_reference"))
      effects.emplace_back(effect, cast<SymbolRefAttr>(ref), resource);
    else
      effects.emplace_back(effect, resource);
  }
}

void SideEffectWithRegionOp::getEffects(
    SmallVectorImpl<TestEffects::EffectInstance> &effects) {
  testSideEffectOpGetEffect(getOperation(), effects);
}

//===----------------------------------------------------------------------===//
// StringAttrPrettyNameOp
//===----------------------------------------------------------------------===//

// This op has fancy handling of its SSA result name.
ParseResult StringAttrPrettyNameOp::parse(OpAsmParser &parser,
                                          OperationState &result) {
  // Add the result types.
  for (size_t i = 0, e = parser.getNumResults(); i != e; ++i)
    result.addTypes(parser.getBuilder().getIntegerType(32));

  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
    return failure();

  // If the attribute dictionary contains no 'names' attribute, infer it from
  // the SSA name (if specified).
  bool hadNames = llvm::any_of(result.attributes, [](NamedAttribute attr) {
    return attr.getName() == "names";
  });

  // If there was no name specified, check to see if there was a useful name
  // specified in the asm file.
  if (hadNames || parser.getNumResults() == 0)
    return success();

  SmallVector<StringRef, 4> names;
  auto *context = result.getContext();

  for (size_t i = 0, e = parser.getNumResults(); i != e; ++i) {
    auto resultName = parser.getResultName(i);
    StringRef nameStr;
    if (!resultName.first.empty() && !isdigit(resultName.first[0]))
      nameStr = resultName.first;

    names.push_back(nameStr);
  }

  auto namesAttr = parser.getBuilder().getStrArrayAttr(names);
  result.attributes.push_back({StringAttr::get(context, "names"), namesAttr});
  return success();
}

void StringAttrPrettyNameOp::print(OpAsmPrinter &p) {
  // Note that we only need to print the "name" attribute if the asmprinter
  // result name disagrees with it.  This can happen in strange cases, e.g.
  // when there are conflicts.
  bool namesDisagree = getNames().size() != getNumResults();

  SmallString<32> resultNameStr;
  for (size_t i = 0, e = getNumResults(); i != e && !namesDisagree; ++i) {
    resultNameStr.clear();
    llvm::raw_svector_ostream tmpStream(resultNameStr);
    p.printOperand(getResult(i), tmpStream);

    auto expectedName = dyn_cast<StringAttr>(getNames()[i]);
    if (!expectedName ||
        tmpStream.str().drop_front() != expectedName.getValue()) {
      namesDisagree = true;
    }
  }

  if (namesDisagree) {
    SmallVector<NamedAttribute> attrs((*this)->getDiscardableAttrs());
    attrs.emplace_back(getNamesAttrName(), getNamesAttr());
    llvm::sort(attrs);
    p.printOptionalAttrDictWithKeyword(attrs);
  } else {
    p.printOptionalAttrDictWithKeyword(
        (*this)->getDiscardableAttrDictionary().getValue(), {"names"});
  }
}

// We set the SSA name in the asm syntax to the contents of the name
// attribute.
void StringAttrPrettyNameOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {

  auto value = getNames();
  for (size_t i = 0, e = value.size(); i != e; ++i)
    if (auto str = dyn_cast<StringAttr>(value[i]))
      if (!str.getValue().empty())
        setNameFn(getResult(i), str.getValue());
}

//===----------------------------------------------------------------------===//
// CustomResultsNameOp
//===----------------------------------------------------------------------===//

void CustomResultsNameOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  ArrayAttr value = getNames();
  for (size_t i = 0, e = value.size(); i != e; ++i)
    if (auto str = dyn_cast<StringAttr>(value[i]))
      if (!str.empty())
        setNameFn(getResult(i), str.getValue());
}

//===----------------------------------------------------------------------===//
// ResultNameFromTypeOp
//===----------------------------------------------------------------------===//

void ResultNameFromTypeOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  auto result = getResult();
  auto setResultNameFn = [&](::llvm::StringRef name) {
    setNameFn(result, name);
  };
  auto opAsmTypeInterface =
      ::mlir::cast<::mlir::OpAsmTypeInterface>(result.getType());
  opAsmTypeInterface.getAsmName(setResultNameFn);
}

//===----------------------------------------------------------------------===//
// BlockArgumentNameFromTypeOp
//===----------------------------------------------------------------------===//

void BlockArgumentNameFromTypeOp::getAsmBlockArgumentNames(
    ::mlir::Region &region, ::mlir::OpAsmSetValueNameFn setNameFn) {
  for (auto &block : region) {
    for (auto arg : block.getArguments()) {
      if (auto opAsmTypeInterface =
              ::mlir::dyn_cast<::mlir::OpAsmTypeInterface>(arg.getType())) {
        auto setArgNameFn = [&](StringRef name) { setNameFn(arg, name); };
        opAsmTypeInterface.getAsmName(setArgNameFn);
      }
    }
  }
}

//===----------------------------------------------------------------------===//
// ResultTypeWithTraitOp
//===----------------------------------------------------------------------===//

LogicalResult ResultTypeWithTraitOp::verify() {
  if ((*this)->getResultTypes()[0].hasTrait<TypeTrait::TestTypeTrait>())
    return success();
  return emitError("result type should have trait 'TestTypeTrait'");
}

//===----------------------------------------------------------------------===//
// AttrWithTraitOp
//===----------------------------------------------------------------------===//

LogicalResult AttrWithTraitOp::verify() {
  if (getAttr().hasTrait<AttributeTrait::TestAttrTrait>())
    return success();
  return emitError("'attr' attribute should have trait 'TestAttrTrait'");
}

//===----------------------------------------------------------------------===//
// RegionIfOp
//===----------------------------------------------------------------------===//

void RegionIfOp::print(OpAsmPrinter &p) {
  p << " ";
  p.printOperands(getOperands());
  p << ": " << getOperandTypes();
  p.printArrowTypeList(getResultTypes());
  p << " then ";
  p.printRegion(getThenRegion(),
                /*printEntryBlockArgs=*/true,
                /*printBlockTerminators=*/true);
  p << " else ";
  p.printRegion(getElseRegion(),
                /*printEntryBlockArgs=*/true,
                /*printBlockTerminators=*/true);
  p << " join ";
  p.printRegion(getJoinRegion(),
                /*printEntryBlockArgs=*/true,
                /*printBlockTerminators=*/true);
}

ParseResult RegionIfOp::parse(OpAsmParser &parser, OperationState &result) {
  SmallVector<OpAsmParser::UnresolvedOperand, 2> operandInfos;
  SmallVector<Type, 2> operandTypes;

  result.regions.reserve(3);
  Region *thenRegion = result.addRegion();
  Region *elseRegion = result.addRegion();
  Region *joinRegion = result.addRegion();

  // Parse operand, type and arrow type lists.
  if (parser.parseOperandList(operandInfos) ||
      parser.parseColonTypeList(operandTypes) ||
      parser.parseArrowTypeList(result.types))
    return failure();

  // Parse all attached regions.
  if (parser.parseKeyword("then") || parser.parseRegion(*thenRegion, {}, {}) ||
      parser.parseKeyword("else") || parser.parseRegion(*elseRegion, {}, {}) ||
      parser.parseKeyword("join") || parser.parseRegion(*joinRegion, {}, {}))
    return failure();

  return parser.resolveOperands(operandInfos, operandTypes,
                                parser.getCurrentLocation(), result.operands);
}

OperandRange RegionIfOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  assert(llvm::is_contained({&getThenRegion(), &getElseRegion()},
                            successor.getSuccessor()) &&
         "invalid region index");
  return getOperands();
}

void RegionIfOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // We always branch to the join region.
  if (!point.isParent()) {
    if (point.getTerminatorPredecessorOrNull()->getParentRegion() !=
        &getJoinRegion())
      regions.push_back(RegionSuccessor(&getJoinRegion()));
    else
      regions.push_back(RegionSuccessor(getOperation()));
    return;
  }

  // The then and else regions are the entry regions of this op.
  regions.push_back(RegionSuccessor(&getThenRegion()));
  regions.push_back(RegionSuccessor(&getElseRegion()));
}

ValueRange RegionIfOp::getSuccessorInputs(RegionSuccessor successor) {
  if (successor.isOperation())
    return getResults();
  if (successor == &getThenRegion())
    return getThenArgs();
  if (successor == &getElseRegion())
    return getElseArgs();
  if (successor == &getJoinRegion())
    return getJoinArgs();
  llvm_unreachable("invalid region successor");
}

void RegionIfOp::getRegionInvocationBounds(
    ArrayRef<Attribute> operands,
    SmallVectorImpl<InvocationBounds> &invocationBounds) {
  // Each region is invoked at most once.
  invocationBounds.assign(/*NumElts=*/3, /*Elt=*/{0, 1});
}

//===----------------------------------------------------------------------===//
// AnyCondOp
//===----------------------------------------------------------------------===//

void AnyCondOp::getSuccessorRegions(RegionBranchPoint point,
                                    SmallVectorImpl<RegionSuccessor> &regions) {
  // The parent op branches into the only region, and the region branches back
  // to the parent op.
  if (point.isParent())
    regions.emplace_back(&getRegion());
  else
    regions.push_back(RegionSuccessor(getOperation()));
}

ValueRange AnyCondOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getResults()) : ValueRange();
}

void AnyCondOp::getRegionInvocationBounds(
    ArrayRef<Attribute> operands,
    SmallVectorImpl<InvocationBounds> &invocationBounds) {
  invocationBounds.emplace_back(1, 1);
}

//===----------------------------------------------------------------------===//
// SingleBlockImplicitTerminatorOp
//===----------------------------------------------------------------------===//

/// Testing the correctness of some traits.
static_assert(
    llvm::is_detected<OpTrait::has_implicit_terminator_t,
                      SingleBlockImplicitTerminatorOp>::value,
    "has_implicit_terminator_t does not match SingleBlockImplicitTerminatorOp");
static_assert(OpTrait::hasSingleBlockImplicitTerminator<
                  SingleBlockImplicitTerminatorOp>::value,
              "hasSingleBlockImplicitTerminator does not match "
              "SingleBlockImplicitTerminatorOp");

//===----------------------------------------------------------------------===//
// SingleNoTerminatorCustomAsmOp
//===----------------------------------------------------------------------===//

ParseResult SingleNoTerminatorCustomAsmOp::parse(OpAsmParser &parser,
                                                 OperationState &state) {
  Region *body = state.addRegion();
  if (parser.parseRegion(*body, /*arguments=*/{}, /*argTypes=*/{}))
    return failure();
  return success();
}

void SingleNoTerminatorCustomAsmOp::print(OpAsmPrinter &printer) {
  printer.printRegion(
      getRegion(), /*printEntryBlockArgs=*/false,
      // This op has a single block without terminators. But explicitly mark
      // as not printing block terminators for testing.
      /*printBlockTerminators=*/false);
}

//===----------------------------------------------------------------------===//
// TestVerifiersOp
//===----------------------------------------------------------------------===//

LogicalResult TestVerifiersOp::verify() {
  if (!getRegion().hasOneBlock())
    return emitOpError("`hasOneBlock` trait hasn't been verified");

  Operation *definingOp = getInput().getDefiningOp();
  if (definingOp && failed(mlir::verify(definingOp)))
    return emitOpError("operand hasn't been verified");

  // Avoid using `emitRemark(msg)` since that will trigger an infinite verifier
  // loop.
  mlir::emitRemark(getLoc(), "success run of verifier");

  return success();
}

LogicalResult TestVerifiersOp::verifyRegions() {
  if (!getRegion().hasOneBlock())
    return emitOpError("`hasOneBlock` trait hasn't been verified");

  for (Block &block : getRegion())
    for (Operation &op : block)
      if (failed(mlir::verify(&op)))
        return emitOpError("nested op hasn't been verified");

  // Avoid using `emitRemark(msg)` since that will trigger an infinite verifier
  // loop.
  mlir::emitRemark(getLoc(), "success run of region verifier");

  return success();
}

//===----------------------------------------------------------------------===//
// Test InferIntRangeInterface
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// TestWithBoundsOp
//===----------------------------------------------------------------------===//

LogicalResult TestWithBoundsOp::verify() {
  Type type = getElementTypeOrSelf(getResult().getType());
  unsigned expectedWidth = 0;
  if (type.isIndex())
    expectedWidth = IndexType::kInternalStorageBitWidth;
  else if (auto intTy = llvm::dyn_cast<IntegerType>(type))
    expectedWidth = intTy.getWidth();
  if (expectedWidth != 0 && getUmin().getBitWidth() != expectedWidth)
    return emitOpError("bound attribute width (")
           << getUmin().getBitWidth() << ") does not match result type width ("
           << expectedWidth << ")";
  return success();
}

void TestWithBoundsOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                         SetIntRangeFn setResultRanges) {
  setResultRanges(getResult(), {getUmin(), getUmax(), getSmin(), getSmax()});
}

//===----------------------------------------------------------------------===//
// TestWithoutBoundsOp
//===----------------------------------------------------------------------===//

void TestWithoutBoundsOp::inferResultRangesFromOptional(
    ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRanges) {
  // Mimic ops with uninitialized range.
  setResultRanges(getResult(), IntegerValueRange{});
}

//===----------------------------------------------------------------------===//
// TestWithBoundsRegionOp
//===----------------------------------------------------------------------===//

ParseResult TestWithBoundsRegionOp::parse(OpAsmParser &parser,
                                          OperationState &result) {
  if (parser.parseOptionalAttrDict(result.attributes))
    return failure();

  // Parse the input argument
  OpAsmParser::Argument argInfo;
  if (failed(parser.parseArgument(argInfo, true)))
    return failure();

  // Parse the body region, and reuse the operand info as the argument info.
  Region *body = result.addRegion();
  return parser.parseRegion(*body, argInfo, /*enableNameShadowing=*/false);
}

void TestWithBoundsRegionOp::print(OpAsmPrinter &p) {
  SmallVector<NamedAttribute> attrs((*this)->getDiscardableAttrs());
  attrs.emplace_back(getUminAttrName(), getUminAttr());
  attrs.emplace_back(getUmaxAttrName(), getUmaxAttr());
  attrs.emplace_back(getSminAttrName(), getSminAttr());
  attrs.emplace_back(getSmaxAttrName(), getSmaxAttr());
  llvm::sort(attrs);
  p.printOptionalAttrDict(attrs);
  p << ' ';
  p.printRegionArgument(getRegion().getArgument(0), /*argAttrs=*/{},
                        /*omitType=*/false);
  p << ' ';
  p.printRegion(getRegion(), /*printEntryBlockArgs=*/false);
}

void TestWithBoundsRegionOp::inferResultRanges(
    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
  Value arg = getRegion().getArgument(0);
  setResultRanges(arg, {getUmin(), getUmax(), getSmin(), getSmax()});
}

//===----------------------------------------------------------------------===//
// TestIncrementOp
//===----------------------------------------------------------------------===//

void TestIncrementOp::inferResultRanges(ArrayRef<ConstantIntRanges> argRanges,
                                        SetIntRangeFn setResultRanges) {
  const ConstantIntRanges &range = argRanges[0];
  APInt one(range.umin().getBitWidth(), 1);
  setResultRanges(getResult(),
                  {range.umin().uadd_sat(one), range.umax().uadd_sat(one),
                   range.smin().sadd_sat(one), range.smax().sadd_sat(one)});
}

//===----------------------------------------------------------------------===//
// TestReflectBoundsOp
//===----------------------------------------------------------------------===//

void TestReflectBoundsOp::inferResultRanges(
    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
  const ConstantIntRanges &range = argRanges[0];
  MLIRContext *ctx = getContext();
  Builder b(ctx);
  Type sIntTy, uIntTy;
  // For plain `IntegerType`s, we can derive the appropriate signed and unsigned
  // Types for the Attributes.
  Type type = getElementTypeOrSelf(getType());
  if (auto intTy = llvm::dyn_cast<IntegerType>(type)) {
    unsigned bitwidth = intTy.getWidth();
    sIntTy = b.getIntegerType(bitwidth, /*isSigned=*/true);
    uIntTy = b.getIntegerType(bitwidth, /*isSigned=*/false);
  } else {
    sIntTy = uIntTy = type;
  }

  setUminAttr(b.getIntegerAttr(uIntTy, range.umin()));
  setUmaxAttr(b.getIntegerAttr(uIntTy, range.umax()));
  setSminAttr(b.getIntegerAttr(sIntTy, range.smin()));
  setSmaxAttr(b.getIntegerAttr(sIntTy, range.smax()));
  setResultRanges(getResult(), range);
}

//===----------------------------------------------------------------------===//
// ConversionFuncOp
//===----------------------------------------------------------------------===//

ParseResult ConversionFuncOp::parse(OpAsmParser &parser,
                                    OperationState &result) {
  auto buildFuncType =
      [](Builder &builder, ArrayRef<Type> argTypes, ArrayRef<Type> results,
         function_interface_impl::VariadicFlag,
         std::string &) { return builder.getFunctionType(argTypes, results); };

  return function_interface_impl::parseFunctionOp(
      parser, result, /*allowVariadic=*/false,
      getFunctionTypeAttrName(result.name), buildFuncType,
      getArgAttrsAttrName(result.name), getResAttrsAttrName(result.name));
}

void ConversionFuncOp::print(OpAsmPrinter &p) {
  function_interface_impl::printFunctionOp(
      p, *this, /*isVariadic=*/false, getFunctionTypeAttrName(),
      getArgAttrsAttrName(), getResAttrsAttrName());
}

//===----------------------------------------------------------------------===//
// TestValueWithBoundsOp
//===----------------------------------------------------------------------===//

void TestValueWithBoundsOp::populateBoundsForIndexValue(
    Value v, ValueBoundsConstraintSet &cstr) {
  cstr.bound(v) >= getMin().getSExtValue();
  cstr.bound(v) <= getMax().getSExtValue();
}

//===----------------------------------------------------------------------===//
// ReifyBoundOp
//===----------------------------------------------------------------------===//

mlir::presburger::BoundType ReifyBoundOp::getBoundType() {
  if (getType() == "EQ")
    return mlir::presburger::BoundType::EQ;
  if (getType() == "LB")
    return mlir::presburger::BoundType::LB;
  if (getType() == "UB")
    return mlir::presburger::BoundType::UB;
  llvm_unreachable("invalid bound type");
}

LogicalResult ReifyBoundOp::verify() {
  if (getType() != "EQ" && getType() != "LB" && getType() != "UB")
    return emitOpError("invalid bound type '")
           << getType() << "', expected 'EQ', 'LB', or 'UB'";
  if (isa<ShapedType>(getVar().getType())) {
    if (!getDim().has_value())
      return emitOpError("expected 'dim' attribute for shaped type variable");
  } else if (getVar().getType().isIndex() || getVar().getType().isInteger()) {
    if (getVar().getType().isInteger() && !getAllowIntegerType())
      return emitOpError("integer variable requires 'allow_integer_type'");
    if (getDim().has_value())
      return emitOpError(
          "unexpected 'dim' attribute for index/integer variable");
  } else {
    return emitOpError(
        "expected index-typed/integer-typed variable or shape type variable");
  }
  if (getConstant() && getScalable())
    return emitOpError("'scalable' and 'constant' are mutually exlusive");
  if (getScalable() != getVscaleMin().has_value())
    return emitOpError("expected 'vscale_min' if and only if 'scalable'");
  if (getScalable() != getVscaleMax().has_value())
    return emitOpError("expected 'vscale_min' if and only if 'scalable'");
  return success();
}

ValueBoundsConstraintSet::Variable ReifyBoundOp::getVariable() {
  if (getDim().has_value())
    return ValueBoundsConstraintSet::Variable(getVar(), *getDim());
  return ValueBoundsConstraintSet::Variable(getVar());
}

//===----------------------------------------------------------------------===//
// CompareOp
//===----------------------------------------------------------------------===//

ValueBoundsConstraintSet::ComparisonOperator
CompareOp::getComparisonOperator() {
  if (getCmp() == "EQ")
    return ValueBoundsConstraintSet::ComparisonOperator::EQ;
  if (getCmp() == "LT")
    return ValueBoundsConstraintSet::ComparisonOperator::LT;
  if (getCmp() == "LE")
    return ValueBoundsConstraintSet::ComparisonOperator::LE;
  if (getCmp() == "GT")
    return ValueBoundsConstraintSet::ComparisonOperator::GT;
  if (getCmp() == "GE")
    return ValueBoundsConstraintSet::ComparisonOperator::GE;
  llvm_unreachable("invalid comparison operator");
}

mlir::ValueBoundsConstraintSet::Variable CompareOp::getLhs() {
  if (!getLhsMap())
    return ValueBoundsConstraintSet::Variable(getVarOperands()[0]);
  SmallVector<Value> mapOperands(
      getVarOperands().slice(0, getLhsMap()->getNumInputs()));
  return ValueBoundsConstraintSet::Variable(*getLhsMap(), mapOperands);
}

mlir::ValueBoundsConstraintSet::Variable CompareOp::getRhs() {
  int64_t rhsOperandsBegin = getLhsMap() ? getLhsMap()->getNumInputs() : 1;
  if (!getRhsMap())
    return ValueBoundsConstraintSet::Variable(
        getVarOperands()[rhsOperandsBegin]);
  SmallVector<Value> mapOperands(
      getVarOperands().slice(rhsOperandsBegin, getRhsMap()->getNumInputs()));
  return ValueBoundsConstraintSet::Variable(*getRhsMap(), mapOperands);
}

LogicalResult CompareOp::verify() {
  if (getCompose() && (getLhsMap() || getRhsMap()))
    return emitOpError(
        "'compose' not supported when 'lhs_map' or 'rhs_map' is present");
  int64_t expectedNumOperands = getLhsMap() ? getLhsMap()->getNumInputs() : 1;
  expectedNumOperands += getRhsMap() ? getRhsMap()->getNumInputs() : 1;
  if (getVarOperands().size() != size_t(expectedNumOperands))
    return emitOpError("expected ")
           << expectedNumOperands << " operands, but got "
           << getVarOperands().size();
  return success();
}

//===----------------------------------------------------------------------===//
// TestOpInPlaceSelfFold
//===----------------------------------------------------------------------===//

OpFoldResult TestOpInPlaceSelfFold::fold(FoldAdaptor adaptor) {
  if (!getFolded()) {
    // The folder adds the "folded" if not present.
    setFolded(true);
    return getResult();
  }
  return {};
}

//===----------------------------------------------------------------------===//
// TestOpFoldWithFoldAdaptor
//===----------------------------------------------------------------------===//

OpFoldResult TestOpFoldWithFoldAdaptor::fold(FoldAdaptor adaptor) {
  int64_t sum = 0;
  if (auto value = dyn_cast_or_null<IntegerAttr>(adaptor.getOp()))
    sum += value.getValue().getSExtValue();

  for (Attribute attr : adaptor.getVariadic())
    if (auto value = dyn_cast_or_null<IntegerAttr>(attr))
      sum += 2 * value.getValue().getSExtValue();

  for (ArrayRef<Attribute> attrs : adaptor.getVarOfVar())
    for (Attribute attr : attrs)
      if (auto value = dyn_cast_or_null<IntegerAttr>(attr))
        sum += 3 * value.getValue().getSExtValue();

  sum += 4 * std::distance(adaptor.getBody().begin(), adaptor.getBody().end());

  return IntegerAttr::get(getType(), sum);
}

//===----------------------------------------------------------------------===//
// OpWithInferTypeAdaptorInterfaceOp
//===----------------------------------------------------------------------===//

LogicalResult OpWithInferTypeAdaptorInterfaceOp::inferReturnTypes(
    MLIRContext *, std::optional<Location> location,
    OpWithInferTypeAdaptorInterfaceOp::Adaptor adaptor,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  if (adaptor.getX().getType() != adaptor.getY().getType()) {
    return emitOptionalError(location, "operand type mismatch ",
                             adaptor.getX().getType(), " vs ",
                             adaptor.getY().getType());
  }
  inferredReturnTypes.assign({adaptor.getX().getType()});
  return success();
}

//===----------------------------------------------------------------------===//
// OpWithRefineTypeInterfaceOp
//===----------------------------------------------------------------------===//

// TODO: We should be able to only define either inferReturnType or
// refineReturnType, currently only refineReturnType can be omitted.
LogicalResult OpWithRefineTypeInterfaceOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &returnTypes) {
  returnTypes.clear();
  return OpWithRefineTypeInterfaceOp::refineReturnTypes(
      context, location, operands, attributes, properties, regions,
      returnTypes);
}

LogicalResult OpWithRefineTypeInterfaceOp::refineReturnTypes(
    MLIRContext *, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &returnTypes) {
  if (operands[0].getType() != operands[1].getType()) {
    return emitOptionalError(location, "operand type mismatch ",
                             operands[0].getType(), " vs ",
                             operands[1].getType());
  }
  // TODO: Add helper to make this more concise to write.
  if (returnTypes.empty())
    returnTypes.resize(1, nullptr);
  if (returnTypes[0] && returnTypes[0] != operands[0].getType())
    return emitOptionalError(location,
                             "required first operand and result to match");
  returnTypes[0] = operands[0].getType();
  return success();
}

//===----------------------------------------------------------------------===//
// TilingNoDpsOp
//===----------------------------------------------------------------------===//

SmallVector<Range> TilingNoDpsOp::getIterationDomain(OpBuilder &builder) {
  return {};
}

SmallVector<utils::IteratorType> TilingNoDpsOp::getLoopIteratorTypes() {
  return {};
}

FailureOr<TilingResult> TilingNoDpsOp::getTiledImplementation(
    OpBuilder &builder, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, ArrayRef<InnerTileAlignment>) {
  return getTiledImplementation(builder, offsets, sizes);
}

FailureOr<TilingResult>
TilingNoDpsOp::getTiledImplementation(OpBuilder &builder,
                                      ArrayRef<OpFoldResult> offsets,
                                      ArrayRef<OpFoldResult> sizes) {
  return failure();
}

LogicalResult TilingNoDpsOp::getResultTilePosition(
    OpBuilder &builder, unsigned resultNumber, ArrayRef<OpFoldResult> offsets,
    ArrayRef<OpFoldResult> sizes, SmallVector<OpFoldResult> &resultOffsets,
    SmallVector<OpFoldResult> &resultSizes) {
  return failure();
}

//===----------------------------------------------------------------------===//
// OpWithShapedTypeInferTypeAdaptorInterfaceOp
//===----------------------------------------------------------------------===//

LogicalResult
OpWithShapedTypeInferTypeAdaptorInterfaceOp::inferReturnTypeComponents(
    MLIRContext *context, std::optional<Location> location,
    OpWithShapedTypeInferTypeAdaptorInterfaceOp::Adaptor adaptor,
    SmallVectorImpl<ShapedTypeComponents> &inferredReturnShapes) {
  // Create return type consisting of the last element of the first operand.
  auto operandType = adaptor.getOperand1().getType();
  auto sval = dyn_cast<ShapedType>(operandType);
  if (!sval)
    return emitOptionalError(location, "only shaped type operands allowed");
  int64_t dim = sval.hasRank() ? sval.getShape().front() : ShapedType::kDynamic;
  auto type = IntegerType::get(context, 17);

  Attribute encoding;
  if (auto rankedTy = dyn_cast<RankedTensorType>(sval))
    encoding = rankedTy.getEncoding();
  inferredReturnShapes.push_back(ShapedTypeComponents({dim}, type, encoding));
  return success();
}

LogicalResult
OpWithShapedTypeInferTypeAdaptorInterfaceOp::reifyReturnTypeShapes(
    OpBuilder &builder, ValueRange operands,
    llvm::SmallVectorImpl<Value> &shapes) {
  shapes = SmallVector<Value, 1>{
      builder.createOrFold<tensor::DimOp>(getLoc(), operands.front(), 0)};
  return success();
}

//===----------------------------------------------------------------------===//
// TestOpWithPropertiesAndInferredType
//===----------------------------------------------------------------------===//

LogicalResult TestOpWithPropertiesAndInferredType::inferReturnTypes(
    MLIRContext *context, std::optional<Location>, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {

  Adaptor adaptor(operands, attributes, properties, regions);
  inferredReturnTypes.push_back(IntegerType::get(
      context, adaptor.getLhs() + adaptor.getProperties().rhs));
  return success();
}

//===----------------------------------------------------------------------===//
// SegmentedRegionBranchOp / SegmentedRegionBranchTerminatorOp
//===----------------------------------------------------------------------===//

void SegmentedRegionBranchOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns, MLIRContext *context) {
  populateRegionBranchOpInterfaceCanonicalizationPatterns(patterns,
                                                          getOperationName());
}

void SegmentedRegionBranchOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  regions.emplace_back(&getFirstRegion());
  regions.emplace_back(&getSecondRegion());
  if (!point.isParent())
    regions.emplace_back(getOperation());
}

ValueRange
SegmentedRegionBranchOp::getSuccessorInputs(RegionSuccessor successor) {
  if (successor.isOperation())
    return {};
  return successor.getSuccessor()->front().getArguments();
}

OperandRange
SegmentedRegionBranchOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  return successor.getSuccessor() == &getFirstRegion() ? getFirst()
                                                       : getSecond();
}

MutableOperandRange
SegmentedRegionBranchTerminatorOp::getMutableSuccessorOperands(
    RegionSuccessor successor) {
  if (successor.isOperation())
    return MutableOperandRange(getOperation(), getNumOperands(), 0);
  auto parent = cast<SegmentedRegionBranchOp>((*this)->getParentOp());
  return successor.getSuccessor() == &parent.getFirstRegion()
             ? getFirstMutable()
             : getSecondMutable();
}

//===----------------------------------------------------------------------===//
// IsolatedRegionBranchOp
//===----------------------------------------------------------------------===//

void IsolatedRegionBranchOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  if (!point.isParent()) {
    regions.emplace_back(getOperation());
    return;
  }
  for (Region &region : getBranches())
    regions.emplace_back(&region);
}

ValueRange
IsolatedRegionBranchOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation()
             ? ValueRange(getOutputs())
             : ValueRange(successor.getSuccessor()->getArguments());
}

OperandRange
IsolatedRegionBranchOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  return getInputs();
}

void IsolatedRegionBranchOp::getCanonicalizationPatterns(
    RewritePatternSet &patterns, MLIRContext *context) {
  populateRegionBranchOpInterfaceCanonicalizationPatterns(patterns,
                                                          getOperationName());
}

MutableOperandRange
IsolatedRegionYieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
  return getValuesMutable();
}

//===----------------------------------------------------------------------===//
// LoopBlockOp
//===----------------------------------------------------------------------===//

void LoopBlockOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  regions.emplace_back(&getBody());
  if (point.isParent())
    return;

  regions.push_back(RegionSuccessor(getOperation()));
}

ValueRange LoopBlockOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getOperation()->getResults())
                                 : ValueRange(getBody().getArguments());
}

OperandRange LoopBlockOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  assert(successor.getSuccessor() == &getBody());
  return MutableOperandRange(getInitMutable());
}

//===----------------------------------------------------------------------===//
// LoopBlockTerminatorOp
//===----------------------------------------------------------------------===//

MutableOperandRange
LoopBlockTerminatorOp::getMutableSuccessorOperands(RegionSuccessor successor) {
  if (successor.isOperation())
    return getExitArgMutable();
  return getNextIterArgMutable();
}

//===----------------------------------------------------------------------===//
// LoopWithExtraResultOp / LoopWithExtraResultYieldOp
//===----------------------------------------------------------------------===//

void LoopWithExtraResultOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Parent always enters the body; the body can loop back or exit to parent.
  regions.emplace_back(&getBody());
  if (!point.isParent())
    regions.push_back(RegionSuccessor(getOperation()));
}

ValueRange
LoopWithExtraResultOp::getSuccessorInputs(RegionSuccessor successor) {
  // When branching to the parent, only iterResult (result #1) is a successor
  // input; extraResult (#0) is not. Similaly when branching to the body, only
  // body block arg #1 (iterArg) is a successor input; arg #0 is not.
  if (successor.isOperation())
    return getResults().drop_front(1);
  return getBody().getArguments().drop_front(1);
}

OperandRange LoopWithExtraResultOp::getEntrySuccessorOperands(RegionSuccessor) {
  return MutableOperandRange(getInitMutable());
}

MutableOperandRange
LoopWithExtraResultYieldOp::getMutableSuccessorOperands(RegionSuccessor) {
  return getIterArgMutable();
}

//===----------------------------------------------------------------------===//
// TestCrashingReturnOp
//===----------------------------------------------------------------------===//

LogicalResult TestCrashingReturnOp::verify() {
  if (!getValid())
    return emitOpError("requires 'valid' unit attribute");
  return success();
}

MutableOperandRange
TestCrashingReturnOp::getMutableSuccessorOperands(RegionSuccessor successor) {
  if (!getValid())
    llvm::report_fatal_error("getMutableSuccessorOperands called on unverified "
                             "test.crashing_return");
  return getArgsMutable();
}

//===----------------------------------------------------------------------===//
// TestReturnWithIgnoredValueOp
//===----------------------------------------------------------------------===//

MutableOperandRange TestReturnWithIgnoredValueOp::getMutableSuccessorOperands(
    RegionSuccessor /*successor*/) {
  return getValuesMutable();
}

//===----------------------------------------------------------------------===//
// SwitchWithNoBreakOp
//===----------------------------------------------------------------------===//

void TestNoTerminatorOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {}

//===----------------------------------------------------------------------===//
// Test InferIntRangeInterface
//===----------------------------------------------------------------------===//

OpFoldResult ManualCppOpWithFold::fold(ArrayRef<Attribute> attributes) {
  // Just a simple fold for testing purposes that reads an operands constant
  // value and returns it.
  if (!attributes.empty())
    return attributes.front();
  return nullptr;
}

//===----------------------------------------------------------------------===//
// Tensor/Buffer Ops
//===----------------------------------------------------------------------===//

void ReadBufferOp::getEffects(
    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
        &effects) {
  // The buffer operand is read.
  effects.emplace_back(MemoryEffects::Read::get(), &getBufferMutable(),
                       SideEffects::DefaultResource::get());
  // The buffer contents are dumped.
  effects.emplace_back(MemoryEffects::Write::get(),
                       SideEffects::DefaultResource::get());
}

//===----------------------------------------------------------------------===//
// Test Dataflow
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// TestCallAndStoreOp
//===----------------------------------------------------------------------===//

CallInterfaceCallable TestCallAndStoreOp::getCallableForCallee() {
  return getCallee();
}

void TestCallAndStoreOp::setCalleeFromCallable(CallInterfaceCallable callee) {
  setCalleeAttr(cast<SymbolRefAttr>(callee));
}

Operation::operand_range TestCallAndStoreOp::getArgOperands() {
  return getCalleeOperands();
}

MutableOperandRange TestCallAndStoreOp::getArgOperandsMutable() {
  return getCalleeOperandsMutable();
}

//===----------------------------------------------------------------------===//
// TestCallOnDeviceOp
//===----------------------------------------------------------------------===//

CallInterfaceCallable TestCallOnDeviceOp::getCallableForCallee() {
  return getCallee();
}

void TestCallOnDeviceOp::setCalleeFromCallable(CallInterfaceCallable callee) {
  setCalleeAttr(cast<SymbolRefAttr>(callee));
}

Operation::operand_range TestCallOnDeviceOp::getArgOperands() {
  return getForwardedOperands();
}

MutableOperandRange TestCallOnDeviceOp::getArgOperandsMutable() {
  return getForwardedOperandsMutable();
}

//===----------------------------------------------------------------------===//
// TestCallAndProduceOp
//===----------------------------------------------------------------------===//

CallInterfaceCallable TestCallAndProduceOp::getCallableForCallee() {
  return getCallee();
}

void TestCallAndProduceOp::setCalleeFromCallable(CallInterfaceCallable callee) {
  setCalleeAttr(cast<SymbolRefAttr>(callee));
}

Operation::operand_range TestCallAndProduceOp::getArgOperands() {
  return getForwardedOperands();
}

MutableOperandRange TestCallAndProduceOp::getArgOperandsMutable() {
  return getForwardedOperandsMutable();
}

Operation::result_range TestCallAndProduceOp::getForwardedResults() {
  // The first result (`produced_status`) is produced by this operation and is
  // not forwarded from the callee.
  return getForwarded();
}

LogicalResult
TestCallAndProduceOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  return call_interface_impl::verifyCallOpInterface(*this, symbolTable);
}

//===----------------------------------------------------------------------===//
// TestStoreWithARegion
//===----------------------------------------------------------------------===//

void TestStoreWithARegion::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  if (point.isParent())
    regions.emplace_back(&getBody());
  else
    regions.push_back(RegionSuccessor(getOperation()));
}

ValueRange TestStoreWithARegion::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getOperation()->getResults())
                                 : ValueRange(getBody().front().getArguments());
}

//===----------------------------------------------------------------------===//
// TestStoreWithALoopRegion
//===----------------------------------------------------------------------===//

void TestStoreWithALoopRegion::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Both the operation itself and the region may be branching into the body or
  // back into the operation itself. It is possible for the operation not to
  // enter the body.
  regions.emplace_back(&getBody());
  regions.push_back(RegionSuccessor(getOperation()));
}

ValueRange
TestStoreWithALoopRegion::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getOperation()->getResults())
                                 : ValueRange(getBody().front().getArguments());
}

//===----------------------------------------------------------------------===//
// TestRegionTypesCompatOp
//===----------------------------------------------------------------------===//

void TestRegionTypesCompatOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  if (point.isParent())
    regions.emplace_back(&getBody());
  else
    regions.push_back(RegionSuccessor(getOperation()));
}

OperandRange
TestRegionTypesCompatOp::getEntrySuccessorOperands(RegionSuccessor) {
  return getEntries();
}

ValueRange
TestRegionTypesCompatOp::getSuccessorInputs(RegionSuccessor successor) {
  if (successor.isOperation())
    return getResults();
  return getBody().getArguments();
}

bool TestRegionTypesCompatOp::areTypesCompatible(Type lhs, Type rhs) {
  return lhs == rhs || (isa<IntegerType>(lhs) && isa<IntegerType>(rhs));
}

//===----------------------------------------------------------------------===//
// TestLoopTypesCompatOp
//===----------------------------------------------------------------------===//

void TestLoopTypesCompatOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  regions.emplace_back(&getBody());
  if (!point.isParent())
    regions.push_back(RegionSuccessor(getOperation()));
}

OperandRange TestLoopTypesCompatOp::getEntrySuccessorOperands(RegionSuccessor) {
  return getInitArgs();
}

ValueRange
TestLoopTypesCompatOp::getSuccessorInputs(RegionSuccessor successor) {
  if (successor.isOperation())
    return getResults();
  return getBody().getArguments();
}

MutableArrayRef<OpOperand> TestLoopTypesCompatOp::getInitsMutable() {
  return getInitArgsMutable();
}

Block::BlockArgListType TestLoopTypesCompatOp::getRegionIterArgs() {
  return getBody().getArguments();
}

std::optional<MutableArrayRef<OpOperand>>
TestLoopTypesCompatOp::getYieldedValuesMutable() {
  return cast<TestTypesCompatYieldOp>(getBody().front().getTerminator())
      .getArgsMutable();
}

std::optional<ResultRange> TestLoopTypesCompatOp::getLoopResults() {
  return getResults();
}

SmallVector<Region *> TestLoopTypesCompatOp::getLoopRegions() {
  return {&getBody()};
}

bool TestLoopTypesCompatOp::areTypesCompatible(Type lhs, Type rhs) {
  return lhs == rhs || (isa<IntegerType>(lhs) && isa<IntegerType>(rhs));
}

void TestRegionTypeChangerOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  if (point.isParent())
    regions.emplace_back(&getBody());
  else
    regions.push_back(RegionSuccessor(getOperation()));
}

OperandRange
TestRegionTypeChangerOp::getEntrySuccessorOperands(RegionSuccessor) {
  return getEntries();
}

ValueRange
TestRegionTypeChangerOp::getSuccessorInputs(RegionSuccessor successor) {
  if (successor.isOperation())
    return getResults();
  return getBody().getArguments();
}

bool TestRegionTypeChangerOp::areTypesCompatible(Type lhs, Type rhs) {
  return true;
}

//===----------------------------------------------------------------------===//
// TestVersionedOpA
//===----------------------------------------------------------------------===//

LogicalResult
TestVersionedOpA::readProperties(mlir::DialectBytecodeReader &reader,
                                 mlir::OperationState &state) {
  auto &prop = state.getOrAddProperties<Properties>();
  if (mlir::failed(reader.readAttribute(prop.dims)))
    return mlir::failure();

  // Check if we have a version. If not, assume we are parsing the current
  // version.
  auto maybeVersion = reader.getDialectVersion<test::TestDialect>();
  if (succeeded(maybeVersion)) {
    // If version is less than 2.0, there is no additional attribute to parse.
    // We can materialize missing properties post parsing before verification.
    const auto *version =
        reinterpret_cast<const TestDialectVersion *>(*maybeVersion);
    if ((version->major_ < 2)) {
      return success();
    }
  }

  if (mlir::failed(reader.readAttribute(prop.modifier)))
    return mlir::failure();
  return mlir::success();
}

void TestVersionedOpA::writeProperties(mlir::DialectBytecodeWriter &writer) {
  auto &prop = getProperties();
  writer.writeAttribute(prop.dims);

  auto maybeVersion = writer.getDialectVersion<test::TestDialect>();
  if (succeeded(maybeVersion)) {
    // If version is less than 2.0, there is no additional attribute to write.
    const auto *version =
        reinterpret_cast<const TestDialectVersion *>(*maybeVersion);
    if ((version->major_ < 2)) {
      llvm::outs() << "downgrading op properties...\n";
      return;
    }
  }
  writer.writeAttribute(prop.modifier);
}

//===----------------------------------------------------------------------===//
// TestOpWithVersionedProperties
//===----------------------------------------------------------------------===//

llvm::LogicalResult TestOpWithVersionedProperties::readFromMlirBytecode(
    mlir::DialectBytecodeReader &reader, test::VersionedProperties &prop) {
  uint64_t value1, value2 = 0;
  if (failed(reader.readVarInt(value1)))
    return failure();

  // Check if we have a version. If not, assume we are parsing the current
  // version.
  auto maybeVersion = reader.getDialectVersion<test::TestDialect>();
  bool needToParseAnotherInt = true;
  if (succeeded(maybeVersion)) {
    // If version is less than 2.0, there is no additional attribute to parse.
    // We can materialize missing properties post parsing before verification.
    const auto *version =
        reinterpret_cast<const TestDialectVersion *>(*maybeVersion);
    if ((version->major_ < 2))
      needToParseAnotherInt = false;
  }
  if (needToParseAnotherInt && failed(reader.readVarInt(value2)))
    return failure();

  prop.value1 = value1;
  prop.value2 = value2;
  return success();
}

void TestOpWithVersionedProperties::writeToMlirBytecode(
    mlir::DialectBytecodeWriter &writer,
    const test::VersionedProperties &prop) {
  writer.writeVarInt(prop.value1);
  writer.writeVarInt(prop.value2);
}

//===----------------------------------------------------------------------===//
// TestMultiSlotAlloca
//===----------------------------------------------------------------------===//

llvm::SmallVector<MemorySlot> TestMultiSlotAlloca::getPromotableSlots() {
  SmallVector<MemorySlot> slots;
  for (Value result : getResults()) {
    slots.push_back(MemorySlot{
        result, cast<MemRefType>(result.getType()).getElementType()});
  }
  return slots;
}

Value TestMultiSlotAlloca::getDefaultValue(const MemorySlot &slot,
                                           OpBuilder &builder) {
  return TestOpConstant::create(builder, getLoc(), slot.valueType,
                                builder.getI32IntegerAttr(42));
}

void TestMultiSlotAlloca::handleBlockArgument(const MemorySlot &slot,
                                              BlockArgument argument,
                                              OpBuilder &builder) {
  // Not relevant for testing.
}

/// Creates a new TestMultiSlotAlloca operation, just without the `slot`.
static std::optional<TestMultiSlotAlloca>
createNewMultiAllocaWithoutSlot(const MemorySlot &slot, OpBuilder &builder,
                                TestMultiSlotAlloca oldOp) {

  if (oldOp.getNumResults() == 1) {
    oldOp.erase();
    return std::nullopt;
  }

  SmallVector<Type> newTypes;
  SmallVector<Value> remainingValues;

  for (Value oldResult : oldOp.getResults()) {
    if (oldResult == slot.ptr)
      continue;
    remainingValues.push_back(oldResult);
    newTypes.push_back(oldResult.getType());
  }

  OpBuilder::InsertionGuard guard(builder);
  builder.setInsertionPoint(oldOp);
  auto replacement =
      TestMultiSlotAlloca::create(builder, oldOp->getLoc(), newTypes);
  for (auto [oldResult, newResult] :
       llvm::zip_equal(remainingValues, replacement.getResults()))
    oldResult.replaceAllUsesWith(newResult);

  oldOp.erase();
  return replacement;
}

std::optional<PromotableAllocationOpInterface>
TestMultiSlotAlloca::handlePromotionComplete(const MemorySlot &slot,
                                             Value defaultValue,
                                             OpBuilder &builder) {
  if (defaultValue && defaultValue.use_empty())
    defaultValue.getDefiningOp()->erase();
  return createNewMultiAllocaWithoutSlot(slot, builder, *this);
}

SmallVector<DestructurableMemorySlot>
TestMultiSlotAlloca::getDestructurableSlots() {
  SmallVector<DestructurableMemorySlot> slots;
  for (Value result : getResults()) {
    auto memrefType = cast<MemRefType>(result.getType());
    auto destructurable = dyn_cast<DestructurableTypeInterface>(memrefType);
    if (!destructurable)
      continue;

    std::optional<DenseMap<Attribute, Type>> destructuredType =
        destructurable.getSubelementIndexMap();
    if (!destructuredType)
      continue;
    slots.emplace_back(
        DestructurableMemorySlot{{result, memrefType}, *destructuredType});
  }
  return slots;
}

DenseMap<Attribute, MemorySlot> TestMultiSlotAlloca::destructure(
    const DestructurableMemorySlot &slot,
    const SmallPtrSetImpl<Attribute> &usedIndices, OpBuilder &builder,
    SmallVectorImpl<DestructurableAllocationOpInterface> &newAllocators) {
  OpBuilder::InsertionGuard guard(builder);
  builder.setInsertionPointAfter(*this);

  DenseMap<Attribute, MemorySlot> slotMap;

  for (Attribute usedIndex : usedIndices) {
    Type elemType = slot.subelementTypes.lookup(usedIndex);
    MemRefType elemPtr = MemRefType::get({}, elemType);
    auto subAlloca = TestMultiSlotAlloca::create(builder, getLoc(), elemPtr);
    newAllocators.push_back(subAlloca);
    slotMap.try_emplace<MemorySlot>(usedIndex,
                                    {subAlloca.getResult(0), elemType});
  }

  return slotMap;
}

std::optional<DestructurableAllocationOpInterface>
TestMultiSlotAlloca::handleDestructuringComplete(
    const DestructurableMemorySlot &slot, OpBuilder &builder) {
  return createNewMultiAllocaWithoutSlot(slot, builder, *this);
}

//===----------------------------------------------------------------------===//
// TestTransparentAlias
//===----------------------------------------------------------------------===//

void TestTransparentAlias::getPromotableSlotAliases(
    OpOperand &aliasedSlotPointerOperand, const MemorySlot & /*parentSlot*/,
    SmallVectorImpl<MemorySlot> &newMemorySlots) {
  if (aliasedSlotPointerOperand.get() != getSource())
    return;
  Type elemType = cast<MemRefType>(getResult().getType()).getElementType();
  newMemorySlots.push_back(MemorySlot{getResult(), elemType});
}

bool TestTransparentAlias::canUsesBeRemoved(
    const SmallPtrSetImpl<OpOperand *> &blockingUses,
    SmallVectorImpl<OpOperand *> &newBlockingUses,
    const DataLayout &dataLayout) {
  for (OpOperand &use : getResult().getUses())
    newBlockingUses.push_back(&use);
  return true;
}

DeletionKind TestTransparentAlias::removeBlockingUses(
    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
  return DeletionKind::Delete;
}

//===----------------------------------------------------------------------===//
// TestTransparentCastAlias
//===----------------------------------------------------------------------===//

void TestTransparentCastAlias::getPromotableSlotAliases(
    OpOperand &aliasedSlotPointerOperand, const MemorySlot & /*parentSlot*/,
    SmallVectorImpl<MemorySlot> &newMemorySlots) {
  if (aliasedSlotPointerOperand.get() != getSource())
    return;
  Type elemType = cast<MemRefType>(getResult().getType()).getElementType();
  newMemorySlots.push_back(MemorySlot{getResult(), elemType});
}

bool TestTransparentCastAlias::canUsesBeRemoved(
    const SmallPtrSetImpl<OpOperand *> &blockingUses,
    SmallVectorImpl<OpOperand *> &newBlockingUses,
    const DataLayout &dataLayout) {
  for (OpOperand &use : getResult().getUses())
    newBlockingUses.push_back(&use);
  return true;
}

DeletionKind TestTransparentCastAlias::removeBlockingUses(
    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
  return DeletionKind::Delete;
}

Value TestTransparentCastAlias::projectSlotValueToAliasValue(
    OpOperand & /*aliasedSlotPointerOperand*/,
    const MemorySlot & /*parentSlot*/, const MemorySlot &aliasSlot,
    Value slotValue, OpBuilder &builder) {
  if (slotValue.getType() == aliasSlot.valueType)
    return slotValue;
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            aliasSlot.valueType, slotValue)
      .getResult(0);
}

Value TestTransparentCastAlias::projectAliasValueToSlotValue(
    OpOperand & /*aliasedSlotPointerOperand*/, const MemorySlot &parentSlot,
    const MemorySlot & /*aliasSlot*/, Value aliasValue, Value /*reachingDef*/,
    OpBuilder &builder) {
  if (aliasValue.getType() == parentSlot.valueType)
    return aliasValue;
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            parentSlot.valueType, aliasValue)
      .getResult(0);
}

//===----------------------------------------------------------------------===//
// TestTransparentDualAlias
//===----------------------------------------------------------------------===//

void TestTransparentDualAlias::getPromotableSlotAliases(
    OpOperand &aliasedSlotPointerOperand, const MemorySlot & /*parentSlot*/,
    SmallVectorImpl<MemorySlot> &newMemorySlots) {
  if (aliasedSlotPointerOperand.get() != getSource())
    return;
  // Expose both results as aliases of the same parent at their own
  // signedness (same bit width as the parent's signless i32).
  newMemorySlots.push_back(MemorySlot{
      getResultSigned(),
      cast<MemRefType>(getResultSigned().getType()).getElementType()});
  newMemorySlots.push_back(MemorySlot{
      getResultUnsigned(),
      cast<MemRefType>(getResultUnsigned().getType()).getElementType()});
}

bool TestTransparentDualAlias::canUsesBeRemoved(
    const SmallPtrSetImpl<OpOperand *> &blockingUses,
    SmallVectorImpl<OpOperand *> &newBlockingUses,
    const DataLayout &dataLayout) {
  for (Value result : getResults())
    for (OpOperand &use : result.getUses())
      newBlockingUses.push_back(&use);
  return true;
}

DeletionKind TestTransparentDualAlias::removeBlockingUses(
    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
  return DeletionKind::Delete;
}

Value TestTransparentDualAlias::projectSlotValueToAliasValue(
    OpOperand & /*aliasedSlotPointerOperand*/,
    const MemorySlot & /*parentSlot*/, const MemorySlot &aliasSlot,
    Value slotValue, OpBuilder &builder) {
  if (slotValue.getType() == aliasSlot.valueType)
    return slotValue;
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            aliasSlot.valueType, slotValue)
      .getResult(0);
}

Value TestTransparentDualAlias::projectAliasValueToSlotValue(
    OpOperand & /*aliasedSlotPointerOperand*/, const MemorySlot &parentSlot,
    const MemorySlot & /*aliasSlot*/, Value aliasValue, Value /*reachingDef*/,
    OpBuilder &builder) {
  if (aliasValue.getType() == parentSlot.valueType)
    return aliasValue;
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            parentSlot.valueType, aliasValue)
      .getResult(0);
}

//===----------------------------------------------------------------------===//
// TestPartialAlias
//===----------------------------------------------------------------------===//

void TestPartialAlias::getPromotableSlotAliases(
    OpOperand &aliasedSlotPointerOperand, const MemorySlot & /*parentSlot*/,
    SmallVectorImpl<MemorySlot> &newMemorySlots) {
  if (aliasedSlotPointerOperand.get() != getSource())
    return;
  newMemorySlots.push_back(MemorySlot{
      getResult(), cast<MemRefType>(getResult().getType()).getElementType()});
}

bool TestPartialAlias::canUsesBeRemoved(
    const SmallPtrSetImpl<OpOperand *> &blockingUses,
    SmallVectorImpl<OpOperand *> &newBlockingUses,
    const DataLayout &dataLayout) {
  for (OpOperand &use : getResult().getUses())
    newBlockingUses.push_back(&use);
  return true;
}

DeletionKind TestPartialAlias::removeBlockingUses(
    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
  return DeletionKind::Delete;
}

Value TestPartialAlias::projectSlotValueToAliasValue(
    OpOperand & /*aliasedSlotPointerOperand*/,
    const MemorySlot & /*parentSlot*/, const MemorySlot &aliasSlot,
    Value slotValue, OpBuilder &builder) {
  // Sub-value extraction: 1-input cast.
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            aliasSlot.valueType, slotValue)
      .getResult(0);
}

Value TestPartialAlias::projectAliasValueToSlotValue(
    OpOperand & /*aliasedSlotPointerOperand*/, const MemorySlot &parentSlot,
    const MemorySlot & /*aliasSlot*/, Value aliasValue, Value reachingDef,
    OpBuilder &builder) {
  // Sub-value insertion into the current reaching definition: emit a 2-input
  // cast taking both the new alias value and the existing parent value.
  return UnrealizedConversionCastOp::create(builder, getLoc(),
                                            parentSlot.valueType,
                                            ValueRange{aliasValue, reachingDef})
      .getResult(0);
}

//===----------------------------------------------------------------------===//
// TestSlotTracker
//===----------------------------------------------------------------------===//

bool TestSlotTracker::canUsesBeRemoved(
    const SmallPtrSetImpl<OpOperand *> &blockingUses,
    SmallVectorImpl<OpOperand *> &newBlockingUses,
    const DataLayout &dataLayout) {
  if (blockingUses.size() != 1)
    return false;
  return (*blockingUses.begin())->get() == getSource();
}

DeletionKind TestSlotTracker::removeBlockingUses(
    const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
  return DeletionKind::Delete;
}

bool TestSlotTracker::requiresReplacedValues() { return true; }

void TestSlotTracker::visitReplacedValues(
    ArrayRef<std::pair<Operation *, Value>> definitions, OpBuilder &builder) {
  for (auto [op, value] : definitions) {
    builder.setInsertionPointAfter(op);
    TestTrackedValue::create(builder, getLoc(), value, getNameAttr());
  }
}

::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
    ::mlir::RewriterBase &rewriter,
    const ::mlir::bufferization::BufferizationOptions &options,
    ::mlir::bufferization::BufferizationState &state) {
  auto buffer =
      mlir::bufferization::getBuffer(rewriter, getInput(), options, state);
  if (mlir::failed(buffer))
    return failure();

  const auto bufferizedOutType =
      mlir::bufferization::getBufferType(getOutput(), options, state);
  if (mlir::failed(bufferizedOutType))
    return failure();

  // replace op with memref analogy
  auto dummyMemrefOp = test::TestDummyMemrefOp::create(
      rewriter, getLoc(), *bufferizedOutType, *buffer);

  mlir::bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),
                                                     dummyMemrefOp.getResult());

  return mlir::success();
}

mlir::FailureOr<mlir::bufferization::BufferLikeType>
test::TestDummyTensorOp::getBufferType(
    mlir::Value value, const mlir::bufferization::BufferizationOptions &options,
    const mlir::bufferization::BufferizationState &,
    llvm::SmallVector<::mlir::Value> &) {
  // let unknown type converter handle the types - assume test operations work
  // with test-one-shot-module-bufferize pass.
  return options.unknownTypeConverterFn(
      cast<mlir::bufferization::TensorLikeType>(value.getType()), nullptr,
      options);
}

::mlir::LogicalResult test::TestCreateTensorOp::bufferize(
    ::mlir::RewriterBase &rewriter,
    const ::mlir::bufferization::BufferizationOptions &options,
    ::mlir::bufferization::BufferizationState &state) {
  // Note: mlir::bufferization::getBufferType() would internally call
  // TestCreateTensorOp::getBufferType()
  const auto bufferizedOutType =
      mlir::bufferization::getBufferType(getOutput(), options, state);
  if (mlir::failed(bufferizedOutType))
    return failure();

  // replace op with memref analogy
  auto createMemrefOp =
      test::TestCreateMemrefOp::create(rewriter, getLoc(), *bufferizedOutType);

  mlir::bufferization::replaceOpWithBufferizedValues(
      rewriter, getOperation(), createMemrefOp.getResult());

  return mlir::success();
}

mlir::FailureOr<mlir::bufferization::BufferLikeType>
test::TestCreateTensorOp::getBufferType(
    mlir::Value value, const mlir::bufferization::BufferizationOptions &options,
    const mlir::bufferization::BufferizationState &,
    llvm::SmallVector<::mlir::Value> &) {
  // let unknown type converter handle the types - assume test operations work
  // with test-one-shot-module-bufferize pass.
  return options.unknownTypeConverterFn(
      cast<mlir::bufferization::TensorLikeType>(value.getType()), nullptr,
      options);
}

::mlir::LogicalResult test::TestTensorWithFutureLayoutOp::bufferize(
    ::mlir::RewriterBase &rewriter,
    const ::mlir::bufferization::BufferizationOptions &options,
    ::mlir::bufferization::BufferizationState &state) {
  const auto bufferizedOutType =
      mlir::bufferization::getBufferType(getOutput(), options, state);
  if (mlir::failed(bufferizedOutType))
    return failure();

  auto createMemrefOp =
      test::TestCreateMemrefOp::create(rewriter, getLoc(), *bufferizedOutType);
  mlir::bufferization::replaceOpWithBufferizedValues(
      rewriter, getOperation(), createMemrefOp.getResult());
  return mlir::success();
}

mlir::FailureOr<mlir::bufferization::BufferLikeType>
test::TestTensorWithFutureLayoutOp::getBufferType(
    mlir::Value value, const mlir::bufferization::BufferizationOptions &,
    const mlir::bufferization::BufferizationState &,
    llvm::SmallVector<::mlir::Value> &) {
  return llvm::TypeSwitch<mlir::Type,
                          mlir::FailureOr<mlir::bufferization::BufferLikeType>>(
             value.getType())
      .Case([&](RankedTensorType tensorType) {
        // Set the memref layout to the op's 'layout' attribute, ignoring
        // any pre-existing tensor encoding. This is what lets two
        // `test.tensor_with_layout` ops produce *bufferized* memrefs with
        // different layouts while keeping their *tensor* result types identical
        // -- which is required to construct SCF iter_arg/branch mismatches that
        // the verifier still accepts.
        auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
        return cast<bufferization::BufferLikeType>(MemRefType::get(
            tensorType.getShape(), tensorType.getElementType(), layout));
      })
      .Case([&](TestTensorType tensorType) {
        auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
        return cast<bufferization::BufferLikeType>(
            TestMemrefType::get(tensorType.getContext(), tensorType.getShape(),
                                tensorType.getElementType(), layout));
      })
      .Default([&](Type) { return emitError("unknown type"); });
}

LogicalResult test::TestForceNewLayoutOp::bufferize(
    RewriterBase &rewriter, const bufferization::BufferizationOptions &options,
    bufferization::BufferizationState &state) {
  auto buffer = bufferization::getBuffer(rewriter, getInput(), options, state);
  if (failed(buffer))
    return failure();

  const auto bufferizedOutType =
      bufferization::getBufferType(getOutput(), options, state);
  if (failed(bufferizedOutType))
    return failure();

  auto dummyMemrefOp = test::TestDummyMemrefOp::create(
      rewriter, getLoc(), *bufferizedOutType, *buffer);
  bufferization::replaceOpWithBufferizedValues(rewriter, getOperation(),
                                               dummyMemrefOp.getResult());
  return success();
}

FailureOr<bufferization::BufferLikeType>
test::TestForceNewLayoutOp::getBufferType(
    mlir::Value value, const bufferization::BufferizationOptions &options,
    const bufferization::BufferizationState &,
    llvm::SmallVector<::mlir::Value> &) {
  return llvm::TypeSwitch<mlir::Type,
                          mlir::FailureOr<mlir::bufferization::BufferLikeType>>(
             value.getType())
      .Case([&](RankedTensorType tensorType) {
        // Set the memref layout to the op's 'layout' attribute, ignoring any
        // pre-existing tensor encoding.
        auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
        return cast<bufferization::BufferLikeType>(MemRefType::get(
            tensorType.getShape(), tensorType.getElementType(), layout));
      })
      .Case([&](TestTensorType tensorType) {
        auto layout = cast<MemRefLayoutAttrInterface>(getLayout());
        return cast<bufferization::BufferLikeType>(
            TestMemrefType::get(tensorType.getContext(), tensorType.getShape(),
                                tensorType.getElementType(), layout));
      })
      .Default([&](Type) { return emitError("unknown type"); });
}

// Define a custom builder for ManyRegionsOp declared in TestOps.td.
//  OpBuilder<(ins "::std::unique_ptr<::mlir::Region>":$firstRegion,
//                 "::std::unique_ptr<::mlir::Region>":$secondRegion)>
void test::ManyRegionsOp::build(
    mlir::OpBuilder &builder, mlir::OperationState &state,
    llvm::SmallVectorImpl<std::unique_ptr<mlir::Region>> &&regions) {
  for (auto &&regionPtr : std::move(regions))
    state.addRegion(std::move(regionPtr));
  ManyRegionsOp::build(builder, state, {}, regions.size());
}