Swizzle Node

MaterialGraphProject/Assets/UnityShaderEditor/Editor/Data/Nodes/Channel/SwizzleNode.cs

-using UnityEditor.Graphing;
-using System.Linq;
+using UnityEditor.Graphing;
+using UnityEditor.ShaderGraph.Drawing.Controls;
+using UnityEngine;
-    /* [Title("Channel/Swizzle")]
-     public class SwizzleNode : Function1Input
-     {
-         public enum SwizzleChannel
-         {
-             R = 0,
-             G = 1,
-             B = 2,
-             A = 3,
-         }
+    [Title("Channel/Swizzle")]
+    public class SwizzleNode : AbstractMaterialNode, IGeneratesBodyCode, IGeneratesFunction
+    {
+        public SwizzleNode()
+        {
+            name = "Swizzle";
+            UpdateNodeAfterDeserialization();
+        }
-         [SerializeField]
-         private SwizzleChannel[] m_SwizzleChannels = new SwizzleChannel[4] { SwizzleChannel.R, SwizzleChannel.G, SwizzleChannel.B, SwizzleChannel.A };
+        const int InputSlotId = 0;
+        const int OutputSlotId = 1;
+        const string kInputSlotName = "In";
+        const string kOutputSlotName = "Out";
-         public SwizzleChannel[] swizzleChannels
-         {
-             get { return m_SwizzleChannels; }
-             set
-             {
-                 if (m_SwizzleChannels == value)
-                     return;
+        public override bool hasPreview
+        {
+            get { return true; }
+        }
-                 m_SwizzleChannels = value;
-                 if (onModified != null)
-                 {
-                     onModified(this, ModificationScope.Graph);
-                 }
-             }
-         }
+        string GetFunctionName()
+        {
+            return "Unity_Swizzle_" + precision + "_" + GuidEncoder.Encode(guid);
+        }
-         public SwizzleNode()
-         {
-             name = "Swizzle";
-         }
+        public sealed override void UpdateNodeAfterDeserialization()
+        {
+            AddSlot(new DynamicVectorMaterialSlot(InputSlotId, kInputSlotName, kInputSlotName, SlotType.Input, Vector4.zero));
+            AddSlot(new Vector4MaterialSlot(OutputSlotId, kOutputSlotName, kOutputSlotName, SlotType.Output, Vector4.zero));
+            RemoveSlotsNameNotMatching(new[] { InputSlotId, OutputSlotId });
+        }
-         protected override string GetFunctionName()
-         {
-             return "";
-         }
+        static Dictionary<TextureChannel, string> m_ComponentList = new Dictionary<TextureChannel, string>
+        {
+            {TextureChannel.Red, ".r" },
+            {TextureChannel.Green, ".g" },
+            {TextureChannel.Blue, ".b" },
+            {TextureChannel.Alpha, ".a" },
+        };
-         protected override string GetFunctionCallBody(string inputValue)
-         {
-             string[] channelNames = { "r", "g", "b", "a" };
-             var inputSlot = FindInputSlot<MaterialSlot>(InputSlotId);
-             var outputSlot = FindOutputSlot<MaterialSlot>(OutputSlotId);
+        [SerializeField]
+        private TextureChannel m_RedChannel;
-             if (inputSlot == null)
-                 return "1.0";
-             if (outputSlot == null)
-                 return "1.0";
+        [ChannelEnumControl("Red Out")]
+        public TextureChannel redChannel
+        {
+            get { return m_RedChannel; }
+            set
+            {
+                if (m_RedChannel == value)
+                    return;
-             int numInputChannels = (int)SlotValueHelper.GetChannelCount(inputSlot.concreteValueType);
-             int numOutputChannels = (int)SlotValueHelper.GetChannelCount(outputSlot.concreteValueType);
+                m_RedChannel = value;
+                if (onModified != null)
+                {
+                    onModified(this, ModificationScope.Graph);
+                }
+            }
+        }
-             //int numInputChannels = (int)inputSlot.concreteValueType;
-             int numOutputChannels = (int)outputSlot.concreteValueType;
-             if (owner.GetEdges(inputSlot.slotReference).ToList().Count() == 0)
-                 numInputChannels = 0;
-             if (owner.GetEdges(outputSlot.slotReference).ToList().Count() == 0)
-                 numOutputChannels = 0;
+        [SerializeField]
+        private TextureChannel m_GreenChannel;
-             if (numOutputChannels == 0)
-             {
-                 return "1.0";
-             }
+        [ChannelEnumControl("Green Out")]
+        public TextureChannel greenChannel
+        {
+            get { return m_GreenChannel; }
+            set
+            {
+                if (m_GreenChannel == value)
+                    return;
-             string outputString = precision + "4(";
-             //float4(1.0,1.0,1.0,1.0)
-             if (numInputChannels == 0)
-             {
-                 outputString += "1.0, 1.0, 1.0, 1.0).";
-             }
-             else
-             {
-                 //float4(arg1.
-                 outputString += inputValue + ".";
+                m_GreenChannel = value;
+                if (onModified != null)
+                {
+                    onModified(this, ModificationScope.Graph);
+                }
+            }
+        }
+
+        [SerializeField]
+        private TextureChannel m_BlueChannel;
+
+        [ChannelEnumControl("Blue Out")]
+        public TextureChannel blueChannel
+        {
+            get { return m_BlueChannel; }
+            set
+            {
+                if (m_BlueChannel == value)
+                    return;
+
+                m_BlueChannel = value;
+                if (onModified != null)
+                {
+                    onModified(this, ModificationScope.Graph);
+                }
+            }
+        }
+
+        [SerializeField]
+        private TextureChannel m_AlphaChannel;
+
+        [ChannelEnumControl("Alpha Out")]
+        public TextureChannel alphaChannel
+        {
+            get { return m_AlphaChannel; }
+            set
+            {
+                if (m_AlphaChannel == value)
+                    return;
+
+                m_AlphaChannel = value;
+                if (onModified != null)
+                {
+                    onModified(this, ModificationScope.Graph);
+                }
+            }
+        }
+
+        void ValidateChannelCount()
+        {
+            int channelCount = (int)SlotValueHelper.GetChannelCount(FindInputSlot<MaterialSlot>(InputSlotId).concreteValueType);
+            if ((int)redChannel >= channelCount)
+                redChannel = TextureChannel.Red;
+            if ((int)greenChannel >= channelCount)
+                greenChannel = TextureChannel.Red;
+            if ((int)blueChannel >= channelCount)
+                blueChannel = TextureChannel.Red;
+            if ((int)alphaChannel >= channelCount)
+                alphaChannel = TextureChannel.Red;
+        }
+
+        string GetFunctionPrototype(string inArg, string outArg)
+        {
+            return string.Format("void {0} ({1} {2}, out {3} {4})", GetFunctionName(), 
+                ConvertConcreteSlotValueTypeToString(precision, FindInputSlot<MaterialSlot>(InputSlotId).concreteValueType), inArg, 
+                ConvertConcreteSlotValueTypeToString(precision, FindOutputSlot<MaterialSlot>(OutputSlotId).concreteValueType), outArg);
+        }
-                 //float4(arg1.xy
-                 int i = 0;
-                 for (; i < numInputChannels; ++i)
-                 {
-                     int channel = (int)m_SwizzleChannels[i];
-                     outputString += channelNames[channel];
-                 }
+        public void GenerateNodeCode(ShaderGenerator visitor, GenerationMode generationMode)
+        {
+            ValidateChannelCount();
+            string inputValue = GetSlotValue(InputSlotId, generationMode);
+            string outputValue = GetSlotValue(OutputSlotId, generationMode);
+            visitor.AddShaderChunk(string.Format("{0} {1};", ConvertConcreteSlotValueTypeToString(precision, FindOutputSlot<MaterialSlot>(OutputSlotId).concreteValueType), GetVariableNameForSlot(OutputSlotId)), true);
+            visitor.AddShaderChunk(GetFunctionCallBody(inputValue, outputValue), true);
+        }
+
+        string GetFunctionCallBody(string inputValue, string outputValue)
+        {
+            return GetFunctionName() + " (" + inputValue + ", " + outputValue + ");";
+        }
+
+        public void GenerateNodeFunction(ShaderGenerator visitor, GenerationMode generationMode)
+        {
+            ValidateChannelCount();
+            var outputString = new ShaderGenerator();
+            outputString.AddShaderChunk(GetFunctionPrototype("In", "Out"), true);
+            outputString.AddShaderChunk("{", true);
+            outputString.Indent();
+
+            outputString.AddShaderChunk(string.Format("Out = {0} ({1}, {2}, {3}, {4});",
+                ConvertConcreteSlotValueTypeToString(precision, FindOutputSlot<MaterialSlot>(OutputSlotId).concreteValueType),
+                kInputSlotName + m_ComponentList[m_RedChannel],
+                kInputSlotName + m_ComponentList[m_GreenChannel],
+                kInputSlotName + m_ComponentList[m_BlueChannel],
+                kInputSlotName + m_ComponentList[m_AlphaChannel]), true);
-                 //float4(arg1.xy, 1.0, 1.0)
-                 for (; i < 4; i++)
-                 {
-                     outputString += ", 1.0";
-                 }
-                 outputString += ").";
-             }
+            outputString.Deindent();
+            outputString.AddShaderChunk("}", true);
-             //float4(arg1.xy, 1.0, 1.0).rg
-             for (int j = 0; j < numOutputChannels; ++j)
-             {
-                 outputString += channelNames[j];
-             }
-             return outputString;
-         }
-     }*/
-}
+            visitor.AddShaderChunk(outputString.GetShaderString(0), true);
+        }
+    }
+}