// 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/period_self.h"

#include "toolchain/base/kind_switch.h"
#include "toolchain/check/control_flow.h"
#include "toolchain/check/convert.h"
#include "toolchain/check/facet_type.h"
#include "toolchain/check/generic.h"
#include "toolchain/check/inst.h"
#include "toolchain/check/subst.h"
#include "toolchain/check/type.h"
#include "toolchain/check/type_completion.h"
#include "toolchain/sem_ir/typed_insts.h"

namespace Carbon::Check {

auto MakePeriodSelfFacetValue(Context& context, SemIR::LocId loc_id,
                              SemIR::TypeId self_type_id) -> SemIR::InstId {
  CARBON_CHECK(self_type_id == SemIR::ErrorInst::TypeId ||
               context.types().Is<SemIR::FacetType>(self_type_id));
  auto entity_name_id = context.entity_names().AddCanonical({
      .name_id = SemIR::NameId::PeriodSelf,
      .parent_scope_id = context.scope_stack().PeekNameScopeId(),
  });
  auto inst_id = AddInst<SemIR::SymbolicBinding>(
      context, loc_id,
      {
          .type_id = self_type_id,
          .entity_name_id = entity_name_id,
          // `None` because there is no equivalent non-symbolic value.
          .value_id = SemIR::InstId::None,
      });
  auto existing = context.scope_stack().LookupOrAddName(
      SemIR::NameId::PeriodSelf, inst_id, ScopeIndex::None,
      IsCurrentPositionReachable(context));
  // Shouldn't have any names in newly created scope.
  CARBON_CHECK(!existing.has_value());
  return inst_id;
}

class SubstPeriodSelfCallbacks : public SubstInstCallbacks {
 public:
  explicit SubstPeriodSelfCallbacks(
      Context* context, SemIR::LocId loc_id,
      SemIR::ConstantId period_self_replacement_id,
      SubstPeriodSelfBehaviour behaviour, SubstPeriodSelfRebuildInst rebuild)
      : SubstInstCallbacks(context),
        loc_id_(loc_id),
        period_self_replacement_id_(period_self_replacement_id),
        behaviour_(behaviour),
        rebuild_callback_(rebuild) {}

  virtual ~SubstPeriodSelfCallbacks() {
    CARBON_CHECK(designator_states_.empty());
  }

  auto Subst(SemIR::InstId& inst_id) -> SubstResult override {
    // FacetTypes are concrete even if they have `.Self` inside them, but we
    // don't recurse into FacetTypes, so we can use this as a base case. This
    // avoids infinite recursion on TypeType and ErrorInst.
    if (context().constant_values().Get(inst_id).is_concrete()) {
      return FullySubstituted;
    }
    // Don't recurse into nested facet types, even if they are symbolic. Leave
    // their `.Self` as is.
    if (context().insts().Is<SemIR::FacetType>(inst_id)) {
      return FullySubstituted;
    }

    // Look for implicit use of `.Self` in designators: `.X` is really
    // `.Self.X`.
    if (auto access =
            context().insts().TryGetAs<SemIR::ImplWitnessAccess>(inst_id)) {
      if (auto witness = context().insts().TryGetAs<SemIR::LookupImplWitness>(
              access->witness_id)) {
        // Canonicalization not necessary; we are working with the constant
        // value already, and the query self in a witness is already
        // canonicalized.
        if (IsPeriodSelf(context(), witness->query_self_inst_id,
                         /*canonicalize=*/false)) {
          // We are entering a designator. We watch for the witness next.
          designator_states_.push_back(WitnessNext);
          return SubstOperands;
        }
      }
    }

    if (GetDesignatorState() == WitnessNext) {
      if (auto witness =
              context().insts().TryGetAs<SemIR::LookupImplWitness>(inst_id)) {
        // The query self comes next, before the specific interface. This
        // deepnds on the order of the operands in the LookupImplWitness, and
        // that Subst visits them in top to bottom order.
        designator_states_.back() = WitnessSelfNext;
        return SubstOperands;
      }
    }

    // Canonicalization not necessary; we are working with the constant
    // value already, and the query self in a witness is already
    // canonicalized.
    if (IsPeriodSelf(context(), inst_id,
                     /*canonicalize=*/false)) {
      bool is_implicit_self_in_desigator = false;
      if (GetDesignatorState() == WitnessSelfNext) {
        is_implicit_self_in_desigator = true;
        designator_states_.back() = RebuildNext;
      }

      switch (behaviour_) {
        case SubstPeriodSelfBehaviour::All:
          inst_id = GetReplacement(inst_id);
          break;
        case SubstPeriodSelfBehaviour::ImplicitOnly:
          if (is_implicit_self_in_desigator) {
            inst_id = GetReplacement(inst_id);
          }
          break;
        case SubstPeriodSelfBehaviour::ExplicitOnly:
          if (!is_implicit_self_in_desigator) {
            inst_id = GetReplacement(inst_id);
          }
          break;
      }
      return FullySubstituted;
    }

    return SubstOperands;
  }

