Addressing pull request e3eb1cc

bigorangemachine · Dec 8, 2014 · 423cf6e · 423cf6e
1 parent 1217de4
commit 423cf6e
Showing 1 changed file with 48 additions and 13 deletions.
diff --git a/src/KRSVesselShot.cs b/src/KRSVesselShot.cs
@@ -101,14 +101,15 @@ public KRSVesselShot()
             LoadShaders();
             UpdateShipBounds();
 
-            GameEvents.onPartAttach.Add(PartAttached);
-            GameEvents.onPartRemove.Add(PartRemoved);
+
+            GameEvents.onPartAttach.Add(PartModified);
+            GameEvents.onPartRemove.Add(PartModified);
         }
 
         ~KRSVesselShot()
         {
-            GameEvents.onPartAttach.Remove(PartAttached);
-            GameEvents.onPartRemove.Remove(PartRemoved);
+            GameEvents.onPartAttach.Remove(PartModified);
+            GameEvents.onPartRemove.Remove(PartModified);
         }
 
         // Sets up Orthographic and Perspective camera.
@@ -167,28 +168,52 @@ private void ReplacePartShaders(Part part)
             var model = part.transform.Find("model");
             if (!model) return;
 
-            foreach (var r in model.GetComponentsInChildren<MeshRenderer>())
+            Dictionary<MeshRenderer, Shader> MeshRendererLibrary = new Dictionary<MeshRenderer,Shader>();
+
+            foreach (MeshRenderer mr in model.GetComponentsInChildren<MeshRenderer>())
             {
                 Material mat;
-                if (Materials.TryGetValue(r.material.shader.name, out mat))
+                if (Materials.TryGetValue(mr.material.shader.name, out mat))
                 {
-                    r.material.shader = mat.shader;
+                    if (!MeshRendererLibrary.ContainsKey(mr))
+                    {
+                        MeshRendererLibrary.Add(mr, mr.material.shader);
+                    }
+                    mr.material.shader = mat.shader;
                 }
                 else
                 {
-                    MonoBehaviour.print("[Warning] " + this.GetType().Name + "No replacement for " + r.material.shader + " in " + part + "/*/" + r);
+                    MonoBehaviour.print("[Warning] " + this.GetType().Name + "No replacement for " + mr.material.shader + " in " + part + "/*/" + mr);
                 }
             }
+            if (!PartShaderLibrary.ContainsKey(part))
+            {
+                PartShaderLibrary.Add(part, MeshRendererLibrary);
+            }
         }
+
+        Dictionary<Part,Dictionary<MeshRenderer, Shader>> PartShaderLibrary = new Dictionary<Part,Dictionary<MeshRenderer,Shader>>();
 
-        private void PartAttached(GameEvents.HostTargetAction<Part, Part> data)
+        private void RestorePartShaders(Part part)
         {
-            ReplacePartShaders(data.host);
-            ReplacePartShaders(data.target);
-            UpdateShipBounds();
+            var model = part.transform.Find("model");
+            if (!model) return;
+
+            Dictionary<MeshRenderer,Shader> MeshRendererLibrary;
+            if (PartShaderLibrary.TryGetValue(part, out MeshRendererLibrary))
+            {
+                foreach (MeshRenderer mr in model.GetComponentsInChildren<MeshRenderer>())
+                {
+                    Shader OldShader;
+                    if (MeshRendererLibrary.TryGetValue(mr, out OldShader))
+                    {
+                        mr.material.shader = OldShader;
+                    }
+                }
+            }            
         }
 
-        private void PartRemoved(GameEvents.HostTargetAction<Part, Part> data)
+        private void PartModified(GameEvents.HostTargetAction<Part, Part> data)
         {
             UpdateShipBounds();
         }
@@ -226,6 +251,11 @@ public Vector3 GetShipSize()
 
         public void GenTexture(Vector3 direction, int imageWidth = -1, int imageHeight = -1)
         {
+            foreach (Part p in EditorLogic.fetch.ship)
+            {
+                ReplacePartShaders(p);
+            }
+
             var minusDir = -direction;
             this.Camera.clearFlags = CameraClearFlags.SolidColor;
             if(this.Effects["Blue Print"].Enabled){
@@ -340,6 +370,11 @@ public void GenTexture(Vector3 direction, int imageWidth = -1, int imageHeight =
                     }
                 }
             }
+
+            foreach (Part p in EditorLogic.fetch.ship)
+            {
+                RestorePartShaders(p);
+            }
         }
 
         private void SaveTexture(String fileName)