Fix nvmap object ref count error

src/core/hle/service/nvdrv/devices/nvmap.cpp

@@ -69,13 +69,41 @@ VAddr nvmap::GetObjectAddress(u32 handle) const {
     return object->addr;
 }
 
+std::shared_ptr<nvmap::Object> nvmap::GetObject(u32 handle) const {
+    auto itr = handles.find(handle);
+    if (itr != handles.end()) {
+        return itr->second;
+    }
+    return {};
+}
+
+u32 nvmap::IncrementObjectRefCount(u32 handle) {
+    auto object = GetObject(handle);
+    ASSERT(object);
+    auto ref_count = ++object->ref_count;
+    return ref_count;
+}
+
+u32 nvmap::DecrementObjectRefCount(u32 handle) {
+    auto itr = handles.find(handle);
+    if (itr != handles.end()) {
+        auto ref_count = --itr->second->ref_count;
+        if (ref_count == 0) {
+            handles.erase(itr);
+        }
+        return ref_count;
+    }
+
+    UNREACHABLE();
+}
+
 u32 nvmap::CreateObject(u32 size) {
     // Create a new nvmap object and obtain a handle to it.
     auto object = std::make_shared<Object>();
     object->id = next_id++;
     object->size = size;
     object->status = Object::Status::Created;
-    object->refcount = 1;
+    object->ref_count = 1;
 
     const u32 handle = next_handle++;
 
@@ -183,7 +211,7 @@ NvResult nvmap::IocFromId(const std::vector<u8>& input, std::vector<u8>& output)
         return NvResult::BadValue;
     }
 
-    itr->second->refcount++;
+    itr->second->ref_count++;
 
     // Return the existing handle instead of creating a new one.
     params.handle = itr->first;
@@ -250,7 +278,7 @@ NvResult nvmap::IocFree(const std::vector<u8>& input, std::vector<u8>& output) {
         LOG_ERROR(Service_NVDRV, "Object does not exist, handle={:08X}", params.handle);
         return NvResult::BadValue;
     }
