TypeVarInferencePass.cpp

Go to the documentation of this file.

//===-- TypeVarInferencePass.cpp -------------------------------*- C++ -*-===//
//
// Part of the LLZK Project, under the Apache License v2.0.
// See LICENSE.txt for license information.
// Copyright 2026 Project LLZK
// SPDX-License-Identifier: Apache-2.0
//
//===----------------------------------------------------------------------===//
//===----------------------------------------------------------------------===//
 
#include "llzk/Dialect/Array/IR/Ops.h"
#include "llzk/Dialect/Array/Util/ArrayTypeHelper.h"
#include "llzk/Dialect/Function/IR/Ops.h"
#include "llzk/Dialect/POD/IR/Types.h"
#include "llzk/Dialect/Polymorphic/IR/Ops.h"
#include "llzk/Dialect/Polymorphic/Transforms/TransformationPasses.h"
#include "llzk/Dialect/Struct/IR/Ops.h"
#include "llzk/Dialect/Verif/IR/Ops.h"
#include "llzk/Util/Constants.h"
#include "llzk/Util/SymbolHelper.h"
#include "llzk/Util/SymbolLookup.h"
#include "llzk/Util/SymbolTableLLZK.h"
#include "llzk/Util/TypeHelper.h"
#include "llzk/Util/Walk.h"
 
#include <mlir/IR/BuiltinOps.h>
#include <mlir/IR/PatternMatch.h>
 
#include <llvm/ADT/DenseMap.h>
#include <llvm/ADT/DenseSet.h>
#include <llvm/ADT/SmallVector.h>
#include <llvm/ADT/StringMap.h>
#include <llvm/Support/raw_ostream.h>
 
#include <memory>
#include <vector>
 
// Include the generated base pass class definitions.
namespace llzk::polymorphic {
#define GEN_PASS_DEF_TYPEVARINFERENCEPASS
#include "llzk/Dialect/Polymorphic/Transforms/TransformationPasses.h.inc"
} // namespace llzk::polymorphic
 
#include "SharedImpl.h"
 
#define DEBUG_TYPE "llzk-infer-tvar"
 
using namespace mlir;
using namespace llzk;
using namespace llzk::array;
using namespace llzk::component;
using namespace llzk::function;
using namespace llzk::pod;
using namespace llzk::polymorphic;
using namespace llzk::polymorphic::detail;
 
namespace {

struct InferredType {
  Type type;
  Location loc;
};

using SpecializedCallableCloneCache = DenseMap<Operation *, llvm::StringMap<SymbolRefAttr>>;
using SpecializedTemplateCloneCache = DenseMap<Operation *, llvm::StringMap<StringAttr>>;

static SmallVector<StringAttr> getStringPieces(SymbolRefAttr ref) {
  return llvm::to_vector(llvm::map_range(getPieces(ref), [](FlatSymbolRefAttr piece) {
    return piece.getAttr();
  }));
}

static StringAttr getFlatSymbolName(Attribute attr) {
  auto symRef = llvm::dyn_cast_if_present<SymbolRefAttr>(attr);
  if (!symRef || !symRef.getNestedReferences().empty()) {
    return {};
  }
  return symRef.getRootReference();
}

static bool isTypeVarParam(TemplateParamOp op) {
  std::optional<Type> declaredTy = op.getTypeOpt();
  if (!declaredTy) {
    return false;
  }
  auto tvarTy = llvm::dyn_cast<TypeVarType>(*declaredTy);
  return tvarTy && tvarTy.getRefName() == op.getName();
}

static bool isCurrentTemplateTypeVarParamSymbol(StringAttr symbolName, Operation *op) {
  TemplateOp templateOp = getParentOfType<TemplateOp>(op);
  if (!templateOp) {
    return false;
  }
  TemplateParamOp paramOp = templateOp.getConstNamed<TemplateParamOp>(symbolName);
  return paramOp && isTypeVarParam(paramOp);
}

template <typename ConverterT>
static bool
convertTypeRange(TypeRange oldTypes, SmallVectorImpl<Type> &newTypes, ConverterT &converter) {
  bool changed = false;
  newTypes.reserve(oldTypes.size());
  for (Type oldTy : oldTypes) {
    Type newTy = converter.convertType(oldTy);
    newTypes.push_back(newTy);
    changed |= newTy != oldTy;
  }
  return changed;
}

template <typename ConverterT>
static Type convertPodType(PodType podTy, MLIRContext *ctx, ConverterT &converter) {
  SmallVector<RecordAttr> newRecords;
  bool changed = false;
  for (RecordAttr record : podTy.getRecords()) {
    Type newRecordTy = converter.convertType(record.getType());
    newRecords.push_back(RecordAttr::get(ctx, record.getName(), newRecordTy));
    changed |= newRecordTy != record.getType();
  }
  return changed ? PodType::get(ctx, newRecords) : podTy;
}

template <typename ConvertElementFn>
static Type convertArrayElementType(ArrayType arrTy, ConvertElementFn convertElement) {
  Type newElemTy = convertElement(arrTy.getElementType());
  if (newElemTy == arrTy.getElementType()) {
    return arrTy;
  }
  return llzk::polymorphic::detail::flattenInstantiatedArrayType(arrTy, newElemTy);
}

template <typename ConverterT>
static Type convertFunctionType(FunctionType funcTy, ConverterT &converter) {
  SmallVector<Type> newInputs;
  SmallVector<Type> newResults;
  bool changed = convertTypeRange(funcTy.getInputs(), newInputs, converter);
  changed |= convertTypeRange(funcTy.getResults(), newResults, converter);
  return changed ? FunctionType::get(funcTy.getContext(), newInputs, newResults) : funcTy;
}

template <typename ConvertTypeFn, typename ConvertNestedAttrFn>
static Attribute convertTypeOrArrayAttr(
    Attribute attr, MLIRContext *ctx, ConvertTypeFn convertType,
    ConvertNestedAttrFn convertNestedAttr
) {
  if (!attr) {
    return attr;
  }
  if (auto tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
    Type newTy = convertType(tyAttr.getValue());
    return newTy == tyAttr.getValue() ? attr : TypeAttr::get(newTy);
  }
  if (auto arrAttr = llvm::dyn_cast<ArrayAttr>(attr)) {
    SmallVector<Attribute> newAttrs;
    bool changed = false;
    for (Attribute nested : arrAttr.getValue()) {
      Attribute newNested = convertNestedAttr(nested);
      newAttrs.push_back(newNested);
      changed |= newNested != nested;
    }
    return changed ? ArrayAttr::get(ctx, newAttrs) : attr;
  }
  return attr;
}

static bool templateArgUnifiesWithType(Attribute attr, Type expectedTy) {
  Attribute expectedAttr = TypeAttr::get(expectedTy);
  return typeParamsUnify(ArrayRef<Attribute> {attr}, ArrayRef<Attribute> {expectedAttr});
}

class TypeVarReplacementConverter {
  MLIRContext *ctx_;
  /// Fully-qualified symbol path of the template currently being rewritten.
  SmallVector<StringAttr> templatePath_;
  SmallVector<StringAttr> oldParamOrder_;
  DenseSet<StringAttr> removedParams_;
  DenseMap<StringAttr, Type> replacements_;
  bool trimResolvedParams_;
 
public:
  /// Create a converter for a single template rewrite.
  ///
  /// `oldParamOrder` is captured before erasing any `poly.param` operations, so
  /// call-site template argument lists can be trimmed using their original
  /// positional layout.
  TypeVarReplacementConverter(
      MLIRContext *c, ArrayRef<StringAttr> templatePath, ArrayRef<StringAttr> oldParamOrder,
      const DenseMap<StringAttr, Type> &replacements, bool trimResolvedParams = true
  )
      : ctx_(c), templatePath_(templatePath), oldParamOrder_(oldParamOrder),
        replacements_(replacements), trimResolvedParams_(trimResolvedParams) {
    for (const auto &entry : replacements_) {
      removedParams_.insert(entry.first);
    }
  }


  /// or unsupported types are returned unchanged.
  Type convertType(Type ty) const {
    DenseSet<StringAttr> resolvingParams;
    return convertType(ty, resolvingParams);
  }
 
private:
  Type convertType(Type ty, DenseSet<StringAttr> &resolvingParams) const {
    if (!ty) {
      return ty;
    }
    if (auto tvarTy = llvm::dyn_cast<TypeVarType>(ty)) {
      StringAttr paramName = tvarTy.getNameRef().getAttr();
      auto it = replacements_.find(paramName);
      if (it == replacements_.end()) {
        return ty;
      }
      if (!resolvingParams.insert(paramName).second) {
        return ty;
      }
      Type replacement = convertType(it->second, resolvingParams);
      resolvingParams.erase(paramName);
      return replacement;
    }
    if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
      return convertArrayElementType(arrTy, [this, &resolvingParams](Type elemTy) {
        return convertType(elemTy, resolvingParams);
      });
    }
    if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
      return convertStructType(structTy, resolvingParams);
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      return convertPodType(podTy, ctx_, *this);
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      return convertFunctionType(funcTy, *this);
    }
    return ty;
  }
 
public:
  Attribute convertAttr(Attribute attr) const {
    DenseSet<StringAttr> resolvingParams;
    return convertAttr(attr, resolvingParams);
  }

  LogicalResult validateOperation(Operation *op) const {
    if (auto createOp = llvm::dyn_cast<CreateArrayOp>(op)) {
      if (failed(validateCreateArrayOp(createOp))) {
        return failure();
      }
    }
    if (auto func = llvm::dyn_cast<FuncDefOp>(op)) {
      if (failed(validateType(func.getFunctionType(), op))) {
        return failure();
      }
    }
    for (Region &region : op->getRegions()) {
      for (Block &block : region.getBlocks()) {
        for (Type argTy : block.getArgumentTypes()) {
          if (failed(validateType(argTy, op))) {
            return failure();
          }
        }
      }
    }
    for (Type resultTy : op->getResultTypes()) {
      if (failed(validateType(resultTy, op))) {
        return failure();
      }
    }
    for (NamedAttribute attr : op->getAttrs()) {
      if (failed(validateAttr(attr.getValue(), op))) {
        return failure();
      }
    }
    return success();
  }
 
private:
  LogicalResult validateCreateArrayOp(CreateArrayOp createOp) const {
    if (createOp.getElements().empty()) {
      return success();
    }
 
    ArrayType oldResultTy = createOp.getType();
    auto newResultTy = llvm::dyn_cast<ArrayType>(convertType(oldResultTy));
    auto newElemTy = llvm::dyn_cast<ArrayType>(convertType(oldResultTy.getElementType()));
    if (!newResultTy || !newElemTy || newResultTy == oldResultTy || oldResultTy.hasStaticShape()) {
      return success();
    }
 
    return createOp.emitError()
           << "cannot rewrite initialized array.new with non-static initializer shape "
           << oldResultTy;
  }

  LogicalResult validateType(Type ty, Operation *diagnosticOp) const {
    if (!ty) {
      return success();
    }
    if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
      return validateType(arrTy.getElementType(), diagnosticOp);
    }
    if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
      return validateStructType(structTy, diagnosticOp);
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      for (RecordAttr record : podTy.getRecords()) {
        if (failed(validateType(record.getType(), diagnosticOp))) {
          return failure();
        }
      }
      return success();
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      for (Type inputTy : funcTy.getInputs()) {
        if (failed(validateType(inputTy, diagnosticOp))) {
          return failure();
        }
      }
      for (Type resultTy : funcTy.getResults()) {
        if (failed(validateType(resultTy, diagnosticOp))) {
          return failure();
        }
      }
    }
    return success();
  }

  LogicalResult validateAttr(Attribute attr, Operation *diagnosticOp) const {
    if (!attr) {
      return success();
    }
    if (auto tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      return validateType(tyAttr.getValue(), diagnosticOp);
    }
    if (auto arrAttr = llvm::dyn_cast<ArrayAttr>(attr)) {
      for (Attribute nested : arrAttr.getValue()) {
        if (failed(validateAttr(nested, diagnosticOp))) {
          return failure();
        }
      }
    }
    return success();
  }

  Attribute convertAttr(Attribute attr, DenseSet<StringAttr> &resolvingParams) const {
    return convertTypeOrArrayAttr(attr, ctx_, [this, &resolvingParams](Type ty) {
      return convertType(ty, resolvingParams);
    }, [this, &resolvingParams](Attribute nested) {
      return convertTemplateArgAttr(nested, resolvingParams);
    });
  }
 
public:
  FailureOr<ArrayAttr> convertTemplateParams(
      ArrayAttr params, Operation *diagnosticOp, bool resolveTemplateSymbolArgs = true,
      bool allowCurrentTemplateTypeVarSymbols = false
  ) const {
    if (!params) {
      return ArrayAttr();
    }
    if (params.size() != oldParamOrder_.size()) {
      return params;
    }
    SmallVector<Attribute> kept;
    for (auto [paramName, attr] : llvm::zip_equal(oldParamOrder_, params.getValue())) {
      if (removedParams_.contains(paramName)) {
        if (failed(checkRemovedTemplateParam(
                paramName, attr, diagnosticOp, resolveTemplateSymbolArgs,
                allowCurrentTemplateTypeVarSymbols
            ))) {
          return failure();
        }
        continue;
      }
      kept.push_back(resolveTemplateSymbolArgs ? convertTemplateArgAttr(attr) : attr);
    }
    return kept.empty() ? ArrayAttr() : ArrayAttr::get(ctx_, kept);
  }

  bool hasWildcardForRemovedParam(ArrayAttr params) const {
    if (!params || params.size() != oldParamOrder_.size()) {
      return false;
    }
    for (auto [paramName, attr] : llvm::zip_equal(oldParamOrder_, params.getValue())) {
      if (!removedParams_.contains(paramName)) {
        continue;
      }
      if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr); intAttr && isDynamic(intAttr)) {
        return true;
      }
    }
    return false;
  }
 
private:
  Attribute convertTemplateArgAttr(Attribute attr) const {
    DenseSet<StringAttr> resolvingParams;
    return convertTemplateArgAttr(attr, resolvingParams);
  }

  Attribute convertTemplateArgAttr(Attribute attr, DenseSet<StringAttr> &resolvingParams) const {
    if (StringAttr symbolName = getFlatSymbolName(attr)) {
      auto it = replacements_.find(symbolName);
      if (it != replacements_.end()) {
        if (!resolvingParams.insert(symbolName).second) {
          return attr;
        }
        Type replacement = convertType(it->second, resolvingParams);
        resolvingParams.erase(symbolName);
        return TypeAttr::get(replacement);
      }
    }
    return convertAttr(attr, resolvingParams);
  }

  LogicalResult checkRemovedTemplateParam(
      StringAttr paramName, Attribute attr, Operation *diagnosticOp, bool resolveTemplateSymbolArgs,
      bool allowCurrentTemplateTypeVarSymbols = false
  ) const {
    auto replacementIt = replacements_.find(paramName);
    assert(replacementIt != replacements_.end() && "removed parameter must have a replacement");
    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr); intAttr && isDynamic(intAttr)) {
      return success();
    }
    Attribute convertedAttr = resolveTemplateSymbolArgs ? convertTemplateArgAttr(attr) : attr;
    if (StringAttr symbolName = getFlatSymbolName(convertedAttr)) {
      if (allowCurrentTemplateTypeVarSymbols &&
          isCurrentTemplateTypeVarParamSymbol(symbolName, diagnosticOp)) {
        return success();
      }
      return emitRemovedTemplateParamMismatch(paramName, attr, replacementIt->second, diagnosticOp);
    }
    if (templateArgUnifiesWithType(convertedAttr, convertType(replacementIt->second))) {
      return success();
    }
 
    return emitRemovedTemplateParamMismatch(paramName, attr, replacementIt->second, diagnosticOp);
  }

  LogicalResult emitRemovedTemplateParamMismatch(
      StringAttr paramName, Attribute attr, Type replacementTy, Operation *diagnosticOp
  ) const {
    InFlightDiagnostic diag = diagnosticOp->emitError()
                              << "explicit template argument for inferred parameter @"
                              << paramName.getValue() << " must match inferred type "
                              << replacementTy << ", but found ";
    if (auto typeAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      diag << typeAttr.getValue();
    } else {
      diag << attr;
    }
    return diag;
  }

  bool isOwnedStructType(StructType structTy) const {
    SmallVector<StringAttr> structPath = getStringPieces(structTy.getNameRef());
    return structPath.size() > templatePath_.size() &&
           std::equal(templatePath_.begin(), templatePath_.end(), structPath.begin());
  }

  StructType convertStructType(StructType structTy, DenseSet<StringAttr> &resolvingParams) const {
    ArrayAttr params = structTy.getParams();
    if (!params) {
      return structTy;
    }
 
    SmallVector<Attribute> newParams;
    bool changed = false;
    bool removeOwnedParams = isOwnedStructType(structTy) && params.size() == oldParamOrder_.size();
    for (auto [index, attr] : llvm::enumerate(params.getValue())) {
      if (trimResolvedParams_ && removeOwnedParams &&
          removedParams_.contains(oldParamOrder_[index])) {
        changed = true;
        continue;
      }
      Attribute newAttr = convertTemplateArgAttr(attr, resolvingParams);
      newParams.push_back(newAttr);
      changed |= newAttr != attr;
    }
    return changed ? getStructTypeWithParams(structTy.getNameRef(), ctx_, newParams) : structTy;
  }

  LogicalResult validateStructType(StructType structTy, Operation *diagnosticOp) const {
    ArrayAttr params = structTy.getParams();
    if (!params) {
      return success();
    }
    bool removeOwnedParams = isOwnedStructType(structTy) && params.size() == oldParamOrder_.size();
    for (auto indexedAttr : llvm::enumerate(params.getValue())) {
      unsigned index = indexedAttr.index();
      Attribute attr = indexedAttr.value();
      if (trimResolvedParams_ && removeOwnedParams &&
          removedParams_.contains(oldParamOrder_[index])) {
        if (failed(checkRemovedTemplateParam(
                oldParamOrder_[index], attr, diagnosticOp, /*resolveTemplateSymbolArgs=*/true,
                /*allowCurrentTemplateTypeVarSymbols=*/false
            ))) {
          return failure();
        }
        continue;
      }
      if (failed(validateAttr(attr, diagnosticOp))) {
        return failure();
      }
    }
    return success();
  }
};

