From c3981d0e163393e2cde74db3ab8ab10ca22e3f3e Mon Sep 17 00:00:00 2001
From: Alex Barney <thealexbarney@gmail.com>
Date: Mon, 8 Jul 2024 12:29:03 -0700
Subject: [PATCH] Implement StorageDuplicator

---
 .../FsSrv/Impl/SaveDataTransferUtility.cs     | 73 ++++++++++++++++---
 src/LibHac/FsSystem/PooledBuffer.cs           |  8 +-
 2 files changed, 69 insertions(+), 12 deletions(-)

diff --git a/src/LibHac/FsSrv/Impl/SaveDataTransferUtility.cs b/src/LibHac/FsSrv/Impl/SaveDataTransferUtility.cs
index a14bead8..91763492 100644
--- a/src/LibHac/FsSrv/Impl/SaveDataTransferUtility.cs
+++ b/src/LibHac/FsSrv/Impl/SaveDataTransferUtility.cs
@@ -3,6 +3,7 @@
 using System;
 using LibHac.Common;
 using LibHac.Common.FixedArrays;
+using LibHac.Diag;
 using LibHac.Fs;
 using LibHac.Fs.Fsa;
 using LibHac.FsSystem;
@@ -36,38 +37,92 @@ public class StorageDuplicator : IDisposable
     private SharedRef<IStorage> _destinationFileStorage;
     private SharedRef<IFileSystem> _destinationFileSystem;
     private long _offset;
-    private long _restSize;
+    private long _remainingSize;
 
     public StorageDuplicator(ref readonly SharedRef<IStorage> sourceFileStorage,
         ref readonly SharedRef<IStorage> destinationFileStorage,
         ref readonly SharedRef<IFileSystem> destinationFileSystem)
     {
-        throw new NotImplementedException();
+        _sourceFileStorage = SharedRef<IStorage>.CreateCopy(in sourceFileStorage);
+        _destinationFileStorage = SharedRef<IStorage>.CreateCopy(in destinationFileStorage);
+        _destinationFileSystem = SharedRef<IFileSystem>.CreateCopy(in destinationFileSystem);
+
+        _offset = 0;
+        _remainingSize = -1;
     }
 
     public void Dispose()
     {
-        throw new NotImplementedException();
+        _sourceFileStorage.Destroy();
+        _destinationFileStorage.Destroy();
+        _destinationFileSystem.Destroy();
     }
 
     public Result Initialize()
     {
-        throw new NotImplementedException();
+        Result res = _sourceFileStorage.Get.GetSize(out long sourceSize);
+        if (res.IsFailure()) return res.Miss();
+
+        res = _destinationFileStorage.Get.GetSize(out long destinationSize);
+        if (res.IsFailure()) return res.Miss();
+
+        if (sourceSize != destinationSize)
+            return ResultFs.SaveDataPorterSaveDataModified.Log();
+
+        _remainingSize = destinationSize;
+        return Result.Success;
     }
 
     public Result FinalizeObject()
     {
-        throw new NotImplementedException();
+        Result res = _destinationFileStorage.Get.Flush();
+        if (res.IsFailure()) return res.Miss();
+
+        res = _destinationFileSystem.Get.Commit();
+        if (res.IsFailure()) return res.Miss();
+
+        return Result.Success;
     }
 
-    public Result ProcessDuplication(out long outRestSize, Span<byte> workBuffer, long sizeToProcess)
+    public Result ProcessDuplication(out long outRemainingSize, Span<byte> workBuffer, long sizeToProcess)
     {
-        throw new NotImplementedException();
+        UnsafeHelpers.SkipParamInit(out outRemainingSize);
+        Assert.SdkGreaterEqual(_remainingSize, 0);
+
+        long remainingSizeToProcess = Math.Min(sizeToProcess, _remainingSize);
+
+        while (remainingSizeToProcess > 0)
+        {
+            int processSize = (int)Math.Min(workBuffer.Length, Math.Min(sizeToProcess, remainingSizeToProcess));
+            Span<byte> buffer = workBuffer.Slice(0, processSize);
+
+            Result res = _sourceFileStorage.Get.Read(_offset, buffer);
+            if (res.IsFailure()) return res.Miss();
+
+            res = _destinationFileStorage.Get.Write(_offset, buffer);
+            if (res.IsFailure()) return res.Miss();
+
+            _offset += processSize;
+            remainingSizeToProcess -= processSize;
+            _remainingSize -= processSize;
+        }
+
+        outRemainingSize = _remainingSize;
+        return Result.Success;
     }
 
-    public Result ProcessDuplication(out long outRestSize, long sizeToProcess)
+    public Result ProcessDuplication(out long outRemainingSize, long sizeToProcess)
     {
-        throw new NotImplementedException();
+        const int bufferSize = 0x8000;
+
+        UnsafeHelpers.SkipParamInit(out outRemainingSize);
+        Assert.SdkGreaterEqual(_remainingSize, 0);
+
+        using var buffer = new PooledBuffer();
+        Result res = buffer.Allocate(bufferSize, bufferSize);
+        if (res.IsFailure()) return res.Miss();
+
+        return ProcessDuplication(out outRemainingSize, buffer.GetBuffer(), sizeToProcess).Ret();
     }
 }
 
diff --git a/src/LibHac/FsSystem/PooledBuffer.cs b/src/LibHac/FsSystem/PooledBuffer.cs
index 7f6733ed..59d586f9 100644
--- a/src/LibHac/FsSystem/PooledBuffer.cs
+++ b/src/LibHac/FsSystem/PooledBuffer.cs
@@ -137,10 +137,10 @@ public struct PooledBuffer : IDisposable
         return enableLargeCapacity ? HeapAllocatableSizeMaxForLarge : HeapAllocatableSizeMax;
     }
 
-    public void Allocate(int idealSize, int requiredSize) => AllocateCore(idealSize, requiredSize, false);
-    public void AllocateParticularlyLarge(int idealSize, int requiredSize) => AllocateCore(idealSize, requiredSize, true);
+    public Result Allocate(int idealSize, int requiredSize) => AllocateCore(idealSize, requiredSize, false).Ret();
+    public Result AllocateParticularlyLarge(int idealSize, int requiredSize) => AllocateCore(idealSize, requiredSize, true).Ret();
 
-    private void AllocateCore(int idealSize, int requiredSize, bool enableLargeCapacity)
+    private Result AllocateCore(int idealSize, int requiredSize, bool enableLargeCapacity)
     {
         Assert.SdkRequiresNull(_array);
 
@@ -160,6 +160,8 @@ public struct PooledBuffer : IDisposable
         }
 
         _length = _array.Length;
+
+        return Result.Success;
     }
 
     public void Deallocate()