-    if (!itr->second->refcount) {
+    if (!itr->second->ref_count) {
         LOG_ERROR(
             Service_NVDRV,
             "There is no references to this object. The object is already freed. handle={:08X}",
@@ -258,11 +286,9 @@ NvResult nvmap::IocFree(const std::vector<u8>& input, std::vector<u8>& output) {
         return NvResult::BadValue;
     }
 
-    itr->second->refcount--;
-
     params.size = itr->second->size;
 
-    if (itr->second->refcount == 0) {
+    if (DecrementObjectRefCount(params.handle) == 0) {
         params.flags = Freed;
         // The address of the nvmap is written to the output if we're finally freeing it, otherwise
         // 0 is written.
@@ -272,8 +298,6 @@ NvResult nvmap::IocFree(const std::vector<u8>& input, std::vector<u8>& output) {
         params.address = 0;
     }
 
-    handles.erase(params.handle);
-
     std::memcpy(output.data(), &params, sizeof(params));
     return NvResult::Success;
 }

src/core/hle/service/nvdrv/devices/nvmap.h

@@ -41,17 +41,14 @@ public:
         u8 kind;
         VAddr addr;
         Status status;
-        u32 refcount;
+        u32 ref_count;
         u32 dma_map_addr;
     };
 
-    std::shared_ptr<Object> GetObject(u32 handle) const {
-        auto itr = handles.find(handle);
-        if (itr != handles.end()) {
-            return itr->second;
-        }
-        return {};
-    }
+    std::shared_ptr<Object> GetObject(u32 handle) const;
+
+    u32 IncrementObjectRefCount(u32 handle);
+    u32 DecrementObjectRefCount(u32 handle);
 
 private:
     /// Id to use for the next handle that is created.

src/core/hle/service/nvflinger/buffer_queue_consumer.cpp

@@ -5,10 +5,12 @@
 // https://cs.android.com/android/platform/superproject/+/android-5.1.1_r38:frameworks/native/libs/gui/BufferQueueConsumer.cpp
 
 #include "common/logging/log.h"
+#include "core/hle/service/nvdrv/devices/nvmap.h"
 #include "core/hle/service/nvflinger/buffer_item.h"
 #include "core/hle/service/nvflinger/buffer_queue_consumer.h"
 #include "core/hle/service/nvflinger/buffer_queue_core.h"
 #include "core/hle/service/nvflinger/producer_listener.h"
+#include "core/hle/service/nvflinger/ui/graphic_buffer.h"
 
 namespace Service::android {
 
@@ -17,6 +19,11 @@ BufferQueueConsumer::BufferQueueConsumer(std::shared_ptr<BufferQueueCore> core_)
 
 BufferQueueConsumer::~BufferQueueConsumer() = default;
 
+void BufferQueueConsumer::SetNVMapInstance(
+    std::shared_ptr<Service::Nvidia::Devices::nvmap> instance) {
+    nvmap = instance;
+}
+
 Status BufferQueueConsumer::AcquireBuffer(BufferItem* out_buffer,
                                           std::chrono::nanoseconds expected_present) {
     std::scoped_lock lock{core->mutex};
@@ -133,6 +140,8 @@ Status BufferQueueConsumer::ReleaseBuffer(s32 slot, u64 frame_number, const Fenc
 
         slots[slot].buffer_state = BufferState::Free;
 
+        nvmap->DecrementObjectRefCount(slots[slot].graphic_buffer->BufferId());
+
         listener = core->connected_producer_listener;
 
         LOG_DEBUG(Service_NVFlinger, "releasing slot {}", slot);

src/core/hle/service/nvflinger/buffer_queue_consumer.h

@@ -13,6 +13,10 @@
 #include "core/hle/service/nvflinger/buffer_queue_defs.h"
 #include "core/hle/service/nvflinger/status.h"
 
+namespace Service::Nvidia::Devices {
+class nvmap;
+} // namespace Service::Nvidia::Devices
+
 namespace Service::android {
 
 class BufferItem;
@@ -24,6 +28,8 @@ public:
     explicit BufferQueueConsumer(std::shared_ptr<BufferQueueCore> core_);
     ~BufferQueueConsumer();
 
+    void SetNVMapInstance(std::shared_ptr<Service::Nvidia::Devices::nvmap> instance);
+
     Status AcquireBuffer(BufferItem* out_buffer, std::chrono::nanoseconds expected_present);
     Status ReleaseBuffer(s32 slot, u64 frame_number, const Fence& release_fence);
     Status Connect(std::shared_ptr<IConsumerListener> consumer_listener, bool controlled_by_app);
@@ -32,6 +38,7 @@ public:
 private:
     std::shared_ptr<BufferQueueCore> core;
     BufferQueueDefs::SlotsType& slots;
+    std::shared_ptr<Service::Nvidia::Devices::nvmap> nvmap;
 };
 
 } // namespace Service::android

src/core/hle/service/nvflinger/buffer_queue_producer.cpp

@@ -14,7 +14,7 @@
 #include "core/hle/kernel/k_writable_event.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/service/kernel_helpers.h"
-#include "core/hle/service/nvdrv/nvdrv.h"
+#include "core/hle/service/nvdrv/devices/nvmap.h"
 #include "core/hle/service/nvflinger/buffer_queue_core.h"
 #include "core/hle/service/nvflinger/buffer_queue_producer.h"
 #include "core/hle/service/nvflinger/consumer_listener.h"
@@ -530,6 +530,8 @@ Status BufferQueueProducer::QueueBuffer(s32 slot, const QueueBufferInput& input,
         item.is_droppable = core->dequeue_buffer_cannot_block || async;
         item.swap_interval = swap_interval;
 
+        nvmap->IncrementObjectRefCount(item.graphic_buffer->BufferId());
+
         sticky_transform = sticky_transform_;
 
         if (core->queue.empty()) {
@@ -920,4 +922,9 @@ Kernel::KReadableEvent& BufferQueueProducer::GetNativeHandle() {
     return buffer_wait_event->GetReadableEvent();
 }
 
+void BufferQueueProducer::SetNVMapInstance(
+    std::shared_ptr<Service::Nvidia::Devices::nvmap> instance) {
+    nvmap = instance;
+}
+
 } // namespace Service::android

src/core/hle/service/nvflinger/buffer_queue_producer.h

@@ -31,6 +31,10 @@ namespace Service::KernelHelpers {
 class ServiceContext;
 } // namespace Service::KernelHelpers
 
+namespace Service::Nvidia::Devices {
+class nvmap;
+} // namespace Service::Nvidia::Devices
+
 namespace Service::android {
 
 class BufferQueueCore;
@@ -46,6 +50,8 @@ public:
 
     Kernel::KReadableEvent& GetNativeHandle() override;
 
+    void SetNVMapInstance(std::shared_ptr<Service::Nvidia::Devices::nvmap> instance);
+
 public:
     Status RequestBuffer(s32 slot, std::shared_ptr<GraphicBuffer>* buf);
     Status SetBufferCount(s32 buffer_count);
@@ -70,6 +76,7 @@ private:
 
     Kernel::KEvent* buffer_wait_event{};
     Service::KernelHelpers::ServiceContext& service_context;
+    std::shared_ptr<Service::Nvidia::Devices::nvmap> nvmap;
 
     std::shared_ptr<BufferQueueCore> core;
     BufferQueueDefs::SlotsType& slots;

src/core/hle/service/nvflinger/nvflinger.cpp

@@ -14,6 +14,7 @@
 #include "core/core_timing.h"
 #include "core/hle/kernel/k_readable_event.h"
 #include "core/hle/service/nvdrv/devices/nvdisp_disp0.h"
+#include "core/hle/service/nvdrv/devices/nvmap.h"
 #include "core/hle/service/nvdrv/nvdrv.h"
 #include "core/hle/service/nvflinger/buffer_item_consumer.h"
 #include "core/hle/service/nvflinger/buffer_queue_core.h"
@@ -132,7 +133,9 @@ std::optional<u64> NVFlinger::CreateLayer(u64 display_id) {
 
 void NVFlinger::CreateLayerAtId(VI::Display& display, u64 layer_id) {
     const auto buffer_id = next_buffer_queue_id++;
-    display.CreateLayer(layer_id, buffer_id);
+    auto nvmap = nvdrv->GetDevice<Nvidia::Devices::nvmap>("/dev/nvmap");
+    ASSERT(nvmap);
+    display.CreateLayer(layer_id, buffer_id, nvmap);
 }
 
 void NVFlinger::CloseLayer(u64 layer_id) {

src/core/hle/service/vi/display/vi_display.cpp

@@ -63,10 +63,13 @@ void Display::SignalVSyncEvent() {
     vsync_event->GetWritableEvent().Signal();
 }
 
-void Display::CreateLayer(u64 layer_id, u32 binder_id) {
+void Display::CreateLayer(u64 layer_id, u32 binder_id,
+                          std::shared_ptr<Service::Nvidia::Devices::nvmap> instance) {
     ASSERT_MSG(layers.empty(), "Only one layer is supported per display at the moment");
 
     auto [core, producer, consumer] = CreateBufferQueue(service_context);
+    producer->SetNVMapInstance(instance);
+    consumer->SetNVMapInstance(instance);
 
     auto buffer_item_consumer = std::make_shared<android::BufferItemConsumer>(std::move(consumer));
     buffer_item_consumer->Connect(false);

src/core/hle/service/vi/display/vi_display.h

@@ -26,6 +26,10 @@ namespace Service::NVFlinger {
 class HosBinderDriverServer;
 }
 
+namespace Service::Nvidia::Devices {
+class nvmap;
+}
+
 namespace Service::VI {
 
 class Layer;
@@ -84,7 +88,8 @@ public:
     /// @param layer_id The ID to assign to the created layer.
     /// @param binder_id The ID assigned to the buffer queue.
     ///
-    void CreateLayer(u64 layer_id, u32 binder_id);
+    void CreateLayer(u64 layer_id, u32 binder_id,
+                     std::shared_ptr<Service::Nvidia::Devices::nvmap> instance);
 
     /// Closes and removes a layer from this display with the given ID.
     ///