/* Copyright 2019 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.
==============================================================================*/
#include <initializer_list>
#include <memory>
#include <string>
#include <vector>

#include "absl/memory/memory.h"
#include "absl/strings/match.h"
#include "tensorflow/lite/core/interpreter.h"
#include "tensorflow/lite/core/model.h"
#include "tensorflow/lite/experimental/resource/lookup_interfaces.h"
#include "tensorflow/lite/kernels/internal/tensor_ctypes.h"
#include "tensorflow/lite/kernels/test_util.h"
#include "tensorflow/lite/testing/util.h"

namespace tflite {

// Forward declaration for op kernels.
namespace ops {
namespace builtin {

TfLiteRegistration* Register_HASHTABLE();
TfLiteRegistration* Register_HASHTABLE_FIND();
TfLiteRegistration* Register_HASHTABLE_IMPORT();
TfLiteRegistration* Register_HASHTABLE_SIZE();

}  // namespace builtin
}  // namespace ops

namespace {

using ::testing::ElementsAreArray;

typedef enum {
  kResourceTensorId = 0,
  kKeyTensorId = 1,
  kValueTensorId = 2,
  kQueryTensorId = 3,
  kResultTensorId = 4,
  kSizeTensorId = 5,
  kDefaultValueTensorId = 6,
  kResourceTwoTensorId = 7,
  kKeyTwoTensorId = 8,
  kValueTwoTensorId = 9,
  kQueryTwoTensorId = 10,
  kResultTwoTensorId = 11,
  kSizeTwoTensorId = 12,
  kDefaultValueTwoTensorId = 13,
} TensorIds;

template <typename T>
void SetTensorData(Interpreter* interpreter, int tensorId,
                   std::vector<T> data) {
  auto* tensor = interpreter->tensor(tensorId);
  auto* tensor_data = GetTensorData<T>(tensor);
  int i = 0;
  for (auto item : data) {
    tensor_data[i++] = item;
  }
}

template <>
void SetTensorData(Interpreter* interpreter, int tensorId,
                   std::vector<std::string> data) {
  auto* tensor = interpreter->tensor(tensorId);
  DynamicBuffer buf;
  for (auto item : data) {
    buf.AddString(item.c_str(), item.length());
  }
  buf.WriteToTensorAsVector(tensor);
}

TensorType ConvertTfLiteType(TfLiteType type) {
  // Currently, hashtable kernels support INT64 and STRING types only.
  switch (type) {
    case kTfLiteInt64:
      return TensorType_INT64;
    case kTfLiteString:
      return TensorType_STRING;
    default:
      CHECK(false);  // Not reached.
      return TensorType_MIN;
  }
}

// HashtableGraph generates a graph with hash table ops. This class can create
// the following scenarios:
//
// - Default graph: One hash table resource with import, lookup, and size ops.
// - Graph without any import node
// - Graph with two import nodes
// - Graph has two hash table resources.
//
template <typename KeyType, typename ValueType>
class HashtableGraph {
 public:
  HashtableGraph(TfLiteType key_type, TfLiteType value_type)
      : key_type_(key_type), value_type_(value_type) {
    interpreter_ = std::make_unique<Interpreter>(&error_reporter_);
    InitOpRegistrations();
  }
  ~HashtableGraph() {}

  void BuildDefaultGraph() {
    TfLiteHashtableParams* hashtable_params = GetHashtableParams();

    int node_index;
    // Hash table node.
    interpreter_->AddNodeWithParameters(
        {}, {kResourceTensorId}, nullptr, 0,
        reinterpret_cast<void*>(hashtable_params), hashtable_registration_,
        &node_index);

    // Hash table import node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kKeyTensorId, kValueTensorId}, {}, nullptr, 0,
        nullptr, hashtable_import_registration_, &node_index);

    // Hash table lookup node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kQueryTensorId, kDefaultValueTensorId},
        {kResultTensorId}, nullptr, 0, nullptr, hashtable_find_registration_,
        &node_index);

