/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.

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

#ifndef TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_
#define TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_

#include "tensorflow/lite/micro/compatibility.h"
#include "tensorflow/lite/micro/micro_allocator.h"
#include "tensorflow/lite/micro/recording_simple_memory_allocator.h"

namespace tflite {

// List of buckets currently recorded by this class. Each type keeps a list of
// allocated information during model initialization.
enum class RecordedAllocationType {
  kTfLiteTensorArray,
  kTfLiteTensorArrayQuantizationData,
  kTfLiteTensorVariableBufferData,
  kNodeAndRegistrationArray,
  kOpData,
};

// Container for holding information about allocation recordings by a given
// type. Each recording contains the number of bytes requested, the actual bytes
// allocated (can defer from requested by alignment), and the number of items
// allocated.
typedef struct RecordedAllocation {
  RecordedAllocation() : requested_bytes(0), used_bytes(0), count(0) {}
  size_t requested_bytes;
  size_t used_bytes;
  size_t count;
} RecordedAllocation;

// Utility subclass of MicroAllocator that records all allocations
// inside the arena. A summary of allocations can be logged through the
// ErrorReporter by invoking LogAllocations(). This special allocator requires
// an instance of RecordingSimpleMemoryAllocator to capture allocations in the
// head and tail. Arena allocation recording can be retrieved by type through
// the GetRecordedAllocation() function. This class should only be used for
// auditing memory usage or integration testing.
class RecordingMicroAllocator : public MicroAllocator {
 public:
  static RecordingMicroAllocator* Create(uint8_t* tensor_arena,
                                         size_t arena_size,
                                         ErrorReporter* error_reporter);

  // Returns the recorded allocations information for a given allocation type.
  RecordedAllocation GetRecordedAllocation(
      RecordedAllocationType allocation_type) const;

  const RecordingSimpleMemoryAllocator* GetSimpleMemoryAllocator() const;

  // Logs out through the ErrorReporter all allocation recordings by type
  // defined in RecordedAllocationType.
  void PrintAllocations() const;

 protected:
  TfLiteStatus AllocateTfLiteTensorArray(TfLiteContext* context,
                                         const SubGraph* subgraph) override;
  TfLiteStatus PopulateTfLiteTensorArrayFromFlatbuffer(
      const Model* model, TfLiteContext* context,
      const SubGraph* subgraph) override;
  TfLiteStatus AllocateNodeAndRegistrations(
      const SubGraph* subgraph,
      NodeAndRegistration** node_and_registrations) override;
  TfLiteStatus PrepareNodeAndRegistrationDataFromFlatbuffer(
      const Model* model, const SubGraph* subgraph,
      const MicroOpResolver& op_resolver,
      NodeAndRegistration* node_and_registrations) override;
  TfLiteStatus AllocateVariables(TfLiteContext* context,
                                 const SubGraph* subgraph) override;

  void SnapshotAllocationUsage(RecordedAllocation& recorded_allocation);
  void RecordAllocationUsage(RecordedAllocation& recorded_allocation);

 private:
  RecordingMicroAllocator(RecordingSimpleMemoryAllocator* memory_allocator,
                          ErrorReporter* error_reporter);

  void PrintRecordedAllocation(RecordedAllocationType allocation_type,
                               const char* allocation_name,
                               const char* allocation_description) const;

  const RecordingSimpleMemoryAllocator* recording_memory_allocator_;

  RecordedAllocation recorded_tflite_tensor_array_data_;
  RecordedAllocation recorded_tflite_tensor_array_quantization_data_;
  RecordedAllocation recorded_tflite_tensor_variable_buffer_data_;
  RecordedAllocation recorded_node_and_registration_array_data_;
  RecordedAllocation recorded_op_data_;

  TF_LITE_REMOVE_VIRTUAL_DELETE
};

}  // namespace tflite

#endif  // TENSORFLOW_LITE_MICRO_RECORDING_MICRO_ALLOCATOR_H_