template <typename ConverterT, typename SetSignatureFn>
static void updateCallableSignature(
    FunctionType oldFuncTy, Region &body, ConverterT &converter, SetSignatureFn setSignature
) {
  Type converted = converter.convertType(oldFuncTy);
  auto newFuncTy = llvm::cast<FunctionType>(converted);
  if (oldFuncTy == newFuncTy) {
    return;
  }
 
  setSignature(newFuncTy);
  if (body.empty()) {
    return;
  }
 
  Block &entryBlock = body.front();
  assert(entryBlock.getNumArguments() == newFuncTy.getNumInputs());
  for (auto [arg, newTy] : llvm::zip_equal(entryBlock.getArguments(), newFuncTy.getInputs())) {
    arg.setType(newTy);
  }
}

template <typename ConverterT>
static void updateFuncSignature(FuncDefOp func, ConverterT &converter) {
  updateCallableSignature(
      func.getFunctionType(), func.getFunctionBody(), converter,
      [func](FunctionType newFuncTy) mutable { func.setType(newFuncTy); }
  );
}

template <typename ConverterT>
static void updateContractSignature(verif::ContractOp contract, ConverterT &converter) {
  updateCallableSignature(
      contract.getFunctionType(), contract.getBody(), converter,
      [contract](FunctionType newFuncTy) mutable { contract.setFunctionType(newFuncTy); }
  );
}

template <typename ConverterT> static bool convertCallCallee(CallOp callOp, ConverterT &converter) {
  if constexpr (requires { converter.convertCallCallee(callOp); }) {
    SymbolRefAttr oldCallee = callOp.getCalleeAttr();
    SymbolRefAttr newCallee = converter.convertCallCallee(callOp);
    if (newCallee != oldCallee) {
      callOp.setCalleeAttr(newCallee);
      return true;
    }
  }
  return false;
}

template <typename ConverterT>
static bool convertContractTarget(verif::ContractOp contract, ConverterT &converter) {
  if constexpr (requires { converter.convertContractTarget(contract); }) {
    SymbolRefAttr oldTarget = contract.getTargetAttr();
    SymbolRefAttr newTarget = converter.convertContractTarget(contract);
    if (newTarget != oldTarget) {
      contract.setTargetAttr(newTarget);
      return true;
    }
  }
  return false;
}

template <typename ConverterT> static bool converterFailed(ConverterT &converter) {
  if constexpr (requires { converter.hadFailure(); }) {
    return converter.hadFailure();
  }
  return false;
}

template <typename ConverterT>
static void startConverterOperation(ConverterT &converter, Operation *op) {
  if constexpr (requires { converter.startOperation(op); }) {
    converter.startOperation(op);
  }
}

template <typename ConverterT>
static FailureOr<bool> convertOperationTypes(Operation *op, ConverterT &converter) {
  if (auto createOp = llvm::dyn_cast<CreateArrayOp>(op)) {
    ArrayType oldResultTy = createOp.getType();
    auto newResultTy = llvm::dyn_cast<ArrayType>(converter.convertType(oldResultTy));
    auto newElemTy = llvm::dyn_cast<ArrayType>(converter.convertType(oldResultTy.getElementType()));
    if (converterFailed(converter)) {
      return failure();
    }
    if (newResultTy && newElemTy && newResultTy != oldResultTy && !createOp.getElements().empty() &&
        oldResultTy.hasStaticShape()) {
      // Validate every init value in the `array.new` before creating replacement ops to avoid
      // dangling IR if one of the init values is invalid.
      SmallVector<Type> newElementValueTypes;
      newElementValueTypes.reserve(createOp.getElements().size());
      for (auto [index, element] : llvm::enumerate(createOp.getElements())) {
        Type newElementValueTy = converter.convertType(element.getType());
        if (converterFailed(converter)) {
          return failure();
        }
        if (newElementValueTy != newElemTy) {
          createOp.emitError() << "cannot rewrite initialized array.new: initializer " << index
                               << " converts to " << newElementValueTy << ", but expected "
                               << newElemTy;
          return failure();
        }
        newElementValueTypes.push_back(newElementValueTy);
      }
 
      OpBuilder builder(createOp);
      Location loc = createOp.getLoc();
      CreateArrayOp newCreate = builder.create<CreateArrayOp>(loc, newResultTy);
      ArrayIndexGen idxGen = ArrayIndexGen::from(oldResultTy);
      for (auto [index, element] : llvm::enumerate(createOp.getElements())) {
        Type newElementValueTy = newElementValueTypes[index];
        if (element.getType() != newElementValueTy) {
          element.setType(newElementValueTy);
        }
        std::optional<SmallVector<Value>> indices =
            idxGen.delinearize(checkedCast<int64_t>(index), loc, builder);
        assert(indices && "static array initializer index should delinearize");
        builder.create<InsertArrayOp>(loc, newCreate.getResult(), ValueRange(*indices), element);
      }
      createOp.getResult().replaceAllUsesWith(newCreate.getResult());
      createOp.erase();
      return true;
    }
  }
 
  if (auto readOp = llvm::dyn_cast<ReadArrayOp>(op)) {
    Type newResultTy = converter.convertType(readOp.getResult().getType());
    if (converterFailed(converter)) {
      return failure();
    }
    if (auto newArrayTy = llvm::dyn_cast<ArrayType>(newResultTy)) {
      OpBuilder builder(readOp);
      auto extractOp = builder.create<ExtractArrayOp>(
          readOp.getLoc(), newArrayTy, readOp.getArrRef(), readOp.getIndices()
      );
      readOp.getResult().replaceAllUsesWith(extractOp.getResult());
      readOp.erase();
      return true;
    }
  }
 
  if (auto writeOp = llvm::dyn_cast<WriteArrayOp>(op)) {
    Type newRvalueTy = converter.convertType(writeOp.getRvalue().getType());
    if (converterFailed(converter)) {
      return failure();
    }
    if (llvm::isa<ArrayType>(newRvalueTy)) {
      OpBuilder builder(writeOp);
      builder.create<InsertArrayOp>(
          writeOp.getLoc(), writeOp.getArrRef(), writeOp.getIndices(), writeOp.getRvalue()
      );
      writeOp.erase();
      return true;
    }
  }
 
  bool changed = false;
 
  if (auto func = llvm::dyn_cast<FuncDefOp>(op)) {
    FunctionType oldFuncTy = func.getFunctionType();
    updateFuncSignature(func, converter);
    if (converterFailed(converter)) {
      return failure();
    }
    changed |= oldFuncTy != func.getFunctionType();
  }
  if (auto contract = llvm::dyn_cast<verif::ContractOp>(op)) {
    FunctionType oldFuncTy = contract.getFunctionType();
    updateContractSignature(contract, converter);
    if (converterFailed(converter)) {
      return failure();
    }
    changed |= oldFuncTy != contract.getFunctionType();
    changed |= convertContractTarget(contract, converter);
    if (converterFailed(converter)) {
      return failure();
    }
  }
 
  for (Region &region : op->getRegions()) {
    for (Block &block : region.getBlocks()) {
      for (BlockArgument arg : block.getArguments()) {
        Type newTy = converter.convertType(arg.getType());
        if (converterFailed(converter)) {
          return failure();
        }
        if (newTy != arg.getType()) {
          arg.setType(newTy);
          changed = true;
        }
      }
    }
  }
 
  for (Value result : op->getResults()) {
    Type newTy = converter.convertType(result.getType());
    if (converterFailed(converter)) {
      return failure();
    }
    if (newTy != result.getType()) {
      result.setType(newTy);
      changed = true;
    }
  }
 
  if (auto callOp = llvm::dyn_cast<CallOp>(op)) {
    changed |= convertCallCallee(callOp, converter);
    if (converterFailed(converter)) {
      return failure();
    }
  }
 
  SmallVector<NamedAttribute> newAttrs;
  bool attrsChanged = false;
  newAttrs.reserve(op->getAttrs().size());
  for (NamedAttribute attr : op->getAttrs()) {
    Attribute newAttr = converter.convertAttr(attr.getValue());
    if (converterFailed(converter)) {
      return failure();
    }
    newAttrs.emplace_back(attr.getName(), newAttr);
    attrsChanged |= newAttr != attr.getValue();
  }
  if (attrsChanged) {
    op->setAttrs(DictionaryAttr::get(op->getContext(), newAttrs));
    changed = true;
  }
 
  return changed;
}

template <typename ConverterT>
static FailureOr<bool> convertOperationTypesInAndTrack(Operation *root, ConverterT &converter) {
  bool changed = false;
  WalkResult res = root->walk([&converter, &changed](Operation *op) {
    startConverterOperation(converter, op);
    FailureOr<bool> opChanged = convertOperationTypes(op, converter);
    if (failed(opChanged)) {
      return WalkResult::interrupt();
    }
    changed |= *opChanged;
    return WalkResult::advance();
  });
  if (res.wasInterrupted()) {
    return failure();
  }
  return changed;
}

template <typename ConverterT>
static LogicalResult convertOperationTypesIn(Operation *root, ConverterT &converter) {
  return failure(failed(convertOperationTypesInAndTrack(root, converter)));
}

static void removeIdentityCasts(Operation *root) {
  for (UnifiableCastOp castOp : walkCollect<UnifiableCastOp>(*root)) {
    if (castOp.getInput().getType() != castOp.getResult().getType()) {
      continue;
    }
    castOp.getResult().replaceAllUsesWith(castOp.getInput());
    castOp.erase();
  }
}

template <typename ConverterT>
static LogicalResult convertTemplateExprTypesIn(TemplateOp templateOp, ConverterT &converter) {
  for (TemplateExprOp expr : templateOp.getConstOps<TemplateExprOp>()) {
    if (failed(convertOperationTypesIn(expr.getOperation(), converter))) {
      return failure();
    }
    removeIdentityCasts(expr.getOperation());
  }
  return success();
}

static std::string buildSpecializedTemplateCloneCacheKey(
    StringRef templateName, ArrayRef<StringAttr> oldParamOrder,
    const DenseMap<Attribute, Attribute> &paramNameToConcrete
) {
  std::string key;
  llvm::raw_string_ostream os(key);
  os << templateName.size() << ':' << templateName;
  for (StringAttr paramName : oldParamOrder) {
    os << '|';
    os << paramName.getValue().size() << ':' << paramName.getValue() << '=';
    auto concreteIt = paramNameToConcrete.find(FlatSymbolRefAttr::get(paramName));
    if (concreteIt == paramNameToConcrete.end()) {
      os << '_';
      continue;
    }
 
    std::string attrText;
    llvm::raw_string_ostream attrOs(attrText);
    concreteIt->second.print(attrOs);
    os << attrText.size() << ':' << attrText;
  }
  return key;
}

static bool canCopyTemplateExprToSpecialization(
    TemplateExprOp expr, const DenseSet<StringAttr> &availableNames
) {
  bool canCopy = true;
  expr->walk([&](ConstReadOp readOp) {
    if (!availableNames.contains(readOp.getConstNameAttr().getAttr())) {
      canCopy = false;
      return WalkResult::interrupt();
    }
    return WalkResult::advance();
  });
  return canCopy;
}

static void copyPreservedTemplateExprs(
    TemplateOp parentTemplate, Block &newTemplateBody, ArrayRef<Attribute> remainingNames
) {
  DenseSet<StringAttr> availableNames;
  for (Attribute name : remainingNames) {
    FlatSymbolRefAttr nameSym = llvm::cast<FlatSymbolRefAttr>(name);
    availableNames.insert(nameSym.getAttr());
  }
 
  for (TemplateExprOp expr : parentTemplate.getConstOps<TemplateExprOp>()) {
    if (canCopyTemplateExprToSpecialization(expr, availableNames)) {
      newTemplateBody.push_back(expr->clone());
    }
  }
}

static FailureOr<TemplateOp> getOrCreateSpecializedTemplateClone(
    TemplateOp parentTemplate, ArrayRef<StringAttr> oldParamOrder,
    const DenseMap<Attribute, Attribute> &paramNameToConcrete, ArrayAttr callParams,
    SymbolTableCollection &tables, llvm::StringMap<StringAttr> &templateClones,
    InstantiationLayout &layout
) {
  FailureOr<InstantiationLayout> layoutResult =
      buildInstantiationLayout(parentTemplate, callParams, paramNameToConcrete);
  if (failed(layoutResult)) {
    return failure();
  }
  layout = std::move(*layoutResult);
  std::string cacheKey = buildSpecializedTemplateCloneCacheKey(
      parentTemplate.getSymName(), oldParamOrder, paramNameToConcrete
  );
 
  ModuleOp parentModule = getParentOfType<ModuleOp>(parentTemplate);
  if (!parentModule) {
    return failure();
  }
  SymbolTable &moduleSymbols = tables.getSymbolTable(parentModule);
 
  auto cachedName = templateClones.find(cacheKey);
  if (cachedName != templateClones.end()) {
    Operation *existing = moduleSymbols.lookup(cachedName->second);
    if (auto existingTemplate = llvm::dyn_cast_or_null<TemplateOp>(existing)) {
      return existingTemplate;
    }
    return failure();
  }
 
  TemplateOp newTemplate = parentTemplate.cloneWithoutRegions();
  newTemplate.setSymName(layout.templateNameWithAttrs);
  setInstantiationNamePattern(
      newTemplate, layout.remainingNames.empty() ? ArrayAttr() : layout.namePattern
  );
  assert(newTemplate->getNumRegions() > 0 && "region exists");
  newTemplate.getBodyRegion().emplaceBlock();
  Block &newTemplateBody = newTemplate.getBodyRegion().front();
  SymbolTable &parentTemplateSymbols = tables.getSymbolTable(parentTemplate);
  for (Attribute name : layout.remainingNames) {
    FlatSymbolRefAttr nameSym = llvm::cast<FlatSymbolRefAttr>(name);
    Operation *paramOp = parentTemplateSymbols.lookup(nameSym.getAttr());
    assert(paramOp && "symbol must exist");
    newTemplateBody.push_back(paramOp->clone());
  }
  copyPreservedTemplateExprs(parentTemplate, newTemplateBody, layout.remainingNames);
 
  moduleSymbols.insert(newTemplate, Block::iterator(parentTemplate));
  templateClones.try_emplace(cacheKey, newTemplate.getSymNameAttr());
  return newTemplate;
}

class ConcreteStructInstantiationConverter {
  struct StructInstantiationTypes {
    StructType localType;
    StructType remoteType;
  };

  MLIRContext *ctx_;
  ModuleOp module_;
  SymbolTableCollection &tables_;
  DenseMap<StructType, StructInstantiationTypes> instantiations_;
  SpecializedTemplateCloneCache templateClones_;
  DenseSet<SymbolRefAttr> instantiatedCloneNames_;
  DenseMap<StructType, StructType> activeLocalStructReplacements_;
  DenseMap<StringAttr, Type> activeTypeReplacements_;
  bool hasFailure = false;
 
public:
  ConcreteStructInstantiationConverter(MLIRContext *c, ModuleOp m, SymbolTableCollection &t)
      : ctx_(c), module_(m), tables_(t) {}

  bool hadFailure() const { return hasFailure; }

