From 04b123f16ec54ed9a614ee5f045b78d8feec9ce8 Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Wed, 23 Oct 2019 15:44:00 -0500
Subject: [PATCH] Adjust opened SaveDataInfoReaders if the collection is
 modified

---
 src/LibHac/Fs/Shim/SaveDataManagement.cs      |  7 ++-
 src/LibHac/FsService/ISaveDataInfoReader.cs   |  2 +-
 src/LibHac/FsService/SaveDataIndexer.cs       | 56 ++++++++++++++++++-
 .../FsService/SaveDataInfoFilterReader.cs     |  5 ++
 4 files changed, 67 insertions(+), 3 deletions(-)

diff --git a/src/LibHac/Fs/Shim/SaveDataManagement.cs b/src/LibHac/Fs/Shim/SaveDataManagement.cs
index eb5ed709..12aca8b2 100644
--- a/src/LibHac/Fs/Shim/SaveDataManagement.cs
+++ b/src/LibHac/Fs/Shim/SaveDataManagement.cs
@@ -227,7 +227,7 @@ namespace LibHac.Fs.Shim
         }
     }
 
-    public struct SaveDataIterator
+    public struct SaveDataIterator : IDisposable
     {
         private FileSystemClient FsClient { get; }
         private ISaveDataInfoReader Reader { get; }
@@ -259,5 +259,10 @@ namespace LibHac.Fs.Shim
 
             return rc;
         }
+
+        public void Dispose()
+        {
+            Reader?.Dispose();
+        }
     }
 }
diff --git a/src/LibHac/FsService/ISaveDataInfoReader.cs b/src/LibHac/FsService/ISaveDataInfoReader.cs
index d0ece142..2dc3abb5 100644
--- a/src/LibHac/FsService/ISaveDataInfoReader.cs
+++ b/src/LibHac/FsService/ISaveDataInfoReader.cs
@@ -2,7 +2,7 @@
 
 namespace LibHac.FsService
 {
-    public interface ISaveDataInfoReader
+    public interface ISaveDataInfoReader : IDisposable
     {
         Result ReadSaveDataInfo(out long readCount, Span<byte> saveDataInfoBuffer);
     }
diff --git a/src/LibHac/FsService/SaveDataIndexer.cs b/src/LibHac/FsService/SaveDataIndexer.cs
index 909f57fb..345d09fa 100644
--- a/src/LibHac/FsService/SaveDataIndexer.cs
+++ b/src/LibHac/FsService/SaveDataIndexer.cs
@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Diagnostics;
+using System.Linq;
 using System.Runtime.CompilerServices;
 using System.Runtime.InteropServices;
 using LibHac.Common;
@@ -26,6 +27,7 @@ namespace LibHac.FsService
         private bool IsKvdbLoaded { get; set; }
         private ulong LastPublishedId { get; set; }
         private int Version { get; set; }
+        private List<SaveDataInfoReader> OpenedReaders { get; } = new List<SaveDataInfoReader>();
 
         public SaveDataIndexer(FileSystemClient fsClient, string mountName, SaveDataSpaceId spaceId, ulong saveDataId)
         {
@@ -356,6 +358,8 @@ namespace LibHac.FsService
 
                 var reader = new SaveDataInfoReader(this);
 
+                OpenedReaders.Add(reader);
+
                 infoReader = reader;
 
                 return Result.Success;
@@ -490,10 +494,55 @@ namespace LibHac.FsService
 
         private Result AdjustOpenedInfoReaders(ref SaveDataAttribute key)
         {
-            // todo
+            // If a new key is added or removed during iteration of the list,
+            // make sure the current item of the iterator remains the same
+
+            // Todo: A more efficient way of doing this
+            List<SaveDataAttribute> list = KvDatabase.ToList().Select(x => x.key).ToList();
+
+            int index = list.BinarySearch(key);
+
+            bool keyWasAdded = index >= 0;
+
+            if (!keyWasAdded)
+            {
+                // If the item was not found, List<T>.BinarySearch returns a negative number that
+                // is the bitwise complement of the index of the next element that is larger than the item
+                index = ~index;
+            }
+
+            foreach (SaveDataInfoReader reader in OpenedReaders)
+            {
+                if (keyWasAdded)
+                {
+                    // New key was inserted before the iterator's position
+                    // increment the position to compensate
+                    if(reader.Position >= index)
+                    {
+                        reader.Position++;
+                    }
+                }
+                else
+                {
+                    // The position should be decremented if the iterator's position is
+                    // after the key that came directly after the deleted key
+                    if (reader.Position > index)
+                    {
+                        reader.Position--;
+                    }
+                }
+            }
+
             return Result.Success;
         }
 
+        private void CloseReader(SaveDataInfoReader reader)
+        {
+            bool wasRemoved = OpenedReaders.Remove(reader);
+
+            Debug.Assert(wasRemoved);
+        }
+
         private ref struct Mounter
         {
             private FileSystemClient FsClient { get; set; }
@@ -602,6 +651,11 @@ namespace LibHac.FsService
                     return Result.Success;
                 }
             }
+
+            public void Dispose()
+            {
+                Indexer?.CloseReader(this);
+            }
         }
     }
 }
diff --git a/src/LibHac/FsService/SaveDataInfoFilterReader.cs b/src/LibHac/FsService/SaveDataInfoFilterReader.cs
index 312c00ad..af38cfc4 100644
--- a/src/LibHac/FsService/SaveDataInfoFilterReader.cs
+++ b/src/LibHac/FsService/SaveDataInfoFilterReader.cs
@@ -48,6 +48,11 @@ namespace LibHac.FsService
 
             return Result.Success;
         }
+
+        public void Dispose()
+        {
+            Reader?.Dispose();
+        }
     }
 
     [StructLayout(LayoutKind.Explicit, Size = 0x50)]