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

#include "llvm/ADT/STLExtras.h"
#include "toolchain/sem_ir/expr_info.h"
#include "toolchain/sem_ir/inst.h"
#include "toolchain/sem_ir/type_info.h"
#include "toolchain/sem_ir/typed_insts.h"

namespace Carbon::Lower {

static auto GetPointeeType(FunctionContext::TypeInFile type)
    -> FunctionContext::TypeInFile {
  return {.file = type.file,
          .type_id = type.file->GetPointeeType(type.type_id)};
}

// Given an index within a SemIR aggregate type, returns the corresponding index
// of the element within the LLVM type suitable for use with the getelementptr
// instruction.
static auto GetElementIndex(FunctionContext::TypeInFile type,
                            SemIR::ElementIndex idx) -> unsigned int {
  auto type_inst = type.file->types().GetAsInst(type.type_id);

  if (auto custom_layout_type = type_inst.TryAs<SemIR::CustomLayoutType>()) {
    // For custom layout types, we form an array of i8 as the LLVM type, so the
    // offset in the type is the getelementptr index.
    // TODO: This offset might not fit into an `unsigned int`.
    auto offset = type.file->custom_layouts().Get(
        custom_layout_type
            ->layout_id)[SemIR::CustomLayoutId::FirstFieldIndex + idx.index];
    CARBON_CHECK(offset == offset.AlignedTo(SemIR::ObjectSize::Bytes(1)),
                 "Element offset not byte-aligned: {0}", offset);
    return offset.bytes();
  }

  // For now, struct and tuple types map directly into LLVM struct types with
  // identical field numbering.
  CARBON_CHECK((type_inst.IsOneOf<SemIR::StructType, SemIR::TupleType>()),
               "Indexing unexpected aggregate type {0}", type_inst);
  return idx.index;
}

// Extracts an element of an aggregate, such as a struct, tuple, or class, by
// index. Depending on the expression category and value representation of the
// aggregate input, this will either produce a value or a reference.
auto GetAggregateElement(FunctionContext& context, SemIR::InstId aggr_inst_id,
                         SemIR::ElementIndex idx, SemIR::InstId result_inst_id,
                         llvm::Twine name) -> llvm::Value* {
  auto* aggr_value = context.GetValue(aggr_inst_id);

  switch (SemIR::GetExprCategory(context.sem_ir(), aggr_inst_id)) {
    case SemIR::ExprCategory::RefTagged:
    case SemIR::ExprCategory::Error:
    case SemIR::ExprCategory::NotExpr:
    case SemIR::ExprCategory::Pattern:
    case SemIR::ExprCategory::ReprInitializing:
    case SemIR::ExprCategory::InPlaceInitializing:
    case SemIR::ExprCategory::Mixed:
    case SemIR::ExprCategory::Dependent:
      CARBON_FATAL(
          "Unexpected expression category for aggregate access into {0}",
          context.sem_ir().insts().Get(aggr_inst_id));

    case SemIR::ExprCategory::Value: {
      auto aggr_type = context.GetTypeIdOfInst(aggr_inst_id);
      auto value_repr = context.GetValueRepr(aggr_type);
      CARBON_CHECK(
          value_repr.repr.aggregate_kind != SemIR::ValueRepr::NotAggregate,
          "aggregate type should have aggregate value representation");
      switch (value_repr.repr.kind) {
        case SemIR::ValueRepr::Unknown:
          CARBON_FATAL("Lowering access to incomplete aggregate type");
        case SemIR::ValueRepr::Dependent:
          CARBON_FATAL("Lowering access to dependent aggregate type");
        case SemIR::ValueRepr::None:
          return aggr_value;
        case SemIR::ValueRepr::Copy:
          // We are holding the values of the aggregate directly, elementwise.
          return context.builder().CreateExtractValue(
              aggr_value, GetElementIndex(value_repr.type(), idx), name);
        case SemIR::ValueRepr::Pointer: {
          // The value representation is a pointer to an aggregate that we want
          // to index into.
          auto value_rep_type = GetPointeeType(value_repr.type());
          auto* value_type = context.GetType(value_rep_type);
          auto* elem_ptr = context.builder().CreateStructGEP(
              value_type, aggr_value, GetElementIndex(value_rep_type, idx),
              name);

          if (!value_repr.repr.elements_are_values()) {
            // `elem_ptr` points to an object representation, which is our
            // result.
            return elem_ptr;
          }

          // `elem_ptr` points to a value representation. Load it.
          auto result_type = context.GetTypeIdOfInst(result_inst_id);
          auto result_value_type = context.GetValueRepr(result_type).type();
          return context.LoadObject(result_value_type, elem_ptr,
                                    name + ".load");
        }
        case SemIR::ValueRepr::Custom:
          CARBON_FATAL(
              "Aggregate should never have custom value representation");
      }
    }

    case SemIR::ExprCategory::DurableRef:
    case SemIR::ExprCategory::EphemeralRef: {
      // Just locate the aggregate element.
      auto aggr_type = context.GetTypeIdOfInst(aggr_inst_id);
      auto object_repr = FunctionContext::TypeInFile{
          .file = aggr_type.file,
          .type_id = aggr_type.file->types().GetObjectRepr(aggr_type.type_id)};
      return context.builder().CreateStructGEP(
          context.GetType(object_repr), aggr_value,
          GetElementIndex(object_repr, idx), name);
    }
  }
}

auto EmitAggregateValueRepr(FunctionContext& context,
                            SemIR::InstId value_inst_id,
                            SemIR::InstBlockId refs_id) -> llvm::Value* {
  auto type = context.GetTypeIdOfInst(value_inst_id);
  auto value_repr = context.GetValueRepr(type);
  auto value_type = value_repr.type();
  switch (value_repr.repr.kind) {
    case SemIR::ValueRepr::Unknown:
      CARBON_FATAL("Lowering value of incomplete aggregate type");

    case SemIR::ValueRepr::Dependent:
      CARBON_FATAL("Lowering value of dependent aggregate type");

    case SemIR::ValueRepr::None:
      // TODO: Add a helper to get a "no value representation" value.
      return llvm::PoisonValue::get(context.GetType(value_type));

    case SemIR::ValueRepr::Copy: {
      auto refs = context.sem_ir().inst_blocks().Get(refs_id);
      CARBON_CHECK(
          refs.size() == 1,
          "Unexpected size for aggregate with by-copy value representation");
      // TODO: Remove the LLVM StructType wrapper in this case, so we don't
      // need this `insert_value` wrapping.
      return context.builder().CreateInsertValue(
          llvm::PoisonValue::get(context.GetType(value_type)),
          context.GetValue(refs[0]), {0});
    }

    case SemIR::ValueRepr::Pointer: {
      auto* llvm_value_rep_type = context.GetType(GetPointeeType(value_type));

      // Write the value representation to a local alloca so we can produce a
      // pointer to it as the value representation of the struct or tuple.
      auto* alloca = context.builder().CreateAlloca(llvm_value_rep_type);
      for (auto [i, ref_id] :
           llvm::enumerate(context.sem_ir().inst_blocks().Get(refs_id))) {
        context.StoreObject(
            context.GetValueRepr(context.GetTypeIdOfInst(ref_id)).type(),
            context.GetValue(ref_id),
            context.builder().CreateStructGEP(llvm_value_rep_type, alloca, i));
      }
      return alloca;
    }

    case SemIR::ValueRepr::Custom:
      CARBON_FATAL("Aggregate should never have custom value representation");
  }
}

auto EmitAggregateInitializer(FunctionContext& context,
                              SemIR::InstId init_inst_id,
                              SemIR::InstBlockId refs_id, llvm::Twine name)
    -> llvm::Value* {
  auto type = context.GetTypeIdOfInst(init_inst_id);
  auto* llvm_type = context.GetType(type);
  auto refs = context.sem_ir().inst_blocks().Get(refs_id);

  switch (context.GetInitRepr(type).kind) {
    case SemIR::InitRepr::None: {
      // TODO: Add a helper to poison a value slot.
      return llvm::PoisonValue::get(llvm_type);
    }

    case SemIR::InitRepr::InPlace: {
      // Finish initialization of constant fields. We will have skipped this
      // when emitting the initializers because they have constant values.
      //
      // TODO: This emits the initializers for constant fields after all
      // initialization of non-constant fields. This may be observable in some
      // ways such as under a debugger in a debug build. It would be preferable
      // to initialize the constant portions of the aggregate first, but this
      // will likely need a change to the SemIR representation.
      //
      // TODO: If most of the bytes of the result have known constant values,
      // it'd be nice to emit a memcpy from a constant followed by the
      // non-constant initialization.
      for (auto [i, ref_id] : llvm::enumerate(refs)) {
        if (context.sem_ir().constant_values().Get(ref_id).is_constant()) {
          auto dest_id =
              SemIR::FindStorageArgForInitializer(context.sem_ir(), ref_id);
          auto src_id = ref_id;
          auto storage_type = context.GetTypeIdOfInst(dest_id);
          context.InitializeStorage(storage_type, dest_id, src_id);
        }
      }
      // TODO: Add a helper to poison a value slot.
      return llvm::PoisonValue::get(llvm_type);
    }

    case SemIR::InitRepr::ByCopy: {
      auto refs = context.sem_ir().inst_blocks().Get(refs_id);
      CARBON_CHECK(
          refs.size() == 1,
          "Unexpected size for aggregate with by-copy value representation");
      // TODO: Remove the LLVM StructType wrapper in this case, so we don't
      // need this `insert_value` wrapping.
      return context.builder().CreateInsertValue(
          llvm::PoisonValue::get(llvm_type), context.GetValue(refs[0]), {0},
          name);
    }

    case SemIR::InitRepr::Abstract:
      CARBON_FATAL("Lowering aggregate initialization of abstract type {0}",
                   type.file->types().GetAsInst(type.type_id));

    case SemIR::InitRepr::Incomplete:
      CARBON_FATAL("Lowering aggregate initialization of incomplete type {0}",
                   type.file->types().GetAsInst(type.type_id));

    case SemIR::InitRepr::Dependent:
      CARBON_FATAL("Lowering aggregate initialization of dependent type {0}",
                   type.file->types().GetAsInst(type.type_id));
  }
}

}  // namespace Carbon::Lower