  auto Rebuild(SemIR::InstId orig_inst_id, SemIR::Inst new_inst)
      -> SemIR::InstId override {
    TryPopDesignatorState(orig_inst_id);
    if (rebuild_callback_) {
      if (auto inst_id = rebuild_callback_(new_inst); inst_id.has_value()) {
        return inst_id;
      }
    }
    return RebuildNewInst(SemIR::LocId(orig_inst_id), new_inst);
  }

  auto ReuseUnchanged(SemIR::InstId orig_inst_id) -> SemIR::InstId override {
    TryPopDesignatorState(orig_inst_id);
    return orig_inst_id;
  }

 private:
  enum DesignatorState {
    // We are not inside a designator.
    None,
    // We are looking for the LookupImplWitness of a designator next.
    WitnessNext,
    // We are looking for `.Self`, the next one will be the query self of the
    // LookupImplWitness of a designator. This is the implicit use of `.Self`.
    WitnessSelfNext,
    // We have seen the query self of the designator, and are now looking for
    // the designator to be rebuilt with the replacement. Any other `.Self` that
    // we find are explicit uses of `.Self`, such as in the designator's
    // specific interface.
    RebuildNext,
  };

  auto GetReplacement(SemIR::InstId period_self) -> SemIR::InstId {
    auto period_self_type_id = context().insts().Get(period_self).type_id();
    CARBON_CHECK(context().types().Is<SemIR::FacetType>(period_self_type_id));

    auto replacement_self_inst_id =
        context().constant_values().GetInstId(period_self_replacement_id_);
    auto replacement_type_id =
        context().insts().Get(replacement_self_inst_id).type_id();
    CARBON_CHECK(context().types().IsFacetType(replacement_type_id));

    // If the replacement has the same type as `.Self`, use it directly.
    if (replacement_type_id == period_self_type_id) {
      return replacement_self_inst_id;
    }

    // If we have already converted the replacement to the type of `.Self`, use
    // our previous conversion.
    if (period_self_type_id == cached_replacement_type_id_) {
      return cached_replacement_id_;
    }

    // Convert the replacement facet to the type of `.Self`.
    cached_replacement_id_ = ConvertReplacement(
        replacement_self_inst_id, replacement_type_id, period_self_type_id);
    cached_replacement_type_id_ = period_self_type_id;
    return cached_replacement_id_;
  }