  Type convertType(Type ty) {
    if (!ty || hasFailure) {
      return ty;
    }
    if (auto tvarTy = llvm::dyn_cast<TypeVarType>(ty)) {
      auto it = activeTypeReplacements_.find(tvarTy.getNameRef().getAttr());
      return it == activeTypeReplacements_.end() ? ty : it->second;
    }
    if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
      return convertArrayElementType(arrTy, [this](Type elemTy) { return convertType(elemTy); });
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      return convertPodType(podTy, ctx_, *this);
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      return convertFunctionType(funcTy, *this);
    }
    if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
      return convertStructType(structTy);
    }
    return ty;
  }

  Attribute convertAttr(Attribute attr) {
    if (hasFailure) {
      return attr;
    }
    return convertTypeOrArrayAttr(attr, ctx_, [this](Type ty) {
      return convertType(ty);
    }, [this](Attribute nested) { return convertTemplateArgAttr(nested); });
  }

  SymbolRefAttr convertCallCallee(CallOp callOp) {
    SymbolRefAttr callee = callOp.getCalleeAttr();
    if (!callee || callee.getNestedReferences().empty()) {
      return callee;
    }
 
    StructType targetStructTy = getStructFunctionTargetType(callOp);
    if (!targetStructTy) {
      return callee;
    }
    auto convertedStructTy = llvm::dyn_cast<StructType>(convertType(targetStructTy));
    if (!convertedStructTy) {
      return callee;
    }
 
    SymbolRefAttr convertedStructName = convertedStructTy.getNameRef();
    if (convertedStructName == getPrefixAsSymbolRefAttr(callee)) {
      return callee;
    }
    SmallVector<FlatSymbolRefAttr> pieces = getPieces(convertedStructName);
    pieces.push_back(FlatSymbolRefAttr::get(callee.getLeafReference()));
    return asSymbolRefAttr(pieces);
  }

  SymbolRefAttr convertContractTarget(verif::ContractOp contract) {
    SymbolRefAttr target = contract.getTargetAttr();
    if (!target || failed(contract.getStructTarget(tables_))) {
      return target;
    }
 
    FunctionType contractTy = contract.getFunctionType();
    if (contractTy.getNumInputs() == 0) {
      return target;
    }
    auto selfTy = llvm::dyn_cast<StructType>(contractTy.getInput(0));
    if (!selfTy || selfTy.getNameRef() == target) {
      return target;
    }
    return instantiatedCloneNames_.contains(selfTy.getNameRef()) ? selfTy.getNameRef() : target;
  }
 
private:
  static StructType getStructFunctionTargetType(CallOp callOp) {
    StringAttr calleeLeaf = callOp.getCallee().getLeafReference();
    if (calleeLeaf == FUNC_NAME_CONSTRAIN) {
      return dyn_cast<StructType>(callOp.getSelfValueFromConstrain().getType());
    }
    if (calleeLeaf == FUNC_NAME_COMPUTE || calleeLeaf == FUNC_NAME_PRODUCT) {
      return callOp.getSingleResultTypeOfWitnessGen();
    }
    return {};
  }

  Attribute convertTemplateArgAttr(Attribute attr) {
    if (StringAttr symbolName = getFlatSymbolName(attr)) {
      auto it = activeTypeReplacements_.find(symbolName);
      if (it != activeTypeReplacements_.end()) {
        return TypeAttr::get(it->second);
      }
    }
    return convertAttr(attr);
  }

  StructType convertStructType(StructType structTy) {
    ArrayAttr params = structTy.getParams();
    if (isNullOrEmpty(params)) {
      return structTy;
    }
 
    SmallVector<Attribute> newParams;
    bool changed = false;
    for (Attribute attr : params.getValue()) {
      Attribute newAttr = convertTemplateArgAttr(attr);
      newParams.push_back(newAttr);
      changed |= newAttr != attr;
    }
    StructType convertedTy =
        changed ? StructType::get(structTy.getNameRef(), ArrayAttr::get(ctx_, newParams))
                : structTy;
 
    if (auto it = activeLocalStructReplacements_.find(convertedTy);
        it != activeLocalStructReplacements_.end()) {
      return it->second;
    }
    if (instantiatedCloneNames_.contains(convertedTy.getNameRef())) {
      return convertedTy;
    }
 
    if (llvm::any_of(newParams, [](Attribute attr) {
      return !isConcreteStructParamAttr(attr, /*allowStructParams=*/false);
    })) {
      return convertedTy;
    }
 
    FailureOr<StructType> cloneTy = getOrCreateStructClone(convertedTy, newParams);
    if (failed(cloneTy)) {
      hasFailure = true;
      return convertedTy;
    }
    return *cloneTy;
  }

  FailureOr<StructType>
  getOrCreateStructClone(StructType concreteStructTy, ArrayRef<Attribute> concreteParams) {
    if (auto it = instantiations_.find(concreteStructTy); it != instantiations_.end()) {
      return it->second.remoteType;
    }
 
    FailureOr<SymbolLookupResult<StructDefOp>> lookup =
        concreteStructTy.getDefinition(tables_, module_, /*emitError=*/false);
    if (failed(lookup)) {
      return failure();
    }
 
    StructDefOp origStruct = lookup->get();
    TemplateOp parentTemplate = getParentOfType<TemplateOp>(origStruct);
    if (!parentTemplate) {
      return failure();
    }
    StructType typeAtDef = origStruct.getType();
    ArrayAttr paramNames = typeAtDef.getParams();
    if (!paramNames || paramNames.size() != concreteParams.size()) {
      return failure();
    }
 
    DenseMap<StringAttr, Type> typeReplacements;
    DenseMap<Attribute, Attribute> paramNameToConcrete;
    SmallVector<StringAttr> oldParamOrder;
    oldParamOrder.reserve(paramNames.size());
    SmallVector<Attribute> convertedSourceParams;
    convertedSourceParams.reserve(concreteParams.size());
    for (auto [paramName, concreteAttr] : llvm::zip_equal(paramNames.getValue(), concreteParams)) {
      auto paramSym = llvm::dyn_cast<FlatSymbolRefAttr>(paramName);
      if (!paramSym) {
        return failure();
      }
      oldParamOrder.push_back(paramSym.getAttr());
      auto concreteType = llvm::dyn_cast<TypeAttr>(concreteAttr);
      if (concreteType) {
        paramNameToConcrete.try_emplace(FlatSymbolRefAttr::get(paramSym.getAttr()), concreteAttr);
        typeReplacements.try_emplace(paramSym.getAttr(), concreteType.getValue());
        convertedSourceParams.push_back(concreteType);
      } else {
        convertedSourceParams.push_back(paramName);
      }
    }
    if (paramNameToConcrete.empty()) {
      return concreteStructTy;
    }
 
    InstantiationLayout layout;
    ArrayAttr concreteParamArray = ArrayAttr::get(ctx_, concreteParams);
    FailureOr<TemplateOp> newTemplate = getOrCreateSpecializedTemplateClone(
        parentTemplate, oldParamOrder, paramNameToConcrete, concreteParamArray, tables_,
        templateClones_[parentTemplate.getOperation()], layout
    );
    if (failed(newTemplate)) {
      return failure();
    }
 
    SymbolTable &templateSymbols = tables_.getSymbolTable(*newTemplate);
    StructDefOp clone = origStruct.clone();
    templateSymbols.insert(clone);
    StructType localTy =
        getStructTypeWithParams(clone.getFullyQualifiedName(), clone.getType().getParams());
    StructType remoteTy =
        getStructTypeWithParams(clone.getFullyQualifiedName(), layout.rewrittenCallParams);
    instantiations_.try_emplace(concreteStructTy, StructInstantiationTypes {localTy, remoteTy});
    instantiatedCloneNames_.insert(clone.getFullyQualifiedName());
 
    DenseMap<StringAttr, Type> previousReplacements = std::move(activeTypeReplacements_);
    DenseMap<StructType, StructType> previousLocalStructReplacements =
        std::move(activeLocalStructReplacements_);
    activeTypeReplacements_ = std::move(typeReplacements);
    activeLocalStructReplacements_.clear();
    activeLocalStructReplacements_.try_emplace(concreteStructTy, localTy);
    activeLocalStructReplacements_.try_emplace(
        StructType::get(typeAtDef.getNameRef(), ArrayAttr::get(ctx_, convertedSourceParams)),
        localTy
    );
    activeLocalStructReplacements_.try_emplace(
        StructType::get(localTy.getNameRef(), ArrayAttr::get(ctx_, convertedSourceParams)), localTy
    );
    activeLocalStructReplacements_.try_emplace(
        StructType::get(localTy.getNameRef(), ArrayAttr::get(ctx_, concreteParams)), localTy
    );
    SymbolRefAttr cloneNameRef = clone.getFullyQualifiedName();
    if (failed(convertTemplateExprTypesIn(*newTemplate, *this))) {
      activeLocalStructReplacements_ = std::move(previousLocalStructReplacements);
      activeTypeReplacements_ = std::move(previousReplacements);
      instantiations_.erase(concreteStructTy);
      instantiatedCloneNames_.erase(cloneNameRef);
      clone.erase();
      return failure();
    }
    if (failed(convertOperationTypesIn(clone.getOperation(), *this))) {
      activeLocalStructReplacements_ = std::move(previousLocalStructReplacements);
      activeTypeReplacements_ = std::move(previousReplacements);
      instantiations_.erase(concreteStructTy);
      instantiatedCloneNames_.erase(cloneNameRef);
      clone.erase();
      return failure();
    }
    removeIdentityCasts(clone.getOperation());
    activeLocalStructReplacements_ = std::move(previousLocalStructReplacements);
    activeTypeReplacements_ = std::move(previousReplacements);
    return remoteTy;
  }
};

struct TemplateInferenceInfo {
  TemplateOp templateOp;
  SmallVector<StringAttr> templatePath;
  SmallVector<StringAttr> oldParamOrder;
  DenseMap<StringAttr, TemplateParamOp> typeVarParams;
  DenseMap<StringAttr, InferredType> replacements;
  DenseMap<StringAttr, InferredType> templateScopeReplacements;
  DenseMap<Operation *, DenseMap<StringAttr, InferredType>> functionReplacements;
};

static DenseMap<Attribute, Attribute>
buildParamNameToCallArg(ArrayAttr callParams, ArrayRef<StringAttr> oldParamOrder) {
  DenseMap<Attribute, Attribute> paramNameToCallArg;
  if (!callParams || callParams.size() != oldParamOrder.size()) {
    return paramNameToCallArg;
  }
  for (auto [paramName, attr] : llvm::zip_equal(oldParamOrder, callParams.getValue())) {
    paramNameToCallArg.try_emplace(FlatSymbolRefAttr::get(paramName), attr);
  }
  return paramNameToCallArg;
}

static Type
substituteExplicitCallTypeArgs(Type ty, ArrayAttr callParams, ArrayRef<StringAttr> oldParamOrder) {
  if (!callParams || callParams.size() != oldParamOrder.size()) {
    return ty;
  }
 
  DenseMap<StringAttr, Type> callTypeArgs;
  for (auto [paramName, attr] : llvm::zip_equal(oldParamOrder, callParams.getValue())) {
    if (auto tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      callTypeArgs.try_emplace(paramName, tyAttr.getValue());
    }
  }
  TypeVarReplacementConverter converter(
      ty.getContext(), ArrayRef<StringAttr> {}, oldParamOrder, callTypeArgs,
      /*trimResolvedParams=*/false
  );
  return converter.convertType(ty);
}

class ExplicitCallTemplateParamSubstituter {
  const DenseMap<Attribute, Attribute> &paramNameToCallArg_;
  DenseSet<Attribute> resolvingParams_;
 
public:
  explicit ExplicitCallTemplateParamSubstituter(
      const DenseMap<Attribute, Attribute> &paramNameToCallArg
  )
      : paramNameToCallArg_(paramNameToCallArg) {}
 
  Type substituteType(Type ty);
  Attribute substituteAttr(Attribute attr);
};

Attribute ExplicitCallTemplateParamSubstituter::substituteAttr(Attribute attr) {
  if (!attr) {
    return attr;
  }
  auto it = paramNameToCallArg_.find(attr);
  if (it != paramNameToCallArg_.end()) {
    return it->second;
  }
  if (auto tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
    Type newTy = substituteType(tyAttr.getValue());
    return newTy == tyAttr.getValue() ? attr : TypeAttr::get(newTy);
  }
  if (auto arrAttr = llvm::dyn_cast<ArrayAttr>(attr)) {
    SmallVector<Attribute> newAttrs;
    bool changed = false;
    for (Attribute nested : arrAttr.getValue()) {
      Attribute newNested = substituteAttr(nested);
      newAttrs.push_back(newNested);
      changed |= newNested != nested;
    }
    return changed ? ArrayAttr::get(attr.getContext(), newAttrs) : attr;
  }
  return attr;
}

Type ExplicitCallTemplateParamSubstituter::substituteType(Type ty) {
  if (!ty) {
    return ty;
  }
  if (auto tvarTy = llvm::dyn_cast<TypeVarType>(ty)) {
    Attribute paramRef = tvarTy.getNameRef();
    auto it = paramNameToCallArg_.find(paramRef);
    if (it == paramNameToCallArg_.end()) {
      return ty;
    }
    auto tyAttr = llvm::dyn_cast<TypeAttr>(it->second);
    if (!tyAttr || tyAttr.getValue() == ty) {
      return ty;
    }
    if (!resolvingParams_.insert(paramRef).second) {
      return ty;
    }
    Type replacement = substituteType(tyAttr.getValue());
    resolvingParams_.erase(paramRef);
    return replacement;
  }
  if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
    SmallVector<Attribute> newDims;
    bool changed = false;
    for (Attribute dim : arrTy.getDimensionSizes()) {
      Attribute newDim = substituteAttr(dim);
      newDims.push_back(newDim);
      changed |= newDim != dim;
    }
    Type newElemTy = substituteType(arrTy.getElementType());
    if (!changed && newElemTy == arrTy.getElementType()) {
      return arrTy;
    }
    return flattenInstantiatedArrayType(
        arrTy.cloneWith(arrTy.getElementType(), newDims), newElemTy
    );
  }
  if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
    ArrayAttr params = structTy.getParams();
    if (!params) {
      return structTy;
    }
    SmallVector<Attribute> newParams;
    bool changed = false;
    for (Attribute param : params.getValue()) {
      Attribute newParam = substituteAttr(param);
      newParams.push_back(newParam);
      changed |= newParam != param;
    }
    return changed ? getStructTypeWithParams(structTy.getNameRef(), ty.getContext(), newParams)
                   : structTy;
  }
  if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
    SmallVector<RecordAttr> newRecords;
    bool changed = false;
    for (RecordAttr record : podTy.getRecords()) {
      Type newRecordTy = substituteType(record.getType());
      newRecords.push_back(RecordAttr::get(ty.getContext(), record.getName(), newRecordTy));
      changed |= newRecordTy != record.getType();
    }
    return changed ? PodType::get(ty.getContext(), newRecords) : podTy;
  }
  if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
    SmallVector<Type> newInputs;
    SmallVector<Type> newResults;
    bool changed = false;
    newInputs.reserve(funcTy.getNumInputs());
    newResults.reserve(funcTy.getNumResults());
    for (Type inputTy : funcTy.getInputs()) {
      Type newInputTy = substituteType(inputTy);
      newInputs.push_back(newInputTy);
      changed |= newInputTy != inputTy;
    }
    for (Type resultTy : funcTy.getResults()) {
      Type newResultTy = substituteType(resultTy);
      newResults.push_back(newResultTy);
      changed |= newResultTy != resultTy;
    }
    return changed ? FunctionType::get(ty.getContext(), newInputs, newResults) : funcTy;
  }
  return ty;
}

static Type substituteExplicitCallTemplateParams(
    Type ty, ArrayAttr callParams, ArrayRef<StringAttr> oldParamOrder
) {
  auto paramNameToCallArg = buildParamNameToCallArg(callParams, oldParamOrder);
  ExplicitCallTemplateParamSubstituter substituter(paramNameToCallArg);
  return substituter.substituteType(ty);
}

static bool symbolMatchesParam(SymbolRefAttr symRef, StringAttr paramName) {
  return symRef && symRef.getNestedReferences().empty() && symRef.getRootReference() == paramName;
}

class ParamMentionChecker {
  StringAttr paramName_;
 
public:
  explicit ParamMentionChecker(StringAttr paramName) : paramName_(paramName) {}

  bool typeMentions(Type ty) const {
    if (!ty) {
      return false;
    }
    if (auto tvarTy = llvm::dyn_cast<TypeVarType>(ty)) {
      return tvarTy.getNameRef().getAttr() == paramName_;
    }
    if (auto arrayTy = llvm::dyn_cast<ArrayType>(ty)) {
      return typeMentions(arrayTy.getElementType()) ||
             llvm::any_of(arrayTy.getDimensionSizes(), [this](Attribute dim) {
        return attrMentions(dim, /*allowSymbolRefs=*/true);
      });
    }
    if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
      ArrayAttr params = structTy.getParams();
      return params && attrMentions(params, /*allowSymbolRefs=*/true);
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      return llvm::any_of(podTy.getRecords(), [this](RecordAttr record) {
        return typeMentions(record.getType());
      });
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      return llvm::any_of(funcTy.getInputs(), [this](Type input) {
        return typeMentions(input);
      }) || llvm::any_of(funcTy.getResults(), [this](Type result) { return typeMentions(result); });
    }
    return false;
  }

  bool attrMentions(Attribute attr, bool allowSymbolRefs) const {
    if (!attr) {
      return false;
    }
    if (auto typeAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      return typeMentions(typeAttr.getValue());
    }
    if (auto arrayAttr = llvm::dyn_cast<ArrayAttr>(attr)) {
      return llvm::any_of(arrayAttr, [this](Attribute nested) {
        return attrMentions(nested, /*allowSymbolRefs=*/true);
      });
    }
    return allowSymbolRefs && symbolMatchesParam(llvm::dyn_cast<SymbolRefAttr>(attr), paramName_);
  }
};

