From d7dc8ebb679a973e1ae126a0867c858e90bff86f Mon Sep 17 00:00:00 2001
From: Alexandre Mutel <alexandre_mutel@live.com>
Date: Sat, 17 Sep 2022 10:39:10 +0200
Subject: [PATCH] Add `VirtualArena.Reset(VirtualArenaResetKind resetKind)` to
 allow to decommit partial committed memory

---
 src/Varena.Tests/TestVirtualArena.cs | 14 +++++++++
 src/Varena/VirtualArena.cs           | 43 ++++++++++++++++++++++++++--
 src/Varena/VirtualArenaResetKind.cs  | 26 +++++++++++++++++
 3 files changed, 80 insertions(+), 3 deletions(-)
 create mode 100644 src/Varena/VirtualArenaResetKind.cs

diff --git a/src/Varena.Tests/TestVirtualArena.cs b/src/Varena.Tests/TestVirtualArena.cs
index d7de39e..e7216b6 100644
--- a/src/Varena.Tests/TestVirtualArena.cs
+++ b/src/Varena.Tests/TestVirtualArena.cs
@@ -23,6 +23,20 @@ public void TestCreateBufferReset()
         var span2 = buffer.AllocateRange(1024);
         Assert.AreEqual((nuint)1024, buffer.AllocatedBytes);
         Assert.AreEqual((nuint)1 << 16, buffer.CommittedBytes);
+        buffer.Reset(VirtualArenaResetKind.KeepMinimalCommitted);
+        Assert.AreEqual((nuint)1 << 16, buffer.CommittedBytes);
+        buffer.Reset(VirtualArenaResetKind.KeepAllCommitted);
+        Assert.AreEqual((nuint)1 << 16, buffer.CommittedBytes);
+
+        // Allocate a bit more than one commit block
+        buffer.AllocateRange(1024 + (1 << 16));
+        Assert.AreEqual((nuint)2 << 16, buffer.CommittedBytes);
+        buffer.Reset(VirtualArenaResetKind.KeepAllCommitted);
+        Assert.AreEqual((nuint)2 << 16, buffer.CommittedBytes);
+        buffer.Reset(VirtualArenaResetKind.KeepMinimalCommitted);
+        Assert.AreEqual((nuint)1 << 16, buffer.CommittedBytes);
+
+        Assert.Throws<ArgumentOutOfRangeException>(() => buffer.Reset((VirtualArenaResetKind)123));
     }
 
     [Test]
diff --git a/src/Varena/VirtualArena.cs b/src/Varena/VirtualArena.cs
index 713a7e2..9072be7 100644
--- a/src/Varena/VirtualArena.cs
+++ b/src/Varena/VirtualArena.cs
@@ -110,12 +110,49 @@ public void Dispose()
     /// <exception cref="VirtualMemoryException">If this method fails to uncommit the memory used by this arena.</exception>
     public void Reset()
     {
-        if (CapacityInBytes > 0 && !_handler.TryUnCommit(_range))
+        Reset(VirtualArenaResetKind.Default);
+    }
+
+    /// <summary>
+    /// Resets the memory allocated by this arena with the specified reset options.
+    /// </summary>
+    /// <param name="resetKind">The kind of reset.</param>
+    /// <exception cref="VirtualMemoryException">If this method fails to uncommit the memory used by this arena.</exception>
+    public void Reset(VirtualArenaResetKind resetKind)
+    {
+        // We decommit only if we have something to decommit.
+        if (CommittedBytes > 0)
         {
-            throw new VirtualMemoryException($"Unable to un-commit the memory range for the arena `{Name}`");
+            bool decommitSuccess = true;
+            switch (resetKind)
+            {
+                case VirtualArenaResetKind.Default:
+                    // We decommit the entire committed bytes
+                    decommitSuccess = _handler.TryUnCommit(new VirtualMemoryRange(this.BaseAddress, _committedBytes));
+                    _committedBytes = 0;
+                    break;
+                case VirtualArenaResetKind.KeepAllCommitted:
+                    break;
+                case VirtualArenaResetKind.KeepMinimalCommitted:
+                    // We decommit only entire committed bytes
+                    var newCommittedBytes = this.CommitPageSizeMultiplier * this._handler.PageSize;
+                    var bytesToDecommit = _committedBytes - newCommittedBytes;
+                    if (bytesToDecommit > 0)
+                    {
+                        decommitSuccess = _handler.TryUnCommit(new VirtualMemoryRange((nint)this.BaseAddress + (nint)(newCommittedBytes), bytesToDecommit));
+                    }
+                    _committedBytes = newCommittedBytes;
+                    break;
+                default:
+                    throw new ArgumentOutOfRangeException(nameof(resetKind), resetKind, null);
+            }
+
+            if (!decommitSuccess)
+            {
+                throw new VirtualMemoryException($"Unable to un-commit the memory range for the arena `{Name}`");
+            }
         }
 
-        _committedBytes = 0;
         _allocatedBytes = 0;
         ResetImpl();
     }
diff --git a/src/Varena/VirtualArenaResetKind.cs b/src/Varena/VirtualArenaResetKind.cs
new file mode 100644
index 0000000..eb32b05
--- /dev/null
+++ b/src/Varena/VirtualArenaResetKind.cs
@@ -0,0 +1,26 @@
+// Copyright (c) Alexandre Mutel. All rights reserved.
+// Licensed under the BSD-Clause 2 license.
+// See license.txt file in the project root for full license information.
+
+namespace Varena;
+
+/// <summary>
+/// Type of reset for <see cref="VirtualArena.Reset(VirtualArenaResetKind)"/>.
+/// </summary>
+public enum VirtualArenaResetKind {
+
+    /// <summary>
+    /// All committed bytes are uncommitted.
+    /// </summary>
+    Default,
+
+    /// <summary>
+    /// Keep all committed bytes.
+    /// </summary>
+    KeepAllCommitted,
+
+    /// <summary>
+    /// Keep minimal committed bytes which is the committed memory page size multiplier.
+    /// </summary>
+    KeepMinimalCommitted,
+}
\ No newline at end of file