#!/usr/bin/env ruby

# Copyright (c) 2021-2024 Huawei Device Co., Ltd.
# 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
#
# http://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_relative 'common.irt'

$available_regs = $temps_mask + :arg0 + :arg1
# there are 2 temp registers used in the case of arm64 (instead of 3 for
# x86_64 and arm32) that's why one callee-saved register is added to the mask
if Options.arch == :arm64
  $available_regs = $available_regs + :callee1
end

AllocateObjectTlabValidation = {
  spills_count_max: { default: 2, arm32: 9999 },  # TODO(msherstennikov): set to 0 once regalloc supports smart temps
  code_size_max: { arm64: 116, x86_64: 477, arm32: 9999 }
  # TODO(msherstennikov): revert back code size values, once regalloc supports smart temps
  # code_size_max: { arm64: 100, x86_64: 125, arm32: 9999 }
}

AllocateArrayTlabValidation = {
  spills_count_max: { default: 2, arm32: 9999 },  # TODO(msherstennikov): set to 0 once regalloc supports smart temps
  code_size_max: { arm64: 136, x86_64: 476, arm32: 9999 }
  # TODO(msherstennikov): revert back code size values, once regalloc supports smart temps
  # code_size_max: { arm64: 128, x86_64: 145, arm32: 9999 }
}

function(:AllocateObjectTlab,
         params: {klass: 'ref', size: 'word'},
         regmap: $full_regmap,
         mode: [:FastPath],
         regalloc_set: $available_regs,
         validate: AllocateObjectTlabValidation) {

  if Options.arch == :arm32
    Intrinsic(:UNREACHABLE).Terminator.void
    ReturnVoid().void
    next
  end

  # Load pointer to the TLAB from TLS
  tlab_ptr := LoadI(%tr).Imm(Constants::TLAB_OFFSET).ptr

  # Load pointer to the start address of free memory in the TLAB
  start := LoadI(tlab_ptr).Imm(Constants::TLAB_CUR_FREE_POSITION_OFFSET).ptr

  # Load pointer to the end address of free memory in the TLAB
  tls_end := LoadI(tlab_ptr).Imm(Constants::TLAB_MEMORY_END_ADDR_OFFSET).ptr

  tls_size := Sub(tls_end, start).word

  If(tls_size, size).LT.Unlikely.b {
    ep_offset = get_entrypoint_offset("CREATE_OBJECT_BY_CLASS_SLOW_PATH")
    Intrinsic(:SLOW_PATH_ENTRY, klass, size).AddImm(ep_offset).MethodAsImm("CreateObjectByClassBridge").Terminator.ptr
    Intrinsic(:UNREACHABLE).Terminator.void if defines.DEBUG
  } Else {
    Intrinsic(:WRITE_TLAB_STATS_SAFE, start, size, Cast(-1).u64).void if defines.DEBUG
    if defines.__SANITIZE_ADDRESS__ || defines.__SANITIZE_THREAD__
      call_runtime_save_all(Constants::ANNOTATE_SANITIZERS_NO_BRIDGE, start, size).void
    end
    new_start := Add(start, size).ptr
    StoreI(start, klass).Imm(Constants::OBJECT_CLASS_OFFSET).ref
    addr := Add(tlab_ptr, Constants::TLAB_CUR_FREE_POSITION_OFFSET).ptr
    StoreI(addr, new_start).Imm(0).Volatile.ptr
    Return(start).ptr
  }
}

# Allow three temps, two arguments registers and one callee saved register
$available_regs = $available_regs + :callee0

def AllocateArrayTlab(element_size)
  function("AllocateArrayTlab#{element_size}".to_sym,
           params: {klass: 'ref', elements_num: 'word'},
           regmap: $full_regmap,
           mode: [:FastPath],
           regalloc_set: $available_regs,
           validate: AllocateArrayTlabValidation) {
    if Options.arch == :arm32
      ReturnVoid().void
      next
    end
    elements_num := And(elements_num, "0x00000000ffffffff").word
    if element_size == 8
      size := elements_num
    elsif element_size == 16
      size := Shl(elements_num, 1).word
    elsif element_size == 32
      size := Shl(elements_num, 2).word
    elsif element_size == 64
      size := Shl(elements_num, 3).word
    else
      raise "Wrong element size #{element_size}"
    end
    # Add sizeof(Array) and do align
    size := Add(size, "DEFAULT_ALIGNMENT_IN_BYTES - 1 + CORETYPES_ARRAY_CLASS_SIZE").word
    size := And(size, "(~(DEFAULT_ALIGNMENT_IN_BYTES - 1))").word

    # Load pointer to the TLAB from TLS
    tlab_ptr := LoadI(%tr).Imm(Constants::TLAB_OFFSET).ptr

    # Load pointer to the start address of free memory in the TLAB
    start := LoadI(tlab_ptr).Imm(Constants::TLAB_CUR_FREE_POSITION_OFFSET).ptr

    # Load pointer to the end address of free memory in the TLAB
    tls_end := LoadI(tlab_ptr).Imm(Constants::TLAB_MEMORY_END_ADDR_OFFSET).ptr

    tls_size := Sub(tls_end, start).word

    If(tls_size, size).LT.Unlikely.b {
      ep_offset = get_entrypoint_offset("CREATE_ARRAY_SLOW_PATH")
      Intrinsic(:SLOW_PATH_ENTRY, klass, elements_num).AddImm(ep_offset).MethodAsImm("CreateArrayBridge").Terminator.ptr
      Intrinsic(:UNREACHABLE).Terminator.void if defines.DEBUG
    } Else {
      Intrinsic(:WRITE_TLAB_STATS_SAFE, start, size, Cast(-1).u64).void if defines.DEBUG
      if defines.__SANITIZE_ADDRESS__ || defines.__SANITIZE_THREAD__
        call_runtime_save_all(Constants::ANNOTATE_SANITIZERS_NO_BRIDGE, start, size).void
      end
      new_start := Add(start, size).ptr
      StoreI(start, klass).Imm(Constants::OBJECT_CLASS_OFFSET).ref
      StoreI(start, elements_num).Imm(Constants::ARRAY_LENGTH_OFFSET).word
      addr := Add(tlab_ptr, Constants::TLAB_CUR_FREE_POSITION_OFFSET).ptr
      StoreI(addr, new_start).Imm(0).Volatile.ptr
      Return(start).ptr
    }
  }
end

AllocateArrayTlab(8)
AllocateArrayTlab(16)
AllocateArrayTlab(32)
AllocateArrayTlab(64)