LibHac/tests/LibHac.Tests/Fs/StorageTester.cs

using System;
using System.IO;
using System.Linq;
using LibHac.Fs;

namespace LibHac.Tests.Fs
{
    public class StorageTester
    {
        private Random _random;
        private byte[][] _backingArrays;
        private byte[][] _buffers;
        private int _size;

        private int[] _frequentAccessOffsets;
        private int _lastAccessEnd;
        private int _totalAccessCount;
        private Configuration _config;

        public class Configuration
        {
            public Entry[] Entries { get; set; }
            public int[] SizeClassProbs { get; set; }
            public int[] SizeClassMaxSizes { get; set; }
            public int[] TaskProbs { get; set; }
            public int[] AccessTypeProbs { get; set; }
            public ulong RngSeed { get; set; }
            public int FrequentAccessBlockCount { get; set; }
        }

        public StorageTester(Configuration config)
        {
            Entry[] entries = config.Entries;

            if (entries.Length < 2)
            {
                throw new ArgumentException("At least 2 storage entries must be provided", nameof(config.Entries));
            }

            if (entries.Select(x => x.BackingArray.Length).Distinct().Count() != 1)
            {
                throw new ArgumentException("All storages must have the same size.", nameof(config.Entries));
            }

            if (entries[0].BackingArray.Length == 0)
            {
                throw new ArgumentException("The storage size must be greater than 0.", nameof(config.Entries));
            }

            _config = config;
            _random = new Random(config.RngSeed);

            _backingArrays = entries.Select(x => x.BackingArray).ToArray();

            _buffers = new byte[entries.Length][];
            for (int i = 0; i < entries.Length; i++)
            {
                _buffers[i] = new byte[config.SizeClassMaxSizes[^1]];
            }

            _size = entries[0].BackingArray.Length;
            _lastAccessEnd = 0;

            _frequentAccessOffsets = new int[config.FrequentAccessBlockCount];
            for (int i = 0; i < _frequentAccessOffsets.Length; i++)
            {
                _frequentAccessOffsets[i] = ChooseOffset(AccessType.Random);
            }
        }

        //public StorageTester(ulong rngSeed, int frequentAccessBlockCount, params Entry[] entries)
        //{
        //    if (entries.Length < 2)
        //    {
        //        throw new ArgumentException("At least 2 storage entries must be provided", nameof(entries));
        //    }

        //    if (entries.Select(x => x.BackingArray.Length).Distinct().Count() != 1)
        //    {
        //        throw new ArgumentException("All storages must have the same size.", nameof(entries));
        //    }

        //    if (entries[0].BackingArray.Length == 0)
        //    {
        //        throw new ArgumentException("The storage size must be greater than 0.", nameof(entries));
        //    }

        //    _random = new Random(rngSeed);

        //    _entries = entries;
        //    _backingArrays = entries.Select(x => x.BackingArray).ToArray();

        //    _buffers = new byte[entries.Length][];
        //    for (int i = 0; i < entries.Length; i++)
        //    {
        //        _buffers[i] = new byte[SizeClassMaxSizes[^1]];
        //    }

        //    _size = _entries[0].BackingArray.Length;
        //    _lastAccessEnd = 0;

        //    _frequentAccessOffsets = new int[frequentAccessBlockCount];
        //    for (int i = 0; i < _frequentAccessOffsets.Length; i++)
        //    {
        //        _frequentAccessOffsets[i] = ChooseOffset(AccessType.Random);
        //    }
        //}

        public void Run(long accessCount)
        {
            long endCount = _totalAccessCount + accessCount;

            while (_totalAccessCount < endCount)
            {
                Task task = ChooseTask();
                switch (task)
                {
                    case Task.Read:
                        RunRead();
                        break;
                    case Task.Write:
                        RunWrite();
                        break;
                    case Task.Flush:
                        RunFlush();
                        break;
                }

                _totalAccessCount++;
            }
        }

        private void RunRead()
        {
            int sizeClass = ChooseSizeClass();
            AccessType accessType = ChooseAccessType();
            int offset = ChooseOffset(accessType);
            int size = ChooseSize(offset, sizeClass);

            for (int i = 0; i < _config.Entries.Length; i++)
            {
                Entry entry = _config.Entries[i];
                entry.Storage.Read(offset, _buffers[i].AsSpan(0, size)).ThrowIfFailure();
            }

            if (!CompareBuffers(_buffers, size))
            {
                throw new InvalidDataException($"Read: Offset {offset}; Size {size}");
            }
        }

        private void RunWrite()
        {
            int sizeClass = ChooseSizeClass();
            AccessType accessType = ChooseAccessType();
            int offset = ChooseOffset(accessType);
            int size = ChooseSize(offset, sizeClass);

            Span<byte> buffer = _buffers[0].AsSpan(0, size);
            _random.NextBytes(buffer);

            for (int i = 0; i < _config.Entries.Length; i++)
            {
                Entry entry = _config.Entries[i];
                entry.Storage.Write(offset, buffer).ThrowIfFailure();
            }
        }

        private void RunFlush()
        {
            foreach (Entry entry in _config.Entries)
            {
                entry.Storage.Flush().ThrowIfFailure();
            }

            if (!CompareBuffers(_backingArrays, _size))
            {
                throw new InvalidDataException("Flush");
            }
        }

        private Task ChooseTask() => (Task)ChooseProb(_config.TaskProbs);
        private int ChooseSizeClass() => ChooseProb(_config.SizeClassProbs);
        private AccessType ChooseAccessType() => (AccessType)ChooseProb(_config.AccessTypeProbs);

        private int ChooseOffset(AccessType type) => type switch
        {
            AccessType.Random => _random.Next(0, _size),
            AccessType.Sequential => _lastAccessEnd == _size ? 0 : _lastAccessEnd,
            AccessType.FrequentBlock => _frequentAccessOffsets[_random.Next(0, _frequentAccessOffsets.Length)],
            _ => 0
        };

        private int ChooseSize(int offset, int sizeClass)
        {
            int availableSize = Math.Max(0, _size - offset);
            int randSize = _random.Next(0, _config.SizeClassMaxSizes[sizeClass]);
            return Math.Min(availableSize, randSize);
        }

        private int ChooseProb(int[] weights)
        {
            int total = 0;
            foreach (int weight in weights)
            {
                total += weight;
            }

            int rand = _random.Next(0, total);
            int currentThreshold = 0;

            for (int i = 0; i < weights.Length; i++)
            {
                currentThreshold += weights[i];

                if (rand < currentThreshold)
                    return i;
            }

            return 0;
        }

        private bool CompareBuffers(byte[][] buffers, int size)
        {
            Span<byte> baseBuffer = buffers[0].AsSpan(0, size);

            for (int i = 1; i < buffers.Length; i++)
            {
                Span<byte> testBuffer = buffers[i].AsSpan(0, size);
                if (!baseBuffer.SequenceEqual(testBuffer))
                {
                    return false;
                }
            }

            return true;
        }

        public readonly struct Entry
        {
            public readonly IStorage Storage;
            public readonly byte[] BackingArray;

            public Entry(IStorage storage, byte[] backingArray)
            {
                Storage = storage;
                BackingArray = backingArray;
            }
        }

        private enum Task
        {
            Read = 0,
            Write = 1,
            Flush = 2
        }

        private enum AccessType
        {
            Random = 0,
            Sequential = 1,
            FrequentBlock = 2
        }
    }
}