Search window respects shader stage

ShaderGraph/com.unity.shadergraph/Editor/Data/Graphs/MaterialSlot.cs

 
         public bool IsCompatibleWith(MaterialSlot otherSlot)
         {
-            if(stageCapability != ShaderStageCapability.All && otherSlot.stageCapability != ShaderStageCapability.All)
-            {
-                if(stageCapability != otherSlot.stageCapability)
-                    return false;
-            }
-
             return otherSlot != null
                 && otherSlot.owner != owner
                 && otherSlot.isInputSlot != isInputSlot
+        }
+
+        public bool IsCompatibleStageWith(MaterialSlot otherSlot)
+        {
+            var candidateStage = otherSlot.stageCapability;
+            return stageCapability == ShaderStageCapability.All || candidateStage == stageCapability;
         }
 
         public virtual string GetDefaultValue(GenerationMode generationMode)

ShaderGraph/com.unity.shadergraph/Editor/Drawing/SearchWindowProvider.cs

                 return !materialSlot.IsCompatibleWith(connectedSlot);
             });
 
+            m_Slots.RemoveAll(slot =>
+            {
+                var materialSlot = (MaterialSlot)slot;
+                return !materialSlot.IsCompatibleStageWith(connectedSlot);
+            });
+
             if (hasSingleSlot && m_Slots.Count == 1)
             {
                 nodeEntries.Add(new NodeEntry