From da62e92784fbae15942ae4be57fb7eb3b4a9b5b5 Mon Sep 17 00:00:00 2001
From: ameerj <52414509+ameerj@users.noreply.github.com>
Date: Thu, 6 May 2021 11:20:52 -0400
Subject: [PATCH] nvflinger: Create layers when they are queried but not found

Fixes Shantae softlock on boot.
---
 src/core/hle/service/nvflinger/nvflinger.cpp | 30 +++++++++++++++++---
 src/core/hle/service/nvflinger/nvflinger.h   | 10 ++++++-
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/src/core/hle/service/nvflinger/nvflinger.cpp b/src/core/hle/service/nvflinger/nvflinger.cpp
index c43593e7f..810312dc4 100644
--- a/src/core/hle/service/nvflinger/nvflinger.cpp
+++ b/src/core/hle/service/nvflinger/nvflinger.cpp
@@ -139,11 +139,15 @@ std::optional<u64> NVFlinger::CreateLayer(u64 display_id) {
     }
 
     const u64 layer_id = next_layer_id++;
+    CreateLayerAtId(*display, layer_id);
+    return layer_id;
+}
+
+void NVFlinger::CreateLayerAtId(VI::Display& display, u64 layer_id) {
     const u32 buffer_queue_id = next_buffer_queue_id++;
     buffer_queues.emplace_back(
         std::make_unique<BufferQueue>(system.Kernel(), buffer_queue_id, layer_id));
-    display->CreateLayer(layer_id, *buffer_queues.back());
-    return layer_id;
+    display.CreateLayer(layer_id, *buffer_queues.back());
 }
 
 void NVFlinger::CloseLayer(u64 layer_id) {
@@ -154,9 +158,9 @@ void NVFlinger::CloseLayer(u64 layer_id) {
     }
 }
 
-std::optional<u32> NVFlinger::FindBufferQueueId(u64 display_id, u64 layer_id) const {
+std::optional<u32> NVFlinger::FindBufferQueueId(u64 display_id, u64 layer_id) {
     const auto lock_guard = Lock();
-    const auto* const layer = FindLayer(display_id, layer_id);
+    const auto* const layer = FindOrCreateLayer(display_id, layer_id);
 
     if (layer == nullptr) {
         return std::nullopt;
@@ -232,6 +236,24 @@ const VI::Layer* NVFlinger::FindLayer(u64 display_id, u64 layer_id) const {
     return display->FindLayer(layer_id);
 }
 
+VI::Layer* NVFlinger::FindOrCreateLayer(u64 display_id, u64 layer_id) {
+    auto* const display = FindDisplay(display_id);
+
+    if (display == nullptr) {
+        return nullptr;
+    }
+
+    auto* layer = display->FindLayer(layer_id);
+
+    if (layer == nullptr) {
+        LOG_DEBUG(Service, "Layer at id {} not found. Trying to create it.", layer_id);
+        CreateLayerAtId(*display, layer_id);
+        return display->FindLayer(layer_id);
+    }
+
+    return layer;
+}
+
 void NVFlinger::Compose() {
     for (auto& display : displays) {
         // Trigger vsync for this display at the end of drawing
diff --git a/src/core/hle/service/nvflinger/nvflinger.h b/src/core/hle/service/nvflinger/nvflinger.h
index 6fe2c7f2a..ebc82c688 100644
--- a/src/core/hle/service/nvflinger/nvflinger.h
+++ b/src/core/hle/service/nvflinger/nvflinger.h
@@ -67,7 +67,7 @@ public:
     /// Finds the buffer queue ID of the specified layer in the specified display.
     ///
     /// If an invalid display ID or layer ID is provided, then an empty optional is returned.
-    [[nodiscard]] std::optional<u32> FindBufferQueueId(u64 display_id, u64 layer_id) const;
+    [[nodiscard]] std::optional<u32> FindBufferQueueId(u64 display_id, u64 layer_id);
 
     /// Gets the vsync event for the specified display.
     ///
@@ -100,6 +100,14 @@ private:
     /// Finds the layer identified by the specified ID in the desired display.
     [[nodiscard]] const VI::Layer* FindLayer(u64 display_id, u64 layer_id) const;
 
+    /// Finds the layer identified by the specified ID in the desired display,
+    /// or creates the layer if it is not found.
+    /// To be used when the system expects the specified ID to already exist.
+    [[nodiscard]] VI::Layer* FindOrCreateLayer(u64 display_id, u64 layer_id);
+
+    /// Creates a layer with the specified layer ID in the desired display.
+    void CreateLayerAtId(VI::Display& display, u64 layer_id);
+
     static void VSyncThread(NVFlinger& nv_flinger);
 
     void SplitVSync();