rar5 improve memory usage

use ArrayPool for stream buffer use stackalloc for methods on file decompression code path
adamhathcock · Mar 14, 2024 · d2d4755 · d2d4755
1 parent 4afc7ae
commit d2d4755
Show file tree

Hide file tree

Showing 6 changed files with 74 additions and 30 deletions.
diff --git a/src/SharpCompress/BufferPool.cs b/src/SharpCompress/BufferPool.cs
@@ -0,0 +1,33 @@
+using System.Buffers;
+
+namespace SharpCompress;
+
+internal static class BufferPool
+{
+    /// <summary>
+    ///     Gets a buffer.
+    /// </summary>
+    /// <param name="bufferSize">Size of the buffer.</param>
+    /// <returns></returns>
+    public static byte[] Rent(int bufferSize)
+    {
+#if NETCOREAPP || NETSTANDARD2_1_OR_GREATER
+        return ArrayPool<byte>.Shared.Rent(bufferSize);
+#else
+        return new byte[bufferSize];
+#endif
+    }
+
+    /// <summary>
+    ///     Returns the buffer.
+    /// </summary>
+    /// <param name="buffer">The buffer.</param>
+    public static void Return(byte[] buffer)
+    {
+#if NETCOREAPP || NETSTANDARD2_1_OR_GREATER
+        ArrayPool<byte>.Shared.Return(buffer);
+#else
+        // no-op
+#endif
+    }
+}
diff --git a/src/SharpCompress/Compressors/Rar/RarStream.cs b/src/SharpCompress/Compressors/Rar/RarStream.cs
@@ -14,7 +14,7 @@ internal class RarStream : Stream
 
     private bool fetch;
 
-    private byte[] tmpBuffer = new byte[65536];
+    private byte[] tmpBuffer = BufferPool.Rent(65536);
     private int tmpOffset;
     private int tmpCount;
 
@@ -40,6 +40,11 @@ protected override void Dispose(bool disposing)
     {
         if (!isDisposed)
         {
+            if (disposing)
+            {
+                BufferPool.Return(this.tmpBuffer);
+                this.tmpBuffer = null;
+            }
             isDisposed = true;
             base.Dispose(disposing);
             readStream.Dispose();
@@ -118,16 +123,7 @@ public override void Write(byte[] buffer, int offset, int count)
         }
         if (count > 0)
         {
-            if (tmpBuffer.Length < tmpCount + count)
-            {
-                var newBuffer = new byte[
-                    tmpBuffer.Length * 2 > tmpCount + count
-                        ? tmpBuffer.Length * 2
-                        : tmpCount + count
-                ];
-                Buffer.BlockCopy(tmpBuffer, 0, newBuffer, 0, tmpCount);
-                tmpBuffer = newBuffer;
-            }
+            EnsureBufferCapacity(count);
             Buffer.BlockCopy(buffer, offset, tmpBuffer, tmpCount, count);
             tmpCount += count;
             tmpOffset = 0;
@@ -138,4 +134,20 @@ public override void Write(byte[] buffer, int offset, int count)
             unpack.Suspended = false;
         }
     }
+
+    private void EnsureBufferCapacity(int count)
+    {
+        if (this.tmpBuffer.Length < this.tmpCount + count)
+        {
+            var newLength =
+                this.tmpBuffer.Length * 2 > this.tmpCount + count
+                    ? this.tmpBuffer.Length * 2
+                    : this.tmpCount + count;
+            var newBuffer = BufferPool.Rent(newLength);
+            Buffer.BlockCopy(this.tmpBuffer, 0, newBuffer, 0, this.tmpCount);
+            var oldBuffer = this.tmpBuffer;
+            this.tmpBuffer = newBuffer;
+            BufferPool.Return(oldBuffer);
+        }
+    }
 }
diff --git a/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.cs b/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.cs
@@ -1,4 +1,4 @@
-using System;
+using System;
 using System.IO;
 using SharpCompress.Common.Rar.Headers;
 #if !Rar2017_64bit
@@ -67,15 +67,15 @@ public void DoUnpack()
 
     private void UnstoreFile()
     {
-        var b = new byte[0x10000];
+        Span<byte> b = stackalloc byte[(int)Math.Min(0x10000, DestUnpSize)];
         do
         {
-            var n = readStream.Read(b, 0, (int)Math.Min(b.Length, DestUnpSize));
+            var n = readStream.Read(b);
             if (n == 0)
             {
                 break;
             }
-            writeStream.Write(b, 0, n);
+            writeStream.Write(b.Slice(0, n));
             DestUnpSize -= n;
         } while (!Suspended);
     }

