From c28128c7dd4acb6bcc6a945200b058d44b142814 Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Sat, 14 Aug 2021 19:49:51 -0700
Subject: [PATCH] Add SharedRef<T> and WeakRef<T>

---
 src/LibHac/Common/SharedRef.cs                | 397 ++++++++++++++++++
 src/LibHac/Common/UniqueRef.cs                |  38 +-
 src/LibHac/Fs/FileStorageBasedFileSystem.cs   |   2 +-
 src/LibHac/Fs/Fsa/DirectoryAccessor.cs        |   2 +-
 src/LibHac/Fs/Fsa/FileAccessor.cs             |   2 +-
 .../FsSrv/Impl/FileSystemInterfaceAdapter.cs  |   4 +-
 .../FsSrv/Impl/IResultConvertFileSystem.cs    |   4 +-
 src/LibHac/FsSrv/Impl/OpenCountFileSystem.cs  |   2 +-
 src/LibHac/FsSystem/AesXtsDirectory.cs        |   2 +-
 src/LibHac/FsSystem/AesXtsFile.cs             |   2 +-
 .../FsSystem/ConcatenationFileSystem.cs       |   2 +-
 .../FsSystem/DirectorySaveDataFileSystem.cs   |   6 +-
 src/LibHac/FsSystem/Utility12.cs              |   2 +-
 13 files changed, 436 insertions(+), 29 deletions(-)
 create mode 100644 src/LibHac/Common/SharedRef.cs

