From 08a384f7c26b84510d754fd24acbe1e94c5cbb73 Mon Sep 17 00:00:00 2001
From: Jack Zhou <zhoushunjie@baidu.com>
Date: Sun, 6 Nov 2022 20:17:35 +0800
Subject: [PATCH] [Other]Fix the fd tensor copy assignment (#506)

Fix the fd tensor copy assignment
---
 fastdeploy/core/fd_tensor.cc | 22 ++++-----
 tests/core/test_fd_tensor.cc | 89 ++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 11 deletions(-)
 create mode 100644 tests/core/test_fd_tensor.cc

diff --git a/fastdeploy/core/fd_tensor.cc b/fastdeploy/core/fd_tensor.cc
index e98a81e1b..8b739d844 100644
--- a/fastdeploy/core/fd_tensor.cc
+++ b/fastdeploy/core/fd_tensor.cc
@@ -89,8 +89,9 @@ void FDTensor::Squeeze(int64_t axis) {
   size_t ndim = shape.size();
   FDASSERT(axis >= 0 && axis < ndim,
            "The allowed 'axis' must be in range of (0, %lu)!", ndim);
-  FDASSERT(shape[axis]==1,
-           "The No.%ld dimension of shape should be 1, but it is %ld!", (long)axis, (long)shape[axis]);
+  FDASSERT(shape[axis] == 1,
+           "The No.%ld dimension of shape should be 1, but it is %ld!",
+           (long)axis, (long)shape[axis]);
   shape.erase(shape.begin() + axis);
 }
 
@@ -220,9 +221,9 @@ bool FDTensor::ReallocFn(size_t nbytes) {
       return buffer_ != nullptr;
 #else
       FDASSERT(false,
-              "The FastDeploy FDTensor allocator didn't compile under "
-              "-DWITH_GPU=ON,"
-              "so this is an unexpected problem happend.");
+               "The FastDeploy FDTensor allocator didn't compile under "
+               "-DWITH_GPU=ON,"
+               "so this is an unexpected problem happend.");
 #endif
     }
     buffer_ = realloc(buffer_, nbytes);
@@ -316,16 +317,15 @@ FDTensor& FDTensor::operator=(const FDTensor& other) {
     if (other.buffer_ == nullptr) {
       FreeFn();
       buffer_ = nullptr;
+      shape = other.shape;
+      name = other.name;
+      dtype = other.dtype;
+      device = other.device;
     } else {
-      Resize(other.shape);
+      Resize(other.shape, other.dtype, other.name, other.device);
       size_t nbytes = Nbytes();
       CopyBuffer(buffer_, other.buffer_, nbytes);
     }
-
-    shape = other.shape;
-    name = other.name;
-    dtype = other.dtype;
-    device = other.device;
     external_data_ptr = other.external_data_ptr;
   }
   return *this;
diff --git a/tests/core/test_fd_tensor.cc b/tests/core/test_fd_tensor.cc
new file mode 100644
index 000000000..ad4d639e4
--- /dev/null
+++ b/tests/core/test_fd_tensor.cc
@@ -0,0 +1,89 @@
+// Copyright (c) 2022 PaddlePaddle 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 <array>
+#include <cstring>
+#include <vector>
+#include "fastdeploy/core/fd_tensor.h"
+#include "gtest/gtest.h"
+#include "gtest_utils.h"
+
+namespace fastdeploy {
+
+TEST(fastdeploy, fd_tensor_constructor) {
+  CheckShape check_shape;
+  CheckData check_data;
+
+  FDTensor tensor1;
+  check_shape(tensor1.shape, {0});
+  ASSERT_EQ(tensor1.name, "");
+  ASSERT_EQ(tensor1.dtype, FDDataType::INT8);
+  ASSERT_EQ(tensor1.device, Device::CPU);
+
+  std::vector<int> inputs = {2, 4, 3, 7, 1, 5};
+  tensor1.SetExternalData({2, 3}, FDDataType::INT32, inputs.data());
+  ASSERT_EQ(tensor1.dtype, FDDataType::INT32);
+
+  FDTensor tensor2(tensor1);
+  check_shape(tensor1.shape, {2, 3});
+  ASSERT_EQ(tensor2.name, "");
+  ASSERT_EQ(tensor2.dtype, FDDataType::INT32);
+  ASSERT_EQ(tensor2.device, Device::CPU);
+
+  FDTensor tensor3;
+  tensor3.Resize({2, 3}, FDDataType::INT32, "tensor3");
+  check_shape(tensor3.shape, {2, 3});
+  ASSERT_EQ(tensor3.Nbytes(), 24);
+
+  // Copy constructor
+  FDTensor tensor4(tensor3);
+  check_shape(tensor4.shape, {2, 3});
+  ASSERT_EQ(tensor3.Nbytes(), tensor4.Nbytes());
+  check_data(reinterpret_cast<int*>(tensor3.Data()),
+             reinterpret_cast<int*>(tensor4.Data()), tensor4.Numel());
+
+  // Move constructor
+  ASSERT_NE(tensor1.external_data_ptr, nullptr);
+  FDTensor tensor5(std::move(tensor1));
+  ASSERT_EQ(tensor1.external_data_ptr, nullptr);
+  ASSERT_EQ(tensor5.external_data_ptr, inputs.data());
+  check_shape(tensor5.shape, {2, 3});
+}
+
+TEST(fastdeploy, fd_tensor_assignment) {
+  CheckShape check_shape;
+  CheckData check_data;
+
+  FDTensor tensor1("T1");
+  std::vector<int> inputs = {2, 4, 3, 7, 1, 5};
+  tensor1.SetExternalData({2, 3}, FDDataType::INT32, inputs.data());
+
+  FDTensor tensor2;
+  tensor2 = tensor1;
+  ASSERT_EQ(tensor2.name, "T1");
+  ASSERT_EQ(tensor2.dtype, FDDataType::INT32);
+  ASSERT_EQ(tensor2.device, Device::CPU);
+  ASSERT_EQ(tensor2.Data(), inputs.data());
+  check_shape(tensor2.shape, {2, 3});
+
+  FDTensor tensor3;
+  tensor3 = std::move(tensor1);
+  ASSERT_EQ(tensor3.name, "T1");
+  ASSERT_EQ(tensor3.dtype, FDDataType::INT32);
+  ASSERT_EQ(tensor3.device, Device::CPU);
+  ASSERT_EQ(tensor3.Data(), inputs.data());
+  ASSERT_EQ(tensor1.Data(), nullptr);
+}
+
+}  // namespace fastdeploy
\ No newline at end of file