// Copyright 2024 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#ifndef SERVICES_WEBNN_COREML_GRAPH_BUILDER_COREML_H_
#define SERVICES_WEBNN_COREML_GRAPH_BUILDER_COREML_H_

#include <cstdint>
#include <memory>
#include <optional>
#include <string_view>
#include <variant>

#include "base/containers/flat_map.h"
#include "base/files/file.h"
#include "base/files/file_path.h"
#include "base/memory/raw_ptr.h"
#include "base/memory/stack_allocated.h"
#include "base/numerics/checked_math.h"
#include "base/time/time.h"
#include "base/types/expected.h"
#include "services/webnn/public/cpp/context_properties.h"
#include "services/webnn/public/cpp/operand_descriptor.h"
#include "services/webnn/public/cpp/webnn_types.h"
#include "services/webnn/public/mojom/webnn_context_provider.mojom.h"
#include "services/webnn/public/mojom/webnn_error.mojom-forward.h"
#include "services/webnn/public/mojom/webnn_graph.mojom.h"
#include "third_party/abseil-cpp/absl/container/flat_hash_map.h"
#include "third_party/coremltools/mlmodel/format/MIL.pb.h"
#include "third_party/coremltools/mlmodel/format/Model.pb.h"

namespace webnn {

class WebNNConstantOperand;

namespace coreml {

struct Float16 {
  uint16_t data;
};

namespace internal {
// Supported tensor types for immediate values. The list can be expanded as
// needed.
template <typename T, typename... U>
concept IsAnyOf = (std::same_as<T, U> || ...);
template <typename T>
concept IsSupportedTensorType =
    IsAnyOf<T, Float16, float, int32_t, int8_t, uint8_t, char, bool>;
}  // namespace internal

inline constexpr char kPlaceholderInputName[] = "placeholder";

// Get name identifiers used in CoreML model files for input or output operands.
std::string GetCoreMLNameFromInput(std::string_view input_name,
                                   OperandId operand_id);
std::string GetCoreMLNameFromOutput(std::string_view output_name,
                                    OperandId operand_id);

// Reads the WebNN graph from the mojom::GraphInfo to
// produce CoreML model and serializes to provided `working_directory`.
// There is nothing macOS-specific in this class.
//
// The instances of the class may not be allocated on the heap, but as a member
// variable of a non-stack-allocated class and be single-use per conversion.
class GraphBuilderCoreml {
  STACK_ALLOCATED();

 public:
  // Tracks Operand information during graph building, so that
  // future operations can look them up based on operand id.
  // When an operation is decomposed, additional `OperandInfo` entities are
  // created to represent intermediate layers.
  struct OperandInfo {
    OperandInfo();
    OperandInfo(std::string name,
                base::span<const uint32_t> dimensions,
                CoreML::Specification::MILSpec::DataType mil_data_type);
    OperandInfo(OperandInfo&);
    OperandInfo(OperandInfo&&);
    ~OperandInfo();

    // Identifier for this operand in coreml model file.
    std::string coreml_name;
    // Due to the limitations of CoreML not supporting 0D input at model
    // entry point, model 0D inputs are splitted into two nodes, with the
    // external facing node that's casted to 1D array and internal node that
    // preserves the 0D shape.
    std::string external_coreml_name;
    std::vector<uint32_t> dimensions;
    CoreML::Specification::MILSpec::DataType mil_data_type;
  };

  struct Result {
    explicit Result(base::FilePath ml_package_dir);
    Result(const Result&) = delete;
    Result& operator=(const Result&) = delete;
    ~Result();

    const base::FilePath& GetModelFilePath();

    [[nodiscard]] const OperandInfo& GetOperandInfo(OperandId operand_id) const;

    const base::FilePath ml_package_dir;
    // `std::unique_ptr` is used for values to provide pointer stabiliy for
    // `GetOperandInfo`.
    absl::flat_hash_map<OperandId, std::unique_ptr<OperandInfo>>
        id_to_operand_info_map;
  };

  // Factory method that creates a GraphBuilderCoreml, builds and serializes the
  // CoreML model to the `working_directory`. This expects the
  // `working_directory` to be an empty directory.
  //
  // Returns unexpected if it fails.
  [[nodiscard]] static base::expected<std::unique_ptr<Result>, mojom::ErrorPtr>
  CreateAndBuild(
      const mojom::GraphInfo& graph_info,
      ContextProperties context_properties,
      mojom::Device device,
      const base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>&
          constant_operands,
      const base::FilePath& working_directory);

  static ContextProperties GetContextProperties();