static bool operationMentionsParam(Operation *op, StringAttr paramName) {
  ParamMentionChecker mentions(paramName);
  if (llvm::any_of(op->getResultTypes(), [&mentions](Type ty) {
    return mentions.typeMentions(ty);
  })) {
    return true;
  }
  for (Region &region : op->getRegions()) {
    for (Block &block : region.getBlocks()) {
      if (llvm::any_of(block.getArgumentTypes(), [&mentions](Type ty) {
        return mentions.typeMentions(ty);
      })) {
        return true;
      }
    }
  }
  return llvm::any_of(op->getAttrs(), [&mentions](NamedAttribute attr) {
    return mentions.attrMentions(attr.getValue(), /*allowSymbolRefs=*/false);
  });
}

static SmallVector<StringAttr> getMentionedTypeVarParams(
    Attribute attr, const DenseMap<StringAttr, TemplateParamOp> &typeVarParams
) {
  SmallVector<StringAttr> mentionedParams;
  for (const auto &entry : typeVarParams) {
    if (ParamMentionChecker(entry.first).attrMentions(attr, /*allowSymbolRefs=*/true)) {
      mentionedParams.push_back(entry.first);
    }
  }
  return mentionedParams;
}

static inline bool attrMentionsAnyTypeVarParam(
    Attribute attr, const DenseMap<StringAttr, TemplateParamOp> &typeVarParams
) {
  return !getMentionedTypeVarParams(attr, typeVarParams).empty();
}

static bool funcMentionsParam(FuncDefOp func, StringAttr paramName) {
  if (operationMentionsParam(func.getOperation(), paramName)) {
    return true;
  }
  return walkContains<Operation *>(func, [paramName](Operation *op) {
    return operationMentionsParam(op, paramName);
  });
}

static bool contractMentionsParam(verif::ContractOp contract, StringAttr paramName) {
  if (operationMentionsParam(contract.getOperation(), paramName)) {
    return true;
  }
  return walkContains<Operation *>(contract, [paramName](Operation *op) {
    return operationMentionsParam(op, paramName);
  });
}

static inline bool targetMentionsParam(FuncDefOp func, StringAttr paramName) {
  return funcMentionsParam(func, paramName);
}

static inline bool targetMentionsParam(verif::ContractOp contract, StringAttr paramName) {
  return contractMentionsParam(contract, paramName);
}

static inline bool exprMentionsParam(TemplateExprOp expr, StringAttr paramName) {
  return walkContains<Operation *>(expr, [paramName](Operation *op) {
    return operationMentionsParam(op, paramName);
  });
}

static bool structUseCoveredByFunctionProof(
    StructDefOp structOp, StringAttr paramName, Type replacementType,
    const DenseMap<Operation *, DenseMap<StringAttr, InferredType>> &functionReplacements
) {
  bool sawMention = false;
  bool missingProof = false;
  structOp.walk([&](FuncDefOp func) {
    if (func->getParentOfType<StructDefOp>() != structOp || !funcMentionsParam(func, paramName)) {
      return;
    }
    sawMention = true;
    auto funcIt = functionReplacements.find(func.getOperation());
    if (funcIt == functionReplacements.end()) {
      missingProof = true;
      return;
    }
    auto replacementIt = funcIt->second.find(paramName);
    if (replacementIt == funcIt->second.end() || replacementIt->second.type != replacementType) {
      missingProof = true;
    }
  });
  return sawMention && !missingProof;
}

static bool hasUncoveredNonFunctionMention(
    TemplateOp templateOp, StringAttr paramName, Type replacementType,
    const DenseMap<Operation *, DenseMap<StringAttr, InferredType>> &functionReplacements
) {
  WalkResult result =
      templateOp.walk([paramName, replacementType, &functionReplacements](Operation *op) {
    if (llvm::isa<FuncDefOp, TemplateExprOp, TemplateParamOp, verif::ContractOp>(op) ||
        hasParentThatIsa<FuncDefOp, TemplateExprOp, verif::ContractOp>(op)) {
      return WalkResult::advance();
    }
    if (!operationMentionsParam(op, paramName)) {
      return WalkResult::advance();
    }
    if (StructDefOp structOp = op->getParentOfType<StructDefOp>()) {
      if (structUseCoveredByFunctionProof(
              structOp, paramName, replacementType, functionReplacements
          )) {
        return WalkResult::advance();
      }
    }
    return WalkResult::interrupt();
  });
  return result.wasInterrupted();
}

class TypeVarInferenceCollector {
  TemplateInferenceInfo &inferenceInfo;
  DenseMap<StringAttr, InferredType> &replacements;
  DenseMap<Value, InferredType> byValue;
  DenseMap<StringAttr, DenseSet<StringAttr>> paramRelations;
  bool changedInIteration = false;
 
public:
  TypeVarInferenceCollector(
      TemplateInferenceInfo &info, DenseMap<StringAttr, InferredType> &scopeReplacements
  )
      : inferenceInfo(info), replacements(scopeReplacements) {}

  LogicalResult collect(Operation *root) {
    do {
      changedInIteration = false;
      auto result = root->walk([this](UnifiableCastOp castOp) -> WalkResult {
        return collectTypePairInferences(
            castOp.getInput().getType(), castOp.getResult().getType(), castOp.getInput(),
            castOp.getResult(), castOp.getLoc()
        );
      });
      if (result.wasInterrupted()) {
        return failure();
      }
    } while (changedInIteration);
    return success();
  }

  LogicalResult collectTypeInferences(Type lhs, Type rhs, Location loc) {
    do {
      changedInIteration = false;
      if (failed(collectTypePairInferences(lhs, rhs, Value(), Value(), loc))) {
        return failure();
      }
    } while (changedInIteration);
    return success();
  }

  LogicalResult collectStructTemplateParamInferences(
      Operation *root, ModuleOp module, DenseMap<Operation *, TemplateInferenceInfo *> &infos,
      SymbolTableCollection &tables
  ) {
    auto result = root->walk([&](Operation *op) -> WalkResult {
      return collectOperationStructTemplateParamInferences(op, module, infos, tables);
    });
    return failure(result.wasInterrupted());
  }

  LogicalResult collectOperationStructTemplateParamInferences(
      Operation *op, ModuleOp module, DenseMap<Operation *, TemplateInferenceInfo *> &infos,
      SymbolTableCollection &tables
  ) {
    Location loc = op->getLoc();
 
    if (auto func = llvm::dyn_cast<FuncDefOp>(op)) {
      if (failed(collectStructTemplateParamInferences(
              func.getFunctionType(), loc, module, infos, tables
          ))) {
        return failure();
      }
    }
 
    for (Region &region : op->getRegions()) {
      for (Block &block : region.getBlocks()) {
        for (Type argTy : block.getArgumentTypes()) {
          if (failed(collectStructTemplateParamInferences(argTy, loc, module, infos, tables))) {
            return failure();
          }
        }
      }
    }
 
    for (Type resultTy : op->getResultTypes()) {
      if (failed(collectStructTemplateParamInferences(resultTy, loc, module, infos, tables))) {
        return failure();
      }
    }
 
    if (auto callOp = llvm::dyn_cast<CallOp>(op)) {
      if (failed(collectCallableTemplateParamInferences(callOp, infos, tables))) {
        return failure();
      }
    }
    if (auto includeOp = llvm::dyn_cast<verif::IncludeOp>(op)) {
      if (failed(collectIncludeTemplateParamInferences(includeOp, infos, tables))) {
        return failure();
      }
    }
 
    for (NamedAttribute attr : op->getAttrs()) {
      if (failed(
              collectStructTemplateParamInferences(attr.getValue(), loc, module, infos, tables)
          )) {
        return failure();
      }
    }
    return success();
  }
 
private:
  LogicalResult collectStructTemplateParamInferences(
      Attribute attr, Location loc, ModuleOp module,
      DenseMap<Operation *, TemplateInferenceInfo *> &infos, SymbolTableCollection &tables
  ) {
    if (!attr) {
      return success();
    }
    if (auto tyAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      return collectStructTemplateParamInferences(tyAttr.getValue(), loc, module, infos, tables);
    }
    if (auto arrAttr = llvm::dyn_cast<ArrayAttr>(attr)) {
      for (Attribute nested : arrAttr.getValue()) {
        if (failed(collectStructTemplateParamInferences(nested, loc, module, infos, tables))) {
          return failure();
        }
      }
    }
    return success();
  }

  LogicalResult collectStructTemplateParamInferences(
      Type ty, Location loc, ModuleOp module, DenseMap<Operation *, TemplateInferenceInfo *> &infos,
      SymbolTableCollection &tables
  ) {
    if (!ty) {
      return success();
    }
    if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
      return collectStructTemplateParamInferences(
          arrTy.getElementType(), loc, module, infos, tables
      );
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      for (RecordAttr record : podTy.getRecords()) {
        if (failed(
                collectStructTemplateParamInferences(record.getType(), loc, module, infos, tables)
            )) {
          return failure();
        }
      }
      return success();
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      for (Type inputTy : funcTy.getInputs()) {
        if (failed(collectStructTemplateParamInferences(inputTy, loc, module, infos, tables))) {
          return failure();
        }
      }
      for (Type resultTy : funcTy.getResults()) {
        if (failed(collectStructTemplateParamInferences(resultTy, loc, module, infos, tables))) {
          return failure();
        }
      }
      return success();
    }
    auto structTy = llvm::dyn_cast<StructType>(ty);
    if (!structTy) {
      return success();
    }
 
    ArrayAttr params = structTy.getParams();
    if (!params) {
      return success();
    }
    for (Attribute attr : params.getValue()) {
      if (failed(collectStructTemplateParamInferences(attr, loc, module, infos, tables))) {
        return failure();
      }
    }
 
    FailureOr<SymbolLookupResult<StructDefOp>> lookup =
        structTy.getDefinition(tables, module, /*reportMissing=*/false);
    if (failed(lookup)) {
      return success();
    }
    TemplateOp parentTemplate = getParentOfType<TemplateOp>(lookup->get().getOperation());
    if (!parentTemplate) {
      return success();
    }
    TemplateInferenceInfo *targetInfo = infos.lookup(parentTemplate.getOperation());
    if (!targetInfo || targetInfo->replacements.empty() ||
        params.size() != targetInfo->oldParamOrder.size()) {
      return success();
    }
 