    // Hash table size node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId}, {kSizeTensorId}, nullptr, 0, nullptr,
        hashtable_size_registration_, &node_index);
  }

  void BuildNoImportGraph() {
    TfLiteHashtableParams* hashtable_params = GetHashtableParams();

    int node_index;
    // Hash table node.
    interpreter_->AddNodeWithParameters(
        {}, {kResourceTensorId}, nullptr, 0,
        reinterpret_cast<void*>(hashtable_params), hashtable_registration_,
        &node_index);

    // Hash table lookup node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kQueryTensorId, kDefaultValueTensorId},
        {kResultTensorId}, nullptr, 0, nullptr, hashtable_find_registration_,
        &node_index);

    // Hash table size node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId}, {kSizeTensorId}, nullptr, 0, nullptr,
        hashtable_size_registration_, &node_index);
  }

  void BuildImportTwiceGraph() {
    TfLiteHashtableParams* hashtable_params = GetHashtableParams();

    int node_index;
    // Hash table node.
    interpreter_->AddNodeWithParameters(
        {}, {kResourceTensorId}, nullptr, 0,
        reinterpret_cast<void*>(hashtable_params), hashtable_registration_,
        &node_index);

    // Hash table import node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kKeyTensorId, kValueTensorId}, {}, nullptr, 0,
        nullptr, hashtable_import_registration_, &node_index);

    // Hash table import node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kKeyTensorId, kValueTensorId}, {}, nullptr, 0,
        nullptr, hashtable_import_registration_, &node_index);

    // Hash table lookup node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kQueryTensorId, kDefaultValueTensorId},
        {kResultTensorId}, nullptr, 0, nullptr, hashtable_find_registration_,
        &node_index);

    // Hash table size node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId}, {kSizeTensorId}, nullptr, 0, nullptr,
        hashtable_size_registration_, &node_index);
  }

  void BuildTwoHashtablesGraph() {
    TfLiteHashtableParams* hashtable_params = GetHashtableParams();

    int node_index;
    // Hash table node.
    interpreter_->AddNodeWithParameters(
        {}, {kResourceTensorId}, nullptr, 0,
        reinterpret_cast<void*>(hashtable_params), hashtable_registration_,
        &node_index);

    // Hash table import node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kKeyTensorId, kValueTensorId}, {}, nullptr, 0,
        nullptr, hashtable_import_registration_, &node_index);

    // Hash table lookup node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId, kQueryTensorId, kDefaultValueTensorId},
        {kResultTensorId}, nullptr, 0, nullptr, hashtable_find_registration_,
        &node_index);

    // Hash table size node.
    interpreter_->AddNodeWithParameters(
        {kResourceTensorId}, {kSizeTensorId}, nullptr, 0, nullptr,
        hashtable_size_registration_, &node_index);

    TfLiteHashtableParams* hashtable_two_params = GetHashtableParams();

    // Hash table two node.
    interpreter_->AddNodeWithParameters(
        {}, {kResourceTwoTensorId}, nullptr, 0,
        reinterpret_cast<void*>(hashtable_two_params), hashtable_registration_,
        &node_index);

    // Hash table two import node.
    interpreter_->AddNodeWithParameters(
        {kResourceTwoTensorId, kKeyTwoTensorId, kValueTwoTensorId}, {}, nullptr,
        0, nullptr, hashtable_import_registration_, &node_index);

    // Hash table two lookup node.
    interpreter_->AddNodeWithParameters(
        {kResourceTwoTensorId, kQueryTwoTensorId, kDefaultValueTwoTensorId},
        {kResultTwoTensorId}, nullptr, 0, nullptr, hashtable_find_registration_,
        &node_index);

    // Hash table two size node.
    interpreter_->AddNodeWithParameters(
        {kResourceTwoTensorId}, {kSizeTwoTensorId}, nullptr, 0, nullptr,
        hashtable_size_registration_, &node_index);
  }

  TfLiteStatus Invoke() { return interpreter_->Invoke(); }

  void SetTable(std::initializer_list<KeyType> keys,
                std::initializer_list<ValueType> values) {
    keys_ = std::vector<KeyType>(keys);
    values_ = std::vector<ValueType>(values);
  }

  void SetTableTwo(std::initializer_list<KeyType> keys,
                   std::initializer_list<ValueType> values) {
    keys_two_ = std::vector<KeyType>(keys);
    values_two_ = std::vector<ValueType>(values);
  }

  void SetQuery(std::initializer_list<KeyType> queries,
                ValueType default_value) {
    queries_ = std::vector<KeyType>(queries);
    default_value_ = default_value;
  }

  void SetQueryForTableTwo(std::initializer_list<KeyType> queries,
                           ValueType default_value) {
    queries_two_ = std::vector<KeyType>(queries);
    default_value_two_ = default_value;
  }

  int64_t GetTableSize() {
    auto* size_tensor = interpreter_->tensor(kSizeTensorId);
    auto size_tensor_shape = GetTensorShape(size_tensor);
    return GetTensorData<int64_t>(size_tensor)[0];
  }

  int64_t GetTableTwoSize() {
    auto* size_tensor = interpreter_->tensor(kSizeTwoTensorId);
    auto size_tensor_shape = GetTensorShape(size_tensor);
    return GetTensorData<int64_t>(size_tensor)[0];
  }

  std::vector<ValueType> GetLookupResult() {
    auto* result_tensor = interpreter_->tensor(kResultTensorId);
    auto result_tensor_shape = GetTensorShape(result_tensor);
    auto* result_tensor_data = GetTensorData<ValueType>(result_tensor);

    int size = result_tensor_shape.FlatSize();
    std::vector<ValueType> result;
    for (int i = 0; i < size; ++i) {
      result.push_back(result_tensor_data[i]);
    }
    return result;
  }

  std::vector<ValueType> GetLookupTwoResult() {
    auto* result_tensor = interpreter_->tensor(kResultTwoTensorId);
    auto result_tensor_shape = GetTensorShape(result_tensor);
    auto* result_tensor_data = GetTensorData<ValueType>(result_tensor);

    int size = result_tensor_shape.FlatSize();
    std::vector<ValueType> result;
    for (int i = 0; i < size; ++i) {
      result.push_back(result_tensor_data[i]);
    }
    return result;
  }

  std::vector<std::string> GetStringLookupResult() {
    auto* result_tensor = interpreter_->tensor(kResultTensorId);
    auto result_tensor_shape = GetTensorShape(result_tensor);

    int size = result_tensor_shape.FlatSize();
    std::vector<std::string> result;
    for (int i = 0; i < size; ++i) {
      auto string_ref = GetString(result_tensor, i);
      result.push_back(std::string(string_ref.str, string_ref.len));
    }
    return result;
  }

  void AddTensors(bool table_two_initialization = false) {
    int first_new_tensor_index;
    if (!table_two_initialization) {
      ASSERT_EQ(interpreter_->AddTensors(7, &first_new_tensor_index),
                kTfLiteOk);
      ASSERT_EQ(interpreter_->SetInputs({kResourceTensorId, kKeyTensorId,
                                         kValueTensorId, kQueryTensorId,
                                         kDefaultValueTensorId}),
                kTfLiteOk);
      ASSERT_EQ(interpreter_->SetOutputs({kResultTensorId, kSizeTensorId}),
                kTfLiteOk);
    } else {
      ASSERT_EQ(interpreter_->AddTensors(14, &first_new_tensor_index),
                kTfLiteOk);
      ASSERT_EQ(
          interpreter_->SetInputs(
              {kResourceTensorId, kKeyTensorId, kValueTensorId, kQueryTensorId,
               kDefaultValueTensorId, kResourceTwoTensorId, kKeyTwoTensorId,
               kValueTwoTensorId, kQueryTwoTensorId, kDefaultValueTwoTensorId}),
          kTfLiteOk);
      ASSERT_EQ(
          interpreter_->SetOutputs({kResultTensorId, kSizeTensorId,
                                    kResultTwoTensorId, kSizeTwoTensorId}),
          kTfLiteOk);
    }

    // Resource id tensor.
    interpreter_->SetTensorParametersReadWrite(
        kResourceTensorId, kTfLiteResource, "", {}, TfLiteQuantization());

    // Key tensor for import.
    interpreter_->SetTensorParametersReadWrite(kKeyTensorId, key_type_, "",
                                               {static_cast<int>(keys_.size())},
                                               TfLiteQuantization());

    // Value tensor for import.
    interpreter_->SetTensorParametersReadWrite(
        kValueTensorId, value_type_, "", {static_cast<int>(values_.size())},
        TfLiteQuantization());

    // Query tensor for lookup.
    interpreter_->SetTensorParametersReadWrite(
        kQueryTensorId, key_type_, "", {static_cast<int>(queries_.size())},
        TfLiteQuantization());

    // Result tensor for lookup result.
    interpreter_->SetTensorParametersReadWrite(
        kResultTensorId, value_type_, "", {static_cast<int>(queries_.size())},
        TfLiteQuantization());

    // Result tensor for size calculation.
    interpreter_->SetTensorParametersReadWrite(kSizeTensorId, kTfLiteInt64, "",
                                               {1}, TfLiteQuantization());

    // Default value tensor for lookup.
    interpreter_->SetTensorParametersReadWrite(
        kDefaultValueTensorId, value_type_, "", {1}, TfLiteQuantization());

    if (table_two_initialization) {
      // Resource id tensor.
      interpreter_->SetTensorParametersReadWrite(
          kResourceTwoTensorId, kTfLiteResource, "", {}, TfLiteQuantization());

      // Key tensor for import.
      interpreter_->SetTensorParametersReadWrite(
          kKeyTwoTensorId, key_type_, "", {static_cast<int>(keys_two_.size())},
          TfLiteQuantization());

      // Value tensor for import.
      interpreter_->SetTensorParametersReadWrite(
          kValueTwoTensorId, value_type_, "",
          {static_cast<int>(values_two_.size())}, TfLiteQuantization());

      // Query tensor for lookup.
      interpreter_->SetTensorParametersReadWrite(
          kQueryTwoTensorId, key_type_, "",
          {static_cast<int>(queries_two_.size())}, TfLiteQuantization());

      // Result tensor for lookup result.
      interpreter_->SetTensorParametersReadWrite(
          kResultTwoTensorId, value_type_, "",
          {static_cast<int>(queries_two_.size())}, TfLiteQuantization());

      // Result tensor for size calculation.
      interpreter_->SetTensorParametersReadWrite(kSizeTwoTensorId, kTfLiteInt64,
                                                 "", {1}, TfLiteQuantization());

      // Default value tensor for lookup.
      interpreter_->SetTensorParametersReadWrite(
          kDefaultValueTwoTensorId, value_type_, "", {1}, TfLiteQuantization());
    }
  }

  TfLiteStatus AllocateTensors(bool table_two_initialization = false) {
    if (interpreter_->AllocateTensors() != kTfLiteOk) {
      return kTfLiteError;
    }

    SetTensorData(interpreter_.get(), kKeyTensorId, keys_);
    SetTensorData(interpreter_.get(), kValueTensorId, values_);
    SetTensorData(interpreter_.get(), kQueryTensorId, queries_);
    SetTensorData(interpreter_.get(), kDefaultValueTensorId,
                  std::vector<ValueType>({default_value_}));

    if (table_two_initialization) {
      SetTensorData(interpreter_.get(), kKeyTwoTensorId, keys_two_);
      SetTensorData(interpreter_.get(), kValueTwoTensorId, values_two_);
      SetTensorData(interpreter_.get(), kQueryTwoTensorId, queries_two_);
      SetTensorData(interpreter_.get(), kDefaultValueTwoTensorId,
                    std::vector<ValueType>({default_value_two_}));
    }
    return kTfLiteOk;
  }

  TestErrorReporter* GetErrorReporter() { return &error_reporter_; }

 private:
  void InitOpRegistrations() {
    hashtable_registration_ = tflite::ops::builtin::Register_HASHTABLE();
    ASSERT_NE(hashtable_registration_, nullptr);

    hashtable_find_registration_ =
        tflite::ops::builtin::Register_HASHTABLE_FIND();
    ASSERT_NE(hashtable_find_registration_, nullptr);

    hashtable_import_registration_ =
        tflite::ops::builtin::Register_HASHTABLE_IMPORT();
    ASSERT_NE(hashtable_import_registration_, nullptr);

    hashtable_size_registration_ =
        tflite::ops::builtin::Register_HASHTABLE_SIZE();
    ASSERT_NE(hashtable_size_registration_, nullptr);
  }

  TfLiteHashtableParams* GetHashtableParams() {
    TfLiteHashtableParams* params = reinterpret_cast<TfLiteHashtableParams*>(
        malloc(sizeof(TfLiteHashtableParams)));
    params->table_id = std::rand();
    params->key_dtype = key_type_;
    params->value_dtype = value_type_;
    return params;
  }

  // Tensor types
  TfLiteType key_type_;
  TfLiteType value_type_;

  // Tensor data
  std::vector<KeyType> keys_;
  std::vector<ValueType> values_;
  std::vector<KeyType> queries_;
  ValueType default_value_;

  // Tensor data for table two.
  std::vector<KeyType> keys_two_;
  std::vector<ValueType> values_two_;
  std::vector<KeyType> queries_two_;
  ValueType default_value_two_;

  // Op registrations.
  TfLiteRegistration* hashtable_registration_;
  TfLiteRegistration* hashtable_find_registration_;
  TfLiteRegistration* hashtable_import_registration_;
  TfLiteRegistration* hashtable_size_registration_;

  // Hashtable params.
  TfLiteHashtableParams* hashtable_params_;
  TfLiteHashtableParams* hashtable_two_params_;

  // Interpreter.
  std::unique_ptr<Interpreter> interpreter_;
  TestErrorReporter error_reporter_;
};

