1 files changed, 49 insertions, 13 deletions
diff --git a/neuralnetworks/aidl/utils/src/Utils.cpp b/neuralnetworks/aidl/utils/src/Utils.cpp
index 95516c854b..03407be4ce 100644
--- a/neuralnetworks/aidl/utils/src/Utils.cpp
+++ b/neuralnetworks/aidl/utils/src/Utils.cpp
@@ -16,12 +16,20 @@
 
 #include "Utils.h"
 
+#include <aidl/android/hardware/common/Ashmem.h>
+#include <aidl/android/hardware/common/MappableFile.h>
+#include <aidl/android/hardware/graphics/common/HardwareBuffer.h>
+#include <android/binder_auto_utils.h>
 #include <android/binder_status.h>
 #include <nnapi/Result.h>
+#include <nnapi/SharedMemory.h>
 
 namespace aidl::android::hardware::neuralnetworks::utils {
 namespace {
 
+nn::GeneralResult<ndk::ScopedFileDescriptor> clone(const ndk::ScopedFileDescriptor& fd);
+using utils::clone;
+
 template <typename Type>
 nn::GeneralResult<std::vector<Type>> cloneVec(const std::vector<Type>& arguments) {
     std::vector<Type> clonedObjects;
@@ -37,24 +45,52 @@ nn::GeneralResult<std::vector<Type>> clone(const std::vector<Type>& arguments) {
     return cloneVec(arguments);
 }
 
+nn::GeneralResult<ndk::ScopedFileDescriptor> clone(const ndk::ScopedFileDescriptor& fd) {
+    auto duplicatedFd = NN_TRY(nn::dupFd(fd.get()));
+    return ndk::ScopedFileDescriptor(duplicatedFd.release());
+}
+
+nn::GeneralResult<common::NativeHandle> clone(const common::NativeHandle& handle) {
+    return common::NativeHandle{
+            .fds = NN_TRY(cloneVec(handle.fds)),
+            .ints = handle.ints,
+    };
+}
+
 }  // namespace
 
 nn::GeneralResult<Memory> clone(const Memory& memory) {
-    common::NativeHandle nativeHandle;
-    nativeHandle.ints = memory.handle.ints;
-    nativeHandle.fds.reserve(memory.handle.fds.size());
-    for (const auto& fd : memory.handle.fds) {
-        const int newFd = dup(fd.get());
-        if (newFd < 0) {
-            return NN_ERROR() << "Couldn't dup a file descriptor";
+    switch (memory.getTag()) {
+        case Memory::Tag::ashmem: {
+            const auto& ashmem = memory.get<Memory::Tag::ashmem>();
+            auto handle = common::Ashmem{
+                    .fd = NN_TRY(clone(ashmem.fd)),
+                    .size = ashmem.size,
+            };
+            return Memory::make<Memory::Tag::ashmem>(std::move(handle));
+        }
+        case Memory::Tag::mappableFile: {
+            const auto& memFd = memory.get<Memory::Tag::mappableFile>();
+            auto handle = common::MappableFile{
+                    .length = memFd.length,
+                    .prot = memFd.prot,
+                    .fd = NN_TRY(clone(memFd.fd)),
+                    .offset = memFd.offset,
+            };
+            return Memory::make<Memory::Tag::mappableFile>(std::move(handle));
+        }
+        case Memory::Tag::hardwareBuffer: {
+            const auto& hardwareBuffer = memory.get<Memory::Tag::hardwareBuffer>();
+            auto handle = graphics::common::HardwareBuffer{
+                    .description = hardwareBuffer.description,
+                    .handle = NN_TRY(clone(hardwareBuffer.handle)),
+            };
+            return Memory::make<Memory::Tag::hardwareBuffer>(std::move(handle));
         }
-        nativeHandle.fds.emplace_back(newFd);
     }
-    return Memory{
-            .handle = std::move(nativeHandle),
-            .size = memory.size,
-            .name = memory.name,
-    };
+    return (NN_ERROR() << "Unrecognized Memory::Tag: " << memory.getTag())
+            .
+            operator nn::GeneralResult<Memory>();
 }
 
 nn::GeneralResult<RequestMemoryPool> clone(const RequestMemoryPool& requestPool) {