diff --git a/src/LibHac/Common/SharedRef.cs b/src/LibHac/Common/SharedRef.cs
new file mode 100644
index 00000000..f6062e1f
--- /dev/null
+++ b/src/LibHac/Common/SharedRef.cs
@@ -0,0 +1,397 @@
+using System;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using InlineIL;
+using LibHac.Diag;
+
+namespace LibHac.Common
+{
+    public static class SharedRefExtensions
+    {
+        // ReSharper disable once EntityNameCapturedOnly.Global
+        public static ref SharedRef<T> Ref<T>(this in SharedRef<T> value) where T : class, IDisposable
+        {
+            IL.Emit.Ldarg(nameof(value));
+            IL.Emit.Ret();
+            throw IL.Unreachable();
+        }
+
+        // ReSharper disable once EntityNameCapturedOnly.Global
+        public static ref WeakRef<T> Ref<T>(this in WeakRef<T> value) where T : class, IDisposable
+        {
+            IL.Emit.Ldarg(nameof(value));
+            IL.Emit.Ret();
+            throw IL.Unreachable();
+        }
+    }
+
+    internal class RefCount
+    {
+        private int _count;
+        private int _weakCount;
+        private IDisposable _value;
+
+        public RefCount(IDisposable value)
+        {
+            _count = 1;
+            _weakCount = 1;
+            _value = value;
+        }
+
+        public int UseCount()
+        {
+            return _count;
+        }
+
+        public void Increment()
+        {
+            // This function shouldn't ever be called after the ref count is reduced to zero.
+            // SharedRef clearing their RefCountBase field after decrementing the ref count *should* already
+            // prevent this if everything works properly.
+            Assert.SdkRequiresGreater(_count, 0);
+
+            Interlocked.Increment(ref _count);
+        }
+
+        public void IncrementWeak()
+        {
+            Assert.SdkRequiresGreater(_count, 0);
+
+            Interlocked.Increment(ref _weakCount);
+        }
+
+        public bool IncrementIfNotZero()
+        {
+            int count = Volatile.Read(ref _count);
+
+            while (count != 0)
+            {
+                int oldValue = Interlocked.CompareExchange(ref _count, count + 1, count);
+                if (oldValue == count)
+                {
+                    return true;
+                }
+
+                count = oldValue;
+            }
+
+            return false;
+        }
+
+        public void Decrement()
+        {
+            if (Interlocked.Decrement(ref _count) == 0)
+            {
+                Destroy();
+                DecrementWeak();
+            }
+        }
+
+        public void DecrementWeak()
+        {
+            if (Interlocked.Decrement(ref _weakCount) == 0)
+            {
+                // Deallocate
+            }
+        }
+
+        private void Destroy()
+        {
+            if (_value is not null)
+            {
+                _value.Dispose();
+                _value = null;
+            }
+        }
+    }
+
+    public struct SharedRef<T> : IDisposable where T : class, IDisposable
+    {
+        // SharedRef and WeakRef should share a base type, but struct inheritance doesn't exist in C#.
+        // This presents a problem because C# also doesn't have friend classes, these two types need to
+        // access each other's fields and we'd rather the fields' visibility stay private. Because the
+        // two types have the same layout we can hack around this with some Unsafe.As shenanigans.
+        private T _value;
+        private RefCount _refCount;
+
+        public SharedRef(T value)
+        {
+            _value = value;
+            _refCount = new RefCount(value);
+        }
+
+        [Obsolete("This method should never be manually called. Use the Destroy method instead.", true)]
+        public void Dispose()
+        {
+            // This function shouldn't be called manually and should always be called at the end of a using block.
+            // This means we don't need to clear any fields because we're going out of scope anyway.
+            _refCount?.Decrement();
+        }
+
+        /// <summary>
+        /// Used to manually dispose the <see cref="SharedRef{T}"/> from the Dispose methods of other types.
+        /// </summary>
+        public void Destroy()
+        {
+            Reset();
+        }
+
+        public readonly T Get => _value;
+        public readonly bool HasValue => Get is not null;
+        public readonly int UseCount => _refCount?.UseCount() ?? 0;
+
+        public static SharedRef<T> CreateMove<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            var sharedRef = new SharedRef<T>();
+
+            sharedRef._value = Unsafe.As<TFrom, T>(ref other._value);
+            sharedRef._refCount = other._refCount;
+
+            other._value = null;
+            other._refCount = null;
+
+            return sharedRef;
+        }
+
+        public static SharedRef<T> CreateCopy<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            var sharedRef = new SharedRef<T>();
+
+            sharedRef._value = Unsafe.As<TFrom, T>(ref other._value);
+            sharedRef._refCount = other._refCount;
+
+            sharedRef._refCount?.Increment();
+
+            return sharedRef;
+        }
+
+        public static SharedRef<T> Create<TFrom>(ref WeakRef<TFrom> other) where TFrom : class, T
+        {
+            ref SharedRef<TFrom> otherShared = ref Unsafe.As<WeakRef<TFrom>, SharedRef<TFrom>>(ref other);
+
+            var sharedRef = new SharedRef<T>();
+
+            if (otherShared._refCount is not null && otherShared._refCount.IncrementIfNotZero())
+            {
+                sharedRef._value = Unsafe.As<TFrom, T>(ref otherShared._value);
+                sharedRef._refCount = otherShared._refCount;
+            }
+            else
+            {
+                ThrowBadWeakPtr();
+            }
+
+            return sharedRef;
+        }
+
+        public static SharedRef<T> Create<TFrom>(ref UniqueRef<TFrom> other) where TFrom : class, T
+        {
+            var sharedRef = new SharedRef<T>();
+            TFrom value = other.Get;
+
+            if (value is not null)
+            {
+                sharedRef._value = value;
+                sharedRef._refCount = new RefCount(value);
+                other.Release();
+            }
+            else
+            {
+                sharedRef._value = null;
+                sharedRef._refCount = null;
+            }
+
+            return sharedRef;
+        }
+
+        public void Swap(ref SharedRef<T> other)
+        {
+            (other._value, _value) = (_value, other._value);
+            (other._refCount, _refCount) = (_refCount, other._refCount);
+        }
+
+        public void Reset()
+        {
+            _value = null;
+            RefCount oldRefCount = _refCount;
+            _refCount = null;
+
+            oldRefCount?.Decrement();
+        }
+
+        public void Reset(T value)
+        {
+            _value = value;
+            RefCount oldRefCount = _refCount;
+            _refCount = new RefCount(value);
+
+            oldRefCount?.Decrement();
+        }
+
+        public void SetByMove<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            RefCount oldRefCount = _refCount;
+
+            _value = Unsafe.As<TFrom, T>(ref other._value);
+            _refCount = other._refCount;
+
+            other._value = null;
+            other._refCount = null;
+
+            oldRefCount?.Decrement();
+        }
+
+        public void SetByCopy<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            RefCount oldRefCount = _refCount;
+            RefCount otherRef = other._refCount;
+
+            otherRef?.Increment();
+
+            _value = Unsafe.As<TFrom, T>(ref other._value);
+            _refCount = otherRef;
+
+            oldRefCount?.Decrement();
+        }
+
+        private static void ThrowBadWeakPtr()
+        {
+            throw new ObjectDisposedException(string.Empty, "bad_weak_ptr");
+        }
+    }
+
+    public struct WeakRef<T> : IDisposable where T : class, IDisposable
+    {
+        private T _value;
+        private RefCount _refCount;
+
+        public WeakRef(ref SharedRef<T> other)
+        {
+            this = Create(ref other);
+        }
+
+        [Obsolete("This method should never be manually called. Use the Destroy method instead.", true)]
+        public void Dispose()
+        {
+            _refCount?.DecrementWeak();
+        }
+
+        // A copy of Dispose so we can call it ourselves inside the struct
+        private void DisposeInternal()
+        {
+            _refCount?.DecrementWeak();
+        }
+
+        /// <summary>
+        /// Used to manually dispose the <see cref="WeakRef{T}"/> from the Dispose methods of other types.
+        /// </summary>
+        public void Destroy()
+        {
+            Reset();
+        }
+
+        public readonly int UseCount => _refCount?.UseCount() ?? 0;
+        public readonly bool Expired => UseCount == 0;
+
+        public static WeakRef<T> CreateMove<TFrom>(ref WeakRef<TFrom> other) where TFrom : class, T
+        {
+            var weakRef = new WeakRef<T>();
+
+            weakRef._value = Unsafe.As<TFrom, T>(ref other._value);
+            weakRef._refCount = other._refCount;
+
+            other._value = null;
+            other._refCount = null;
+
+            return weakRef;
+        }
+
+        public static WeakRef<T> CreateCopy<TFrom>(ref WeakRef<TFrom> other) where TFrom : class, T
+        {
+            var weakRef = new WeakRef<T>();
+
+            if (other._refCount is not null)
+            {
+                weakRef._refCount = other._refCount;
+                weakRef._refCount.IncrementWeak();
+
+                if (weakRef._refCount.IncrementIfNotZero())
+                {
+                    weakRef._value = Unsafe.As<TFrom, T>(ref other._value);
+                    weakRef._refCount.Decrement();
+                }
+            }
+
+            return weakRef;
+        }
+
+        public static WeakRef<T> Create<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            ref WeakRef<TFrom> otherWeak = ref Unsafe.As<SharedRef<TFrom>, WeakRef<TFrom>>(ref other);
+
+            var weakRef = new WeakRef<T>();
+
+            if (otherWeak._refCount is not null)
+            {
+                weakRef._value = Unsafe.As<TFrom, T>(ref otherWeak._value);
+                weakRef._refCount = otherWeak._refCount;
+
+                weakRef._refCount.IncrementWeak();
+            }
+            else
+            {
+                weakRef._value = null;
+                weakRef._refCount = null;
+            }
+
+            return weakRef;
+        }
+
+        public void Swap(ref WeakRef<T> other)
+        {
+            (other._value, _value) = (_value, other._value);
+            (other._refCount, _refCount) = (_refCount, other._refCount);
+        }
+
+        public void Reset()
+        {
+            var temp = new WeakRef<T>();
+            Swap(ref temp);
+            temp.DisposeInternal();
+        }
+
+        public void SetMove<TFrom>(ref WeakRef<TFrom> other) where TFrom : class, T
+        {
+            WeakRef<T> temp = CreateMove(ref other);
+            Swap(ref temp);
+            temp.DisposeInternal();
+        }
+
+        public void SetCopy<TFrom>(ref WeakRef<TFrom> other) where TFrom : class, T
+        {
+            WeakRef<T> temp = CreateCopy(ref other);
+            Swap(ref temp);
+            temp.DisposeInternal();
+        }
+
+        public void Set<TFrom>(ref SharedRef<TFrom> other) where TFrom : class, T
+        {
+            WeakRef<T> temp = Create(ref other);
+            Swap(ref temp);
+            temp.DisposeInternal();
+        }
+
+        public readonly SharedRef<T> Lock()
+        {
+            var sharedRef = new SharedRef<T>();
+
+            if (_refCount is not null && _refCount.IncrementIfNotZero())
+            {
+                Unsafe.As<SharedRef<T>, WeakRef<T>>(ref sharedRef)._value = _value;
+                Unsafe.As<SharedRef<T>, WeakRef<T>>(ref sharedRef)._refCount = _refCount;
+            }
+
+            return sharedRef;
+        }
+    }
+}
diff --git a/src/LibHac/Common/UniqueRef.cs b/src/LibHac/Common/UniqueRef.cs
index bd03d535..5ca4f95e 100644
--- a/src/LibHac/Common/UniqueRef.cs
+++ b/src/LibHac/Common/UniqueRef.cs
@@ -19,10 +19,8 @@ namespace LibHac.Common
     {
         private T _value;
 
-        public void Dispose()
-        {
-            Release()?.Dispose();
-        }
+        public readonly T Get => _value;
+        public readonly bool HasValue => Get is not null;
 
         public UniqueRef(T value)
         {
@@ -34,9 +32,29 @@ namespace LibHac.Common
             _value = other.Release();
         }
 
-        public readonly T Get => _value;
+        [Obsolete("This method should never be manually called. Use the Destroy method instead.", true)]
+        public void Dispose()
+        {
+            _value?.Dispose();
+        }
 
-        public readonly bool HasValue => Get is not null;
+        /// <summary>
+        /// Used to manually dispose the <see cref="UniqueRef{T}"/> from the Dispose methods of other types.
+        /// </summary>
+        public void Destroy()
+        {
+            Reset();
+        }
+
+        public static UniqueRef<T> Create<TFrom>(ref UniqueRef<TFrom> other) where TFrom : class, T
+        {
+            return new UniqueRef<T>(other.Release());
+        }
+
+        public void Swap(ref UniqueRef<T> other)
+        {
+            (other._value, _value) = (_value, other._value);
+        }
 
         public void Reset() => Reset(null);
 
@@ -48,14 +66,6 @@ namespace LibHac.Common
             oldValue?.Dispose();
         }
 
