// WARNING! WARNING! WARNING!
// This file is a temporary hack to enable development of pytorch quantization
//
// It's a stub for wrapping arbitrary cpp types in TorchScript. Proper
// implementation (under development) is to use TorchScript custom types.
// In the meantime, we abuse ByteTensor with custom deleter for this purpose.
//
// Template argument <T> has to be registered with CAFFE_KNOWN_TYPE mechanism.

#include "ATen/ATen.h"

namespace at {
namespace cpp_custom_type_hack {

template<typename T>
T& cast(const Tensor& packed) {
  AT_CHECK(
      packed.scalar_type() == kByte, "Expected temporary cpp type wrapper");
  AT_CHECK(
      packed.storage().data_ptr().get_deleter() ==
          caffe2::TypeMeta::Make<T>().deleteFn(),
      "Expected temporary cpp type wrapper of type ",
      caffe2::TypeMeta::TypeName<T>());
  return *reinterpret_cast<T*>(packed.storage().data_ptr().get());
}

template<typename T>
Tensor create(std::unique_ptr<T> ptr) {
  // We store this instance away in a Tensor and register a deleter function
  // so that we do not leak memory. On the other side, we pull out the storage's
  // data_ptr and get the right typed pointer.
  void* raw_ptr = ptr.release();
  at::DataPtr at_ptr(
      raw_ptr,
      raw_ptr,
      caffe2::TypeMeta::Make<T>().deleteFn(),
      at::kCPU);

  // size doesn't really matter, but we can align it to the actual size
  // returning variables because one likely want to use this hack from python
  auto retval = at::empty(
      {sizeof(T)},
      at::device(kCPU).dtype(at::kByte).is_variable(true).requires_grad(false));
  retval.storage().set_data_ptr(std::move(at_ptr));
  return retval;
}
}
}
