From af420b2b27d90dbee470dfa932f12ae4142abd1a Mon Sep 17 00:00:00 2001
From: Emiel Por <epor@stsci.edu>
Date: Sun, 22 Dec 2024 12:22:03 -0800
Subject: [PATCH] Unify memory block API.

---
 catkit_core/CudaSharedMemory.h |  6 ++++--
 catkit_core/Memory.h           | 16 ++++++++++++++++
 catkit_core/SharedMemory.cpp   |  4 ++--
 catkit_core/SharedMemory.h     |  6 ++++--
 4 files changed, 26 insertions(+), 6 deletions(-)
 create mode 100644 catkit_core/Memory.h

diff --git a/catkit_core/CudaSharedMemory.h b/catkit_core/CudaSharedMemory.h
index 98899e2d5..a5e433f59 100644
--- a/catkit_core/CudaSharedMemory.h
+++ b/catkit_core/CudaSharedMemory.h
@@ -1,6 +1,8 @@
 #ifndef CUDA_SHARED_MEMORY_H
 #define CUDA_SHARED_MEMORY_H
 
+#include "Memory.h"
+
 #include <memory>
 
 #ifdef HAVE_CUDA
@@ -11,7 +13,7 @@ typedef cudaIpcMemHandle_t CudaIpcHandle;
 typedef char CudaIpcHandle[64];
 #endif
 
-class CudaSharedMemory
+class CudaSharedMemory : public Memory
 {
 private:
     CudaSharedMemory(const CudaIpcHandle &ipc_handle, void *device_pointer=nullptr);
@@ -22,7 +24,7 @@ class CudaSharedMemory
     static std::shared_ptr<CudaSharedMemory> Create(size_t num_bytes_in_buffer);
     static std::shared_ptr<CudaSharedMemory> Open(const CudaIpcHandle &ipc_handle);
 
-    void *GetAddress();
+    void *GetAddress(std::size_t offset = 0) override;
 };
 
 #endif // CUDA_SHARED_MEMORY_H
diff --git a/catkit_core/Memory.h b/catkit_core/Memory.h
new file mode 100644
index 000000000..9f20e9ed6
--- /dev/null
+++ b/catkit_core/Memory.h
@@ -0,0 +1,16 @@
+#ifndef MEMORY_H
+#define MEMORY_H
+
+#include <cstddef>
+
+class Memory
+{
+public:
+    virtual ~Memory()
+    {
+    }
+
+    virtual void *GetAddress(std::size_t offset = 0) = 0;
+};
+
+#endif // MEMORY_H
diff --git a/catkit_core/SharedMemory.cpp b/catkit_core/SharedMemory.cpp
index a469b9c40..c8b8d356e 100644
--- a/catkit_core/SharedMemory.cpp
+++ b/catkit_core/SharedMemory.cpp
@@ -82,7 +82,7 @@ SharedMemory::SharedMemory(const std::string &id, FileObject file, bool is_owner
 		throw std::runtime_error("Something went wrong while mapping shared memory file.");
 }
 
-void *SharedMemory::GetAddress()
+void *SharedMemory::GetAddress(std::size_t offset)
 {
-	return m_Buffer;
+	return static_cast<char *>(m_Buffer) + offset;
 }
diff --git a/catkit_core/SharedMemory.h b/catkit_core/SharedMemory.h
index aa6764bd9..ca78b51c6 100644
--- a/catkit_core/SharedMemory.h
+++ b/catkit_core/SharedMemory.h
@@ -1,6 +1,8 @@
 #ifndef SHARED_MEMORY_H
 #define SHARED_MEMORY_H
 
+#include "Memory.h"
+
 #include <memory>
 #include <string>
 
@@ -21,7 +23,7 @@
 	typedef int FileObject;
 #endif
 
-class SharedMemory
+class SharedMemory : public Memory
 {
 private:
 	SharedMemory(const std::string &id, FileObject file, bool is_owner);
@@ -32,7 +34,7 @@ class SharedMemory
 	static std::shared_ptr<SharedMemory> Create(const std::string &id, size_t num_bytes_in_buffer);
 	static std::shared_ptr<SharedMemory> Open(const std::string &id);
 
-	void *GetAddress();
+	void *GetAddress(std::size_t offset = 0) override;
 
 private:
 	std::string m_Id;