  auto ConvertReplacement(SemIR::InstId replacement_self_inst_id,
                          SemIR::TypeId replacement_type_id,
                          SemIR::TypeId period_self_type_id) -> SemIR::InstId {
    // TODO: Replace all empty facet types with TypeType.
    if (period_self_type_id == GetEmptyFacetType(context())) {
      // Convert to an empty facet type (representing TypeType); we don't need
      // any witnesses.
      return ConvertToValueOfType(context(), loc_id_, replacement_self_inst_id,
                                  period_self_type_id);
    }

    // We have a facet or a type, but we need more interfaces in the facet type.
    // We will have to synthesize a symbolic witness for each interface.
    //
    // Why is this okay? The type of `.Self` comes from interfaces that are
    // before it (to the left of it) in the facet type. The replacement for
    // `.Self` will have to impl those interfaces in order to match the facet
    // type, so we know that it is valid to construct these witnesses.

    // Make the replacement into a type, which we will need for the FacetValue.
    if (context().types().Is<SemIR::FacetType>(replacement_type_id)) {
      replacement_self_inst_id = context().constant_values().GetInstId(
          EvalOrAddInst<SemIR::FacetAccessType>(
              context(), loc_id_,
              {.type_id = SemIR::TypeType::TypeId,
               .facet_value_inst_id = replacement_self_inst_id}));
    }

    auto identified_period_self_type_id = RequireIdentifiedFacetType(
        context(), loc_id_,
        context().constant_values().Get(replacement_self_inst_id),
        context().types().GetTypeInstId(period_self_type_id),
        [&](auto& /*builder*/) {
          // Given `I where .Self == ()`, the type of `.Self` is `I` and we're
          // replacing `.Self` with some `T` that must also implement `I`.
          // However `I` can be a generic with arbitrary complexity and the
          // replacement with `T` may fail monomorphization.
          //
          // We don't have any better context to add here really, but we
          // need to accept that errors happen rather than CHECKing that
          // they don't.
        });
    if (!identified_period_self_type_id.has_value()) {
      return SemIR::ErrorInst::InstId;
    }
    const auto& identified_period_self_type =
        context().identified_facet_types().Get(identified_period_self_type_id);
    auto required_impls = identified_period_self_type.required_impls();
    llvm::SmallVector<SemIR::InstId> witnesses;
    witnesses.reserve(required_impls.size());
    for (const auto& req : required_impls) {
      witnesses.push_back(context().constant_values().GetInstId(
          EvalOrAddInst<SemIR::LookupImplWitness>(
              context(), loc_id_,
              {.type_id =
                   GetSingletonType(context(), SemIR::WitnessType::TypeInstId),
               .query_self_inst_id =
                   context().constant_values().GetInstId(req.self_facet_value),
               .query_specific_interface_id =
                   context().specific_interfaces().Add(
                       req.specific_interface)})));
    }
    return context().constant_values().GetInstId(
        EvalOrAddInst<SemIR::FacetValue>(
            context(), loc_id_,
            {
                .type_id = period_self_type_id,
                .type_inst_id =
                    context().types().GetAsTypeInstId(replacement_self_inst_id),
                .witnesses_block_id = context().inst_blocks().Add(witnesses),
            }));
  }

  auto GetDesignatorState() const -> DesignatorState {
    return designator_states_.empty() ? DesignatorState::None
                                      : designator_states_.back();
  }

  auto TryPopDesignatorState(SemIR::InstId orig_inst_id) -> void {
    if (GetDesignatorState() == RebuildNext) {
      if (auto access = context().insts().TryGetAs<SemIR::ImplWitnessAccess>(
              orig_inst_id)) {
        if (auto witness = context().insts().TryGetAs<SemIR::LookupImplWitness>(
                access->witness_id)) {
          // Canonicalization not necessary; we are working with the constant
          // value already, and the query self in a witness is already
          // canonicalized.
          if (IsPeriodSelf(context(), witness->query_self_inst_id,
                           /*canonicalize=*/false)) {
            designator_states_.pop_back();
          }
        }
      }
    }
  }

  SemIR::LocId loc_id_;
  SemIR::ConstantId period_self_replacement_id_;
  SubstPeriodSelfBehaviour behaviour_;
  SubstPeriodSelfRebuildInst rebuild_callback_;

  // The last output of GetReplacement().
  SemIR::InstId cached_replacement_id_ = SemIR::InstId::None;
  // The type of the last output of GetReplacement(). If the type of `.Self`
  // matches, we can reuse the `cached_replacement_id_`.
  SemIR::TypeId cached_replacement_type_id_ = SemIR::TypeId::None;