    for (auto [paramName, attr] : llvm::zip_equal(targetInfo->oldParamOrder, params.getValue())) {
      auto replacementIt = targetInfo->replacements.find(paramName);
      if (replacementIt == targetInfo->replacements.end()) {
        continue;
      }
      Type expectedTy = substituteExplicitCallTemplateParams(
          replacementIt->second.type, params, targetInfo->oldParamOrder
      );
      Attribute expectedAttr = TypeAttr::get(expectedTy);
      if (failed(collectTemplateArgInferences(attr, expectedAttr, loc))) {
        return failure();
      }
    }
    return success();
  }

  SmallVector<Attribute>
  getInferredOmittedTemplateArgs(const UnificationMap &unifyResult, StringAttr targetParamName) {
    SmallVector<Attribute> inferredAttrs;
    auto targetRef = FlatSymbolRefAttr::get(targetParamName);
    auto inferredIt = unifyResult.find({targetRef, Side::RHS});
    if (inferredIt != unifyResult.end()) {
      inferredAttrs.push_back(inferredIt->second);
    }
    for (const auto &entry : unifyResult) {
      if (entry.first.second == Side::LHS && entry.second == targetRef) {
        inferredAttrs.push_back(entry.first.first);
      }
    }
    return inferredAttrs;
  }

  template <typename TargetOp>
  TemplateInferenceInfo *
  getTargetTemplateInfo(TargetOp targetOp, DenseMap<Operation *, TemplateInferenceInfo *> &infos) {
    TemplateOp parentTemplate = getParentOfType<TemplateOp>(targetOp.getOperation());
    if (!parentTemplate) {
      return nullptr;
    }
    TemplateInferenceInfo *targetInfo = infos.lookup(parentTemplate.getOperation());
    if (!targetInfo || targetInfo->replacements.empty()) {
      return nullptr;
    }
    return targetInfo;
  }

  template <typename CallableOp>
  LogicalResult collectCallableUseTemplateParamInferences(
      CallableOp callableOp, FunctionType targetSignature, TemplateInferenceInfo &targetInfo
  ) {
    ArrayAttr params = callableOp.getTemplateParamsAttr();
    if (!isNullOrEmpty(params)) {
      if (params.size() != targetInfo.oldParamOrder.size()) {
        return success();
      }
      struct DeferredExplicitArgInference {
        Attribute attr;
        Attribute expectedAttr;
        SmallVector<StringAttr> mentionedParams;
      };
      SmallVector<DeferredExplicitArgInference> deferredExplicitArgInferences;
      bool deferredToSignature = false;
      for (auto [paramName, attr] : llvm::zip_equal(targetInfo.oldParamOrder, params)) {
        auto replacementIt = targetInfo.replacements.find(paramName);
        if (replacementIt == targetInfo.replacements.end()) {
          continue;
        }
        Type expectedTy = substituteExplicitCallTemplateParams(
            replacementIt->second.type, params, targetInfo.oldParamOrder
        );
        Attribute expectedAttr = TypeAttr::get(expectedTy);
        if (attrMentionsAnyTypeVarParam(attr, inferenceInfo.typeVarParams)) {
          deferredToSignature = true;
          SmallVector<StringAttr> mentionedParams =
              getMentionedTypeVarParams(attr, inferenceInfo.typeVarParams);
          if (!ParamMentionChecker(paramName).typeMentions(targetSignature)) {
            if (failed(collectTemplateArgInferences(attr, expectedAttr, callableOp.getLoc()))) {
              return failure();
            }
          } else {
            deferredExplicitArgInferences.push_back(
                {attr, expectedAttr, std::move(mentionedParams)}
            );
          }
          continue;
        }
        if (failed(collectTemplateArgInferences(attr, expectedAttr, callableOp.getLoc()))) {
          return failure();
        }
      }
      if (deferredToSignature) {
        DenseMap<StringAttr, Type> targetReplacements;
        for (const auto &entry : targetInfo.replacements) {
          targetReplacements.try_emplace(entry.first, entry.second.type);
        }
        TypeVarReplacementConverter converter(
            callableOp.getContext(), targetInfo.templatePath, targetInfo.oldParamOrder,
            targetReplacements, /*trimResolvedParams=*/false
        );
        Type rewrittenTy = substituteExplicitCallTemplateParams(
            converter.convertType(targetSignature), params, targetInfo.oldParamOrder
        );
        if (failed(collectTypeInferences(
                callableOp.getTypeSignature(), llvm::cast<FunctionType>(rewrittenTy),
                callableOp.getLoc()
            ))) {
          return failure();
        }
 
        for (const DeferredExplicitArgInference &deferred : deferredExplicitArgInferences) {
          bool signatureCanInferAllParams =
              llvm::all_of(deferred.mentionedParams, [&](StringAttr mentionedParam) {
            return ParamMentionChecker(mentionedParam).typeMentions(callableOp.getTypeSignature());
          });
          bool signatureHasInferredAllParams =
              llvm::all_of(deferred.mentionedParams, [&](StringAttr mentionedParam) {
            return replacements.contains(mentionedParam);
          });
          if (signatureCanInferAllParams && signatureHasInferredAllParams) {
            continue;
          }
          if (failed(collectTemplateArgInferences(
                  deferred.attr, deferred.expectedAttr, callableOp.getLoc()
              ))) {
            return failure();
          }
        }
      }
      return success();
    }
 
    FailureOr<UnificationMap> unifyResult = callableOp.unifyTypeSignature(targetSignature);
    if (failed(unifyResult)) {
      return success();
    }
 
    for (StringAttr paramName : targetInfo.oldParamOrder) {
      auto replacementIt = targetInfo.replacements.find(paramName);
      if (replacementIt == targetInfo.replacements.end()) {
        continue;
      }
      SmallVector<Attribute> inferredAttrs =
          getInferredOmittedTemplateArgs(*unifyResult, paramName);
      if (inferredAttrs.empty()) {
        continue;
      }
 
      Attribute expectedAttr = TypeAttr::get(replacementIt->second.type);
      for (Attribute inferredAttr : inferredAttrs) {
        if (!inferredAttr || !typeParamsUnify({inferredAttr}, {expectedAttr})) {
          InFlightDiagnostic diag = callableOp.emitError()
                                    << "implicit template argument for inferred parameter @"
                                    << paramName.getValue() << " must match inferred type "
                                    << replacementIt->second.type << ", but found ";
          if (auto typeAttr = llvm::dyn_cast_if_present<TypeAttr>(inferredAttr)) {
            diag << typeAttr.getValue();
          } else {
            diag << inferredAttr;
          }
          return diag;
        }
        if (failed(collectTemplateArgInferences(inferredAttr, expectedAttr, callableOp.getLoc()))) {
          return failure();
        }
      }
    }
    return success();
  }

  LogicalResult collectCallableTemplateParamInferences(
      CallOp callOp, DenseMap<Operation *, TemplateInferenceInfo *> &infos,
      SymbolTableCollection &tables
  ) {
    FailureOr<SymbolLookupResult<FuncDefOp>> target = callOp.getCalleeTarget(tables);
    if (failed(target)) {
      return success();
    }
    FuncDefOp targetFunc = target->get();
    TemplateInferenceInfo *targetInfo = getTargetTemplateInfo(targetFunc, infos);
    if (!targetInfo) {
      return success();
    }
    return collectCallableUseTemplateParamInferences(
        callOp, targetFunc.getFunctionType(), *targetInfo
    );
  }

  LogicalResult collectIncludeTemplateParamInferences(
      verif::IncludeOp includeOp, DenseMap<Operation *, TemplateInferenceInfo *> &infos,
      SymbolTableCollection &tables
  ) {
    FailureOr<SymbolLookupResult<verif::ContractOp>> target = includeOp.getCalleeTarget(tables);
    if (failed(target)) {
      return success();
    }
    verif::ContractOp targetContract = target->get();
    TemplateInferenceInfo *targetInfo = getTargetTemplateInfo(targetContract, infos);
    if (!targetInfo) {
      return success();
    }
    return collectCallableUseTemplateParamInferences(
        includeOp, targetContract.getFunctionType(), *targetInfo
    );
  }

  LogicalResult recordInference(StringAttr paramName, Type inferredTy, Value value, Location loc) {
    if (!inferenceInfo.typeVarParams.contains(paramName)) {
      return success();
    }
    if (auto inferredTvar = llvm::dyn_cast<TypeVarType>(inferredTy)) {
      return recordParamRelation(paramName, inferredTvar.getNameRef().getAttr(), loc);
    }
    if (!isTypeVarFreeType(inferredTy)) {
      return success();
    }
 
    auto reportConflict = [&](StringRef kind, Location originalLoc, Type originalTy) {
      InFlightDiagnostic diag = emitError(loc) << "conflicting inferred type for " << kind << " @"
                                               << paramName.getValue() << ": " << originalTy
                                               << " vs " << inferredTy;
      diag.attachNote(originalLoc) << "previous inference here";
      return diag;
    };
 
    auto byParamIt = replacements.find(paramName);
    bool learnedParamInference = false;
    if (byParamIt == replacements.end()) {
      replacements.try_emplace(paramName, InferredType {inferredTy, loc});
      changedInIteration = true;
      learnedParamInference = true;
    } else if (byParamIt->second.type != inferredTy) {
      return reportConflict("template parameter", byParamIt->second.loc, byParamIt->second.type);
    }
 
    if (value) {
      auto byValueIt = byValue.find(value);
      if (byValueIt == byValue.end()) {
        byValue.try_emplace(value, InferredType {inferredTy, loc});
        changedInIteration = true;
      } else if (byValueIt->second.type != inferredTy) {
        return reportConflict("SSA value using", byValueIt->second.loc, byValueIt->second.type);
      }
    }
 
    if (learnedParamInference) {
      auto relatedIt = paramRelations.find(paramName);
      if (relatedIt != paramRelations.end()) {
        for (StringAttr relatedParam : relatedIt->second) {
          if (failed(recordInference(relatedParam, inferredTy, Value(), loc))) {
            return failure();
          }
        }
      }
    }
    return success();
  }

  LogicalResult recordParamRelation(StringAttr lhsParam, StringAttr rhsParam, Location loc) {
    if (lhsParam == rhsParam || !inferenceInfo.typeVarParams.contains(rhsParam)) {
      return success();
    }
 
    bool inserted = paramRelations[lhsParam].insert(rhsParam).second;
    inserted |= paramRelations[rhsParam].insert(lhsParam).second;
    if (!inserted) {
      return success();
    }
 
    changedInIteration = true;
 
    auto lhsIt = replacements.find(lhsParam);
    if (lhsIt != replacements.end() &&
        failed(recordInference(rhsParam, lhsIt->second.type, Value(), loc))) {
      return failure();
    }
    auto rhsIt = replacements.find(rhsParam);
    if (rhsIt != replacements.end() &&
        failed(recordInference(lhsParam, rhsIt->second.type, Value(), loc))) {
      return failure();
    }
    return success();
  }

  LogicalResult
  collectTypePairInferences(Type lhs, Type rhs, Value lhsValue, Value rhsValue, Location loc) {
    if (auto lhsTvar = llvm::dyn_cast<TypeVarType>(lhs)) {
      if (failed(recordInference(lhsTvar.getNameRef().getAttr(), rhs, lhsValue, loc))) {
        return failure();
      }
      if (rhsValue) {
        auto rhsValueIt = byValue.find(rhsValue);
        if (rhsValueIt != byValue.end() &&
            failed(recordInference(
                lhsTvar.getNameRef().getAttr(), rhsValueIt->second.type, lhsValue, loc
            ))) {
          return failure();
        }
      }
    }
    if (auto rhsTvar = llvm::dyn_cast<TypeVarType>(rhs)) {
      if (failed(recordInference(rhsTvar.getNameRef().getAttr(), lhs, rhsValue, loc))) {
        return failure();
      }
      if (lhsValue) {
        auto lhsValueIt = byValue.find(lhsValue);
        if (lhsValueIt != byValue.end() &&
            failed(recordInference(
                rhsTvar.getNameRef().getAttr(), lhsValueIt->second.type, rhsValue, loc
            ))) {
          return failure();
        }
      }
    }
 
    if (auto lhsArr = llvm::dyn_cast<ArrayType>(lhs)) {
      if (auto rhsArr = llvm::dyn_cast<ArrayType>(rhs)) {
        return collectTypePairInferences(
            lhsArr.getElementType(), rhsArr.getElementType(), Value(), Value(), loc
        );
      }
    }
 
    if (auto lhsStruct = llvm::dyn_cast<StructType>(lhs)) {
      if (auto rhsStruct = llvm::dyn_cast<StructType>(rhs)) {
        ArrayRef<Attribute> lhsParams =
            lhsStruct.getParams() ? lhsStruct.getParams().getValue() : ArrayRef<Attribute> {};
        ArrayRef<Attribute> rhsParams =
            rhsStruct.getParams() ? rhsStruct.getParams().getValue() : ArrayRef<Attribute> {};
        if (lhsParams.size() != rhsParams.size()) {
          return success();
        }
        for (auto [lhsAttr, rhsAttr] : llvm::zip_equal(lhsParams, rhsParams)) {
          if (failed(collectTemplateArgInferences(lhsAttr, rhsAttr, loc))) {
            return failure();
          }
        }
      }
    }
 
    if (auto lhsPod = llvm::dyn_cast<PodType>(lhs)) {
      if (auto rhsPod = llvm::dyn_cast<PodType>(rhs)) {
        ArrayRef<RecordAttr> lhsRecords = lhsPod.getRecords();
        ArrayRef<RecordAttr> rhsRecords = rhsPod.getRecords();
        if (lhsRecords.size() != rhsRecords.size()) {
          return success();
        }
        for (auto [lhsRecord, rhsRecord] : llvm::zip_equal(lhsRecords, rhsRecords)) {
          if (lhsRecord.getName() != rhsRecord.getName()) {
            return success();
          }
          if (failed(collectTypePairInferences(
                  lhsRecord.getType(), rhsRecord.getType(), Value(), Value(), loc
              ))) {
            return failure();
          }
        }
      }
    }
 
    if (auto lhsFunc = llvm::dyn_cast<FunctionType>(lhs)) {
      if (auto rhsFunc = llvm::dyn_cast<FunctionType>(rhs)) {
        if (lhsFunc.getNumInputs() != rhsFunc.getNumInputs() ||
            lhsFunc.getNumResults() != rhsFunc.getNumResults()) {
          return success();
        }
        for (auto [lhsInput, rhsInput] :
             llvm::zip_equal(lhsFunc.getInputs(), rhsFunc.getInputs())) {
          if (failed(collectTypePairInferences(lhsInput, rhsInput, Value(), Value(), loc))) {
            return failure();
          }
        }
        for (auto [lhsResult, rhsResult] :
             llvm::zip_equal(lhsFunc.getResults(), rhsFunc.getResults())) {
          if (failed(collectTypePairInferences(lhsResult, rhsResult, Value(), Value(), loc))) {
            return failure();
          }
        }
      }
    }
 
    return success();
  }

  LogicalResult collectTemplateArgInferences(Attribute lhsAttr, Attribute rhsAttr, Location loc) {
    auto lhsTyAttr = llvm::dyn_cast<TypeAttr>(lhsAttr);
    auto rhsTyAttr = llvm::dyn_cast<TypeAttr>(rhsAttr);
    if (lhsTyAttr && rhsTyAttr) {
      return collectTypePairInferences(
          lhsTyAttr.getValue(), rhsTyAttr.getValue(), Value(), Value(), loc
      );
    }
 
    if (StringAttr lhsSymbolName = getFlatSymbolName(lhsAttr)) {
      if (StringAttr rhsSymbolName = getFlatSymbolName(rhsAttr)) {
        return recordParamRelation(lhsSymbolName, rhsSymbolName, loc);
      }
      if (rhsTyAttr && failed(recordInference(lhsSymbolName, rhsTyAttr.getValue(), Value(), loc))) {
        return failure();
      }
    }
    if (StringAttr rhsSymbolName = getFlatSymbolName(rhsAttr)) {
      if (lhsTyAttr && failed(recordInference(rhsSymbolName, lhsTyAttr.getValue(), Value(), loc))) {
        return failure();
      }
    }
    return success();
  }
};

static DenseMap<StringAttr, Type>
getFunctionProofReplacements(const TemplateInferenceInfo &info, FuncDefOp func) {
  DenseMap<StringAttr, Type> replacements;
  auto funcIt = info.functionReplacements.find(func.getOperation());
  if (funcIt == info.functionReplacements.end()) {
    return replacements;
  }
  for (const auto &entry : funcIt->second) {
    replacements.try_emplace(entry.first, entry.second.type);
  }
  return replacements;
}

static std::optional<unsigned>
getParamIndex(ArrayRef<StringAttr> paramOrder, StringAttr paramName) {
  for (auto indexedParam : llvm::enumerate(paramOrder)) {
    if (indexedParam.value() == paramName) {
      return indexedParam.index();
    }
  }
  return std::nullopt;
}

static DenseMap<Attribute, Attribute>
buildParamNameToConcrete(const DenseMap<StringAttr, Type> &replacements) {
  DenseMap<Attribute, Attribute> paramNameToConcrete;
  for (const auto &entry : replacements) {
    paramNameToConcrete.try_emplace(
        FlatSymbolRefAttr::get(entry.first), TypeAttr::get(entry.second)
    );
  }
  return paramNameToConcrete;
}

static bool callParamsMatchReplacements(
    ArrayAttr callParams, ArrayRef<StringAttr> oldParamOrder,
    const DenseMap<StringAttr, Type> &replacements,
    const DenseMap<StringAttr, InferredType> *wildcardAllowedReplacements = nullptr
) {
  if (!callParams || callParams.size() != oldParamOrder.size()) {
    return false;
  }
  for (const auto &entry : replacements) {
    std::optional<unsigned> index = getParamIndex(oldParamOrder, entry.first);
    if (!index) {
      return false;
    }
    Type expectedTy = substituteExplicitCallTypeArgs(entry.second, callParams, oldParamOrder);
    Attribute attr = callParams[*index];
    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr);
        intAttr && isDynamic(intAttr) && wildcardAllowedReplacements &&
        wildcardAllowedReplacements->contains(entry.first)) {
      continue;
    }
    if (!templateArgUnifiesWithType(attr, expectedTy)) {
      return false;
    }
  }
  return true;
}

static LogicalResult diagnoseCallParamsMismatch(
    Operation *callableOp, ArrayAttr callParams, ArrayRef<StringAttr> oldParamOrder,
    const DenseMap<StringAttr, Type> &replacements,
    const DenseMap<StringAttr, InferredType> *wildcardAllowedReplacements, bool explicitCallParams
) {
  if (callParamsMatchReplacements(
          callParams, oldParamOrder, replacements, wildcardAllowedReplacements
      )) {
    return success();
  }
  for (const auto &entry : replacements) {
    std::optional<unsigned> index = getParamIndex(oldParamOrder, entry.first);
    if (!index || !callParams || *index >= callParams.size()) {
      continue;
    }
    Attribute attr = callParams[*index];
    Type expectedTy = substituteExplicitCallTypeArgs(entry.second, callParams, oldParamOrder);
    if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr);
        intAttr && isDynamic(intAttr) && wildcardAllowedReplacements &&
        wildcardAllowedReplacements->contains(entry.first)) {
      continue;
    }
    if (templateArgUnifiesWithType(attr, expectedTy)) {
      continue;
    }
 
    InFlightDiagnostic diag = callableOp->emitError()
                              << (explicitCallParams ? "explicit" : "implicit")
                              << " template argument for inferred parameter @"
                              << entry.first.getValue() << " must match inferred type "
                              << expectedTy << ", but found ";
    if (auto typeAttr = llvm::dyn_cast<TypeAttr>(attr)) {
      diag << typeAttr.getValue();
    } else {
      diag << attr;
    }
    return diag;
  }
  return callableOp->emitError() << (explicitCallParams ? "explicit" : "implicit")
                                 << " template arguments do not match inferred callee types";
}

static ArrayAttr expandCurrentTemplateParamsToOriginalOrder(
    ArrayAttr callParams, const TemplateInferenceInfo &info
) {
  if (!callParams || callParams.size() == info.oldParamOrder.size()) {
    return callParams;
  }
 
  unsigned currentParamCount = llvm::count_if(info.oldParamOrder, [&info](StringAttr paramName) {
    return !info.replacements.contains(paramName);
  });
  if (callParams.size() != currentParamCount) {
    return callParams;
  }
 
  SmallVector<Attribute> expandedParams;
  expandedParams.reserve(info.oldParamOrder.size());
  unsigned currentIndex = 0;
  for (StringAttr paramName : info.oldParamOrder) {
    auto replacementIt = info.replacements.find(paramName);
    if (replacementIt != info.replacements.end()) {
      expandedParams.push_back(TypeAttr::get(replacementIt->second.type));
      continue;
    }
    expandedParams.push_back(callParams[currentIndex++]);
  }
  return ArrayAttr::get(callParams.getContext(), expandedParams);
}

template <typename TargetOp>
static DenseMap<StringAttr, Type> getConcreteCallSiteReplacements(
    ArrayAttr callParams, const TemplateInferenceInfo &info, TargetOp targetOp
) {
  DenseMap<StringAttr, Type> replacements;
  if (!callParams || callParams.size() != info.oldParamOrder.size()) {
    return replacements;
  }
 
  bool sawResidualReplacement = false;
  for (const auto &entry : info.typeVarParams) {
    StringAttr paramName = entry.first;
    if (info.replacements.contains(paramName)) {
      continue;
    }
    if (!targetMentionsParam(targetOp, paramName)) {
      continue;
    }
    std::optional<unsigned> index = getParamIndex(info.oldParamOrder, paramName);
    assert(index && "eligible type-variable parameter must appear in parameter order");
    Attribute attr = callParams[*index];
    auto tyAttr = llvm::dyn_cast<TypeAttr>(attr);
    if (!tyAttr) {
      return DenseMap<StringAttr, Type>();
    }
    replacements.try_emplace(paramName, tyAttr.getValue());
    sawResidualReplacement = true;
  }
  if (!sawResidualReplacement) {
    return DenseMap<StringAttr, Type>();
  }
 
  for (const auto &entry : info.replacements) {
    replacements.try_emplace(entry.first, entry.second.type);
  }
 
  TypeVarReplacementConverter converter(
      info.templatePath.front().getContext(), info.templatePath, info.oldParamOrder, replacements,
      /*trimResolvedParams=*/false
  );
  for (const auto &entry : replacements) {
    if (!isTypeVarFreeType(converter.convertType(entry.second))) {
      return DenseMap<StringAttr, Type>();
    }
  }
  return replacements;
}

static bool isWildcardTemplateArg(Attribute attr) {
  auto intAttr = llvm::dyn_cast_if_present<IntegerAttr>(attr);
  return intAttr && isDynamic(intAttr);
}

static void appendUniqueType(SmallVectorImpl<Type> &types, Type ty) {
  if (!llvm::any_of(types, [ty](Type existing) { return existing == ty; })) {
    types.push_back(ty);
  }
}

static void collectRhsTypeVarCandidates(
    Type lhsTy, Type rhsTy, SymbolRefAttr targetRef, SmallVectorImpl<Type> &candidates
);

static void collectRhsTypeVarCandidates(
    ArrayRef<Attribute> lhsAttrs, ArrayRef<Attribute> rhsAttrs, SymbolRefAttr targetRef,
    SmallVectorImpl<Type> &candidates
) {
  if (lhsAttrs.size() != rhsAttrs.size()) {
    return;
  }
  for (auto [lhsAttr, rhsAttr] : llvm::zip_equal(lhsAttrs, rhsAttrs)) {
    collectRhsTypeVarCandidates(lhsAttr, rhsAttr, targetRef, candidates);
  }
}