-        public void Reset<TFrom>(TFrom value) where TFrom : class, T
-        {
-            T oldValue = _value;
-            _value = value;
-
-            oldValue?.Dispose();
-        }
-
         public void Set(ref UniqueRef<T> other)
         {
             if (Unsafe.AreSame(ref this, ref other))
diff --git a/src/LibHac/Fs/FileStorageBasedFileSystem.cs b/src/LibHac/Fs/FileStorageBasedFileSystem.cs
index dce86c90..84fbb87b 100644
--- a/src/LibHac/Fs/FileStorageBasedFileSystem.cs
+++ b/src/LibHac/Fs/FileStorageBasedFileSystem.cs
@@ -31,7 +31,7 @@ namespace LibHac.Fs
         {
             if (disposing)
             {
-                _baseFile.Dispose();
+                _baseFile.Destroy();
                 _baseFileSystem?.Dispose();
             }
 
diff --git a/src/LibHac/Fs/Fsa/DirectoryAccessor.cs b/src/LibHac/Fs/Fsa/DirectoryAccessor.cs
index ea59e38e..e4a63733 100644
--- a/src/LibHac/Fs/Fsa/DirectoryAccessor.cs
+++ b/src/LibHac/Fs/Fsa/DirectoryAccessor.cs
@@ -21,7 +21,7 @@ namespace LibHac.Fs.Impl
             _directory.Reset();
             _parentFileSystem.NotifyCloseDirectory(this);
 
-            _directory.Dispose();
+            _directory.Destroy();
         }
 
         public FileSystemAccessor GetParent() => _parentFileSystem;
diff --git a/src/LibHac/Fs/Fsa/FileAccessor.cs b/src/LibHac/Fs/Fsa/FileAccessor.cs
index 1d64cdb8..5204fc94 100644
--- a/src/LibHac/Fs/Fsa/FileAccessor.cs
+++ b/src/LibHac/Fs/Fsa/FileAccessor.cs
@@ -52,7 +52,7 @@ namespace LibHac.Fs.Impl
 
             _parentFileSystem?.NotifyCloseFile(this);
 
-            _file.Dispose();
+            _file.Destroy();
         }
 
         public OpenMode GetOpenMode() => _openMode;