// HashtableDefaultGraphTest tests hash table features on a basic graph, created
// by the HashtableGraph class.
template <typename KeyType, typename ValueType>
class HashtableDefaultGraphTest {
 public:
  HashtableDefaultGraphTest(TfLiteType key_type, TfLiteType value_type,
                            std::initializer_list<KeyType> keys,
                            std::initializer_list<ValueType> values,
                            std::initializer_list<KeyType> queries,
                            ValueType default_value, int table_size,
                            std::initializer_list<ValueType> lookup_result) {
    graph_ = std::make_unique<HashtableGraph<KeyType, ValueType>>(key_type,
                                                                  value_type);
    graph_->SetTable(keys, values);
    graph_->SetQuery(queries, default_value);
    graph_->AddTensors();
    graph_->BuildDefaultGraph();

    value_type_ = value_type;
    table_size_ = table_size;
    lookup_result_ = std::vector<ValueType>(lookup_result);
  }

  void Invoke() {
    EXPECT_EQ(graph_->AllocateTensors(), kTfLiteOk);
    EXPECT_EQ(graph_->Invoke(), kTfLiteOk);

    EXPECT_THAT(graph_->GetTableSize(), table_size_);
  }

  void InvokeAndVerifyStringResult() {
    Invoke();
    EXPECT_THAT(graph_->GetStringLookupResult(),
                ElementsAreArray(lookup_result_));
  }

