#include "tensorflow/core/platform/test.h"

namespace mlir {
namespace TFL {
namespace {

TEST(SizeUtilTest, TestConvertsSize) {
  ASSERT_EQ(ConvertToTfliteSize(1), 1);
  ASSERT_EQ(ConvertToTfliteSize(-1), -1);
  ASSERT_EQ(ConvertToTfliteSize(mlir::ShapedType::kDynamic), -1);
}

}  // namespace
}  // namespace TFL

#include "tensorflow/compiler/mlir/lite/utils/size_utils.h"

#include <cstdint>

#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project

namespace mlir {
namespace TFL {

int32_t ConvertToTfliteSize(int64_t size) {
  return mlir::ShapedType::isDynamic(size) ? -1 : static_cast<int32_t>(size);
}

}  // namespace TFL

namespace mlir {
namespace TFL {

// Converts a TF size (64-bit) to TFLite (32-bit) and properly converts TF's
// value for dynamic size (`std::numeric_limits<int64_t>::min()`) to the
// TFLite-specific value.
int32_t ConvertToTfliteSize(int64_t size);

}  // namespace TFL
}  // namespace mlir