diff --git a/src/LibHac/FsSrv/Impl/FileSystemInterfaceAdapter.cs b/src/LibHac/FsSrv/Impl/FileSystemInterfaceAdapter.cs
index 99d429c2..97b24a55 100644
--- a/src/LibHac/FsSrv/Impl/FileSystemInterfaceAdapter.cs
+++ b/src/LibHac/FsSrv/Impl/FileSystemInterfaceAdapter.cs
@@ -37,7 +37,7 @@ namespace LibHac.FsSrv.Impl
 
         public void Dispose()
         {
-            _baseFile.Dispose();
+            _baseFile.Destroy();
             _parentFs?.Dispose();
         }
 
@@ -190,7 +190,7 @@ namespace LibHac.FsSrv.Impl
 
         public void Dispose()
         {
-            _baseDirectory.Dispose();
+            _baseDirectory.Destroy();
             _parentFs?.Dispose();
         }
 
diff --git a/src/LibHac/FsSrv/Impl/IResultConvertFileSystem.cs b/src/LibHac/FsSrv/Impl/IResultConvertFileSystem.cs
index 7af52fda..2a64f8fa 100644
--- a/src/LibHac/FsSrv/Impl/IResultConvertFileSystem.cs
+++ b/src/LibHac/FsSrv/Impl/IResultConvertFileSystem.cs
@@ -17,7 +17,7 @@ namespace LibHac.FsSrv.Impl
 
         public override void Dispose()
         {
-            BaseFile.Dispose();
+            BaseFile.Destroy();
             base.Dispose();
         }
 
