misc: Add strict typing to attachment manager

include/aquamarine/backend/DRM.hpp

@@ -25,9 +25,6 @@ namespace Aquamarine {
         virtual ~CDRMBufferAttachment() {
             ;
         }
-        virtual eAttachmentType type() {
-            return AQ_ATTACHMENT_DRM_BUFFER;
-        }
 
         Hyprutils::Memory::CSharedPointer<CDRMFB> fb;
     };
@@ -40,9 +37,6 @@ namespace Aquamarine {
         virtual ~CDRMBufferUnimportable() {
             ;
         }
-        virtual eAttachmentType type() {
-            return AQ_ATTACHMENT_DRM_KMS_UNIMPORTABLE;
-        }
     };
 
     class CDRMLease {

include/aquamarine/misc/Attachment.hpp

@@ -2,33 +2,49 @@
 
 #include <vector>
 #include <hyprutils/memory/SharedPtr.hpp>
+#include <unordered_map>
+#include <typeindex>
 
 namespace Aquamarine {
-    enum eAttachmentType : uint32_t {
-        AQ_ATTACHMENT_DRM_BUFFER = 0,
-        AQ_ATTACHMENT_DRM_KMS_UNIMPORTABLE,
-        AQ_ATTACHMENT_DRM_RENDERER_DATA,
-    };
-
     class IAttachment {
       public:
         virtual ~IAttachment() {
             ;
         }
-
-        virtual eAttachmentType type() = 0;
     };
 
+    template <typename T>
+    concept AttachmentConcept = std::is_base_of_v<IAttachment, T>;
+
+    // CAttachmentManager is a registry for arbitrary attachment types.
+    // Any type implementing IAttachment can be added, retrieved, and removed from the registry.
+    // However, only one attachment of a given type is permitted.
     class CAttachmentManager {
       public:
-        bool                                           has(eAttachmentType type);
-        Hyprutils::Memory::CSharedPointer<IAttachment> get(eAttachmentType type);
-        void                                           add(Hyprutils::Memory::CSharedPointer<IAttachment> attachment);
-        void                                           remove(Hyprutils::Memory::CSharedPointer<IAttachment> attachment);
-        void                                           removeByType(eAttachmentType type);
-        void                                           clear();
+        template <AttachmentConcept T>
+        bool has() const {
+            return attachments.contains(typeid(T));
+        }
+        template <AttachmentConcept T>
+        Hyprutils::Memory::CSharedPointer<T> get() const {
+            auto it = attachments.find(typeid(T));
+            if (it == attachments.end())
+                return nullptr;
+            // Reinterpret SP<IAttachment> into SP<T>.
+            // This is safe because we looked up this attachment by typeid(T),
+            // so it must be an SP<T>.
+            return *reinterpret_cast<const Hyprutils::Memory::CSharedPointer<T>*>(&it->second);
+        }
+        // Also removes the previous attachment of the same type if one exists
+        void add(Hyprutils::Memory::CSharedPointer<IAttachment> attachment);
+        void remove(Hyprutils::Memory::CSharedPointer<IAttachment> attachment);
+        template <AttachmentConcept T>
+        void removeByType() {
+            attachments.erase(typeid(T));
+        }
+        void clear();
 
       private:
-        std::vector<Hyprutils::Memory::CSharedPointer<IAttachment>> attachments;
+        std::unordered_map<std::type_index, Hyprutils::Memory::CSharedPointer<IAttachment>> attachments;
     };
 };

src/backend/drm/DRM.cpp

@@ -1621,7 +1621,7 @@ bool Aquamarine::CDRMOutput::commitState(bool onlyTest) {
         return false;
     }
 
-    if ((COMMITTED & COutputState::eOutputStateProperties::AQ_OUTPUT_STATE_BUFFER) && STATE.buffer->attachments.has(AQ_ATTACHMENT_DRM_KMS_UNIMPORTABLE)) {
+    if ((COMMITTED & COutputState::eOutputStateProperties::AQ_OUTPUT_STATE_BUFFER) && STATE.buffer->attachments.has<CDRMBufferUnimportable>()) {
         TRACE(backend->backend->log(AQ_LOG_TRACE, "drm: Cannot commit a KMS-unimportable buffer."));
         return false;
     }
@@ -2003,9 +2003,8 @@ SP<CDRMFB> Aquamarine::CDRMFB::create(SP<IBuffer> buffer_, Hyprutils::Memory::CW
     if (isNew)
         *isNew = true;
 
-    if (buffer_->attachments.has(AQ_ATTACHMENT_DRM_BUFFER)) {
-        auto at = (CDRMBufferAttachment*)buffer_->attachments.get(AQ_ATTACHMENT_DRM_BUFFER).get();
-        fb      = at->fb;
+    if (auto at = buffer_->attachments.get<CDRMBufferAttachment>()) {
+        fb = at->fb;
         TRACE(backend_->log(AQ_LOG_TRACE, std::format("drm: CDRMFB: buffer has drmfb attachment with fb {:x}", (uintptr_t)fb.get())));
     }
 
@@ -2036,7 +2035,7 @@ void Aquamarine::CDRMFB::import() {
         return;
     }
 
-    if (buffer->attachments.has(AQ_ATTACHMENT_DRM_KMS_UNIMPORTABLE)) {
+    if (buffer->attachments.has<CDRMBufferUnimportable>()) {
         backend->backend->log(AQ_LOG_ERROR, "drm: Buffer submitted is unimportable");
         return;
     }

src/backend/drm/Renderer.cpp

@@ -698,14 +698,10 @@ constexpr GLenum PIXEL_BUFFER_FORMAT = GL_RGBA;
 
 void             CDRMRenderer::readBuffer(Hyprutils::Memory::CSharedPointer<IBuffer> buf, std::span<uint8_t> out) {
     CEglContextGuard eglContext(*this);
-    auto             hadAttachment = buf->attachments.get(AQ_ATTACHMENT_DRM_RENDERER_DATA);
-    auto             att           = (CDRMRendererBufferAttachment*)hadAttachment.get();
-    if (!hadAttachment) {
-        // should never remove anything, but JIC. We'll leak an EGLImage if this removes anything.
-        buf->attachments.removeByType(AQ_ATTACHMENT_DRM_RENDERER_DATA);
-        auto newAttachment = makeShared<CDRMRendererBufferAttachment>(self, buf, nullptr, 0, 0, SGLTex{}, std::vector<uint8_t>());
-        att                = newAttachment.get();
-        buf->attachments.add(newAttachment);
+    auto             att = buf->attachments.get<CDRMRendererBufferAttachment>();
+    if (!att) {
+        att = makeShared<CDRMRendererBufferAttachment>(self, buf, nullptr, 0, 0, SGLTex{}, std::vector<uint8_t>());
+        buf->attachments.add(att);
     }
 
     auto dma = buf->dmabuf();
@@ -886,31 +882,28 @@ CDRMRenderer::SBlitResult CDRMRenderer::blit(SP<IBuffer> from, SP<IBuffer> to, S
     auto               fromDma = from->dmabuf();
     std::span<uint8_t> intermediateBuf;
     {
-        auto attachment = from->attachments.get(AQ_ATTACHMENT_DRM_RENDERER_DATA);
+        auto attachment = from->attachments.get<CDRMRendererBufferAttachment>();
         if (attachment) {
             TRACE(backend->log(AQ_LOG_TRACE, "EGL (blit): From attachment found"));
-            auto att        = (CDRMRendererBufferAttachment*)attachment.get();
-            fromTex         = att->tex;
-            intermediateBuf = att->intermediateBuf;
+            fromTex         = attachment->tex;
+            intermediateBuf = attachment->intermediateBuf;
         }
 
         if (!fromTex.image && intermediateBuf.empty()) {
             backend->log(AQ_LOG_DEBUG, "EGL (blit): No attachment in from, creating a new image");
             fromTex = glTex(from);
 
-            auto newAttachment = makeShared<CDRMRendererBufferAttachment>(self, from, nullptr, 0, 0, fromTex, std::vector<uint8_t>());
+            attachment = makeShared<CDRMRendererBufferAttachment>(self, from, nullptr, 0, 0, fromTex, std::vector<uint8_t>());
+            from->attachments.add(attachment);
+
             if (!fromTex.image && primaryRenderer) {
                 backend->log(AQ_LOG_DEBUG, "EGL (blit): Failed to create image from source buffer directly, allocating intermediate buffer");
                 static_assert(PIXEL_BUFFER_FORMAT == GL_RGBA); // If the pixel buffer format changes, the below size calculation probably needs to as well.
-                newAttachment->intermediateBuf.resize(fromDma.size.x * fromDma.size.y * 4);
-                intermediateBuf = newAttachment->intermediateBuf;
+                attachment->intermediateBuf.resize(fromDma.size.x * fromDma.size.y * 4);
+                intermediateBuf = attachment->intermediateBuf;
                 fromTex.target  = GL_TEXTURE_2D;
                 GLCALL(glGenTextures(1, &fromTex.texid));
             }
-
-            // should never remove anything, but JIC. We'll leak an EGLImage if this removes anything.
-            from->attachments.removeByType(AQ_ATTACHMENT_DRM_RENDERER_DATA);
-            from->attachments.add(newAttachment);
         }
 
         if (!intermediateBuf.empty() && primaryRenderer) {
@@ -936,13 +929,12 @@ CDRMRenderer::SBlitResult CDRMRenderer::blit(SP<IBuffer> from, SP<IBuffer> to, S
     }
 
     {
-        auto attachment = to->attachments.get(AQ_ATTACHMENT_DRM_RENDERER_DATA);
+        auto attachment = to->attachments.get<CDRMRendererBufferAttachment>();
         if (attachment) {
             TRACE(backend->log(AQ_LOG_TRACE, "EGL (blit): To attachment found"));
-            auto att = (CDRMRendererBufferAttachment*)attachment.get();
-            rboImage = att->eglImage;
-            fboID    = att->fbo;
-            rboID    = att->rbo;
+            rboImage = attachment->eglImage;
+            fboID    = attachment->fbo;
+            rboID    = attachment->rbo;
         }
 
         if (!rboImage) {
@@ -968,8 +960,6 @@ CDRMRenderer::SBlitResult CDRMRenderer::blit(SP<IBuffer> from, SP<IBuffer> to, S
                 return {};
             }
 
-            // should never remove anything, but JIC. We'll leak an RBO and FBO if this removes anything.
-            to->attachments.removeByType(AQ_ATTACHMENT_DRM_RENDERER_DATA);
             to->attachments.add(makeShared<CDRMRendererBufferAttachment>(self, to, rboImage, fboID, rboID, SGLTex{}, std::vector<uint8_t>()));
         }
     }

src/backend/drm/Renderer.hpp

@@ -29,9 +29,6 @@ namespace Aquamarine {
         virtual ~CDRMRendererBufferAttachment() {
             ;
         }
-        virtual eAttachmentType type() {
-            return AQ_ATTACHMENT_DRM_RENDERER_DATA;
-        }
 
         EGLImageKHR                                   eglImage = nullptr;
         GLuint                                        fbo = 0, rbo = 0;

src/misc/Attachment.cpp

@@ -4,32 +4,16 @@ using namespace Aquamarine;
 using namespace Hyprutils::Memory;
 #define SP CSharedPointer
 
-bool Aquamarine::CAttachmentManager::has(eAttachmentType type) {
-    for (auto const& a : attachments) {
-        if (a->type() == type)
-            return true;
-    }
-    return false;
-}
-
-SP<IAttachment> Aquamarine::CAttachmentManager::get(eAttachmentType type) {
-    for (auto const& a : attachments) {
-        if (a->type() == type)
-            return a;
-    }
-    return nullptr;
-}
-
 void Aquamarine::CAttachmentManager::add(SP<IAttachment> attachment) {
-    attachments.emplace_back(attachment);
+    const IAttachment& att   = *attachment;
+    attachments[typeid(att)] = attachment;
 }
 
 void Aquamarine::CAttachmentManager::remove(SP<IAttachment> attachment) {
-    std::erase(attachments, attachment);
-}
-
-void Aquamarine::CAttachmentManager::removeByType(eAttachmentType type) {
-    std::erase_if(attachments, [type](const auto& e) { return e->type() == type; });
+    const IAttachment& att = *attachment;
+    auto               it  = attachments.find(typeid(att));
+    if (it != attachments.end())
+        attachments.erase(it);
 }
 
 void Aquamarine::CAttachmentManager::clear() {