static void collectRhsTypeVarCandidates(
    Attribute lhsAttr, Attribute rhsAttr, SymbolRefAttr targetRef, SmallVectorImpl<Type> &candidates
) {
  if (auto rhsTypeAttr = llvm::dyn_cast_if_present<TypeAttr>(rhsAttr)) {
    if (auto lhsTypeAttr = llvm::dyn_cast_if_present<TypeAttr>(lhsAttr)) {
      collectRhsTypeVarCandidates(
          lhsTypeAttr.getValue(), rhsTypeAttr.getValue(), targetRef, candidates
      );
    }
    return;
  }
 
  auto lhsArray = llvm::dyn_cast_if_present<ArrayAttr>(lhsAttr);
  auto rhsArray = llvm::dyn_cast_if_present<ArrayAttr>(rhsAttr);
  if (!lhsArray || !rhsArray || lhsArray.size() != rhsArray.size()) {
    return;
  }
  collectRhsTypeVarCandidates(lhsArray.getValue(), rhsArray.getValue(), targetRef, candidates);
}

static void collectRhsTypeVarCandidates(
    Type lhsTy, Type rhsTy, SymbolRefAttr targetRef, SmallVectorImpl<Type> &candidates
) {
  if (auto rhsTvar = llvm::dyn_cast<TypeVarType>(rhsTy);
      rhsTvar && rhsTvar.getNameRef() == targetRef) {
    appendUniqueType(candidates, lhsTy);
    return;
  }
 
  if (auto lhsArray = llvm::dyn_cast<ArrayType>(lhsTy)) {
    if (auto rhsArray = llvm::dyn_cast<ArrayType>(rhsTy)) {
      collectRhsTypeVarCandidates(
          lhsArray.getElementType(), rhsArray.getElementType(), targetRef, candidates
      );
      collectRhsTypeVarCandidates(
          lhsArray.getDimensionSizes(), rhsArray.getDimensionSizes(), targetRef, candidates
      );
    }
    return;
  }
 
  if (auto lhsStruct = llvm::dyn_cast<StructType>(lhsTy)) {
    if (auto rhsStruct = llvm::dyn_cast<StructType>(rhsTy)) {
      collectRhsTypeVarCandidates(
          lhsStruct.getParams(), rhsStruct.getParams(), targetRef, candidates
      );
    }
    return;
  }
 
  if (auto lhsPod = llvm::dyn_cast<PodType>(lhsTy)) {
    if (auto rhsPod = llvm::dyn_cast<PodType>(rhsTy);
        rhsPod && lhsPod.getRecords().size() == rhsPod.getRecords().size()) {
      for (auto [lhsRecord, rhsRecord] :
           llvm::zip_equal(lhsPod.getRecords(), rhsPod.getRecords())) {
        collectRhsTypeVarCandidates(
            lhsRecord.getType(), rhsRecord.getType(), targetRef, candidates
        );
      }
    }
    return;
  }
 
  if (auto lhsFunc = llvm::dyn_cast<FunctionType>(lhsTy)) {
    if (auto rhsFunc = llvm::dyn_cast<FunctionType>(rhsTy)) {
      for (auto [lhsInput, rhsInput] : llvm::zip_equal(lhsFunc.getInputs(), rhsFunc.getInputs())) {
        collectRhsTypeVarCandidates(lhsInput, rhsInput, targetRef, candidates);
      }
      for (auto [lhsResult, rhsResult] :
           llvm::zip_equal(lhsFunc.getResults(), rhsFunc.getResults())) {
        collectRhsTypeVarCandidates(lhsResult, rhsResult, targetRef, candidates);
      }
    }
  }
}

static SmallVector<Type>
getConflictingRhsTypeVarCandidates(FunctionType lhs, FunctionType rhs, SymbolRefAttr targetRef) {
  SmallVector<Type> candidates;
  if (lhs.getNumInputs() != rhs.getNumInputs() || lhs.getNumResults() != rhs.getNumResults()) {
    return candidates;
  }
  for (auto [lhsInput, rhsInput] : llvm::zip_equal(lhs.getInputs(), rhs.getInputs())) {
    collectRhsTypeVarCandidates(lhsInput, rhsInput, targetRef, candidates);
  }
  for (auto [lhsResult, rhsResult] : llvm::zip_equal(lhs.getResults(), rhs.getResults())) {
    collectRhsTypeVarCandidates(lhsResult, rhsResult, targetRef, candidates);
  }
  return candidates;
}

template <typename CallableOp, typename TargetOp>
static LogicalResult
emitConflictingInferredTypes(CallableOp callableOp, TargetOp targetOp, StringAttr paramName) {
  InFlightDiagnostic diag = callableOp.emitError()
                            << "conflicting inferred types for @" << paramName.getValue();
  SmallVector<Type> candidates = getConflictingRhsTypeVarCandidates(
      callableOp.getTypeSignature(), targetOp.getFunctionType(), FlatSymbolRefAttr::get(paramName)
  );
  if (candidates.size() >= 2) {
    diag << ": " << candidates.front();
    for (Type candidate : llvm::drop_begin(candidates)) {
      diag << " vs " << candidate;
    }
  }
  return diag;
}

template <typename CallableOp, typename TargetOp>
static FailureOr<Attribute> getInferredRhsTemplateArg(
    CallableOp callableOp, TargetOp targetOp, const UnificationMap &unifyResult,
    StringAttr paramName, StringRef argKind
) {
  auto inferredIt = unifyResult.find({FlatSymbolRefAttr::get(paramName), Side::RHS});
  if (inferredIt == unifyResult.end()) {
    return callableOp.emitError() << "could not infer " << argKind
                                  << " template argument for parameter @" << paramName.getValue()
                                  << " from callee signature";
  }
  if (!inferredIt->second) {
    return emitConflictingInferredTypes(callableOp, targetOp, paramName);
  }
  return inferredIt->second;
}

template <typename CallableOp, typename TargetOp>
static FailureOr<ArrayAttr> getCallSignatureTemplateParams(
    CallableOp callableOp, const TemplateInferenceInfo &info, TargetOp targetOp
) {
  FailureOr<UnificationMap> unifyResult = callableOp.unifyTypeSignature(targetOp.getFunctionType());
  if (failed(unifyResult)) {
    return callableOp.emitError()
           << "could not infer omitted template arguments from callee signature";
  }
 
  SmallVector<Attribute> params;
  params.reserve(info.oldParamOrder.size());
  for (StringAttr paramName : info.oldParamOrder) {
    auto replacementIt = info.replacements.find(paramName);
    if (replacementIt != info.replacements.end()) {
      params.push_back(TypeAttr::get(replacementIt->second.type));
      continue;
    }
    FailureOr<Attribute> inferredAttr =
        getInferredRhsTemplateArg(callableOp, targetOp, *unifyResult, paramName, "omitted");
    if (failed(inferredAttr)) {
      return failure();
    }
    params.push_back(*inferredAttr);
  }
  return ArrayAttr::get(callableOp.getContext(), params);
}

template <typename TargetOp>
static bool hasResidualFunctionTvarWildcard(
    ArrayAttr callParams, const TemplateInferenceInfo &info, TargetOp targetOp
) {
  if (!callParams || callParams.size() != info.oldParamOrder.size()) {
    return false;
  }
  for (const auto &entry : info.typeVarParams) {
    StringAttr paramName = entry.first;
    if (info.replacements.contains(paramName)) {
      continue;
    }
    if (!targetMentionsParam(targetOp, paramName)) {
      continue;
    }
    std::optional<unsigned> index = getParamIndex(info.oldParamOrder, paramName);
    assert(index && "eligible type-variable parameter must appear in parameter order");
    if (isWildcardTemplateArg(callParams[*index])) {
      return true;
    }
  }
  return false;
}

template <typename CallableOp, typename TargetOp>
static FailureOr<ArrayAttr> materializeResidualFunctionTvarWildcards(
    CallableOp callableOp, ArrayAttr callParams, const TemplateInferenceInfo &info,
    TargetOp targetOp
) {
  FailureOr<UnificationMap> unifyResult = callableOp.unifyTypeSignature(targetOp.getFunctionType());
  if (failed(unifyResult)) {
    return callableOp.emitError()
           << "could not infer wildcard template arguments from callee signature";
  }
 
  SmallVector<Attribute> params(callParams.begin(), callParams.end());
  for (const auto &entry : info.typeVarParams) {
    StringAttr paramName = entry.first;
    if (info.replacements.contains(paramName)) {
      continue;
    }
    if (!targetMentionsParam(targetOp, paramName)) {
      continue;
    }
    std::optional<unsigned> index = getParamIndex(info.oldParamOrder, paramName);
    assert(index && "eligible type-variable parameter must appear in parameter order");
    if (!isWildcardTemplateArg(params[*index])) {
      continue;
    }
 
    auto inferredIt = unifyResult->find({FlatSymbolRefAttr::get(paramName), Side::RHS});
    if (inferredIt == unifyResult->end()) {
      auto proofIt = info.functionReplacements.find(targetOp.getOperation());
      if (proofIt == info.functionReplacements.end()) {
        continue;
      }
      auto replacementIt = proofIt->second.find(paramName);
      if (replacementIt == proofIt->second.end()) {
        continue;
      }
      params[*index] = TypeAttr::get(replacementIt->second.type);
      continue;
    }
    if (!inferredIt->second) {
      return emitConflictingInferredTypes(callableOp, targetOp, paramName);
    }
    params[*index] = inferredIt->second;
  }
  return ArrayAttr::get(callableOp.getContext(), params);
}

template <typename TargetOp>
static bool hasResidualFunctionTvar(const TemplateInferenceInfo &info, TargetOp targetOp) {
  return llvm::any_of(info.typeVarParams, [&](const auto &entry) {
    return !info.replacements.contains(entry.first) && targetMentionsParam(targetOp, entry.first);
  });
}

struct CallableSpecializationInputs {
  ArrayAttr params;
  bool explicitParams = false;
  bool paramsChanged = false;
  DenseMap<StringAttr, Type> replacements;
};

template <typename CallableOp, typename TargetOp>
static LogicalResult prepareCallableSpecialization(
    CallableOp callableOp, const TemplateInferenceInfo &info, TargetOp targetOp,
    CallableSpecializationInputs &inputs, bool &shouldSpecialize
) {
  shouldSpecialize = false;
  inputs.params = {};
  inputs.explicitParams = false;
  inputs.paramsChanged = false;
  inputs.replacements.clear();
  inputs.params = callableOp.getTemplateParamsAttr();
  inputs.explicitParams = !isNullOrEmpty(inputs.params);
  if (isNullOrEmpty(inputs.params)) {
    FailureOr<ArrayAttr> inferredParams =
        getCallSignatureTemplateParams(callableOp, info, targetOp);
    if (failed(inferredParams)) {
      return failure();
    }
    inputs.params = *inferredParams;
    inputs.paramsChanged = true;
  } else {
    inputs.params = expandCurrentTemplateParamsToOriginalOrder(inputs.params, info);
    if (hasResidualFunctionTvarWildcard(inputs.params, info, targetOp)) {
      FailureOr<ArrayAttr> materializedParams =
          materializeResidualFunctionTvarWildcards(callableOp, inputs.params, info, targetOp);
      if (failed(materializedParams)) {
        return failure();
      }
      inputs.params = *materializedParams;
      inputs.paramsChanged = true;
    }
  }
 
  inputs.replacements = getConcreteCallSiteReplacements(inputs.params, info, targetOp);
  shouldSpecialize = !inputs.replacements.empty();
  return success();
}

static std::string buildTemplateLocalFunctionCloneCacheKey(
    StringRef functionName, ArrayRef<StringAttr> oldParamOrder,
    const DenseMap<StringAttr, Type> &replacements
) {
  std::string key;
  llvm::raw_string_ostream os(key);
  os << functionName.size() << ':' << functionName;
  for (StringAttr paramName : oldParamOrder) {
    os << '|';
    os << paramName.getValue().size() << ':' << paramName.getValue() << '=';
    auto replacementIt = replacements.find(paramName);
    if (replacementIt == replacements.end()) {
      os << '_';
      continue;
    }
 
    std::string typeText;
    llvm::raw_string_ostream typeOs(typeText);
    replacementIt->second.print(typeOs);
    os << typeText.size() << ':' << typeText;
  }
  return key;
}

static SymbolRefAttr getSpecializedFunctionCloneCallee(
    SymbolRefAttr originalCallee, StringAttr templateName, StringAttr cloneName
) {
  SmallVector<FlatSymbolRefAttr> pieces = getPieces(originalCallee);
  assert(pieces.size() >= 2 && "callee must include at least template and function names");
  pieces.pop_back();
  pieces.pop_back();
  pieces.push_back(FlatSymbolRefAttr::get(templateName));
  pieces.push_back(FlatSymbolRefAttr::get(cloneName));
  return asSymbolRefAttr(pieces);
}

template <typename CallableOp, typename ConfigureCloneFn>
static FailureOr<SymbolRefAttr> getOrCreateSpecializedCallableClone(
    TemplateOp templateOp, CallableOp callable, SymbolRefAttr originalCallee,
    const TemplateInferenceInfo &info, const DenseMap<StringAttr, Type> &replacements,
    SymbolTableCollection &tables, llvm::StringMap<SymbolRefAttr> &cloneCallees,
    llvm::StringMap<StringAttr> &templateClones, InstantiationLayout &layout, ArrayAttr callParams,
    ConfigureCloneFn configureClone
) {
  DenseMap<Attribute, Attribute> paramNameToConcrete = buildParamNameToConcrete(replacements);
  FailureOr<TemplateOp> newTemplate = getOrCreateSpecializedTemplateClone(
      templateOp, info.oldParamOrder, paramNameToConcrete, callParams, tables, templateClones,
      layout
  );
  if (failed(newTemplate)) {
    return failure();
  }
 
  std::string cacheKey = buildTemplateLocalFunctionCloneCacheKey(
      callable.getSymName(), info.oldParamOrder, replacements
  );
  auto cachedCallee = cloneCallees.find(cacheKey);
  if (cachedCallee != cloneCallees.end()) {
    return cachedCallee->second;
  }
 
  SymbolTable &templateSymbols = tables.getSymbolTable(*newTemplate);
 
  auto clone = llvm::cast<CallableOp>(callable.getOperation()->clone());
  configureClone(clone);
  templateSymbols.insert(clone);
 
  TypeVarReplacementConverter converter(
      templateOp.getContext(), info.templatePath, info.oldParamOrder, replacements,
      /*trimResolvedParams=*/false
  );
  if (failed(convertTemplateExprTypesIn(*newTemplate, converter))) {
    clone.erase();
    return failure();
  }
  if (failed(convertOperationTypesIn(clone.getOperation(), converter))) {
    clone.erase();
    return failure();
  }
  removeIdentityCasts(clone.getOperation());
  SymbolRefAttr cloneCallee = getSpecializedFunctionCloneCallee(
      originalCallee, newTemplate->getSymNameAttr(), clone.getSymNameAttr()
  );
  cloneCallees.try_emplace(cacheKey, cloneCallee);
  return cloneCallee;
}

static FailureOr<SymbolRefAttr> getOrCreateSpecializedFunctionClone(
    TemplateOp templateOp, FuncDefOp func, SymbolRefAttr originalCallee,
    const TemplateInferenceInfo &info, const DenseMap<StringAttr, Type> &replacements,
    SymbolTableCollection &tables, llvm::StringMap<SymbolRefAttr> &cloneCallees,
    llvm::StringMap<StringAttr> &templateClones, InstantiationLayout &layout, ArrayAttr callParams
) {
  return getOrCreateSpecializedCallableClone(
      templateOp, func, originalCallee, info, replacements, tables, cloneCallees, templateClones,
      layout, callParams, [](FuncDefOp) {}
  );
}

static FailureOr<SymbolRefAttr> getOrCreateSpecializedContractClone(
    TemplateOp templateOp, verif::ContractOp contract, SymbolRefAttr originalCallee,
    SymbolRefAttr specializedTarget, const TemplateInferenceInfo &info,
    const DenseMap<StringAttr, Type> &replacements, SymbolTableCollection &tables,
    llvm::StringMap<SymbolRefAttr> &cloneCallees, llvm::StringMap<StringAttr> &templateClones,
    InstantiationLayout &layout, ArrayAttr callParams
) {
  return getOrCreateSpecializedCallableClone(
      templateOp, contract, originalCallee, info, replacements, tables, cloneCallees,
      templateClones, layout, callParams,
      [specializedTarget](verif::ContractOp clone) { clone.setTargetAttr(specializedTarget); }
  );
}

