COM3D2.CacheEditMenu/COM3D2.CacheEditMenu/Hooks.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Reflection.Emit;
using System.Security.Cryptography;
using System.Text;
using UnityEngine;
using Debug = UnityEngine.Debug;

namespace COM3D2.CacheEditMenu
{
    internal static class BinaryExtensions
    {
        public static string ReadNullableString(this BinaryReader br)
        {
            return br.ReadBoolean() ? br.ReadString() : null;
        }

        public static void WriteNullableString(this BinaryWriter bw, string value)
        {
            bw.Write(value != null);
            if (value != null)
                bw.Write(value);
        }
    }

    public static class Hooks
    {
        private const int CACHE_VERSION = 1021;

        private static readonly Dictionary<string, MenuInfo> infoCache = new Dictionary<string, MenuInfo>();
        private static BinaryWriter cacheWriter;

        private static readonly Func<Dictionary<int, SceneEdit.SMenuItem>> getIdItemDic = () =>
            new Dictionary<int, SceneEdit.SMenuItem>();

        private static readonly Func<Dictionary<int, string>> getTexFileIDDic = () => new Dictionary<int, string>();
        private static readonly Action<bool> setNowMenuFlg = b => { };
        private static readonly Action<int> setNowStrMenuFileID = b => { };

        static Hooks()
        {
            var quickEditAss =
                TryLoadFrom(Path.Combine(Patcher.Patcher.SybarisPath, "COM3D2.QuickEditStart.Managed.dll"));

            if (quickEditAss == null)
                return;

            var mainType = quickEditAss.GetType("COM3D2.QuickEditStart.Managed.QuickEditStartManaged");

            MakeStaticGetter(mainType.GetField("texFileIDDic", BindingFlags.Public | BindingFlags.Static),
                ref getTexFileIDDic);
            MakeStaticGetter(mainType.GetField("idItemDic", BindingFlags.Public | BindingFlags.Static),
                ref getIdItemDic);
            MakeStaticSetter(mainType.GetField("nowMenuFlg", BindingFlags.NonPublic | BindingFlags.Static),
                ref setNowMenuFlg);
            MakeStaticSetter(mainType.GetField("nowStrMenuFileID", BindingFlags.NonPublic | BindingFlags.Static),
                ref setNowStrMenuFileID);
        }

        private static void MakeStaticGetter<T>(FieldInfo field, ref T result) where T : Delegate
        {
            if (field == null)
                return;

            var dm = new DynamicMethod($"CacheEditMenu_Getter_{field.GetHashCode()}", field.FieldType, null,
                typeof(Hooks), true);

            var il = dm.GetILGenerator();

            il.Emit(OpCodes.Ldsfld, field);
            il.Emit(OpCodes.Ret);

            result = dm.CreateDelegate(typeof(T)) as T;
        }

        private static void MakeStaticSetter<T>(FieldInfo field, ref T result) where T : Delegate
        {
            if (field == null)
                return;

            var dm = new DynamicMethod($"CacheEditMenu_Getter_{field.GetHashCode()}", typeof(void),
                new[] {field.FieldType}, typeof(Hooks), true);

            var il = dm.GetILGenerator();

            il.Emit(OpCodes.Ldarg_0);
            il.Emit(OpCodes.Stsfld, field);
            il.Emit(OpCodes.Ret);

            result = dm.CreateDelegate(typeof(T)) as T;
        }

        private static Assembly TryLoadFrom(string path)
        {
            try
            {
                return Assembly.LoadFrom(path);
            }
            catch (Exception)
            {
                return null;
            }
        }

        public static bool Prefix(ref bool result, SceneEdit.SMenuItem mi, string menuFileName)
        {
            Init();

            if (menuFileName.Contains("_zurashi"))
            {
                result = false;
                return false;
            }

            if (menuFileName.Contains("_mekure"))
            {
                result = false;
                return false;
            }

            menuFileName = Path.GetFileName(menuFileName);
            mi.m_strMenuFileName = menuFileName;
            mi.m_nMenuFileRID = menuFileName.ToLower().GetHashCode();

            if (infoCache.TryGetValue(menuFileName, out var menuInfo))
            {
                menuInfo.CopyTo(mi);
                result = true;
                return false;
            }

            return true;
        }

        public static string texFileName = null;

        public static string CreateTexturePrefix(string fileName)
        {
            texFileName = fileName;
            return fileName;
        }

        public static bool Postfix(bool result, SceneEdit.SMenuItem mi, string menuFileName)
        {
            if (!result)
                return false;

            var menuInfo = new MenuInfo
            {
                mi = mi,
                key = menuFileName
            };
            if (texFileName != null)
                menuInfo.texName = texFileName;
            infoCache[menuFileName] = menuInfo;

            try
            {
                menuInfo.Serialize(cacheWriter);
            }
            catch (Exception e)
            {
                Debug.Log(
                    $"Failed to serialize menu file {menuFileName}: {e.Message}. The cache may be corrupted and will be rebuilt on next game run.");
            }

            cacheWriter.Flush();

            texFileName = null;

            return true;
        }

        private static byte[] HashMenus(string[] menuList)
        {
            var md5 = MD5.Create();

            foreach (var s in menuList)
            {
                var buf = Encoding.Unicode.GetBytes(s);
                md5.TransformBlock(buf, 0, buf.Length, null, 0);
            }

            md5.TransformFinalBlock(new byte[0], 0, 0);

            return md5.Hash;
        }

