text: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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.
   */
  
  /*
   * Copyright (c) 2021, OPEN AI LAB
   * Author: haitao@openailab.com
   */
  
  #include "mem_stat.h"
  
  #include <stdio.h>
  #include <stdlib.h>
  #include <string.h>
  
  #include "api/c_api.h"
  #include "utility/sys_port.h"
  #include "utility/vector.h"
  #include "utility/log.h"
  
  #ifdef CONFIG_MEM_STAT
  
  extern void (*enable_mem_stat)(void);
  extern void (*disable_mem_stat)(void);
  
  struct mem_stat
  {
      int alloc_count;
      int free_count;
      int realloc_count;
      int max_block_size;
      int min_block_size;
      int peak_mem_size;
      int cur_mem_size;
  };
  
  struct block_stat
  {
      void* ptr;
      int size;
  };
  
  static int mem_stat_skipped = 1;
  static struct mem_stat mem_stat;
  static struct vector* block_list;
  
  DECLARE_AUTO_INIT_FUNC(init_mem_stat);
  DECLARE_AUTO_EXIT_FUNC(release_mem_stat);
  
  static int find_block_list(void* ptr)
  {
      int n = get_vector_num(block_list);
      int i;
  
      for (i = 0; i < n; i++)
      {
          struct block_stat* block_stat = (struct block_stat*)get_vector_data(block_list, i);
  
          if (block_stat->ptr == ptr)
              break;
      }
  
      if (i == n)
          return -1;
  
      return i;
  }
  
  static void real_enable_mem_stat(void)
  {
      set_skip_stat(0);
  }
  
  static void real_disable_mem_stat(void)
  {
      set_skip_stat(1);
  }
  
  static void init_mem_stat(void)
  {
      memset(&mem_stat, 0x0, sizeof(mem_stat));
      mem_stat.min_block_size = 1 << 20;
  
      block_list = create_vector(sizeof(struct block_stat), NULL);
  
      enable_mem_stat = real_enable_mem_stat;
      disable_mem_stat = real_disable_mem_stat;
  }
  
  void dump_mem_stat(void)
  {
      TLOG_INFO("memory usage stats:\n");
      TLOG_INFO("\talloc_count: %d\n", mem_stat.alloc_count);
      TLOG_INFO("\tfree_count: %d\n", mem_stat.free_count);
      TLOG_INFO("\trealloc_count: %d\n", mem_stat.realloc_count);
      TLOG_INFO("\tmax_block_size: %d\n", mem_stat.max_block_size);
      TLOG_INFO("\tmin_block_size: %d\n", mem_stat.min_block_size);
      TLOG_INFO("\tpeak_mem_size: %d\n", mem_stat.peak_mem_size);
      TLOG_INFO("\tcur_mem_size: %d\n", mem_stat.cur_mem_size);
  }
  
  static void release_mem_stat(void)
  {
      release_vector(block_list);
      dump_mem_stat();
  }
  
  void set_skip_stat(int skip)
  {
      mem_stat_skipped = skip;
  }
  
  int skip_stat(void)
  {
      return mem_stat_skipped;
  }
  
  void* stat_malloc(int size)
  {
      void* ptr = malloc(size);
  
      if (ptr == NULL)
      {
          TLOG_ERR("cannot alloc size: %d\n", size);
          TLOG_ERR("cur mem size: %d peak mem size: %d\n", mem_stat.cur_mem_size, mem_stat.peak_mem_size);
  
          return NULL;
      }
  
      mem_stat.alloc_count++;
      mem_stat.cur_mem_size += size;
  
      if (mem_stat.cur_mem_size > mem_stat.peak_mem_size)
          mem_stat.peak_mem_size = mem_stat.cur_mem_size;
  
      if (size > mem_stat.max_block_size)
          mem_stat.max_block_size = size;
      if (size < mem_stat.min_block_size)
          mem_stat.min_block_size = size;
  
      mem_stat_skipped = 1;
  
      struct block_stat block_stat;
  
      block_stat.ptr = ptr;
      block_stat.size = size;
  
      push_vector_data(block_list, &block_stat);
  
      mem_stat_skipped = 0;
  
      return ptr;
  }
  
  void stat_free(void* ptr)
  {
      int idx = find_block_list(ptr);
  
      if (idx < 0)
      {
          /* a memory not allocated by us ? */
          free(ptr);
          return;
      }
  
      struct block_stat* block_stat = (struct block_stat*)get_vector_data(block_list, idx);
  
      mem_stat.free_count++;
      mem_stat.cur_mem_size -= block_stat->size;
  
      mem_stat_skipped = 1;
  
      remove_vector_via_index(block_list, idx);
  
      mem_stat_skipped = 0;
  
      free(ptr);
  }
  
  void* stat_realloc(void* ptr, size_t size)
  {
      if (ptr == NULL)
          return stat_malloc(size);
  
      int idx = find_block_list(ptr);
  
      if (idx < 0)
          return realloc(ptr, size);
  
      void* new_ptr = realloc(ptr, size);
  
      struct block_stat* block_stat = (struct block_stat*)get_vector_data(block_list, idx);
  
      if (new_ptr == NULL)
      {
          TLOG_ERR("cannot realloc size: %d --> %d\n", block_stat->size, size);
          TLOG_ERR("cur mem size: %d peak mem size: %d\n", mem_stat.cur_mem_size, mem_stat.peak_mem_size);
          return NULL;
      }
  
      if (size > mem_stat.max_block_size)
          mem_stat.max_block_size = size;
  
      if (size < mem_stat.min_block_size)
          mem_stat.min_block_size = size;
  
      mem_stat.realloc_count++;
      mem_stat.cur_mem_size += size - block_stat->size;
  
      if (mem_stat.cur_mem_size > mem_stat.peak_mem_size)
          mem_stat.peak_mem_size = mem_stat.cur_mem_size;
  
      block_stat->ptr = new_ptr;
      block_stat->size = size;
  
      return new_ptr;
  }
  
  #endif
  