  void InvokeAndVerifyIntResult() {
    Invoke();
    EXPECT_THAT(graph_->GetLookupResult(), ElementsAreArray(lookup_result_));
  }

  void InvokeAndVerifyFloatResult() {
    Invoke();
    EXPECT_THAT(graph_->GetLookupResult(),
                ElementsAreArray(ArrayFloatNear(lookup_result_)));
  }

 private:
  std::unique_ptr<HashtableGraph<KeyType, ValueType>> graph_;

  TfLiteType value_type_;
  int table_size_;
  std::vector<ValueType> lookup_result_;
};

TEST(HashtableOpsTest, TestInt64ToStringHashtable) {
  HashtableDefaultGraphTest<std::int64_t, std::string> t(
      kTfLiteInt64, kTfLiteString,
      /*keys=*/{1, 2, 3}, /*values=*/{"a", "b", "c"}, /*queries=*/{2, 3, 4},
      /*default_value=*/"d", /*table_size=*/3,
      /*lookup_result=*/{"b", "c", "d"});
  t.InvokeAndVerifyStringResult();
}

TEST(HashtableOpsTest, TestStringToInt64Hashtable) {
  HashtableDefaultGraphTest<std::string, int64_t> t(
      kTfLiteString, kTfLiteInt64,
      /*keys=*/{"A", "B", "C"}, /*values=*/{4, 5, 6},
      /*queries=*/{"B", "C", "D"},
      /*default_value=*/-1, /*table_size=*/3, /*lookup_result=*/{5, 6, -1});
  t.InvokeAndVerifyIntResult();
}