        private static bool ReadCache(string cachePath)
        {
            if (!File.Exists(cachePath)) 
                return false;
            using (var br = new BinaryReader(File.Open(cachePath, FileMode.Open, FileAccess.Read)))
            {
                try
                {
                    var cacheVer = br.ReadInt32();

                    if (cacheVer != CACHE_VERSION)
                    {
                        Debug.LogWarning("Old cache version, rebuilding...");
                        return false;
                    }

                    var menuFilesHashSize = br.ReadInt32();
                    var menuFilesHash = br.ReadBytes(menuFilesHashSize);
                    if (!menuFilesHash.SequenceEqual(HashMenus(GameUty.MenuFiles)))
                    {
                        Debug.LogWarning(".menu files changed, rebuilding cache...");
                        return false;
                    }
                            
                    var modMenuFilesHashSize = br.ReadInt32();
                    var modMenuFilesHash = br.ReadBytes(modMenuFilesHashSize);
                    if (!modMenuFilesHash.SequenceEqual(HashMenus(GameUty.ModOnlysMenuFiles)))
                    {
                        Debug.LogWarning("Mod .menu files changed, rebuilding cache...");
                        return false;
                    }
                        
                    while (true)
                    {
                        var menuInfo = new MenuInfo();
                        menuInfo.Deserialize(br);
                        infoCache[menuInfo.key] = menuInfo;
                    }
                }
                catch (FormatException e)
                {
                    Debug.Log($"Failed to deserialize cache because {e.Message}. Rebuilding the cache...");
                    return false;
                }
                catch (EndOfStreamException)
                {
                    // End of file, no need to read more
                }
            }

            return true;
        }

        private static void Init()
        {
            if (cacheWriter != null)
                return;
            var cacheDir = Path.Combine(Patcher.Patcher.SybarisPath, "EditMenuCache");
            var cachePath = Path.Combine(cacheDir, "EditMenuCache.dat");

            Directory.CreateDirectory(cacheDir);

            var rebuildCache = !ReadCache(cachePath);

            var stream = rebuildCache ? File.Create(cachePath) : File.OpenWrite(cachePath);
            cacheWriter = new BinaryWriter(stream);

            if (rebuildCache)
            {
                cacheWriter.Write(CACHE_VERSION);
                var menuHash = HashMenus(GameUty.MenuFiles);
                var modMenuHash = HashMenus(GameUty.ModOnlysMenuFiles);
                
                cacheWriter.Write(menuHash.Length);
                cacheWriter.Write(menuHash);
                
                cacheWriter.Write(modMenuHash.Length);
                cacheWriter.Write(modMenuHash);
                
                foreach (var keyValuePair in infoCache)
                    keyValuePair.Value.Serialize(cacheWriter);
                cacheWriter.Flush();
            }

            cacheWriter.BaseStream.Seek(0, SeekOrigin.End);
        }

        private class MenuInfo
        {
            public string key;
            public SceneEdit.SMenuItem mi;
            public string texName;

            public void Deserialize(BinaryReader br)
            {
                key = br.ReadNullableString();
                texName = br.ReadNullableString();
                mi = new SceneEdit.SMenuItem
                {
                    m_strMenuName = br.ReadNullableString(),
                    m_strInfo = br.ReadNullableString(),
                    m_mpn = (MPN) br.ReadInt32(),
                    m_strCateName = br.ReadNullableString(),
                    m_eColorSetMPN = (MPN) br.ReadInt32(),
                    m_strMenuNameInColorSet = br.ReadNullableString(),
                    m_pcMultiColorID = (MaidParts.PARTS_COLOR) br.ReadInt32(),
                    m_boDelOnly = br.ReadBoolean(),
                    m_fPriority = br.ReadSingle(),
                    m_bMan = br.ReadBoolean()
                };
            }

            public void Serialize(BinaryWriter bw)
            {
                bw.WriteNullableString(key);
                bw.WriteNullableString(texName);
                bw.WriteNullableString(mi.m_strMenuName);
                bw.WriteNullableString(mi.m_strInfo);
                bw.Write((int) mi.m_mpn);
                bw.WriteNullableString(mi.m_strCateName);
                bw.Write((int) mi.m_eColorSetMPN);
                bw.WriteNullableString(mi.m_strMenuNameInColorSet);
                bw.Write((int) mi.m_pcMultiColorID);
                bw.Write(mi.m_boDelOnly);
                bw.Write(mi.m_fPriority);
                bw.Write(mi.m_bMan);
            }

            public void CopyTo(SceneEdit.SMenuItem other)
            {
                other.m_strMenuName = mi.m_strMenuName;
                other.m_strInfo = mi.m_strInfo;
                other.m_mpn = mi.m_mpn;
                other.m_strCateName = mi.m_strCateName;
                other.m_eColorSetMPN = mi.m_eColorSetMPN;
                other.m_strMenuNameInColorSet = mi.m_strMenuNameInColorSet;
                other.m_pcMultiColorID = mi.m_pcMultiColorID;
                other.m_boDelOnly = mi.m_boDelOnly;
                other.m_fPriority = mi.m_fPriority;
                other.m_bMan = mi.m_bMan;

                if (string.IsNullOrEmpty(texName))
                    return;

                setNowMenuFlg(true);

                setNowStrMenuFileID(other.m_nMenuFileRID);
                getIdItemDic()[other.m_nMenuFileRID] = other;
                other.m_texIcon = ImportCM.CreateTexture(texName);

                setNowMenuFlg(false);
            }
        }
    }
}