static LogicalResult removeResolvedParams(TemplateInferenceInfo &info) {
  if (info.replacements.empty()) {
    return success();
  }
 
  DenseMap<StringAttr, Type> replacements;
  for (const auto &entry : info.replacements) {
    replacements.try_emplace(entry.first, entry.second.type);
  }
  DenseMap<Attribute, Attribute> paramNameToConcrete = buildParamNameToConcrete(replacements);
  FailureOr<InstantiationLayout> layout =
      buildInstantiationLayout(info.templateOp, ArrayAttr(), paramNameToConcrete);
  if (failed(layout)) {
    return failure();
  }
 
  setInstantiationNamePattern(info.templateOp, ArrayAttr());
  for (auto paramOp : llvm::make_early_inc_range(info.templateOp.getConstOps<TemplateParamOp>())) {
    auto name = paramOp.getSymNameAttr();
    if (info.replacements.contains(name)) {
      paramOp.erase();
    }
  }
  setInstantiationNamePattern(
      info.templateOp, layout->remainingNames.empty() ? ArrayAttr() : layout->namePattern
  );
  return success();
}

template <typename CallableOp>
static LogicalResult
validateWildcardCallableSignature(CallableOp callableOp, FunctionType targetTy) {
  if (succeeded(callableOp.unifyTypeSignature(targetTy))) {
    return success();
  }
  return callableOp.emitError() << "call signature " << callableOp.getTypeSignature()
                                << " does not match inferred callee signature " << targetTy;
}

static void updateCallableResultTypesIfNeeded(
    CallOp callOp, const TypeVarReplacementConverter &converter, bool &modified
) {
  for (Value result : callOp.getResults()) {
    Type newTy = converter.convertType(result.getType());
    if (newTy != result.getType()) {
      result.setType(newTy);
      modified = true;
    }
  }
}

static void
updateCallableResultTypesIfNeeded(verif::IncludeOp, const TypeVarReplacementConverter &, bool &) {}

template <typename CallableOp>
static FailureOr<bool> updateCallableTemplateParamsFor(
    ModuleOp module, DenseMap<Operation *, const TypeVarReplacementConverter *> &converters,
    SymbolTableCollection &tables
) {
  bool modified = false;
  bool failedConversion = false;
  module.walk([&](CallableOp callableOp) {
    if (failedConversion) {
      return;
    }
    auto target = callableOp.getCalleeTarget(tables);
    if (failed(target)) {
      return;
    }
    auto targetOp = target->get();
    TemplateOp parentTemplate = getParentOfType<TemplateOp>(targetOp.getOperation());
    if (!parentTemplate) {
      return;
    }
    const TypeVarReplacementConverter *converter = converters.lookup(parentTemplate.getOperation());
    if (!converter) {
      return;
    }
 
    TemplateOp callableTemplate = getParentOfType<TemplateOp>(callableOp.getOperation());
    bool resolveTemplateSymbolArgs =
        callableTemplate && callableTemplate.getOperation() == parentTemplate.getOperation();
    ArrayAttr oldParams = callableOp.getTemplateParamsAttr();
    FailureOr<ArrayAttr> newParams = converter->convertTemplateParams(
        oldParams, callableOp, resolveTemplateSymbolArgs,
        /*allowCurrentTemplateTypeVarSymbols=*/true
    );
    if (failed(newParams)) {
      failedConversion = true;
      return;
    }
    if (converter->hasWildcardForRemovedParam(oldParams)) {
      auto inferredTargetTy =
          llvm::cast<FunctionType>(converter->convertType(targetOp.getFunctionType()));
      if (failed(validateWildcardCallableSignature(callableOp, inferredTargetTy))) {
        failedConversion = true;
        return;
      }
    }
    if (oldParams != *newParams) {
      callableOp.setTemplateParamsAttr(*newParams);
      modified = true;
    }
 
    updateCallableResultTypesIfNeeded(callableOp, *converter, modified);
  });
  if (failedConversion) {
    return failure();
  }
  return modified;
}

static FailureOr<bool> updateCallableTemplateParams(
    ModuleOp module, DenseMap<Operation *, const TypeVarReplacementConverter *> &converters
) {
  SymbolTableCollection tables;
  FailureOr<bool> callsModified =
      updateCallableTemplateParamsFor<CallOp>(module, converters, tables);
  if (failed(callsModified)) {
    return failure();
  }
  FailureOr<bool> includesModified =
      updateCallableTemplateParamsFor<verif::IncludeOp>(module, converters, tables);
  if (failed(includesModified)) {
    return failure();
  }
  return *callsModified || *includesModified;
}

class ReferencedStructTemplateParamConverter {
  MLIRContext *ctx_;
  ModuleOp module_;
  SymbolTableCollection &tables_;
  DenseMap<Operation *, const TypeVarReplacementConverter *> &converters_;
  Operation *diagnosticOp_ = nullptr;
  bool hasFailure = false;
 
public:
  ReferencedStructTemplateParamConverter(
      MLIRContext *ctx, ModuleOp module, SymbolTableCollection &tables,
      DenseMap<Operation *, const TypeVarReplacementConverter *> &converters
  )
      : ctx_(ctx), module_(module), tables_(tables), converters_(converters) {}

  void startOperation(Operation *op) {
    diagnosticOp_ = op;
    hasFailure = false;
  }

  bool hadFailure() const { return hasFailure; }

  Type convertType(Type ty) {
    if (!ty || hasFailure) {
      return ty;
    }
    if (auto arrTy = llvm::dyn_cast<ArrayType>(ty)) {
      return convertArrayElementType(arrTy, [this](Type elemTy) { return convertType(elemTy); });
    }
    if (auto structTy = llvm::dyn_cast<StructType>(ty)) {
      return convertStructType(structTy);
    }
    if (auto podTy = llvm::dyn_cast<PodType>(ty)) {
      return convertPodType(podTy, ctx_, *this);
    }
    if (auto funcTy = llvm::dyn_cast<FunctionType>(ty)) {
      return convertFunctionType(funcTy, *this);
    }
    return ty;
  }

  Attribute convertAttr(Attribute attr) {
    if (hasFailure) {
      return attr;
    }
    return convertTypeOrArrayAttr(attr, ctx_, [this](Type ty) {
      return convertType(ty);
    }, [this](Attribute nested) { return convertAttr(nested); });
  }
 
private:
  StructType convertStructType(StructType structTy) {
    ArrayAttr params = structTy.getParams();
    if (!params) {
      return structTy;
    }
 
    SmallVector<Attribute> newParams;
    bool changed = false;
    for (Attribute attr : params.getValue()) {
      Attribute newAttr = convertAttr(attr);
      newParams.push_back(newAttr);
      changed |= newAttr != attr;
      if (hasFailure) {
        return structTy;
      }
    }
    StructType convertedTy =
        changed ? getStructTypeWithParams(structTy.getNameRef(), ctx_, newParams) : structTy;
 
    FailureOr<SymbolLookupResult<StructDefOp>> lookup = lookupTopLevelSymbol<StructDefOp>(
        tables_, convertedTy.getNameRef(), module_, /*reportMissing=*/false
    );
    if (failed(lookup)) {
      return convertedTy;
    }
    TemplateOp parentTemplate = getParentOfType<TemplateOp>(lookup->get().getOperation());
    if (!parentTemplate) {
      return convertedTy;
    }
    const TypeVarReplacementConverter *converter =
        converters_.lookup(parentTemplate.getOperation());
    if (!converter) {
      return convertedTy;
    }
 
    TemplateOp useTemplate = getParentOfType<TemplateOp>(diagnosticOp_);
    bool resolveTemplateSymbolArgs =
        useTemplate && useTemplate.getOperation() == parentTemplate.getOperation();
    FailureOr<ArrayAttr> trimmedParams = converter->convertTemplateParams(
        convertedTy.getParams(), diagnosticOp_, resolveTemplateSymbolArgs,
        /*allowCurrentTemplateTypeVarSymbols=*/false
    );
    if (failed(trimmedParams)) {
      hasFailure = true;
      return convertedTy;
    }
    if (convertedTy.getParams() == *trimmedParams) {
      return convertedTy;
    }
    return getStructTypeWithParams(convertedTy.getNameRef(), *trimmedParams);
  }
};

static FailureOr<bool> updateStructTemplateParams(
    ModuleOp module, DenseMap<Operation *, const TypeVarReplacementConverter *> &converters
) {
  SymbolTableCollection tables;
  ReferencedStructTemplateParamConverter converter(module.getContext(), module, tables, converters);
  return convertOperationTypesInAndTrack(module.getOperation(), converter);
}

static LogicalResult instantiateConcreteStructUses(ModuleOp module) {
  SymbolTableCollection tables;
  ConcreteStructInstantiationConverter converter(module.getContext(), module, tables);
  return convertOperationTypesIn(module.getOperation(), converter);
}

static TemplateOp
getContractTargetTemplate(verif::ContractOp contract, SymbolTableCollection &tables) {
  if (FailureOr<SymbolLookupResult<FuncDefOp>> funcTarget = contract.getFuncTarget(tables);
      succeeded(funcTarget)) {
    return getParentOfType<TemplateOp>(funcTarget->get().getOperation());
  }
  if (FailureOr<SymbolLookupResult<StructDefOp>> structTarget = contract.getStructTarget(tables);
      succeeded(structTarget)) {
    return getParentOfType<TemplateOp>(structTarget->get().getOperation());
  }
  return {};
}

static LogicalResult updateExternalContractTemplateParams(
    ModuleOp module, DenseMap<Operation *, const TypeVarReplacementConverter *> &converters
) {
  bool failedConversion = false;
  SymbolTableCollection tables;
  module.walk([&](verif::ContractOp contract) {
    if (failedConversion) {
      return;
    }
    TemplateOp targetTemplate = getContractTargetTemplate(contract, tables);
    if (!targetTemplate) {
      return;
    }
    if (getParentOfType<TemplateOp>(contract.getOperation()) == targetTemplate) {
      return;
    }
    const TypeVarReplacementConverter *converter = converters.lookup(targetTemplate.getOperation());
    if (!converter) {
      return;
    }
 
    auto validationResult = contract.walk([converter](Operation *op) -> WalkResult {
      return converter->validateOperation(op);
    });
    if (validationResult.wasInterrupted() ||
        failed(convertOperationTypesIn(contract.getOperation(), *converter))) {
      failedConversion = true;
      return;
    }
    removeIdentityCasts(contract.getOperation());
  });
  return failure(failedConversion);
}

static LogicalResult specializeFunctionLocalCallables(
    ModuleOp module, DenseMap<Operation *, const TemplateInferenceInfo *> &infoByTemplate,
    SpecializedCallableCloneCache &functionCloneCache,
    SpecializedCallableCloneCache &contractCloneCache,
    SpecializedTemplateCloneCache &templateCloneCache
) {
  bool failedClone = false;
  SymbolTableCollection tables;
  module.walk([&](CallOp callOp) {
    if (failedClone) {
      return;
    }
    FailureOr<SymbolLookupResult<FuncDefOp>> target = callOp.getCalleeTarget(tables);
    if (failed(target)) {
      return;
    }
    FuncDefOp targetFunc = target->get();
    auto parentTemplate = llvm::dyn_cast_or_null<TemplateOp>(targetFunc->getParentOp());
    if (!parentTemplate) {
      return;
    }
    const TemplateInferenceInfo *info = infoByTemplate.lookup(parentTemplate.getOperation());
    if (!info) {
      return;
    }
    if (!hasResidualFunctionTvar(*info, targetFunc)) {
      return;
    }
 
    CallableSpecializationInputs inputs;
    bool shouldSpecialize = false;
    if (failed(
            prepareCallableSpecialization(callOp, *info, targetFunc, inputs, shouldSpecialize)
        )) {
      failedClone = true;
      return;
    }
    if (!shouldSpecialize) {
      return;
    }
    DenseMap<StringAttr, Type> inferredReplacements =
        getFunctionProofReplacements(*info, targetFunc);
    if (failed(diagnoseCallParamsMismatch(
            callOp.getOperation(), inputs.params, info->oldParamOrder, inferredReplacements,
            &info->replacements, /*explicitCallParams=*/inputs.explicitParams
        ))) {
      failedClone = true;
      return;
    }
 
    InstantiationLayout layout;
    FailureOr<SymbolRefAttr> cloneCallee = getOrCreateSpecializedFunctionClone(
        parentTemplate, targetFunc, callOp.getCalleeAttr(), *info, inputs.replacements, tables,
        functionCloneCache[parentTemplate.getOperation()],
        templateCloneCache[parentTemplate.getOperation()], layout, inputs.params
    );
    if (failed(cloneCallee)) {
      failedClone = true;
      return;
    }
 
    callOp.setCalleeAttr(*cloneCallee);
    callOp.setTemplateParamsAttr(layout.rewrittenCallParams);
  });
  module.walk([&](verif::IncludeOp includeOp) {
    if (failedClone) {
      return;
    }
    FailureOr<SymbolLookupResult<verif::ContractOp>> target = includeOp.getCalleeTarget(tables);
    if (failed(target)) {
      return;
    }
    verif::ContractOp targetContract = target->get();
    auto parentTemplate = llvm::dyn_cast_or_null<TemplateOp>(targetContract->getParentOp());
    if (!parentTemplate) {
      return;
    }
    const TemplateInferenceInfo *info = infoByTemplate.lookup(parentTemplate.getOperation());
    if (!info) {
      return;
    }
    if (!hasResidualFunctionTvar(*info, targetContract)) {
      return;
    }
 
    FailureOr<SymbolLookupResult<FuncDefOp>> targetFuncResult =
        targetContract.getFuncTarget(tables);
    if (failed(targetFuncResult)) {
      return;
    }
    FuncDefOp targetFunc = targetFuncResult->get();
    if (getParentOfType<TemplateOp>(targetFunc.getOperation()) != parentTemplate) {
      return;
    }
 
    CallableSpecializationInputs inputs;
    bool shouldSpecialize = false;
    if (failed(prepareCallableSpecialization(
            includeOp, *info, targetContract, inputs, shouldSpecialize
        ))) {
      failedClone = true;
      return;
    }
    if (!shouldSpecialize) {
      return;
    }
    DenseMap<StringAttr, Type> inferredReplacements =
        getFunctionProofReplacements(*info, targetFunc);
    if (failed(diagnoseCallParamsMismatch(
            includeOp.getOperation(), inputs.params, info->oldParamOrder, inferredReplacements,
            &info->replacements, /*explicitCallParams=*/inputs.explicitParams
        ))) {
      failedClone = true;
      return;
    }
 
    InstantiationLayout targetLayout;
    FailureOr<SymbolRefAttr> specializedTarget = getOrCreateSpecializedFunctionClone(
        parentTemplate, targetFunc, targetContract.getTargetAttr(), *info, inputs.replacements,
        tables, functionCloneCache[parentTemplate.getOperation()],
        templateCloneCache[parentTemplate.getOperation()], targetLayout, inputs.params
    );
    if (failed(specializedTarget)) {
      failedClone = true;
      return;
    }
    InstantiationLayout contractLayout;
    FailureOr<SymbolRefAttr> cloneCallee = getOrCreateSpecializedContractClone(
        parentTemplate, targetContract, includeOp.getCalleeAttr(), *specializedTarget, *info,
        inputs.replacements, tables, contractCloneCache[parentTemplate.getOperation()],
        templateCloneCache[parentTemplate.getOperation()], contractLayout, inputs.params
    );
    if (failed(cloneCallee)) {
      failedClone = true;
      return;
    }
 
    includeOp.setCalleeAttr(*cloneCallee);
    includeOp.setTemplateParamsAttr(contractLayout.rewrittenCallParams);
  });
  return failure(failedClone);
}

static bool sameReplacementTypes(
    const DenseMap<StringAttr, InferredType> &lhs, const DenseMap<StringAttr, InferredType> &rhs
) {
  if (lhs.size() != rhs.size()) {
    return false;
  }
  for (const auto &entry : lhs) {
    auto rhsIt = rhs.find(entry.first);
    if (rhsIt == rhs.end() || rhsIt->second.type != entry.second.type) {
      return false;
    }
  }
  return true;
}

static LogicalResult collectContractTargetFunctionInferences(
    TemplateInferenceInfo &info, verif::ContractOp contract, ModuleOp module,
    DenseMap<Operation *, TemplateInferenceInfo *> *infoByTemplate, SymbolTableCollection &tables,
    bool &changed
) {
  FailureOr<SymbolLookupResult<FuncDefOp>> target = contract.getFuncTarget(tables);
  if (failed(target)) {
    return success();
  }
  FuncDefOp targetFunc = target->get();
  if (getParentOfType<TemplateOp>(targetFunc.getOperation()) != info.templateOp) {
    return success();
  }
 
  DenseMap<StringAttr, InferredType> funcReplacements;
  auto funcIt = info.functionReplacements.find(targetFunc.getOperation());
  if (funcIt != info.functionReplacements.end()) {
    funcReplacements = funcIt->second;
  }
  DenseMap<StringAttr, InferredType> oldFuncReplacements = funcReplacements;
 
  TypeVarInferenceCollector collector(info, funcReplacements);
  if (failed(collector.collect(contract.getOperation()))) {
    return failure();
  }
  if (infoByTemplate && failed(collector.collectStructTemplateParamInferences(
                            contract.getOperation(), module, *infoByTemplate, tables
                        ))) {
    return failure();
  }
 
  if (!sameReplacementTypes(oldFuncReplacements, funcReplacements)) {
    changed = true;
  }
  if (funcReplacements.empty()) {
    info.functionReplacements.erase(targetFunc.getOperation());
  } else {
    info.functionReplacements[targetFunc.getOperation()] = std::move(funcReplacements);
  }
  return success();
}