  llvm::SmallVector<DesignatorState> designator_states_;
};

auto SubstPeriodSelf(Context& context, SemIR::LocId loc_id,
                     SemIR::ConstantId const_id,
                     SemIR::ConstantId period_self_replacement_id,
                     SubstPeriodSelfBehaviour behaviour,
                     SubstPeriodSelfRebuildInst rebuild) -> SemIR::ConstantId {
  // Don't replace `.Self` with itself; that is cyclical.
  //
  // If the types differ, we would try to convert the replacement to a `.Self`
  // of the desired type in `const_id`, which is what we already have, so
  // there's nothing we need to do. But trying to do that conversion recurses
  // when the type of the `.Self` contains a `.Self`.
  if (IsPeriodSelf(context, context.constant_values().GetInstId(
                                period_self_replacement_id))) {
    return const_id;
  }

  SubstPeriodSelfCallbacks callbacks(
      &context, loc_id, period_self_replacement_id, behaviour, rebuild);
  auto subst_id = SubstInst(
      context, context.constant_values().GetInstId(const_id), callbacks);
  return context.constant_values().Get(subst_id);
}

static auto SubstPeriodSelfInSpecific(
    Context& context, SemIR::LocId loc_id, SemIR::SpecificId specific_id,
    SemIR::ConstantId period_self_replacement_id,
    SubstPeriodSelfBehaviour behaviour, SubstPeriodSelfRebuildInst rebuild)
    -> SemIR::SpecificId {
  if (!specific_id.has_value()) {
    return specific_id;
  }

  const auto& specific = context.specifics().Get(specific_id);

  // Substitute into the specific without having to construct a FacetType
  // instruction just to hold the specific interface inside a constant id.
  llvm::SmallVector<SemIR::InstId> args(
      context.inst_blocks().Get(specific.args_id));
  for (auto& arg_id : args) {
    auto const_id = context.constant_values().Get(arg_id);
    const_id = SubstPeriodSelf(context, loc_id, const_id,
                               period_self_replacement_id, behaviour, rebuild);
    arg_id = context.constant_values().GetInstId(const_id);
  }
  return MakeSpecific(context, loc_id, specific.generic_id, args);
}

auto SubstPeriodSelf(Context& context, SemIR::LocId loc_id,
                     SemIR::SpecificInterface interface,
                     SemIR::ConstantId period_self_replacement_id,
                     SubstPeriodSelfBehaviour behaviour,
                     SubstPeriodSelfRebuildInst rebuild)
    -> SemIR::SpecificInterface {
  interface.specific_id =
      SubstPeriodSelfInSpecific(context, loc_id, interface.specific_id,
                                period_self_replacement_id, behaviour, rebuild);
  return interface;
}
auto SubstPeriodSelf(Context& context, SemIR::LocId loc_id,
                     SemIR::SpecificNamedConstraint constraint,
                     SemIR::ConstantId period_self_replacement_id,
                     SubstPeriodSelfBehaviour behaviour,
                     SubstPeriodSelfRebuildInst rebuild)
    -> SemIR::SpecificNamedConstraint {
  constraint.specific_id =
      SubstPeriodSelfInSpecific(context, loc_id, constraint.specific_id,
                                period_self_replacement_id, behaviour, rebuild);
  return constraint;
}

auto SubstPeriodSelfInFacetType(Context& context, SemIR::LocId loc_id,
                                SemIR::InstId self_inst_id,
                                SemIR::TypeInstId facet_type_inst_id)
    -> SemIR::TypeInstId {
  auto canon_facet_type_inst_id =
      context.constant_values().GetConstantInstId(facet_type_inst_id);
  if (canon_facet_type_inst_id == SemIR::ErrorInst::TypeInstId) {
    return SemIR::ErrorInst::TypeInstId;
  }

  auto period_self_replacement_id = context.constant_values().Get(self_inst_id);

  auto orig_facet_type =
      context.insts().GetAs<SemIR::FacetType>(canon_facet_type_inst_id);
  const auto& orig_info =
      context.facet_types().Get(orig_facet_type.facet_type_id);

  auto replace_interface = [&](SemIR::SpecificInterface si) {
    return SubstPeriodSelf(context, loc_id, si, period_self_replacement_id,
                           SubstPeriodSelfBehaviour::All);
  };
  auto replace_constraint = [&](SemIR::SpecificNamedConstraint sc) {
    return SubstPeriodSelf(context, loc_id, sc, period_self_replacement_id,
                           SubstPeriodSelfBehaviour::All);
  };
  auto replace_type_impls_interface =
      [&](SemIR::FacetTypeInfo::TypeImplsInterface impls)
      -> SemIR::FacetTypeInfo::TypeImplsInterface {
    auto self = SubstPeriodSelf(
        context, loc_id, context.constant_values().Get(impls.self_type),
        period_self_replacement_id, SubstPeriodSelfBehaviour::All);
    auto interface = SubstPeriodSelf(context, loc_id, impls.specific_interface,
                                     period_self_replacement_id,
                                     SubstPeriodSelfBehaviour::All);
    return {context.constant_values().GetInstId(self), interface};
  };
  auto replace_type_impls_constraint =
      [&](SemIR::FacetTypeInfo::TypeImplsNamedConstraint impls)
      -> SemIR::FacetTypeInfo::TypeImplsNamedConstraint {
    auto self = SubstPeriodSelf(
        context, loc_id, context.constant_values().Get(impls.self_type),
        period_self_replacement_id, SubstPeriodSelfBehaviour::All);
    auto constraint = SubstPeriodSelf(
        context, loc_id, impls.specific_named_constraint,
        period_self_replacement_id, SubstPeriodSelfBehaviour::All);
    return {context.constant_values().GetInstId(self), constraint};
  };
  auto replace_rewrite = [&](SemIR::FacetTypeInfo::RewriteConstraint r)
      -> SemIR::FacetTypeInfo::RewriteConstraint {
    // Designators in rewrite constraints are left alone so they can be found
    // for rewrite constraint resolution. Using rewrites later requires further
    // substitution of the `.Self` references inside designators.
    //
    // The LHS is always a simple designator since only that is allowed, so we
    // just substitute the RHS.
    auto rhs = SubstPeriodSelf(
        context, loc_id, context.constant_values().Get(r.rhs_id),
        period_self_replacement_id, SubstPeriodSelfBehaviour::ExplicitOnly);
    return {r.lhs_id, context.constant_values().GetInstId(rhs)};
  };

  SemIR::FacetTypeInfo info;
  llvm::append_range(
      info.extend_constraints,
      llvm::map_range(orig_info.extend_constraints, replace_interface));
  llvm::append_range(
      info.extend_named_constraints,
      llvm::map_range(orig_info.extend_named_constraints, replace_constraint));
  llvm::append_range(
      info.self_impls_constraints,
      llvm::map_range(orig_info.self_impls_constraints, replace_interface));
  llvm::append_range(info.self_impls_named_constraints,
                     llvm::map_range(orig_info.self_impls_named_constraints,
                                     replace_constraint));
  llvm::append_range(info.type_impls_interfaces,
                     llvm::map_range(orig_info.type_impls_interfaces,
                                     replace_type_impls_interface));
  llvm::append_range(info.type_impls_named_constraints,
                     llvm::map_range(orig_info.type_impls_named_constraints,
                                     replace_type_impls_constraint));
  llvm::append_range(
      info.rewrite_constraints,
      llvm::map_range(orig_info.rewrite_constraints, replace_rewrite));

  info.Canonicalize();
  if (info == orig_info) {
    // Nothing was substituted, keep the original instruction.
    //
    // It is noteworthy that we keep the non-canonical instruction here, since
    // it may have a symbolic value (which is attached to a generic, and can be
    // updated by specifics). Returning the canonical facet type instruction
    // would lose the attachment to the generic which would be incorrect.
    return facet_type_inst_id;
  }

  return AddTypeInst<SemIR::FacetType>(
      context, loc_id,
      {.type_id = SemIR::TypeType::TypeId,
       .facet_type_id = context.facet_types().Add(info)});
}

auto IsPeriodSelf(Context& context, SemIR::InstId inst_id, bool canonicalize)
    -> bool {
  auto const_inst_id = context.constant_values().GetConstantInstId(inst_id);
  if (!const_inst_id.has_value()) {
    return false;
  }
  auto query_inst_id =
      canonicalize ? GetCanonicalFacetOrTypeValue(context, const_inst_id)
                   : inst_id;
  if (auto bind =
          context.insts().TryGetAs<SemIR::SymbolicBinding>(query_inst_id)) {
    const auto& entity_name = context.entity_names().Get(bind->entity_name_id);
    return entity_name.name_id == SemIR::NameId::PeriodSelf;
  }
  return false;
}

}  // namespace Carbon::Check