TEST(HashtableOpsTest, TestNoImport) {
  HashtableGraph<std::string, std::int64_t> graph(kTfLiteString, kTfLiteInt64);
  graph.SetQuery({"1", "2", "3"}, -1);
  graph.AddTensors();
  graph.BuildNoImportGraph();
  EXPECT_EQ(graph.AllocateTensors(), kTfLiteOk);
  EXPECT_EQ(graph.Invoke(), kTfLiteError);
  EXPECT_TRUE(
      absl::StrContains(graph.GetErrorReporter()->error_messages(),
                        "hashtable need to be initialized before using"));
}

TEST(HashtableOpsTest, TestImportTwice) {
  HashtableGraph<std::string, std::int64_t> graph(kTfLiteString, kTfLiteInt64);
  graph.SetTable({"1", "2", "3"}, {4, 5, 6});
  graph.SetQuery({"2", "3", "4"}, -1);
  graph.AddTensors();
  graph.BuildImportTwiceGraph();
  EXPECT_EQ(graph.AllocateTensors(), kTfLiteOk);
  // The invocation of  thesecond import node will be ignored.
  EXPECT_EQ(graph.Invoke(), kTfLiteOk);

  EXPECT_THAT(graph.GetTableSize(), 3);
  EXPECT_THAT(graph.GetLookupResult(), ElementsAreArray({5, 6, -1}));
}

TEST(HashtableOpsTest, TestTwoHashtables) {
  HashtableGraph<std::string, std::int64_t> graph(kTfLiteString, kTfLiteInt64);
  graph.SetTable({"1", "2", "3"}, {4, 5, 6});
  graph.SetQuery({"2", "3", "4"}, -1);
  graph.SetTableTwo({"-1", "-2", "-3"}, {7, 8, 9});
  graph.SetQueryForTableTwo({"-4", "-2", "-3"}, -2);
  graph.AddTensors(/*table_two_initialization=*/true);
  graph.BuildTwoHashtablesGraph();
  EXPECT_EQ(graph.AllocateTensors(/*table_two_initialization=*/true),
            kTfLiteOk);
  EXPECT_EQ(graph.Invoke(), kTfLiteOk);

  EXPECT_THAT(graph.GetTableSize(), 3);
  EXPECT_THAT(graph.GetTableTwoSize(), 3);
  EXPECT_THAT(graph.GetLookupResult(), ElementsAreArray({5, 6, -1}));
  EXPECT_THAT(graph.GetLookupTwoResult(), ElementsAreArray({-2, 8, 9}));
}