  GraphBuilderCoreml(const GraphBuilderCoreml&) = delete;
  GraphBuilderCoreml& operator=(const GraphBuilderCoreml&) = delete;

  ~GraphBuilderCoreml();

 private:
  class WeightsFileHandle;
  class ScopedWeightItem {
   public:
    ScopedWeightItem(WeightsFileHandle& weights_file_handle,
                     size_t byte_size,
                     uint64_t offset);
    ~ScopedWeightItem();
    ScopedWeightItem(const ScopedWeightItem&) = delete;
    ScopedWeightItem& operator=(const ScopedWeightItem&) = delete;

    base::expected<void, mojom::ErrorPtr> WriteBytes(
        base::span<const uint8_t> bytes);

    base::expected<void, mojom::ErrorPtr> Finalize();
    uint64_t offset() { return offset_; }

   private:
    base::raw_ref<WeightsFileHandle> weights_file_handle_;
    bool has_error_ = false;
    bool finalized_ = false;
    size_t byte_size_;
    size_t size_written_ = 0;
    const uint64_t offset_;
  };

  class WeightsFileHandle {
   public:
    static std::optional<std::unique_ptr<GraphBuilderCoreml::WeightsFileHandle>>
    CreateWeightsHandle(const base::FilePath& weights_file_path);

    WeightsFileHandle(base::File weights_file, uint64_t current_offset);
    WeightsFileHandle(const WeightsFileHandle&) = delete;
    WeightsFileHandle(WeightsFileHandle&&) = delete;
    ~WeightsFileHandle();

    // Write a single weight item.
    base::expected<CoreML::Specification::MILSpec::Value, mojom::ErrorPtr>
    Write(OperandId operand_id,
          const WebNNConstantOperand& constant_operand,
          std::optional<base::span<const uint32_t>> reshape_dimensions =
              std::nullopt);

    base::expected<std::unique_ptr<ScopedWeightItem>, mojom::ErrorPtr>
    CreateScopedWeightItem(OperandDataType data_type, size_t byte_size);

    // Need to be called to update weight count after all weights are written.
    base::expected<void, mojom::ErrorPtr> Finalize();

    size_t GetByteSize(OperandDataType data_type);
    friend class ScopedWeightItem;

   private:
    // `WriteBytes`, `WriteItemFinalize` allows callers
    // to make multiple partial writes then close of the weight item.
    base::expected<void, mojom::ErrorPtr> WriteBytes(
        base::span<const uint8_t> bytes);
    base::expected<void, mojom::ErrorPtr> WeightItemFinalize(size_t byte_size);

    base::File weights_file_;
    uint64_t current_offset_ = 0;
    uint32_t num_of_weights_ = 0;
    base::TimeDelta weights_write_time_;
    bool has_error_ = false;
    bool finalized_ = false;
    // Maps operand IDs to offsets in the weight file.
    base::flat_map<OperandId, uint64_t> constant_offsets_;
  };

  GraphBuilderCoreml(
      const mojom::GraphInfo& graph_info,
      ContextProperties context_properties,
      mojom::Device device,
      const base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>&
          constant_operands,
      base::FilePath ml_package_dir,
      std::unique_ptr<WeightsFileHandle> weights_file_handle);

  [[nodiscard]] base::expected<void, mojom::ErrorPtr> BuildCoreMLModel();

  [[nodiscard]] base::expected<void, mojom::ErrorPtr> SerializeModel();

  // No further methods may be called on this class after calling this method.
  [[nodiscard]] std::unique_ptr<Result> FinishAndTakeResult();

