Worms2server/src/Syroot.Worms/Gen2/Archive.cs

using System;
using System.Collections.Generic;
using System.IO;
using System.Text;
using Syroot.IO;
using Syroot.Worms.Core;

namespace Syroot.Worms.Gen2
{
    /// <summary>
    /// Represents a directory of files stored in a DIR file, mostly used to store graphics files. Due to the nowadays
    /// small size of typical directories, the entries and data are loaded directly into a dictionary to profit from
    /// optimal performance when accessing and manipulating the directory.
    /// Used by W2, WA and WWP. S. https://worms2d.info/Graphics_directory.
    /// </summary>
    public class Archive : Dictionary<string, byte[]>, ILoadableFile, ISaveableFile
    {
        // ---- CONSTANTS ----------------------------------------------------------------------------------------------

        private const int _signature = 0x1A524944; // "DIR", 0x1A
        private const int _tocSignature = 0x0000000A;

        private const int _hashBits = 10;
        private const int _hashSize = 1 << _hashBits;

        // ---- CONSTRUCTORS & DESTRUCTOR ------------------------------------------------------------------------------

        /// <summary>
        /// Initializes a new instance of the <see cref="Archive"/> class.
        /// </summary>
        public Archive()
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Archive"/> class, loading the data from the given
        /// <see cref="Stream"/>.
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to load the data from.</param>
        public Archive(Stream stream)
        {
            Load(stream);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Archive"/> class, loading the data from the given file.
        /// </summary>
        /// <param name="fileName">The name of the file to load the data from.</param>
        public Archive(string fileName)
        {            
            Load(fileName);
        }

        // ---- METHODS (PUBLIC) ---------------------------------------------------------------------------------------

        /// <summary>
        /// Loads the data from the given <see cref="Stream"/>.
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to load the data from.</param>
        public void Load(Stream stream)
        {
            if (!stream.CanSeek)
            {
                throw new ArgumentException("Stream requires to be seekable.", nameof(stream));
            }

            Clear();
            using (BinaryDataReader reader = new BinaryDataReader(stream, Encoding.ASCII, true))
            {
                // Read the header.
                if (reader.ReadInt32() != _signature)
                {
                    throw new InvalidDataException("Invalid DIR file signature.");
                }
                int fileSize = reader.ReadInt32();
                int tocOffset = reader.ReadInt32();

                // Read the table of contents.
                reader.Position = tocOffset;
                int tocSignature = reader.ReadInt32();
                if (tocSignature != _tocSignature)
                {
                    throw new InvalidDataException("Invalid DIR table of contents signature.");
                }
                // Generate a data dictionary out of the hash table and file entries.
                int[] hashTable = reader.ReadInt32s(_hashSize);
                foreach (int entryOffset in hashTable)
                {
                    // If the hash is not 0, it points to a list of files which have a hash being the hash table index.
                    if (entryOffset > 0)
                    {
                        int nextEntryOffset = entryOffset;
                        do
                        {
                            reader.Position = tocOffset + nextEntryOffset;
                            nextEntryOffset = reader.ReadInt32();
                            int offset = reader.ReadInt32();
                            int length = reader.ReadInt32();
                            string name = reader.ReadString(BinaryStringFormat.ZeroTerminated);
                            using (reader.TemporarySeek(offset, SeekOrigin.Begin))
                            {
                                Add(name, reader.ReadBytes(length));
                            }   
                        } while (nextEntryOffset != 0);
                    }
                }
            }
        }

        /// <summary>
        /// Loads the data from the given file.
        /// </summary>
        /// <param name="fileName">The name of the file to load the data from.</param>
        public void Load(string fileName)
        {
            using (FileStream stream = new FileStream(fileName, FileMode.Open, FileAccess.Read, FileShare.Read))
            {
                Load(stream);
            }
        }

        /// <summary>
        /// Saves the data into the given <paramref name="stream"/>.
        /// </summary>
        /// <param name="stream">The <see cref="Stream"/> to save the data in.</param>
        public void Save(Stream stream)
        {
            using (BinaryDataWriter writer = new BinaryDataWriter(stream))
            {
                // Write the header.
                writer.Write(_signature);
                Offset fileSizeOffset = writer.ReserveOffset();
                Offset tocOffset = writer.ReserveOffset();

                // Write the data and build the hash table and file entries.
                List<HashTableEntry>[] hashTable = new List<HashTableEntry>[_hashSize];
                foreach (KeyValuePair<string, byte[]> item in this)
                {
                    HashTableEntry entry = new HashTableEntry()
                    {
                        Name = item.Key,
                        Offset = (int)writer.Position,
                        Length = item.Value.Length
                    };
                    writer.Write(item.Value);
                    
                    int hash = CalculateHash(item.Key);
                    if (hashTable[hash] == null)
                    {
                        hashTable[hash] = new List<HashTableEntry>();
                    }
                    hashTable[hash].Add(entry);
                }

                // Write the hash table and file entries.
                int tocStart = (int)writer.Position;
                int fileEntryOffset = sizeof(int) + _hashSize * sizeof(int);
                tocOffset.Satisfy(tocStart);
                writer.Write(_tocSignature);
                for (int i = 0; i < _hashSize; i++)
                {
                    List<HashTableEntry> entries = hashTable[i];
                    if (entries == null)
                    {
                        writer.Write(0);
                    }
                    else
                    {
                        // Write the entries resolving to the current hash.
                        writer.Write(fileEntryOffset);
                        using (writer.TemporarySeek(tocStart + fileEntryOffset, SeekOrigin.Begin))
                        {
                            for (int j = 0; j < entries.Count; j++)
                            {
                                HashTableEntry entry = entries[j];
                                Offset nextEntryOffset = writer.ReserveOffset();
                                writer.Write(entry.Offset);
                                writer.Write(entry.Length);
                                writer.Write(entry.Name, BinaryStringFormat.ZeroTerminated);
                                writer.Align(4);
                                if (j < entries.Count - 1)
                                {
                                    nextEntryOffset.Satisfy((int)writer.Position - tocStart);
                                }
                            }
                            fileEntryOffset = (int)writer.Position - tocStart;
                        }
                    }
                }

                fileSizeOffset.Satisfy(tocStart + fileEntryOffset - 1);
            }
        }

        /// <summary>
        /// Saves the data in the file with the given <paramref name="fileName"/>.
        /// </summary>
        /// <param name="fileName">The name of the file to save the data in.</param>
        public void Save(string fileName)
        {
            using (FileStream stream = new FileStream(fileName, FileMode.Create, FileAccess.Write, FileShare.None))
            {
                Save(stream);
            }
        }

        // ---- METHODS (PRIVATE) --------------------------------------------------------------------------------------

        private int CalculateHash(string name)
        {
            int hash = 0;
            for (int i = 0; i < name.Length; i++)
            {
                hash = ((hash << 1) % _hashSize) | (hash >> (_hashBits - 1) & 1);
                hash += name[i];
                hash %= _hashSize;
            }
            return hash;
        }

        // ---- STRUCTURES ---------------------------------------------------------------------------------------------

        private struct HashTableEntry
        {
            internal string Name;
            internal int Offset;
            internal int Length;
        }
    }
}