static LogicalResult collectContractTargetStructInferences(
    TemplateInferenceInfo &info, verif::ContractOp contract, ModuleOp module,
    DenseMap<Operation *, TemplateInferenceInfo *> *infoByTemplate, SymbolTableCollection &tables,
    bool &changed
) {
  FailureOr<SymbolLookupResult<StructDefOp>> target = contract.getStructTarget(tables);
  if (failed(target)) {
    return success();
  }
  StructDefOp targetStruct = target->get();
  if (getParentOfType<TemplateOp>(targetStruct.getOperation()) != info.templateOp) {
    return success();
  }
 
  DenseMap<StringAttr, InferredType> oldTemplateScopeReplacements = info.templateScopeReplacements;
 
  TypeVarInferenceCollector collector(info, info.templateScopeReplacements);
  FunctionType contractTy = contract.getFunctionType();
  if (contractTy.getNumInputs() > 0) {
    StructType targetSelfTy = targetStruct.getType();
    if (auto contractSelfTy = llvm::dyn_cast<StructType>(contractTy.getInput(0));
        contractSelfTy && contractSelfTy.getNameRef() == targetSelfTy.getNameRef() &&
        failed(collector.collectTypeInferences(contractSelfTy, targetSelfTy, contract.getLoc()))) {
      return failure();
    }
  }
  if (failed(collector.collect(contract.getOperation()))) {
    return failure();
  }
  if (infoByTemplate && failed(collector.collectStructTemplateParamInferences(
                            contract.getOperation(), module, *infoByTemplate, tables
                        ))) {
    return failure();
  }
 
  if (!sameReplacementTypes(oldTemplateScopeReplacements, info.templateScopeReplacements)) {
    changed = true;
  }
  return success();
}

static LogicalResult recomputeTemplateWideReplacements(TemplateInferenceInfo &info) {
  DenseMap<StringAttr, unsigned> mentionCounts;
  DenseMap<StringAttr, unsigned> proofCounts;
  DenseMap<StringAttr, InferredType> commonReplacements;
  DenseSet<StringAttr> incompatibleReplacements;
 
  for (FuncDefOp func : walkCollect<FuncDefOp>(info.templateOp)) {
    for (const auto &entry : info.typeVarParams) {
      if (funcMentionsParam(func, entry.first)) {
        ++mentionCounts[entry.first];
      }
    }
 
    auto funcIt = info.functionReplacements.find(func.getOperation());
    if (funcIt == info.functionReplacements.end()) {
      continue;
    }
    for (const auto &entry : funcIt->second) {
      ++proofCounts[entry.first];
      auto commonIt = commonReplacements.find(entry.first);
      if (commonIt == commonReplacements.end()) {
        commonReplacements.try_emplace(entry.first, entry.second);
      } else if (commonIt->second.type != entry.second.type) {
        incompatibleReplacements.insert(entry.first);
      }
    }
  }
 
  for (TemplateExprOp expr : walkCollect<TemplateExprOp>(info.templateOp)) {
    for (const auto &entry : info.typeVarParams) {
      if (exprMentionsParam(expr, entry.first)) {
        ++mentionCounts[entry.first];
      }
    }
  }
 
  info.replacements = info.templateScopeReplacements;
  for (const auto &entry : info.templateScopeReplacements) {
    auto commonIt = commonReplacements.find(entry.first);
    if (commonIt != commonReplacements.end() && commonIt->second.type != entry.second.type) {
      InFlightDiagnostic diag = emitError(entry.second.loc)
                                << "conflicting inferred type for template parameter @"
                                << entry.first.getValue() << ": " << commonIt->second.type << " vs "
                                << entry.second.type;
      diag.attachNote(commonIt->second.loc) << "previous function-local inference here";
      return diag;
    }
  }
 
  for (const auto &entry : commonReplacements) {
    StringAttr paramName = entry.first;
    if (info.replacements.contains(paramName)) {
      continue;
    }
    if (!hasUncoveredNonFunctionMention(
            info.templateOp, paramName, entry.second.type, info.functionReplacements
        ) &&
        !incompatibleReplacements.contains(paramName) &&
        proofCounts.lookup(paramName) == mentionCounts.lookup(paramName)) {
      info.replacements.try_emplace(paramName, entry.second);
    }
  }
  return success();
}

static LogicalResult inferStructTemplateParamUses(
    ModuleOp module, MutableArrayRef<TemplateInferenceInfo> templateInfos,
    DenseMap<Operation *, TemplateInferenceInfo *> &infoByTemplate
) {
  bool changed = false;
  do {
    changed = false;
    SymbolTableCollection tables;
    for (TemplateInferenceInfo &info : templateInfos) {
      for (FuncDefOp func : walkCollect<FuncDefOp>(info.templateOp)) {
        DenseMap<StringAttr, InferredType> funcReplacements;
        auto funcIt = info.functionReplacements.find(func.getOperation());
        if (funcIt != info.functionReplacements.end()) {
          funcReplacements = funcIt->second;
        }
        DenseMap<StringAttr, InferredType> oldFuncReplacements = funcReplacements;
 
        TypeVarInferenceCollector collector(info, funcReplacements);
        if (failed(collector.collect(func.getOperation())) ||
            failed(collector.collectStructTemplateParamInferences(
                func.getOperation(), module, infoByTemplate, tables
            ))) {
          return failure();
        }
 
        if (!sameReplacementTypes(oldFuncReplacements, funcReplacements)) {
          changed = true;
        }
        if (funcReplacements.empty()) {
          info.functionReplacements.erase(func.getOperation());
        } else {
          info.functionReplacements[func.getOperation()] = std::move(funcReplacements);
        }
      }
 
      for (verif::ContractOp contract : walkCollect<verif::ContractOp>(info.templateOp)) {
        if (failed(collectContractTargetFunctionInferences(
                info, contract, module, &infoByTemplate, tables, changed
            )) ||
            failed(collectContractTargetStructInferences(
                info, contract, module, &infoByTemplate, tables, changed
            ))) {
          return failure();
        }
      }
 
      for (TemplateExprOp expr : walkCollect<TemplateExprOp>(info.templateOp)) {
        DenseMap<StringAttr, InferredType> oldTemplateScopeReplacements =
            info.templateScopeReplacements;
 
        TypeVarInferenceCollector collector(info, info.templateScopeReplacements);
        if (failed(collector.collect(expr.getOperation())) ||
            failed(collector.collectStructTemplateParamInferences(
                expr.getOperation(), module, infoByTemplate, tables
            ))) {
          return failure();
        }
 
        if (!sameReplacementTypes(oldTemplateScopeReplacements, info.templateScopeReplacements)) {
          changed = true;
        }
      }
 
      DenseMap<StringAttr, InferredType> oldTemplateScopeReplacements =
          info.templateScopeReplacements;
      TypeVarInferenceCollector collector(info, info.templateScopeReplacements);
      auto nonFunctionResult = info.templateOp.walk([&](Operation *op) -> WalkResult {
        if (llvm::isa<FuncDefOp, TemplateExprOp, verif::ContractOp>(op) ||
            hasParentThatIsa<FuncDefOp, TemplateExprOp, verif::ContractOp>(op)) {
          return WalkResult::advance();
        }
        return collector.collectOperationStructTemplateParamInferences(
            op, module, infoByTemplate, tables
        );
      });
      if (nonFunctionResult.wasInterrupted()) {
        return failure();
      }
      if (!sameReplacementTypes(oldTemplateScopeReplacements, info.templateScopeReplacements)) {
        changed = true;
      }
 
      DenseMap<StringAttr, InferredType> oldReplacements = info.replacements;
      if (failed(recomputeTemplateWideReplacements(info))) {
        return failure();
      }
      if (!sameReplacementTypes(oldReplacements, info.replacements)) {
        changed = true;
      }
    }
  } while (changed);
  return success();
}

static LogicalResult collectExternalContractTargetInferences(
    ModuleOp module, MutableArrayRef<TemplateInferenceInfo> templateInfos,
    DenseMap<Operation *, TemplateInferenceInfo *> &infoByTemplate
) {
  bool changed = false;
  do {
    changed = false;
    bool failedCollection = false;
    SymbolTableCollection tables;
    module.walk([&](verif::ContractOp contract) {
      TemplateOp targetTemplate = getContractTargetTemplate(contract, tables);
      if (!targetTemplate) {
        return;
      }
      if (getParentOfType<TemplateOp>(contract.getOperation()) == targetTemplate) {
        return;
      }
      TemplateInferenceInfo *info = infoByTemplate.lookup(targetTemplate.getOperation());
      if (!info) {
        return;
      }
      if (failed(collectContractTargetFunctionInferences(
              *info, contract, module, &infoByTemplate, tables, changed
          )) ||
          failed(collectContractTargetStructInferences(
              *info, contract, module, &infoByTemplate, tables, changed
          ))) {
        failedCollection = true;
      }
    });
    if (failedCollection) {
      return failure();
    }
 
    for (TemplateInferenceInfo &info : templateInfos) {
      DenseMap<StringAttr, InferredType> oldReplacements = info.replacements;
      if (failed(recomputeTemplateWideReplacements(info))) {
        return failure();
      }
      if (!sameReplacementTypes(oldReplacements, info.replacements)) {
        changed = true;
      }
    }
  } while (changed);
  return success();
}

static FailureOr<TemplateInferenceInfo> buildInfo(TemplateOp templateOp) {
  TemplateInferenceInfo info;
  info.templateOp = templateOp;
  FailureOr<SymbolRefAttr> templatePath =
      getPathFromRoot(llvm::cast<SymbolOpInterface>(templateOp.getOperation()));
  if (failed(templatePath)) {
    return failure();
  }
  info.templatePath = getStringPieces(*templatePath);
  for (TemplateParamOp paramOp : templateOp.getConstOps<TemplateParamOp>()) {
    auto name = paramOp.getSymNameAttr();
    info.oldParamOrder.push_back(name);
    if (isTypeVarParam(paramOp)) {
      info.typeVarParams.try_emplace(name, paramOp);
    }
  }
 
  for (FuncDefOp func : walkCollect<FuncDefOp>(templateOp)) {
    DenseMap<StringAttr, InferredType> funcReplacements;
    if (failed(TypeVarInferenceCollector(info, funcReplacements).collect(func.getOperation()))) {
      return failure();
    }
    if (!funcReplacements.empty()) {
      info.functionReplacements.try_emplace(func.getOperation(), funcReplacements);
    }
  }
 
  for (TemplateExprOp expr : walkCollect<TemplateExprOp>(templateOp)) {
    if (failed(TypeVarInferenceCollector(info, info.templateScopeReplacements)
                   .collect(expr.getOperation()))) {
      return failure();
    }
  }
 
  bool changed = false;
  SymbolTableCollection tables;
  ModuleOp module = templateOp->getParentOfType<ModuleOp>();
  for (verif::ContractOp contract : walkCollect<verif::ContractOp>(templateOp)) {
    if (failed(collectContractTargetFunctionInferences(
            info, contract, module, /*infoByTemplate=*/nullptr, tables, changed
        )) ||
        failed(collectContractTargetStructInferences(
            info, contract, module, /*infoByTemplate=*/nullptr, tables, changed
        ))) {
      return failure();
    }
  }
 
  if (failed(recomputeTemplateWideReplacements(info))) {
    return failure();
  }
  return info;
}

class PassImpl : public llzk::polymorphic::impl::TypeVarInferencePassBase<PassImpl> {
public:
  using Base = TypeVarInferencePassBase<PassImpl>;
  using Base::Base;
 
private:
  void runOnOperation() override {
    ModuleOp module = getOperation();
 
    // Collect all template-local inferences before mutating IR. Conflicts are
    // reported during collection and abort the pass.
    // Note: SmallVector doesn't work because the element size is too large.
    std::vector<TemplateInferenceInfo> templateInfos;
    WalkResult collectResult = module.walk([&templateInfos](TemplateOp templateOp) {
      if (templateOp.getConstOps<TemplateParamOp>().empty()) {
        return WalkResult::advance();
      }
      FailureOr<TemplateInferenceInfo> info = buildInfo(templateOp);
      if (failed(info)) {
        return WalkResult::interrupt();
      }
      templateInfos.push_back(std::move(*info));
      return WalkResult::advance();
    });
    if (collectResult.wasInterrupted()) {
      signalPassFailure();
      return;
    }
    if (templateInfos.empty()) {
      return;
    }
 
    DenseMap<Operation *, TemplateInferenceInfo *> mutableInfoByTemplate;
    for (TemplateInferenceInfo &info : templateInfos) {
      mutableInfoByTemplate.try_emplace(info.templateOp.getOperation(), &info);
    }
    if (failed(
            collectExternalContractTargetInferences(module, templateInfos, mutableInfoByTemplate)
        )) {
      signalPassFailure();
      return;
    }
    if (failed(inferStructTemplateParamUses(module, templateInfos, mutableInfoByTemplate))) {
      signalPassFailure();
      return;
    }
 
    SmallVector<TemplateInferenceInfo *> rewrites;
    DenseMap<Operation *, const TemplateInferenceInfo *> infoByTemplate;
    for (TemplateInferenceInfo &info : templateInfos) {
      infoByTemplate.try_emplace(info.templateOp.getOperation(), &info);
      if (!info.replacements.empty() || !info.functionReplacements.empty()) {
        rewrites.push_back(&info);
      }
    }
 
    // Rewrite each affected template in place, but keep the converters alive so
    // call-site rewriting can use the same positional template-parameter map.
    SmallVector<std::unique_ptr<TypeVarReplacementConverter>> converterStorage;
    DenseMap<Operation *, const TypeVarReplacementConverter *> convertersByTemplate;
    for (TemplateInferenceInfo *info : rewrites) {
      DenseMap<StringAttr, Type> replacements;
      for (const auto &entry : info->replacements) {
        replacements.try_emplace(entry.first, entry.second.type);
      }
      auto converter = std::make_unique<TypeVarReplacementConverter>(
          module.getContext(), info->templatePath, info->oldParamOrder, replacements
      );
      convertersByTemplate.try_emplace(info->templateOp.getOperation(), converter.get());
      converterStorage.push_back(std::move(converter));
    }
 
    // Clone concrete call-site instantiations before trimming template
    // arguments; the clone decision needs the original explicit arguments.
    SpecializedCallableCloneCache functionCloneCache;
    SpecializedCallableCloneCache contractCloneCache;
    SpecializedTemplateCloneCache templateCloneCache;
    if (failed(specializeFunctionLocalCallables(
            module, infoByTemplate, functionCloneCache, contractCloneCache, templateCloneCache
        ))) {
      signalPassFailure();
      return;
    }
 
    // Trim call/include argument lists before rewriting caller bodies, so a
    // caller-local type variable in an explicit argument for an erased callee
    // parameter is not rewritten in the caller's scope before it can be dropped.
    if (failed(updateCallableTemplateParams(module, convertersByTemplate))) {
      signalPassFailure();
      return;
    }
 
    for (TemplateInferenceInfo *info : rewrites) {
      const TypeVarReplacementConverter *converter =
          convertersByTemplate.lookup(info->templateOp.getOperation());
      assert(converter && "rewritten template must have a converter");
      auto validationResult = info->templateOp.walk([converter](Operation *op) -> WalkResult {
        return converter->validateOperation(op);
      });
      if (validationResult.wasInterrupted()) {
        signalPassFailure();
        return;
      }
      if (failed(convertOperationTypesIn(info->templateOp.getOperation(), *converter))) {
        signalPassFailure();
        return;
      }
      removeIdentityCasts(info->templateOp.getOperation());
    }
 
    if (failed(updateExternalContractTemplateParams(module, convertersByTemplate))) {
      signalPassFailure();
      return;
    }
 
    // Calls that still target rewritten templates are updated after all target
    // templates have their converters registered.
    if (failed(updateCallableTemplateParams(module, convertersByTemplate))) {
      signalPassFailure();
      return;
    }
    // Re-run specialization after rewriting because template bodies may now
    // contain concrete forwarded calls into otherwise unconstrained templates.
    if (failed(specializeFunctionLocalCallables(
            module, infoByTemplate, functionCloneCache, contractCloneCache, templateCloneCache
        ))) {
      signalPassFailure();
      return;
    }
    if (failed(updateStructTemplateParams(module, convertersByTemplate))) {
      signalPassFailure();
      return;
    }
 
    // Erase resolved parameters after all argument-list conversion has used
    // the original order, but before concrete struct instantiation verifies
    // rewritten owned-struct arities.
    for (TemplateInferenceInfo *info : rewrites) {
      if (failed(removeResolvedParams(*info))) {
        signalPassFailure();
        return;
      }
    }
 
    if (failed(instantiateConcreteStructUses(module))) {
      signalPassFailure();
      return;
    }
  }
};
 
} // namespace