  // Add input in Model.description and in Program's main function inputs.
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddInput(
      OperandId input_id,
      CoreML::Specification::MILSpec::Function& main_function,
      CoreML::Specification::MILSpec::Block& block);
  void AddPlaceholderInput(
      CoreML::Specification::MILSpec::Function& main_function,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOutput(
      OperandId output_id);

  // Helper function for simple unary operations.
  enum class SupportedDataType { kFloats, kFloatsAndInt32 };

  [[nodiscard]] base::expected<CoreML::Specification::MILSpec::Operation*,
                               mojom::ErrorPtr>
  CreateUnaryOperation(SupportedDataType supported_data_type,
                       std::string_view op_name,
                       OperandId input_operand_id,
                       OperandId output_operand_id,
                       CoreML::Specification::MILSpec::Block& block,
                       std::string_view operand_op_name);
  // TODO: crbug.com/345271830 - remove this after all callers check with
  // `context_properties_.data_type_limits`.
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddUnaryOperation(
      SupportedDataType supported_data_type,
      std::string_view op_name,
      OperandId input_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block,
      std::string_view operand_op_name);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddUnaryOperation(
      std::string_view op_name,
      OperandId input_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  template <typename T>
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddUnaryOperation(
      SupportedDataType supported_data_type,
      std::string_view op_name,
      const T& operation,
      CoreML::Specification::MILSpec::Block& block,
      std::string_view operand_op_name);
  template <typename T>
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddUnaryOperation(
      std::string_view op_name,
      const T& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddUnaryFloatsOperationWithEpsilon(
      std::string_view op_name,
      OperandId input_operand_id,
      OperandId output_operand_id,
      float epsilon,
      CoreML::Specification::MILSpec::Block& block);
  template <typename T>
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddUnaryFloatsOperationWithEpsilon(
      std::string_view op_name,
      const T& operation,
      float epsilon,
      CoreML::Specification::MILSpec::Block& block);

  // Serialization functions for members of the mojom::Operation union. Keep
  // these functions in the same order as in webnn_graph.mojom.
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForArgMinMax(
      const mojom::ArgMinMax& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForBatchNormalization(
      const mojom::BatchNormalization& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForCast(
      OperandId input_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForClamp(
      OperandId input_operand_id,
      OperandId output_operand_id,
      MLNumber min_value,
      MLNumber max_value,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForClamp(
      const mojom::Clamp& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForConcat(
      base::span<const OperandId> input_operand_ids,
      OperandId output_operand_id,
      uint32_t axis,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForConcat(
      const mojom::Concat& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForConv2d(
      const mojom::Conv2d& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForCumulativeSum(const mojom::CumulativeSum& operation,
                               CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForDequantizeLinear(const mojom::DequantizeLinear& operation,
                                  CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForDequantizeLinearConst(
      const mojom::DequantizeLinear& operation,
      size_t axis,
      bool is_scalar_scale,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForDequantizeLinearConstBlockwise(
      const mojom::DequantizeLinear& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForDequantizeLinearEmulate(
      const mojom::DequantizeLinear& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<std::pair<OperandId, OperandId>, mojom::ErrorPtr>
  ExpandForBlockwise(OperandId input_operand_id,
                     OperandId scale_operand_id,
                     OperandId zero_point_operand_id,
                     CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<OperandId, mojom::ErrorPtr>
  ExpandDimForBlockwise(OperandId input_operand_id,
                        size_t repetition_axis,
                        int32_t repetitions,
                        CoreML::Specification::MILSpec::Block& block);

  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForElementwiseBinary(
      std::variant<OperandId, CoreML::Specification::MILSpec::Value>
          lhs_operand,
      std::variant<OperandId, CoreML::Specification::MILSpec::Value>
          rhs_operand,
      OperandId output_operand_id,
      const mojom::ElementWiseBinary::Kind kind,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForElementwiseUnary(mojom::ElementWiseUnary::Kind kind,
                                  OperandId input_operand_id,
                                  OperandId output_operand_id,
                                  CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForElu(
      const mojom::Elu& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForExpand(
      const mojom::Expand& operation,
      CoreML::Specification::MILSpec::Block& block);
  void AddOperationForFill(CoreML::Specification::MILSpec::Value value,
                           OperandId output_operand_id,
                           CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGather(
      const mojom::Gather& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForGatherElements(const mojom::GatherElements& operation,
                                CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGatherND(
      const mojom::GatherND& operation,
      CoreML::Specification::MILSpec::Block& block);

  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGelu(
      const mojom::Gelu& operation,
      CoreML::Specification::MILSpec::Block& block);

  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGemm(
      OperandId a_operand_id,
      OperandId b_operand_id,
      std::optional<OperandId> c_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block,
      bool a_transpose = false,
      bool b_transpose = false,
      float alpha = 1.0f,
      float beta = 1.0f);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGemm(
      const mojom::Gemm& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGru(
      const mojom::Gru& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForGruCell(
      const mojom::GruCell& operation,
      CoreML::Specification::MILSpec::Block& block);
  base::expected<void, mojom::ErrorPtr> AddOperationForGruSingleStep(
      OperandId input_operand_id,
      OperandId hidden_state_operand_id,
      OperandId output_operand_id,
      base::span<const OperandId> weights,
      base::span<const OperandId> recurrent_weights,
      std::optional<base::span<const OperandId>> biases,
      std::optional<base::span<const OperandId>> recurrent_biases,
      uint32_t hidden_size,
      mojom::GruWeightLayout layout,
      mojom::RecurrentNetworkActivation activation,
      mojom::RecurrentNetworkActivation output_activation,
      bool reset_after,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForHardSigmoid(OperandId input_operand_id,
                             float alpha,
                             float beta,
                             OperandId output_operand_id,
                             CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForHardSigmoid(const mojom::HardSigmoid& operation,
                             CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForHardSwish(
      const mojom::HardSwish& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForInstanceNormalization(
      const mojom::InstanceNormalization& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForLayerNormalization(
      const mojom::LayerNormalization& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForLeakyRelu(
      const mojom::LeakyRelu& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForLinear(
      const mojom::Linear& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForLstm(
      const mojom::Lstm& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForLstm(
      OperandId input_operand_id,
      OperandId weight_operand_id,
      OperandId recurrent_weight_operand_id,
      uint32_t hidden_size,
      std::optional<OperandId> bias_operand_id,
      std::optional<OperandId> recurrent_bias_operand_id,
      std::optional<OperandId> peephole_weight_operand_id,
      std::optional<OperandId> initial_hidden_state_operand_id,
      std::optional<OperandId> initial_cell_state_operand_id,
      bool return_sequence,
      mojom::RecurrentNetworkDirection direction,
      mojom::LstmWeightLayout layout,
      base::span<const mojom::RecurrentNetworkActivation> activations,
      base::span<const OperandId> output_operand_ids,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForLstmCell(
      const mojom::LstmCell& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForMatmul(
      OperandId input_x_operand_id,
      OperandId input_y_operand_id,
      bool transpose_x,
      bool transpose_y,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForMatmul(
      const mojom::Matmul& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForPad(
      const mojom::Pad& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForPool2d(
      const mojom::Pool2d& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForQuantizeLinear(const mojom::QuantizeLinear& operation,
                                CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForQuantizeLinearEmulate(
      const mojom::QuantizeLinear& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForReduce(
      const mojom::Reduce& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForResample2d(
      const mojom::Resample2d& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForReshape(
      OperandId input_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForReshape(
      const mojom::Reshape& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForReverse(
      const mojom::Reverse& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForRound(
      OperandId input_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForScatterElements(const mojom::ScatterElements& operation,
                                 CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForScatterND(
      OperandId input_operand_id,
      OperandId indices_operand_id,
      OperandId updates_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForScatterND(
      const mojom::ScatterND& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForSlice(
      OperandId input_operand_id,
      OperandId output_operand_id,
      base::span<const int32_t> beginnings,
      base::span<const int32_t> endings,
      base::span<const int32_t> strides,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForSlice(
      const mojom::Slice& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForSoftmax(
      const mojom::Softmax& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForSplit(
      OperandId input_operand_id,
      base::span<const OperandId> output_operand_ids,
      uint32_t axis,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForSplit(
      const mojom::Split& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForTile(
      OperandId input_operand_id,
      OperandId output_operand_id,
      base::span<const int32_t> repetitions,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForTile(
      const mojom::Tile& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForTranspose(
      OperandId input_operand_id,
      OperandId output_operand_id,
      base::span<const uint32_t> permutation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForTranspose(
      const mojom::Transpose& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForTriangular(
      const mojom::Triangular& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForPrelu(
      const mojom::Prelu& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  AddOperationForPreluEmulate(const mojom::Prelu& operation,
                              CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForWhere(
      const mojom::Where& operation,
      CoreML::Specification::MILSpec::Block& block);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr> AddOperationForWhere(
      OperandId condition_operand_id,
      OperandId true_value_operand_id,
      OperandId false_value_operand_id,
      OperandId output_operand_id,
      CoreML::Specification::MILSpec::Block& block);

  // Sanitize the indices of gather-like operation.
  base::expected<OperandId, mojom::ErrorPtr> SanitizeIndices(
      OperandId input_operand_id,
      OperandId indices_operand_id,
      std::optional<uint32_t> axis,
      CoreML::Specification::MILSpec::Block& block);
  base::expected<OperandId, mojom::ErrorPtr> SanitizeIndicesForDim(
      OperandId indices_operand_id,
      uint32_t axis_dim,
      base::span<const uint32_t> dimensions,
      CoreML::Specification::MILSpec::Block& block);

  // Helpers.
  const mojom::Operand& GetOperand(OperandId operand_id) const;

  [[nodiscard]] const OperandInfo& GetOperandInfo(OperandId operand_id) const;
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  PopulateFeatureDescription(
      OperandId operand_id,
      ::CoreML::Specification::FeatureDescription& feature_description);

  // Accessors for fields declared in `result_`.
  const base::FilePath& ml_package_dir() const {
    return result_->ml_package_dir;
  }
  absl::flat_hash_map<OperandId, std::unique_ptr<OperandInfo>>&
  id_to_operand_info_map() const {
    return result_->id_to_operand_info_map;
  }

  // MILSpec::Program's Function, Block, Operation's inputs/outputs could be
  // defined as NamedValueType.
  void PopulateNamedValueType(
      OperandId operand_id,
      CoreML::Specification::MILSpec::NamedValueType& named_value_type);
  void PopulateNamedValueType(
      std::string_view name,
      CoreML::Specification::MILSpec::DataType mil_data_type,
      base::span<const uint32_t> dimensions,
      CoreML::Specification::MILSpec::NamedValueType& named_value_type);
  void PopulateNamedValueTypeForInput(
      OperandId operand_id,
      CoreML::Specification::MILSpec::NamedValueType& named_value_type);
  // Update the `id_to_op_input_info_map_` to be used by ops later.
  void UpdateCoreMLInputInfoMap(OperandId operand_id);

  std::string GetCoreMLNameFromOperand(OperandId operand_id);
  [[nodiscard]] base::expected<OperandId, mojom::ErrorPtr>
  GenerateInternalOperandInfo(
      CoreML::Specification::MILSpec::DataType mil_data_type,
      base::span<const uint32_t> dimensions);

  [[nodiscard]] base::expected<void, mojom::ErrorPtr> SetInputFromOperand(
      google::protobuf::Map<std::string,
                            CoreML::Specification::MILSpec::Argument>& inputs,
      std::string_view key,
      OperandId operand_id);
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  SetInputFromConstantOperand(
      google::protobuf::Map<std::string,
                            CoreML::Specification::MILSpec::Argument>& inputs,
      std::string_view key,
      OperandId constant_operand_id,
      std::optional<base::span<const uint32_t>> reshaped_dimensions =
          std::nullopt);

  // Helper function to return input[index] using squeeze(slice(input)).
  base::expected<OperandId, mojom::ErrorPtr> SliceFirstDimension(
      OperandId input_operand_id,
      int32_t index,
      CoreML::Specification::MILSpec::Block& block);

  // Split to output operands and squeeze it.
  base::expected<void, mojom::ErrorPtr> SplitAndSqueeze(
      OperandId input_operand_id,
      base::span<OperandId> output_operand_ids,
      int32_t axis,
      CoreML::Specification::MILSpec::Block& block);
  // Set input from a constant operand with an alternative order. The reordered
  // constant won't be re-used across operations.
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  SetInputFromConstantReordered(
      google::protobuf::Map<std::string,
                            CoreML::Specification::MILSpec::Argument>& inputs,
      std::string_view key,
      base::span<const uint8_t> bytes,
      OperandDataType data_type,
      base::span<const uint32_t> dimensions,
      base::span<const std::pair<size_t, size_t>> new_order);
  // Set input from two constants added up with an alternative order. The
  // reordered constant won't be re-used across operations.
  [[nodiscard]] base::expected<void, mojom::ErrorPtr>
  SetInputFromTwoConstantsReordered(
      google::protobuf::Map<std::string,
                            CoreML::Specification::MILSpec::Argument>& inputs,
      std::string_view key,
      base::span<const uint8_t> a_bytes,
      base::span<const uint8_t> b_bytes,
      OperandDataType data_type,
      base::span<const uint32_t> dimensions,
      base::span<const std::pair<size_t, size_t>> new_order);

  // A reference to the WebNN compute graph that `this` instance is
  // converting to CoreML model. The creator of `this` must ensure the
  // GraphInfo reference passed into `CreateAndBuild()` is valid for as long
  // as `this` exists.
  base::raw_ref<const mojom::GraphInfo> graph_info_;

  base::raw_ref<
      const base::flat_map<OperandId, std::unique_ptr<WebNNConstantOperand>>>
      constant_operands_;

  const ContextProperties context_properties_;
  const mojom::Device device_;

  // Used to generate unique names for internal operands generated for WebNN
  // operations that need to be decomposed into multiple CoreML operations.
  base::CheckedNumeric<OperandId::underlying_type> internal_operand_id_;

  CoreML::Specification::Model ml_model_;
  raw_ptr<CoreML::Specification::MILSpec::Program> program_;
  bool support_blockwise_dequantize_ = false;

  std::unique_ptr<WeightsFileHandle> weights_file_handle_;

  std::unique_ptr<Result> result_;
};

}  // namespace coreml
}  // namespace webnn

#endif  // SERVICES_WEBNN_COREML_GRAPH_BUILDER_COREML_H_