feat: allow icon shaders to be replaced

Source/IconMaterialPatch.cs

@@ -0,0 +1,91 @@
+/*
+This file is part of Shabby.
+
+Shabby is free software: you can redistribute it and/or
+modify it under the terms of the GNU General Public License as published by
+the Free Software Foundation, either version 3 of the License, or
+(at your option) any later version.
+
+Shabby is distributed in the hope that it will be useful,
+but WITHOUT ANY WARRANTY; without even the implied warranty of
+MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+GNU General Public License for more details.
+
+You should have received a copy of the GNU General Public License
+along with Shabby.  If not, see
+<http://www.gnu.org/licenses/>.
+*/
+
+using System.Collections.Generic;
+using System.Linq;
+using System.Reflection;
+using System.Reflection.Emit;
+using HarmonyLib;
+using UnityEngine;
+
+namespace Shabby
+{
+
+[HarmonyPatch(typeof(PartLoader), "SetPartIconMaterials")]
+class SetPartIconMaterialsPatch
+{
+	static MethodInfo mInfo_ShaderFind = AccessTools.Method(typeof(Shader), nameof(Shader.Find));
+	static MethodInfo mInfo_FindOverrideIconShader = AccessTools.Method(typeof(SetPartIconMaterialsPatch), nameof(FindOverrideIconShader));
+
+	static Shader FindOverrideIconShader(Material material)
+	{
+		if (Shabby.iconShaders.TryGetValue(material.shader.name, out var shader)) {
+			Debug.Log($"[Shabby] custom icon shader {material.shader.name} -> {shader.name}");
+			return shader;
+		}
+		return Shabby.FindShader("KSP/ScreenSpaceMask");
+	}
+
+	/// <summary>
+	/// The stock method iterates through every material in the icon prefab and replaces some
+	/// stock shaders with 'ScreenSpaceMask'-prefixed ones. All shaders not explicitly checked,
+	/// including custom shaders, are replaced with 'KSP/ScreenSpaceMask'.
+	/// This transpiler inserts logic to check for additional replacements.
+	/// </summary>
+	static IEnumerable<CodeInstruction> Transpiler(IEnumerable<CodeInstruction> instructions)
+	{
+		var code = instructions.ToList();
+		object loc_material = null;
+
+		for (var i = 0; i < code.Count; ++i) {
+			// Material material = sharedMaterials[j];
+			// IL_002C ldloc.3
+			// IL_002D ldloc.s 4
+			// IL_002F ldelem.ref
+			// IL_0030 stloc.s 6
+			if (loc_material == null
+				&& code[i].opcode == OpCodes.Ldloc_3
+				&& code[i+1].opcode == OpCodes.Ldloc_S
+				&& code[i+2].opcode == OpCodes.Ldelem_Ref
+				&& code[i+3].opcode == OpCodes.Stloc_S) {
+				// Extract the stack index of the material local.
+				loc_material = code[i+3].operand;
+			}
+
+			// material2 = new Material(Shader.Find("KSP/ScreenSpaceMask"));
+			// IL_0191 ldstr "KSP/ScreenSpaceMask"
+			// IL_0196 call class UnityEngine.Shader UnityEngine.Shader::Find(string)
+			// IL_019D newobj instance void UnityEngine.Material::.ctor(class UnityEngine.Shader)
+			// IL_01A2 stloc.s 7
+			if (code[i].Is(OpCodes.Ldstr, "KSP/ScreenSpaceMask") && code[i+1].Calls(mInfo_ShaderFind)) {
+				// Replace the call to Shader.Find with FindOverrideIconShader(material).
+				if (loc_material == null) break;
+				code[i].opcode = OpCodes.Ldloc_S;
+				code[i].operand = loc_material;
+				code[i+1].operand = mInfo_FindOverrideIconShader;
+				Debug.Log("[Shabby] patched part icon shader replacement");
+				return code;
+			}
+		}
+
+		Debug.Log("[Shabby] failed to patch part icon shader replacement");
+		return code;
+	}
+}
+
+}

Source/Makefile

@@ -12,6 +12,7 @@ Shabby_FILES := \
 	assembly/VersionReport.cs	\
 	Shabby.cs					\
 	ShabLoader.cs				\
+	IconMaterialPatch.cs		\
 	MaterialDef.cs				\
 	ModelFilter.cs				\
 	MaterialReplacement.cs		\

Source/MaterialReplacement.cs

@@ -54,7 +54,7 @@ public void ApplyToSharedMaterialIfNotIgnored(Renderer renderer)
 }
 
 [HarmonyPatch(typeof(PartLoader), "CompileModel")]
-public class MaterialReplacementPatch
+class MaterialReplacementPatch
 {
 	static void Postfix(ref GameObject __result, ConfigNode partCfg)
 	{

Source/Shabby.cs

@@ -49,6 +49,8 @@ public class Shabby : MonoBehaviour
 
 		static Dictionary<string, Shader> loadedShaders;
 
+		public static readonly Dictionary<string, Shader> iconShaders = new Dictionary<string, Shader>();
+
 		static readonly Dictionary<string, Replacement> nameReplacements = new Dictionary<string, Replacement>();
 
 		public static void AddShader(Shader shader)
@@ -92,12 +94,21 @@ public static void ModuleManagerPostLoad()
 		{
 			var configNodes = GameDatabase.Instance.GetConfigNodes("SHABBY");
 			foreach (var shabbyNode in configNodes) {
-				var replacementNodes = shabbyNode.GetNodes("REPLACE");
-				foreach (var replacementNode in replacementNodes) {
+				foreach (var replacementNode in shabbyNode.GetNodes("REPLACE")) {
 					Replacement replacement = new Replacement(replacementNode);
-
 					nameReplacements[replacement.name] = replacement;
 				}
+
+				foreach (var iconNode in shabbyNode.GetNodes("ICON_SHADER")) {
+					var shader = iconNode.GetValue("shader");
+					var iconShaderName = iconNode.GetValue("iconShader");
+					var iconShader = FindShader(iconShaderName ?? "");
+					if (string.IsNullOrEmpty(shader) || iconShader == null) {
+						Debug.Log($"[Shabby] invalid icon shader specification {shader} -> {iconShaderName}");
+					} else {
+						iconShaders[shader] = iconShader;
+					}
+				}
 			}
 
 			MaterialDefLibrary.Load();

Source/Shabby.csproj

@@ -52,6 +52,7 @@
     </Reference>
   </ItemGroup>
   <ItemGroup>
+    <Compile Include="IconMaterialPatch.cs" />
     <Compile Include="ModelFilter.cs" />
     <Compile Include="Properties\AssemblyInfo.cs" />
     <Compile Include="Shabby.cs" />