TEST(HashtableOpsTest, TestImportDifferentKeyAndValueSize) {
  HashtableGraph<std::string, std::int64_t> graph(kTfLiteString, kTfLiteInt64);
  graph.SetTable({"1", "2", "3"}, {4, 5});
  graph.SetQuery({"2", "3", "4"}, -1);
  graph.AddTensors();
  graph.BuildDefaultGraph();
  EXPECT_EQ(graph.AllocateTensors(), kTfLiteOk);
  EXPECT_EQ(graph.Invoke(), kTfLiteError);
}

// HashtableOpModel creates a model with one single Hashtable op.
class HashtableOpModel : public SingleOpModel {
 public:
  explicit HashtableOpModel(const int table_id, TensorType key_dtype,
                            TensorType value_dtype) {
    output_ = AddOutput(TensorType_RESOURCE);

    SetBuiltinOp(
        BuiltinOperator_HASHTABLE, BuiltinOptions_HashtableOptions,
        CreateHashtableOptions(builder_, table_id, key_dtype, value_dtype)
            .Union());
    BuildInterpreter({});
  }

  int GetOutput() {
    int* int32_ptr =
        reinterpret_cast<int32_t*>(interpreter_->tensor(0)->data.raw);
    return *int32_ptr;
  }
  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }

  resource::ResourceMap& GetResources() {
    return interpreter_->primary_subgraph().resources();
  }

 private:
  int output_;
};

TEST(HashtableOpsTest, TestHashtable) {
  HashtableOpModel m(/*table_id=*/1, TensorType_INT64, TensorType_STRING);
  EXPECT_EQ(m.GetResources().size(), 0);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  auto& resources = m.GetResources();
  EXPECT_EQ(resources.size(), 1);
  int resource_id = m.GetOutput();
  EXPECT_NE(resource_id, 0);
  auto* hashtable = resource::GetHashtableResource(&resources, resource_id);
  EXPECT_TRUE(hashtable != nullptr);
  EXPECT_TRUE(hashtable->GetKeyType() == kTfLiteInt64);
  EXPECT_TRUE(hashtable->GetValueType() == kTfLiteString);
}

template <typename T>
TfLiteTensor CreateTensor(TfLiteType type, const std::vector<T>& vec) {
  TfLiteTensor tensor = {};
  TfLiteIntArray* dims = TfLiteIntArrayCreate(1);
  dims->data[0] = vec.size();
  tensor.dims = dims;
  tensor.name = "";
  tensor.params = {};
  tensor.quantization = {kTfLiteNoQuantization, nullptr};
  tensor.is_variable = false;
  tensor.allocation_type = kTfLiteDynamic;
  tensor.allocation = nullptr;
  tensor.type = type;
  tensor.bytes = sizeof(T) * vec.size();
  T* data = static_cast<T*>(malloc(sizeof(T) * vec.size()));
  for (int i = 0; i < vec.size(); ++i) {
    data[i] = vec[i];
  }
  tensor.data.raw = reinterpret_cast<char*>(data);
  return tensor;
}

template <>
TfLiteTensor CreateTensor(TfLiteType type,
                          const std::vector<std::string>& vec) {
  TfLiteTensor tensor = {};
  TfLiteIntArray* dims = TfLiteIntArrayCreate(1);
  dims->data[0] = vec.size();
  tensor.dims = dims;
  tensor.name = "";
  tensor.params = {};
  tensor.quantization = {kTfLiteNoQuantization, nullptr};
  tensor.is_variable = false;
  tensor.allocation_type = kTfLiteDynamic;
  tensor.allocation = nullptr;
  tensor.type = type;
  DynamicBuffer buf;
  for (std::string str : vec) {
    buf.AddString(str.c_str(), str.size());
  }
  buf.WriteToTensor(&tensor, nullptr);
  return tensor;
}

template <typename KeyType, typename ValueType>
void InitHashtableResource(resource::ResourceMap* resources, int resource_id,
                           TfLiteType key_type, TfLiteType value_type,
                           std::initializer_list<KeyType> keys,
                           std::initializer_list<ValueType> values) {
  resource::CreateHashtableResourceIfNotAvailable(resources, resource_id,
                                                  key_type, value_type);
  auto lookup = resource::GetHashtableResource(resources, resource_id);

  TfLiteContext context;
  TfLiteTensor key_tensor = CreateTensor<KeyType>(key_type, keys);
  TfLiteTensor value_tensor = CreateTensor<ValueType>(value_type, values);
  lookup->Import(&context, &key_tensor, &value_tensor);
  TfLiteTensorFree(&key_tensor);
  TfLiteTensorFree(&value_tensor);
}

// BaseHashtableOpModel is a base class for creating a model with any single
// hashtable op node, which takes a hash table resource as an input.
class BaseHashtableOpModel : public SingleOpModel {
 public:
  BaseHashtableOpModel() {}

