// 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.

#include "third_party/mlir_edge/iree/vm/bytecode_reader.h"

#include "third_party/mlir_edge/iree/base/shape.h"
#include "third_party/mlir_edge/iree/base/status.h"
#include "third_party/mlir_edge/iree/hal/heap_buffer.h"

namespace iree {
namespace vm {

namespace {
using ::iree::hal::BufferView;
}  // namespace

StatusOr<const uint8_t*> BytecodeReader::AdvanceOffset() {
  *stack_frame_->mutable_offset() = offset();
  // TODO(benvanik): make a flag and/or remove.
  DVLOG(1) << "dispatch(" << stack_frame_->function().name() << "@" << offset()
           << "): " << int(*bytecode_pc_);
  for (int i = 0; i < locals_.size(); ++i) {
    DVLOG(1) << "local[" << i << "] " << locals_[i].DebugStringShort();
  }

  if (breakpoint_table_) {
    auto it = breakpoint_table_->find(offset());
    if (it != breakpoint_table_->end()) {
      // Breakpoint hit!
      RETURN_IF_ERROR(it->second(*stack_));
    }
  }

  return bytecode_pc_++;
}

Status BytecodeReader::SkipLocals(int count) {
  size_t stride = sizeof(uint16_t) * count;
  if (bytecode_pc_ + stride >= bytecode_limit_) {
    return OutOfRangeErrorBuilder(ABSL_LOC) << "Bytecode underflow";
  }
  bytecode_pc_ += stride;
  return OkStatus();
}

Status BytecodeReader::ReadShape(Shape* out_shape) {
  ASSIGN_OR_RETURN(auto shape_dims, ReadIndexList());
  *out_shape = Shape(shape_dims);
  return OkStatus();
}

StatusOr<Shape> BytecodeReader::ReadShapePieces() {
  // TODO(benvanik): rewrite to be faster (multiple offsets to walk both lists).
  ASSIGN_OR_RETURN(auto shape_dims, ReadIndexList());
  if (shape_dims.size() >= kMaxRank) {
    return UnimplementedErrorBuilder(ABSL_LOC)
           << "Shapes limited to rank " << kMaxRank << " right now";
  }
  int expected_dynamic_dims = 0;
  for (int i = 0; i < shape_dims.size(); ++i) {
    if (shape_dims[i] == -1) {
      ++expected_dynamic_dims;
    }
  }

  Shape shape(shape_dims);
  ASSIGN_OR_RETURN(int dynamic_dims, ReadCount());
  if (dynamic_dims != expected_dynamic_dims) {
    return InvalidArgumentErrorBuilder(ABSL_LOC)
           << "Expected " << expected_dynamic_dims << " dynamic dims but only "
           << dynamic_dims << " provided";
  } else if (dynamic_dims) {
    for (int i = 0; i < shape_dims.size(); ++i) {
      if (shape_dims[i] != -1) {
        continue;
      }
      // TODO(benvanik): kill this embarrassment.
      ASSIGN_OR_RETURN(auto dims_piece, ReadSlotElements<int32_t>());
      if (dims_piece.size() != 1) {
        return InvalidArgumentErrorBuilder(ABSL_LOC)
               << "Dims piece has rank " << dims_piece.size() << "; must be 1";
      }
      shape[i] = dims_piece[0];
    }
  }
  return shape;
}

StatusOr<Shape> BytecodeReader::ReadShapePieces(size_t* out_element_count) {
  ASSIGN_OR_RETURN(auto shape, ReadShapePieces());
  *out_element_count = shape.element_count();
  return shape;
}

StatusOr<absl::Span<const int32_t>> BytecodeReader::ReadIndexList() {
  ASSIGN_OR_RETURN(int count, ReadCount());
  int stride = count * sizeof(int32_t);
  if (bytecode_pc_ + stride >= bytecode_limit_) {
    return OutOfRangeErrorBuilder(ABSL_LOC) << "Bytecode underflow";
  }
  auto list = absl::Span<const int32_t>(
      reinterpret_cast<const int32_t*>(bytecode_pc_), count);
  bytecode_pc_ += stride;
  return list;
}

Status BytecodeReader::SwitchStackFrame(StackFrame* new_stack_frame) {
  // Flush old state.
  auto* old_stack_frame = stack_frame_;
  if (old_stack_frame) {
    *old_stack_frame->mutable_offset() = offset();
  }

  // Switch the frame. The FiberState holds the full stack, this is just the
  // current one for easy access.
  stack_frame_ = new_stack_frame;

  // Setup state pointers for faster dereferencing.
  const auto& function = new_stack_frame->function();
  const auto& bytecode = *function.def().bytecode();
  bytecode_base_ = bytecode.contents()->Data();
  bytecode_limit_ = bytecode_base_ + bytecode.contents()->size();
  bytecode_pc_ = bytecode_base_ + new_stack_frame->offset();
  locals_ = new_stack_frame->mutable_locals();
  // TODO(benvanik): reimplement breakpoints as bytecode rewriting.
  int function_ordinal = function.module()
                             .function_table()
                             .LookupFunctionOrdinal(function)
                             .ValueOrDie();
  breakpoint_table_ =
      function.module().function_table().GetFunctionBreakpointTable(
          function_ordinal);
  return OkStatus();
}

Status BytecodeReader::CopyInputsAndSwitchStackFrame(
    StackFrame* src_stack_frame, StackFrame* dst_stack_frame) {
  ASSIGN_OR_RETURN(int32_t src_count, ReadCount());
  for (int i = 0; i < src_count; ++i) {
    ASSIGN_OR_RETURN(auto* src_local,
                     ReadLocal(src_stack_frame->mutable_locals()));
    *dst_stack_frame->mutable_local(i) = *src_local;
  }
  return SwitchStackFrame(dst_stack_frame);
}

Status BytecodeReader::CopyResultsAndSwitchStackFrame(
    StackFrame* src_stack_frame, StackFrame* dst_stack_frame) {
  ASSIGN_OR_RETURN(int32_t src_count, ReadCount());
  // TODO(benvanik): avoid vector.
  absl::InlinedVector<BufferView*, 8> src_locals(src_count);
  for (int i = 0; i < src_count; ++i) {
    ASSIGN_OR_RETURN(src_locals[i],
                     ReadLocal(src_stack_frame->mutable_locals()));
  }
  RETURN_IF_ERROR(SwitchStackFrame(dst_stack_frame));
  ASSIGN_OR_RETURN(int32_t dst_count, ReadCount());
  if (src_count != dst_count) {
    return OutOfRangeErrorBuilder(ABSL_LOC)
           << "Src and dst value counts differ: " << src_count << " vs "
           << dst_count;
  }
  for (int i = 0; i < dst_count; ++i) {
    ASSIGN_OR_RETURN(auto* dst_local,
                     ReadLocal(dst_stack_frame->mutable_locals()));
    *dst_local = *src_locals[i];
  }
  return OkStatus();
}

Status BytecodeReader::CopySlots() {
  ASSIGN_OR_RETURN(int32_t count, ReadCount());
  for (int i = 0; i < count; ++i) {
    ASSIGN_OR_RETURN(auto* src_local,
                     ReadLocal(stack_frame_->mutable_locals()));
    ASSIGN_OR_RETURN(auto* dst_local,
                     ReadLocal(stack_frame_->mutable_locals()));
    *dst_local = *src_local;
  }
  return OkStatus();
}

Status BytecodeReader::BranchToOffset(int32_t offset) {
  const uint8_t* new_bytecode_pc = bytecode_base_ + offset;
  if (new_bytecode_pc < bytecode_base_ || new_bytecode_pc > bytecode_limit_) {
    return OutOfRangeErrorBuilder(ABSL_LOC)
           << "Branch target " << offset
           << " is out of bounds of the function bytecode ("
           << static_cast<size_t>(bytecode_limit_ - bytecode_base_)
           << "b total)";
  }
  bytecode_pc_ = new_bytecode_pc;
  return OkStatus();
}

StatusOr<BufferView> BytecodeReader::ReadConstant() {
  BufferView buffer_view;

  // Element type defines the buffer_view size (but we don't really care about
  // the data format).
  ASSIGN_OR_RETURN(auto element_type, ReadType());
  buffer_view.element_size = element_type.element_size();

  // Parse shape - constants always define a full shape.
  RETURN_IF_ERROR(ReadShape(&buffer_view.shape));

  // Read encoding to determine how the constant data is stored in the file.
  ASSIGN_OR_RETURN(auto encoding, ReadValue<ConstantEncoding>());

  // Get buffer for the constant data.
  switch (encoding) {
    case ConstantEncoding::kDense: {
      // Validate we have all constant data present.
      device_size_t serialized_length = buffer_view.byte_length();
      if (bytecode_pc_ + serialized_length >= bytecode_limit_) {
        return OutOfRangeErrorBuilder(ABSL_LOC)
               << "Constant data out of bounds";
      }

      buffer_view.buffer = hal::HeapBuffer::Wrap(
          hal::MemoryType::kHostLocal, hal::BufferUsage::kAll, bytecode_pc_,
          serialized_length);
      bytecode_pc_ += serialized_length;
      break;
    }
    case ConstantEncoding::kSplat: {
      // Validate we have at least one element worth of data in the buffer.
      if (bytecode_pc_ + buffer_view.element_size >= bytecode_limit_) {
        return OutOfRangeErrorBuilder(ABSL_LOC)
               << "Constant data out of bounds";
      }

      // TODO(benvanik): replace with fancy constant pool and such.
      // NOTE: this is not much different than if a alloc_heap+broadcast pair
      // had been in the IR.
      buffer_view.buffer = hal::HeapBuffer::Allocate(
          hal::MemoryType::kHostLocal, hal::BufferUsage::kAll,
          buffer_view.byte_length());
      switch (buffer_view.element_size) {
        case 1: {
          uint8_t value = *reinterpret_cast<const uint8_t*>(bytecode_pc_);
          RETURN_IF_ERROR(buffer_view.buffer->Fill8(value));
          break;
        }
        case 2: {
          uint16_t value = *reinterpret_cast<const uint16_t*>(bytecode_pc_);
          RETURN_IF_ERROR(buffer_view.buffer->Fill16(value));
          break;
        }
        case 4: {
          uint32_t value = *reinterpret_cast<const uint32_t*>(bytecode_pc_);
          RETURN_IF_ERROR(buffer_view.buffer->Fill32(value));
          break;
        }
        case 8: {
          // TODO(benvanik): add Fill64.
          uint64_t value = *reinterpret_cast<const uint64_t*>(bytecode_pc_);
          ASSIGN_OR_RETURN(auto mapping,
                           buffer_view.buffer->MapMemory<uint64_t>(
                               hal::MemoryAccess::kDiscardWrite));
          auto mapped_data = mapping.mutable_contents();
          for (int i = 0; i < mapping.size(); ++i) {
            mapped_data[i] = value;
          }
          break;
        }
        default:
          return UnimplementedErrorBuilder(ABSL_LOC)
                 << "Unimplemented splat element stride "
                 << buffer_view.element_size;
      }
      bytecode_pc_ += buffer_view.element_size;
      break;
    }
    default:
      return UnimplementedErrorBuilder(ABSL_LOC)
             << "Unimplemented constant encoding "
             << static_cast<int>(encoding);
  }

  return buffer_view;
}

}  // namespace vm
}  // namespace iree