@@ -67,7 +67,7 @@ namespace LibHac.FsSrv.Impl
 
         public override void Dispose()
         {
-            BaseDirectory.Dispose();
+            BaseDirectory.Destroy();
             base.Dispose();
         }
 
diff --git a/src/LibHac/FsSrv/Impl/OpenCountFileSystem.cs b/src/LibHac/FsSrv/Impl/OpenCountFileSystem.cs
index c2c178a2..3e15acec 100644
--- a/src/LibHac/FsSrv/Impl/OpenCountFileSystem.cs
+++ b/src/LibHac/FsSrv/Impl/OpenCountFileSystem.cs
@@ -64,7 +64,7 @@ namespace LibHac.FsSrv.Impl
         public override void Dispose()
         {
             _entryCountSemaphore?.Dispose();
-            _mountCountSemaphore.Dispose();
+            _mountCountSemaphore.Destroy();
             base.Dispose();
         }
     }
diff --git a/src/LibHac/FsSystem/AesXtsDirectory.cs b/src/LibHac/FsSystem/AesXtsDirectory.cs
index d61bf631..ab64ead4 100644
--- a/src/LibHac/FsSystem/AesXtsDirectory.cs
+++ b/src/LibHac/FsSystem/AesXtsDirectory.cs
@@ -24,7 +24,7 @@ namespace LibHac.FsSystem
 
         public override void Dispose()
         {
-            _baseDirectory.Dispose();
+            _baseDirectory.Destroy();
             base.Dispose();
         }
 
diff --git a/src/LibHac/FsSystem/AesXtsFile.cs b/src/LibHac/FsSystem/AesXtsFile.cs
index f25379bd..35a4de7b 100644
--- a/src/LibHac/FsSystem/AesXtsFile.cs
+++ b/src/LibHac/FsSystem/AesXtsFile.cs
@@ -126,7 +126,7 @@ namespace LibHac.FsSystem
         {
             BaseStorage.Flush();
             BaseStorage.Dispose();
-            BaseFile.Dispose();
+            BaseFile.Destroy();
 
             base.Dispose();
         }
diff --git a/src/LibHac/FsSystem/ConcatenationFileSystem.cs b/src/LibHac/FsSystem/ConcatenationFileSystem.cs
index e270e307..5478b0d0 100644
--- a/src/LibHac/FsSystem/ConcatenationFileSystem.cs
+++ b/src/LibHac/FsSystem/ConcatenationFileSystem.cs
@@ -393,7 +393,7 @@ namespace LibHac.FsSystem
             public override void Dispose()
             {
                 _path.Dispose();
-                _baseDirectory.Dispose();
+                _baseDirectory.Destroy();
 
                 base.Dispose();
             }
diff --git a/src/LibHac/FsSystem/DirectorySaveDataFileSystem.cs b/src/LibHac/FsSystem/DirectorySaveDataFileSystem.cs
index 5d625960..3f8f2a9d 100644
--- a/src/LibHac/FsSystem/DirectorySaveDataFileSystem.cs
+++ b/src/LibHac/FsSystem/DirectorySaveDataFileSystem.cs
@@ -73,7 +73,7 @@ namespace LibHac.FsSystem
 
             public override void Dispose()
             {
-                _baseFile.Dispose();
+                _baseFile.Destroy();
 
                 if (_mode.HasFlag(OpenMode.Write))
                 {
@@ -178,10 +178,10 @@ namespace LibHac.FsSystem
 
         public override void Dispose()
         {
-            _lockFile.Dispose();
+            _lockFile.Destroy();
 
             _cacheObserver?.Unregister(_spaceId, _saveDataId);
-            _uniqueBaseFs.Dispose();
+            _uniqueBaseFs.Destroy();
             base.Dispose();
         }
 
diff --git a/src/LibHac/FsSystem/Utility12.cs b/src/LibHac/FsSystem/Utility12.cs
index 3e06dbef..1097e8dc 100644
--- a/src/LibHac/FsSystem/Utility12.cs
+++ b/src/LibHac/FsSystem/Utility12.cs
@@ -94,7 +94,7 @@ namespace LibHac.FsSystem
                 rc = directory.Get.Read(out long entriesRead, SpanHelpers.AsSpan(ref dirEntry));
                 if (rc.IsFailure()) return rc;
 
-                directory.Reset(null);
+                directory.Reset();
 
                 if (entriesRead == 0)
                     break;