  void SetResourceId(int resource_id) {
    auto* tensor = interpreter_->tensor(resource_id_);

    size_t bytesRequired = sizeof(int32_t);
    TfLiteTensorRealloc(bytesRequired, tensor);
    tensor->bytes = bytesRequired;

    TfLiteIntArray* outputSize = TfLiteIntArrayCreate(1);
    outputSize->data[0] = 1;
    if (tensor->dims) TfLiteIntArrayFree(tensor->dims);
    tensor->dims = outputSize;

    int32_t* resource_ptr = reinterpret_cast<int32_t*>(tensor->data.raw);
    resource_ptr[0] = resource_id;
  }

  void CreateHashtableResource(int resource_id) {
    auto key_tensor = interpreter_->tensor(keys_);
    auto value_tensor = interpreter_->tensor(values_);

    auto& resources = GetResources();
    resource::CreateHashtableResourceIfNotAvailable(
        &resources, resource_id, key_tensor->type, value_tensor->type);
  }

  template <typename ValueType>
  std::vector<ValueType> GetOutput() {
    return ExtractVector<ValueType>(output_);
  }

  std::vector<int> GetOutputShape() { return GetTensorShape(output_); }

  resource::ResourceMap& GetResources() {
    return interpreter_->primary_subgraph().resources();
  }

 protected:
  int resource_id_;
  int keys_;
  int values_;
  int output_;

  TensorType key_type_;
  TensorType value_type_;
};

// HashtableFindOpModel creates a model with a HashtableLookup op.
template <typename KeyType, typename ValueType>
class HashtableFindOpModel : public BaseHashtableOpModel {
 public:
  HashtableFindOpModel(const TensorType key_type, const TensorType value_type,
                       int lookup_size) {
    key_type_ = key_type;
    value_type_ = value_type;

    resource_id_ = AddInput({TensorType_RESOURCE, {1}});
    lookup_ = AddInput({key_type, {lookup_size}});
    default_value_ = AddInput({value_type, {1}});

    output_ = AddOutput({value_type, {lookup_size}});

    SetBuiltinOp(BuiltinOperator_HASHTABLE_FIND,
                 BuiltinOptions_HashtableFindOptions,
                 CreateHashtableFindOptions(builder_).Union());
    BuildInterpreter(
        {GetShape(resource_id_), GetShape(lookup_), GetShape(default_value_)});
  }

  void SetLookup(const std::vector<KeyType>& data) {
    PopulateTensor(lookup_, data);
  }

  void SetStringLookup(const std::vector<std::string>& data) {
    PopulateStringTensor(lookup_, data);
  }

  void SetDefaultValue(const std::vector<ValueType>& data) {
    PopulateTensor(default_value_, data);
  }

  void SetStringDefaultValue(const std::vector<std::string>& data) {
    PopulateStringTensor(default_value_, data);
  }

 private:
  int lookup_;
  int default_value_;
};

TEST(HashtableOpsTest, TestHashtableLookupStringToInt64) {
  const int kResourceId = 42;
  HashtableFindOpModel<std::string, std::int64_t> m(TensorType_STRING,
                                                    TensorType_INT64, 3);

  m.SetResourceId(kResourceId);
  m.SetStringLookup({"5", "6", "7"});
  m.SetDefaultValue({4});

  InitHashtableResource<std::string, std::int64_t>(
      &m.GetResources(), kResourceId, kTfLiteString, kTfLiteInt64,
      {"4", "5", "6"}, {1, 2, 3});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  EXPECT_THAT(m.GetOutput<std::int64_t>(), ElementsAreArray({2, 3, 4}));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3}));
}

TEST(HashtableOpsTest, TestHashtableLookupInt64ToString) {
  const int kResourceId = 42;
  HashtableFindOpModel<std::int64_t, std::string> m(TensorType_INT64,
                                                    TensorType_STRING, 3);

  m.SetResourceId(kResourceId);
  m.SetLookup({5, 6, 7});
  m.SetStringDefaultValue({"4"});

  InitHashtableResource<std::int64_t, std::string>(
      &m.GetResources(), kResourceId, kTfLiteInt64, kTfLiteString, {4, 5, 6},
      {"1", "2", "3"});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  EXPECT_THAT(m.GetOutput<std::string>(), ElementsAreArray({"2", "3", "4"}));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({3}));
}

