// Copyright 2019 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef THIRD_PARTY_MLIR_EDGE_IREE_VM_BYTECODE_READER_H_
#define THIRD_PARTY_MLIR_EDGE_IREE_VM_BYTECODE_READER_H_

#include "third_party/absl/base/attributes.h"
#include "third_party/absl/container/inlined_vector.h"
#include "third_party/mlir_edge/iree/base/status.h"
#include "third_party/mlir_edge/iree/hal/buffer_view.h"
#include "third_party/mlir_edge/iree/schemas/bytecode/bytecode_v0.h"
#include "third_party/mlir_edge/iree/vm/function.h"
#include "third_party/mlir_edge/iree/vm/stack.h"
#include "third_party/mlir_edge/iree/vm/stack_frame.h"
#include "third_party/mlir_edge/iree/vm/type.h"

namespace iree {
namespace vm {

class BytecodeReader {
 public:
  explicit BytecodeReader(Stack* stack) : stack_(stack) {}

  int offset() const { return static_cast<int>(bytecode_pc_ - bytecode_base_); }

  StatusOr<const uint8_t*> AdvanceOffset();

  Status SwitchStackFrame(StackFrame* new_stack_frame);
  Status BranchToOffset(int32_t offset);

  Status CopyInputsAndSwitchStackFrame(StackFrame* src_stack_frame,
                                       StackFrame* dst_stack_frame);
  Status CopyResultsAndSwitchStackFrame(StackFrame* src_stack_frame,
                                        StackFrame* dst_stack_frame);
  Status CopySlots();

  StatusOr<hal::BufferView> ReadConstant();

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<int> ReadCount() {
    return ReadValue<uint8_t>();
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<const Type> ReadType() {
    ASSIGN_OR_RETURN(uint8_t type_index, ReadValue<uint8_t>());
    return Type::FromTypeIndex(type_index);
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<const Function> ReadFunction() {
    ASSIGN_OR_RETURN(auto value, ReadValue<uint32_t>());
    const auto& module = stack_frame_->module();
    return module.function_table().LookupFunction(value);
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<const ImportFunction*>
  ReadImportFunction() {
    ASSIGN_OR_RETURN(auto value, ReadValue<uint32_t>());
    const auto& module = stack_frame_->module();
    return module.function_table().LookupImport(value);
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<hal::BufferView*> ReadLocal(
      absl::Span<hal::BufferView> locals) {
    ASSIGN_OR_RETURN(auto value, ReadValue<uint16_t>());
    if (value > locals.size()) {
      return OutOfRangeErrorBuilder(ABSL_LOC)
             << "Out of bounds local access " << value << " of "
             << locals.size();
    }
    return &locals[value];
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<hal::BufferView*> ReadLocal() {
    return ReadLocal(locals_);
  }

  Status SkipLocals(int count);

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<uint8_t> ReadUint8_t() {
    return ReadValue<uint8_t>();
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<uint16_t> ReadUint16_t() {
    return ReadValue<uint16_t>();
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<int32_t> ReadInt32() {
    return ReadValue<int32_t>();
  }

  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<uint32_t> ReadBlockOffset() {
    return ReadValue<uint32_t>();
  }

  template <typename T, size_t N = 8>
  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<absl::InlinedVector<T, N>>
  ReadSlotElements() {
    ASSIGN_OR_RETURN(auto* local, ReadLocal(locals_));
    absl::InlinedVector<T, N> result(local->shape.element_count());
    if (sizeof(T) == local->element_size) {
      // Fast(ish) path: requested element size matches the actual element size.
      RETURN_IF_ERROR(
          local->buffer->ReadData(0, result.data(), result.size() * sizeof(T)));
    } else {
      // Slow path: need to convert the data.
      switch (local->element_size) {
        case 4: {
          ASSIGN_OR_RETURN(auto mapping, local->buffer->MapMemory<int32_t>(
                                             hal::MemoryAccess::kRead));
          for (size_t i = 0; i < result.size(); ++i) {
            result[i] = static_cast<T>(mapping[i]);
          }
          break;
        }
        case 8: {
          ASSIGN_OR_RETURN(auto mapping, local->buffer->MapMemory<int64_t>(
                                             hal::MemoryAccess::kRead));
          for (size_t i = 0; i < result.size(); ++i) {
            result[i] = static_cast<T>(mapping[i]);
          }
          break;
        }
        default:
          return UnimplementedErrorBuilder(ABSL_LOC)
                 << "Unsupported local element size: " << local->element_size;
      }
    }
    return result;
  }

  Status ReadShape(Shape* out_shape);

  StatusOr<Shape> ReadShapePieces();
  StatusOr<Shape> ReadShapePieces(size_t* out_element_count);

  StatusOr<absl::Span<const int32_t>> ReadIndexList();

 private:
  template <typename T>
  ABSL_ATTRIBUTE_ALWAYS_INLINE StatusOr<T> ReadValue() {
    // TODO(benvanik): validate bounds.
    T value = *reinterpret_cast<const T*>(bytecode_pc_);
    bytecode_pc_ += sizeof(T);
    return value;
  }

  Stack* stack_ = nullptr;
  StackFrame* stack_frame_ = nullptr;
  const uint8_t* bytecode_base_ = nullptr;
  const uint8_t* bytecode_limit_ = nullptr;
  const uint8_t* bytecode_pc_ = nullptr;
  absl::Span<hal::BufferView> locals_;
  FunctionTable::BreakpointTable* breakpoint_table_ = nullptr;
};

}  // namespace vm
}  // namespace iree

#endif  // THIRD_PARTY_MLIR_EDGE_IREE_VM_BYTECODE_READER_H_