From c1374409285c6b17df0daf37b427792e25b1cd06 Mon Sep 17 00:00:00 2001
From: Keith Colbert <kcolbert@gmail.com>
Date: Thu, 7 Feb 2019 23:09:50 -0500
Subject: [PATCH] Add alias and assign for Ctxt memory allocations

---
 cufhe/include/cufhe.h  |  3 +++
 cufhe/lib/cufhe.cc     | 36 ++++++++++++++++++++++++++++--------
 cufhe/lib/cufhe_gpu.cu | 24 ++++++++++++++++--------
 3 files changed, 47 insertions(+), 16 deletions(-)

diff --git a/cufhe/include/cufhe.h b/cufhe/include/cufhe.h
index 0693dc3..e63a035 100644
--- a/cufhe/include/cufhe.h
+++ b/cufhe/include/cufhe.h
@@ -106,6 +106,9 @@ struct PubKey {
 struct Ctxt {
   Ctxt(bool is_alias = false);
   ~Ctxt();
+  Ctxt(const Ctxt& that) = delete;
+  Ctxt& operator=(const Ctxt& that) = delete;
+  void assign(void* host_ptr, void* device_ptr);
   LWESample* lwe_sample_;
   MemoryDeleter lwe_sample_deleter_;
   LWESample* lwe_sample_device_;
diff --git a/cufhe/lib/cufhe.cc b/cufhe/lib/cufhe.cc
index 797c19b..df81544 100644
--- a/cufhe/lib/cufhe.cc
+++ b/cufhe/lib/cufhe.cc
@@ -279,25 +279,45 @@ PubKey::~PubKey() {
 }
 
 Ctxt::~Ctxt() {
-  if(lwe_sample_ != nullptr)
-  {
-    lwe_sample_deleter_(lwe_sample_->data());
-    lwe_sample_deleter_ = nullptr;
+  if(lwe_sample_ != nullptr) {
+    if(lwe_sample_deleter_ != nullptr) {
+      lwe_sample_deleter_(lwe_sample_->data());
+      lwe_sample_deleter_ = nullptr;
+    }
+
     lwe_sample_->set_data(nullptr);
     delete lwe_sample_;
     lwe_sample_ = nullptr;
   }
 
-  if(lwe_sample_device_ != nullptr)
-  {
-    lwe_sample_device_deleter_(lwe_sample_device_->data());
-    lwe_sample_device_deleter_ = nullptr;
+  if(lwe_sample_device_ != nullptr) {
+    if(lwe_sample_device_deleter_ != nullptr) {
+      lwe_sample_device_deleter_(lwe_sample_device_->data());
+      lwe_sample_device_deleter_ = nullptr;
+    }
+
     lwe_sample_device_->set_data(nullptr);
     delete lwe_sample_device_;
     lwe_sample_device_ = nullptr;
   }
 }
 
+void Ctxt::assign(void* host_ptr, void* device_ptr) {
+  if(lwe_sample_deleter_ != nullptr) {
+    lwe_sample_deleter_(lwe_sample_->data());
+    lwe_sample_deleter_ = nullptr;
+  }
+
+  lwe_sample_->set_data((LWESample::PointerType)host_ptr);
+
+  if(lwe_sample_device_deleter_ != nullptr) {
+    lwe_sample_device_deleter_(lwe_sample_device_->data());
+    lwe_sample_device_deleter_ = nullptr;
+  }
+
+  lwe_sample_device_->set_data((LWESample::PointerType)device_ptr);
+}
+
 void SetSeed(uint32_t seed) {
   srand(seed);
   RandomGeneratorSetSeed(&seed, 1);
diff --git a/cufhe/lib/cufhe_gpu.cu b/cufhe/lib/cufhe_gpu.cu
index 0b9d8ab..568f298 100644
--- a/cufhe/lib/cufhe_gpu.cu
+++ b/cufhe/lib/cufhe_gpu.cu
@@ -32,15 +32,23 @@ Ctxt::Ctxt(bool is_alias) {
   Param* param = GetDefaultParam();
 
   lwe_sample_ = new LWESample(param->lwe_n_);
-  //pair = AllocatorBoth::New(lwe_sample_->SizeMalloc());
-  pair = AllocatorCPU::New(lwe_sample_->SizeMalloc());
-  lwe_sample_->set_data((LWESample::PointerType)pair.first);
-  lwe_sample_deleter_ = pair.second;
-
   lwe_sample_device_ = new LWESample(param->lwe_n_);
-  pair = AllocatorGPU::New(lwe_sample_device_->SizeMalloc());
-  lwe_sample_device_->set_data((LWESample::PointerType)pair.first);
-  lwe_sample_device_deleter_ = pair.second;
+
+  if (is_alias) {
+    lwe_sample_->set_data(nullptr);
+    lwe_sample_deleter_ = nullptr;
+    lwe_sample_device_->set_data(nullptr);
+    lwe_sample_device_deleter_ = nullptr;
+  } else {
+    //pair = AllocatorBoth::New(lwe_sample_->SizeMalloc());
+    pair = AllocatorCPU::New(lwe_sample_->SizeMalloc());
+    lwe_sample_->set_data((LWESample::PointerType)pair.first);
+    lwe_sample_deleter_ = pair.second;
+
+    pair = AllocatorGPU::New(lwe_sample_device_->SizeMalloc());
+    lwe_sample_device_->set_data((LWESample::PointerType)pair.first);
+    lwe_sample_device_deleter_ = pair.second;
+  }
 }
 
 } // namespace cufhe