// HashtableImportOpModel creates a model with a HashtableImport op.
template <typename KeyType, typename ValueType>
class HashtableImportOpModel : public BaseHashtableOpModel {
 public:
  HashtableImportOpModel(const TensorType key_type, const TensorType value_type,
                         int initdata_size) {
    key_type_ = key_type;
    value_type_ = value_type;

    resource_id_ = AddInput({TensorType_RESOURCE, {1}});
    keys_ = AddInput({key_type, {initdata_size}});
    values_ = AddInput({value_type, {initdata_size}});

    SetBuiltinOp(BuiltinOperator_HASHTABLE_IMPORT,
                 BuiltinOptions_HashtableImportOptions,
                 CreateHashtableImportOptions(builder_).Union());
    BuildInterpreter(
        {GetShape(resource_id_), GetShape(keys_), GetShape(values_)});
  }

  void SetKeys(const std::vector<KeyType>& data) {
    PopulateTensor(keys_, data);
  }

  void SetStringKeys(const std::vector<std::string>& data) {
    PopulateStringTensor(keys_, data);
  }

  void SetValues(const std::vector<ValueType>& data) {
    PopulateTensor(values_, data);
  }

  void SetStringValues(const std::vector<std::string>& data) {
    PopulateStringTensor(values_, data);
  }
};

TEST(HashtableOpsTest, TestHashtableImport) {
  const int kResourceId = 42;
  HashtableImportOpModel<std::int64_t, std::string> m(TensorType_INT64,
                                                      TensorType_STRING, 3);
  EXPECT_EQ(m.GetResources().size(), 0);
  m.SetResourceId(kResourceId);
  m.SetKeys({1, 2, 3});
  m.SetStringValues({"1", "2", "3"});
  m.CreateHashtableResource(kResourceId);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  auto& resources = m.GetResources();
  EXPECT_EQ(resources.size(), 1);
  auto* hashtable = resource::GetHashtableResource(&resources, kResourceId);
  EXPECT_TRUE(hashtable != nullptr);
  EXPECT_TRUE(hashtable->GetKeyType() == kTfLiteInt64);
  EXPECT_TRUE(hashtable->GetValueType() == kTfLiteString);

  EXPECT_EQ(hashtable->Size(), 3);
}

TEST(HashtableOpsTest, TestHashtableImportTwice) {
  const int kResourceId = 42;
  HashtableImportOpModel<std::int64_t, std::string> m(TensorType_INT64,
                                                      TensorType_STRING, 3);
  EXPECT_EQ(m.GetResources().size(), 0);
  m.SetResourceId(kResourceId);
  m.SetKeys({1, 2, 3});
  m.SetStringValues({"1", "2", "3"});
  m.CreateHashtableResource(kResourceId);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  auto& resources = m.GetResources();
  EXPECT_EQ(resources.size(), 1);
  auto* hashtable = resource::GetHashtableResource(&resources, kResourceId);
  EXPECT_TRUE(hashtable != nullptr);
  EXPECT_TRUE(hashtable->GetKeyType() == kTfLiteInt64);
  EXPECT_TRUE(hashtable->GetValueType() == kTfLiteString);
  EXPECT_EQ(hashtable->Size(), 3);
}

// HashtableSizeOpModel creates a model with a HashtableSize op.
template <typename KeyType, typename ValueType>
class HashtableSizeOpModel : public BaseHashtableOpModel {
 public:
  HashtableSizeOpModel(const TensorType key_type, const TensorType value_type) {
    key_type_ = key_type;
    value_type_ = value_type;

    resource_id_ = AddInput({TensorType_RESOURCE, {1}});

    output_ = AddOutput({TensorType_INT64, {1}});

    SetBuiltinOp(BuiltinOperator_HASHTABLE_SIZE,
                 BuiltinOptions_HashtableSizeOptions,
                 CreateHashtableSizeOptions(builder_).Union());
    BuildInterpreter({GetShape(resource_id_)});
  }
};

TEST(HashtableOpsTest, TestHashtableSize) {
  const int kResourceId = 42;
  HashtableSizeOpModel<std::string, std::int64_t> m(TensorType_STRING,
                                                    TensorType_INT64);

  m.SetResourceId(kResourceId);

  InitHashtableResource<std::string, std::int64_t>(
      &m.GetResources(), kResourceId, kTfLiteString, kTfLiteInt64,
      {"4", "5", "6"}, {1, 2, 3});
  ASSERT_EQ(m.Invoke(), kTfLiteOk);

  EXPECT_THAT(m.GetOutput<std::int64_t>(), ElementsAreArray({3}));
  EXPECT_THAT(m.GetOutputShape(), ElementsAreArray({1}));
}

TEST(HashtableOpsTest, TestHashtableSizeNonInitialized) {
  const int kResourceId = 42;
  HashtableSizeOpModel<std::string, std::int64_t> m(TensorType_STRING,
                                                    TensorType_INT64);
  m.SetResourceId(kResourceId);

  // Invoke without hash table initialization.
  EXPECT_NE(m.Invoke(), kTfLiteOk);
}

}  // namespace
}  // namespace tflite