// Part of the Carbon Language project, under the Apache License v2.0 with LLVM
// Exceptions. See /LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "toolchain/check/custom_witness.h"

#include "toolchain/base/kind_switch.h"
#include "toolchain/check/facet_type.h"
#include "toolchain/check/function.h"
#include "toolchain/check/generic.h"
#include "toolchain/check/impl.h"
#include "toolchain/check/impl_lookup.h"
#include "toolchain/check/import_ref.h"
#include "toolchain/check/inst.h"
#include "toolchain/check/name_lookup.h"
#include "toolchain/check/type.h"
#include "toolchain/check/type_completion.h"
#include "toolchain/sem_ir/associated_constant.h"
#include "toolchain/sem_ir/builtin_function_kind.h"
#include "toolchain/sem_ir/ids.h"
#include "toolchain/sem_ir/typed_insts.h"

namespace Carbon::Check {

// Given a value whose type `IsFacetTypeOrError`, returns the corresponding
// type.
static auto GetFacetAsType(Context& context,
                           SemIR::ConstantId facet_or_type_const_id)
    -> SemIR::TypeId {
  auto facet_or_type_id =
      context.constant_values().GetInstId(facet_or_type_const_id);
  auto type_type_id = context.insts().Get(facet_or_type_id).type_id();
  CARBON_CHECK(context.types().IsFacetTypeOrError(type_type_id));

  if (context.types().Is<SemIR::FacetType>(type_type_id)) {
    // It's a facet; access its type.
    facet_or_type_id = context.types().GetTypeInstId(
        GetFacetAccessType(context, facet_or_type_id));
  }
  return context.types().GetTypeIdForTypeInstId(facet_or_type_id);
}

// Returns the body for `Destroy.Op`. This will return `None` if using the
// builtin `NoOp` is appropriate.
//
// TODO: This is a placeholder still not actually destroying things, intended to
// maintain mostly-consistent behavior with current logic while working. That
// also means using `self`.
// TODO: This mirrors `TypeCanDestroy` below, think about ways to share what's
// handled.
static auto MakeDestroyOpBody(Context& context, SemIR::LocId loc_id,
                              SemIR::TypeId self_type_id)
    -> SemIR::InstBlockId {
  context.inst_block_stack().Push();
  auto inst = context.types().GetAsInst(self_type_id);

  while (auto class_type = inst.TryAs<SemIR::ClassType>()) {
    // Switch to looking at the object representation.
    auto class_info = context.classes().Get(class_type->class_id);
    CARBON_CHECK(class_info.is_complete());
    inst = context.types().GetAsInst(
        class_info.GetObjectRepr(context.sem_ir(), class_type->specific_id));
  }

  CARBON_KIND_SWITCH(inst) {
    case SemIR::ArrayType::Kind:
    case SemIR::ConstType::Kind:
    case SemIR::MaybeUnformedType::Kind:
    case SemIR::PartialType::Kind:
    case SemIR::StructType::Kind:
    case SemIR::TupleType::Kind:
      // TODO: Implement iterative destruction of types.
      break;
    case SemIR::BoolType::Kind:
    case SemIR::FloatType::Kind:
    case SemIR::IntType::Kind:
    case SemIR::PointerType::Kind:
      // For trivially destructible types, we don't generate anything, so that
      // this can collapse to a noop implementation when possible.
      break;
    case SemIR::ErrorInst::Kind:
      // Errors can't be destroyed, but we'll still try to generate calls for
      // other members.
      break;
    default:
      CARBON_FATAL("Unexpected type for destroy: {0}", inst);
  }

  if (context.inst_block_stack().PeekCurrentBlockContents().empty()) {
    context.inst_block_stack().PopAndDiscard();
    return SemIR::InstBlockId::None;
  }
  AddInst(context, loc_id, SemIR::Return{});
  return context.inst_block_stack().Pop();
}

// Returns a manufactured `Destroy.Op` function with the `self` parameter typed
// to `self_type_id`.
static auto MakeDestroyOpFunction(Context& context, SemIR::LocId loc_id,
                                  SemIR::TypeId self_type_id,
                                  SemIR::NameScopeId parent_scope_id)
    -> SemIR::InstId {
  auto name_id = context.core_identifiers().AddNameId(CoreIdentifier::Op);

  auto [decl_id, function_id] =
      MakeGeneratedFunctionDecl(context, loc_id,
                                {.parent_scope_id = parent_scope_id,
                                 .name_id = name_id,
                                 .self_type_id = self_type_id});

  auto& function = context.functions().Get(function_id);

  auto body_id = MakeDestroyOpBody(context, loc_id, self_type_id);
  if (body_id.has_value()) {
    function.SetCoreWitness();
    function.body_block_ids.push_back(body_id);
  } else {
    function.SetCoreWitness(SemIR::BuiltinFunctionKind::NoOp);
  }

  return decl_id;
}

static auto MakeCustomWitnessConstantInst(
    Context& context, SemIR::LocId loc_id,
    SemIR::SpecificInterfaceId query_specific_interface_id,
    SemIR::InstBlockId associated_entities_block_id) -> SemIR::InstId {
  // The witness is a CustomWitness of the query interface with a table that
  // contains each associated entity.
  auto const_id = EvalOrAddInst<SemIR::CustomWitness>(
      context, loc_id,
      {.type_id = GetSingletonType(context, SemIR::WitnessType::TypeInstId),
       .elements_id = associated_entities_block_id,
       .query_specific_interface_id = query_specific_interface_id});
  return context.constant_values().GetInstId(const_id);
}

struct TypesForSelfFacet {
  // A FacetType that contains only the query interface.
  SemIR::TypeId facet_type_for_query_specific_interface;
  // The query self as a type, which involves a conversion if it was a facet.
  SemIR::TypeId query_self_as_type_id;
};

static auto GetTypesForSelfFacet(
    Context& context, SemIR::LocId loc_id,
    SemIR::ConstantId query_self_const_id,
    SemIR::SpecificInterfaceId query_specific_interface_id)
    -> TypesForSelfFacet {
  const auto query_specific_interface =
      context.specific_interfaces().Get(query_specific_interface_id);

  // The Self facet will have type FacetType, for the query interface.
  auto facet_type_for_query_specific_interface =
      context.types().GetTypeIdForTypeConstantId(
          EvalOrAddInst<SemIR::FacetType>(
              context, loc_id,
              FacetTypeFromInterface(context,
                                     query_specific_interface.interface_id,
                                     query_specific_interface.specific_id)));
  // The Self facet needs to point to a type value. If it's not one already,
  // convert to type.
  auto query_self_as_type_id = GetFacetAsType(context, query_self_const_id);
  return {facet_type_for_query_specific_interface, query_self_as_type_id};
}

// Build a new facet from the query self, using a CustomWitness for the query
// interface with an entry for each associated entity so far.
static auto MakeSelfFacetWithCustomWitness(
    Context& context, SemIR::LocId loc_id, TypesForSelfFacet query_types,
    SemIR::SpecificInterfaceId query_specific_interface_id,
    SemIR::InstBlockId associated_entities_block_id) -> SemIR::ConstantId {
  // We are building a facet value for a single interface, so the witness block
  // is a single witness for that interface.
  auto witnesses_block_id =
      context.inst_blocks().Add({MakeCustomWitnessConstantInst(
          context, loc_id, query_specific_interface_id,
          associated_entities_block_id)});

  return EvalOrAddInst<SemIR::FacetValue>(
      context, loc_id,
      {.type_id = query_types.facet_type_for_query_specific_interface,
       .type_inst_id =
           context.types().GetTypeInstId(query_types.query_self_as_type_id),
       .witnesses_block_id = witnesses_block_id});
}

auto BuildCustomWitness(Context& context, SemIR::LocId loc_id,
                        SemIR::ConstantId query_self_const_id,
                        SemIR::SpecificInterfaceId query_specific_interface_id,
                        llvm::ArrayRef<SemIR::InstId> values) -> SemIR::InstId {
  const auto query_specific_interface =
      context.specific_interfaces().Get(query_specific_interface_id);
  const auto& interface =
      context.interfaces().Get(query_specific_interface.interface_id);
  auto assoc_entities =
      context.inst_blocks().GetOrEmpty(interface.associated_entities_id);
  if (assoc_entities.size() != values.size()) {
    context.TODO(loc_id, ("Unsupported definition of interface " +
                          context.names().GetFormatted(interface.name_id))
                             .str());
    return SemIR::ErrorInst::InstId;
  }

  auto query_types_for_self_facet = GetTypesForSelfFacet(
      context, loc_id, query_self_const_id, query_specific_interface_id);

  // The values that will go in the witness table.
  llvm::SmallVector<SemIR::InstId> entries;

  // Fill in the witness table.
  for (const auto& [assoc_entity_id, value_id] :
       llvm::zip_equal(assoc_entities, values)) {
    LoadImportRef(context, assoc_entity_id);

    // Build a witness with the current contents of the witness table. This will
    // grow as we progress through the impl. In theory this will build O(n^2)
    // table entries, but in practice n <= 2, so that's OK.
    //
    // This is necessary because later associated entities may refer to earlier
    // associated entities in their signatures. In particular, an associated
    // result type may be used as the return type of an associated function.
    auto self_facet = MakeSelfFacetWithCustomWitness(
        context, loc_id, query_types_for_self_facet,
        query_specific_interface_id, context.inst_blocks().Add(entries));
    auto interface_with_self_specific_id = MakeSpecificWithInnerSelf(
        context, loc_id, interface.generic_id, interface.generic_with_self_id,
        query_specific_interface.specific_id, self_facet);

    auto decl_id =
        context.constant_values().GetInstId(SemIR::GetConstantValueInSpecific(
            context.sem_ir(), interface_with_self_specific_id,
            assoc_entity_id));
    CARBON_CHECK(decl_id.has_value(), "Non-constant associated entity");
    auto decl = context.insts().Get(decl_id);
    CARBON_KIND_SWITCH(decl) {
      case CARBON_KIND(SemIR::StructValue struct_value): {
        if (struct_value.type_id == SemIR::ErrorInst::TypeId) {
          return SemIR::ErrorInst::InstId;
        }
        // TODO: If a thunk is needed, this will build a different value each
        // time it's called, so we won't properly deduplicate repeated
        // witnesses.
        entries.push_back(CheckAssociatedFunctionImplementation(
            context,
            context.types().GetAs<SemIR::FunctionType>(struct_value.type_id),
            query_specific_interface.specific_id, value_id,
            /*defer_thunk_definition=*/false));
        break;
      }
      case CARBON_KIND(SemIR::AssociatedConstantDecl decl): {
        if (decl.type_id == SemIR::ErrorInst::TypeId) {
          return SemIR::ErrorInst::InstId;
        }

        // TODO: remove once we have a test-case for all associated constants.
        // Special-case the ones we want to support in this if-statement, until
        // we're able to account for everything.
        if (decl.type_id != SemIR::TypeType::TypeId) {
          context.TODO(loc_id,
                       "Associated constant of type other than `TypeType` in "
                       "synthesized impl");
          return SemIR::ErrorInst::InstId;
        }

        auto type_id = context.insts().Get(value_id).type_id();
        CARBON_CHECK(type_id == SemIR::TypeType::TypeId ||
                     type_id == SemIR::ErrorInst::TypeId);
        auto impl_witness_associated_constant =
            AddInst<SemIR::ImplWitnessAssociatedConstant>(
                context, loc_id, {.type_id = type_id, .inst_id = value_id});
        entries.push_back(impl_witness_associated_constant);
        break;
      }
      default:
        CARBON_CHECK(decl_id == SemIR::ErrorInst::InstId,
                     "Unexpected kind of associated entity {0}", decl);
        return SemIR::ErrorInst::InstId;
    }
  }

  // TODO: Consider building one witness after all associated constants, and
  // then a second after all associated functions, rather than building one in
  // each `StructValue`. Right now the code is written assuming at most one
  // function, though this CHECK can be removed as a temporary workaround.
  auto associated_functions = llvm::count_if(entries, [&](SemIR::InstId id) {
    return context.insts().Get(id).kind() == SemIR::InstKind::FunctionDecl;
  });
  CARBON_CHECK(associated_functions <= 1,
               "TODO: Support multiple associated functions");

  return MakeCustomWitnessConstantInst(context, loc_id,
                                       query_specific_interface_id,
                                       context.inst_blocks().Add(entries));
}

auto GetCoreInterface(Context& context, SemIR::InterfaceId interface_id)
    -> CoreInterface {
  const auto& interface = context.interfaces().Get(interface_id);
  if (!context.name_scopes().IsCorePackage(interface.parent_scope_id) ||
      !interface.name_id.AsIdentifierId().has_value()) {
    return CoreInterface::Unknown;
  }

  constexpr auto CoreIdentifiersToInterfaces = std::array{
      std::pair{CoreIdentifier::Copy, CoreInterface::Copy},
      std::pair{CoreIdentifier::CppUnsafeDeref, CoreInterface::CppUnsafeDeref},
      std::pair{CoreIdentifier::Default, CoreInterface::Default},
      std::pair{CoreIdentifier::Destroy, CoreInterface::Destroy},
      std::pair{CoreIdentifier::IntFitsIn, CoreInterface::IntFitsIn}};

  for (auto [core_identifier, core_interface] : CoreIdentifiersToInterfaces) {
    if (interface.name_id ==
        context.core_identifiers().AddNameId(core_identifier)) {
      return core_interface;
    }
  }
  return CoreInterface::Unknown;
}

// Returns true if the `Self` should impl `Destroy`.
static auto TypeCanDestroy(Context& context,
                           SemIR::ConstantId query_self_const_id,
                           SemIR::InterfaceId destroy_interface_id) -> bool {
  auto inst = context.insts().Get(context.constant_values().GetInstId(
      GetCanonicalFacetOrTypeValue(context, query_self_const_id)));

  // For facet values, look if the FacetType provides the same.
  if (auto facet_type =
          context.types().TryGetAs<SemIR::FacetType>(inst.type_id())) {
    const auto& info = context.facet_types().Get(facet_type->facet_type_id);
    for (auto interface : info.extend_constraints) {
      if (interface.interface_id == destroy_interface_id) {
        return true;
      }
    }
  }

  CARBON_KIND_SWITCH(inst) {
    case CARBON_KIND(SemIR::ClassType class_type): {
      auto class_info = context.classes().Get(class_type.class_id);
      // Incomplete and abstract classes can't be destroyed.
      if (!class_info.is_complete() ||
          class_info.inheritance_kind ==
              SemIR::Class::InheritanceKind::Abstract) {
        return false;
      }

      // `LookupCppImpl` handles C++ types.
      if (context.name_scopes().Get(class_info.scope_id).is_cpp_scope()) {
        return false;
      }

      // TODO: Return false if the object repr doesn't impl `Destroy`.
      return true;
    }
    case SemIR::ArrayType::Kind:
    case SemIR::ConstType::Kind:
    case SemIR::MaybeUnformedType::Kind:
    case SemIR::PartialType::Kind:
    case SemIR::StructType::Kind:
    case SemIR::TupleType::Kind:
      // TODO: Return false for types that indirectly reference a type that
      // doesn't impl `Destroy`.
      return true;
    case SemIR::BoolType::Kind:
    case SemIR::FloatType::Kind:
    case SemIR::IntType::Kind:
    case SemIR::PointerType::Kind:
      // Trivially destructible.
      return true;
    default:
      return false;
  }
}

static auto MakeDestroyWitness(
    Context& context, SemIR::LocId loc_id,
    SemIR::ConstantId query_self_const_id,
    SemIR::SpecificInterfaceId query_specific_interface_id)
    -> std::optional<SemIR::InstId> {
  auto query_specific_interface =
      context.specific_interfaces().Get(query_specific_interface_id);

  if (!TypeCanDestroy(context, query_self_const_id,
                      query_specific_interface.interface_id)) {
    return std::nullopt;
  }

  if (query_self_const_id.is_symbolic()) {
    return SemIR::InstId::None;
  }

  // Mark functions with the interface's scope as a hint to mangling. This does
  // not add them to the scope.
  auto parent_scope_id = context.interfaces()
                             .Get(query_specific_interface.interface_id)
                             .scope_without_self_id;

  auto self_type_id = GetFacetAsType(context, query_self_const_id);
  auto op_id =
      MakeDestroyOpFunction(context, loc_id, self_type_id, parent_scope_id);
  return BuildCustomWitness(context, loc_id, query_self_const_id,
                            query_specific_interface_id, {op_id});
}

static auto MakeIntFitsInWitness(
    Context& context, SemIR::LocId loc_id,
    SemIR::ConstantId query_self_const_id,
    SemIR::SpecificInterfaceId query_specific_interface_id)
    -> std::optional<SemIR::InstId> {
  auto query_specific_interface =
      context.specific_interfaces().Get(query_specific_interface_id);

  auto args_id = query_specific_interface.specific_id;
  if (!args_id.has_value()) {
    return std::nullopt;
  }
  auto args_block_id = context.specifics().Get(args_id).args_id;
  auto args_block = context.inst_blocks().Get(args_block_id);
  if (args_block.size() != 1) {
    return std::nullopt;
  }

  auto dest_const_id = context.constant_values().Get(args_block[0]);
  if (!dest_const_id.is_constant()) {
    return std::nullopt;
  }

  auto src_type_id = GetFacetAsType(context, query_self_const_id);
  auto dest_type_id = GetFacetAsType(context, dest_const_id);

  auto context_fn = [](DiagnosticContextBuilder& /*builder*/) -> void {};
  if (!RequireCompleteType(context, src_type_id, loc_id, context_fn) ||
      !RequireCompleteType(context, dest_type_id, loc_id, context_fn)) {
    return std::nullopt;
  }

  auto src_info = context.types().TryGetIntTypeInfo(src_type_id);
  auto dest_info = context.types().TryGetIntTypeInfo(dest_type_id);

  if (!src_info || !dest_info) {
    return std::nullopt;
  }

  // If the bit width is unknown (e.g., due to symbolic evaluation), we cannot
  // determine whether it fits yet.
  if (src_info->bit_width == IntId::None ||
      dest_info->bit_width == IntId::None) {
    return std::nullopt;
  }

  const auto& src_width = context.ints().Get(src_info->bit_width);
  const auto& dest_width = context.ints().Get(dest_info->bit_width);

  bool fits = false;
  if (src_info->is_signed && !dest_info->is_signed) {
    // Signed -> unsigned: would truncate the sign bit.
    fits = false;
  } else if (src_info->is_signed == dest_info->is_signed) {
    // Signed -> signed or unsigned -> unsigned: allow widening or preserving
    // width.
    fits = src_width.sle(dest_width);
  } else {
    // Unsigned -> signed: strict widening required.
    fits = src_width.slt(dest_width);
  }

  if (!fits) {
    return std::nullopt;
  }

  return BuildCustomWitness(context, loc_id, query_self_const_id,
                            query_specific_interface_id, {});
}

auto LookupCustomWitness(Context& context, SemIR::LocId loc_id,
                         CoreInterface core_interface,
                         SemIR::ConstantId query_self_const_id,
                         SemIR::SpecificInterfaceId query_specific_interface_id)
    -> std::optional<SemIR::InstId> {
  switch (core_interface) {
    case CoreInterface::Destroy:
      return MakeDestroyWitness(context, loc_id, query_self_const_id,
                                query_specific_interface_id);
    case CoreInterface::IntFitsIn:
      return MakeIntFitsInWitness(context, loc_id, query_self_const_id,
                                  query_specific_interface_id);
    case CoreInterface::Copy:
    case CoreInterface::CppUnsafeDeref:
    case CoreInterface::Default:
    case CoreInterface::Unknown:
      // TODO: Handle more interfaces, particularly copy, move, and conversion.
      return std::nullopt;
  }
}

}  // namespace Carbon::Check