diff --git a/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack20_cpp.cs b/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack20_cpp.cs
@@ -373,8 +373,8 @@ private void UnpWriteBuf20()
 
     private bool ReadTables20()
     {
-        var BitLength = new byte[BC20];
-        var Table = new byte[MC20 * 4];
+        Span<byte> BitLength = stackalloc byte[checked((int)BC20)];
+        Span<byte> Table = stackalloc byte[checked((int)MC20 * 4)];
         if (Inp.InAddr > ReadTop - 25)
         {
             if (!UnpReadBuf())
@@ -410,13 +410,13 @@ private bool ReadTables20()
             TableSize = NC20 + DC20 + RC20;
         }
 
-        for (uint I = 0; I < BC20; I++)
+        for (int I = 0; I < checked((int)BC20); I++)
         {
             BitLength[I] = (byte)(Inp.getbits() >> 12);
             Inp.addbits(4);
         }
         MakeDecodeTables(BitLength, 0, BlockTables.BD, BC20);
-        for (uint I = 0; I < TableSize; )
+        for (int I = 0; I < checked((int)TableSize); )
         {
             if (Inp.InAddr > ReadTop - 5)
             {
@@ -487,8 +487,7 @@ private bool ReadTables20()
             MakeDecodeTables(Table, (int)NC20, BlockTables.DD, DC20);
             MakeDecodeTables(Table, (int)(NC20 + DC20), BlockTables.RD, RC20);
         }
-        //x memcpy(UnpOldTable20,Table,sizeof(UnpOldTable20));
-        Array.Copy(Table, UnpOldTable20, UnpOldTable20.Length);
+        Table.CopyTo(this.UnpOldTable20);
         return true;
     }
 

diff --git a/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack50_cpp.cs b/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack50_cpp.cs
@@ -1,4 +1,4 @@
-#nullable disable
+#nullable disable
 
 using System;
 using static SharpCompress.Compressors.Rar.UnpackV2017.PackDef;
@@ -752,8 +752,8 @@ ref UnpackBlockTables Tables
             }
         }
 
-        var BitLength = new byte[BC];
-        for (uint I = 0; I < BC; I++)
+        Span<byte> BitLength = stackalloc byte[checked((int)BC)];
+        for (int I = 0; I < BC; I++)
         {
             uint Length = (byte)(Inp.fgetbits() >> 12);
             Inp.faddbits(4);
@@ -784,9 +784,9 @@ ref UnpackBlockTables Tables
 
         MakeDecodeTables(BitLength, 0, Tables.BD, BC);
 
-        var Table = new byte[HUFF_TABLE_SIZE];
-        const uint TableSize = HUFF_TABLE_SIZE;
-        for (uint I = 0; I < TableSize; )
+        Span<byte> Table = stackalloc byte[checked((int)HUFF_TABLE_SIZE)];
+        const int TableSize = checked((int)HUFF_TABLE_SIZE);
+        for (int I = 0; I < TableSize; )
         {
             if (!Inp.ExternalBuffer && Inp.InAddr > ReadTop - 5)
             {

diff --git a/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack_cpp.cs b/src/SharpCompress/Compressors/Rar/UnpackV2017/Unpack.unpack_cpp.cs
@@ -259,7 +259,7 @@ private void UnpInitData(bool Solid)
     // LengthTable contains the length in bits for every element of alphabet.
     // Dec is the structure to decode Huffman code/
     // Size is size of length table and DecodeNum field in Dec structure,
-    private void MakeDecodeTables(byte[] LengthTable, int offset, DecodeTable Dec, uint Size)
+    private void MakeDecodeTables(Span<byte> LengthTable, int offset, DecodeTable Dec, uint Size)
     {
         // Size of alphabet and DecodePos array.
         Dec.MaxNum = Size;
@@ -269,7 +269,7 @@ private void MakeDecodeTables(byte[] LengthTable, int offset, DecodeTable Dec, u
         //memset(LengthCount,0,sizeof(LengthCount));
         for (size_t I = 0; I < Size; I++)
         {
-            LengthCount[LengthTable[offset + I] & 0xf]++;
+            LengthCount[LengthTable[checked((int)(offset + I))] & 0xf]++;
         }
 
         // We must not calculate the number of zero length codes.
@@ -318,7 +318,7 @@ private void MakeDecodeTables(byte[] LengthTable, int offset, DecodeTable Dec, u
         for (uint I = 0; I < Size; I++)
         {
             // Get the current bit length.
-            var _CurBitLength = (byte)(LengthTable[offset + I] & 0xf);
+            var _CurBitLength = (byte)(LengthTable[checked((int)(offset + I))] & 0xf);
 
             if (_CurBitLength != 0)
             {