1 files changed, 29 insertions, 6 deletions

diff --git a/source/blender/nodes/shader/node_shader_util.cc b/source/blender/nodes/shader/node_shader_util.cc
index 16a4c5bae88..728e2760f9a 100644
--- a/source/blender/nodes/shader/node_shader_util.cc
+++ b/source/blender/nodes/shader/node_shader_util.cc
@@ -162,8 +162,21 @@ static void data_from_gpu_stack_list(ListBase *sockets, bNodeStack **ns, GPUNode
   }
 }
 
-bNode *nodeGetActiveTexture(bNodeTree *ntree)
+bool nodeSupportsActiveFlag(const bNode *node, int sub_activity)
+{
+  BLI_assert(ELEM(sub_activity, NODE_ACTIVE_TEXTURE, NODE_ACTIVE_PAINT_CANVAS));
+  switch (sub_activity) {
+    case NODE_ACTIVE_TEXTURE:
+      return node->typeinfo->nclass == NODE_CLASS_TEXTURE;
+    case NODE_ACTIVE_PAINT_CANVAS:
+      return ELEM(node->type, SH_NODE_TEX_IMAGE, SH_NODE_ATTRIBUTE);
+  }
+  return false;
+}
+
+static bNode *node_get_active(bNodeTree *ntree, int sub_activity)
 {
+  BLI_assert(ELEM(sub_activity, NODE_ACTIVE_TEXTURE, NODE_ACTIVE_PAINT_CANVAS));
   /* this is the node we texture paint and draw in textured draw */
   bNode *inactivenode = nullptr, *activetexnode = nullptr, *activegroup = nullptr;
   bool hasgroup = false;
@@ -173,14 +186,14 @@ bNode *nodeGetActiveTexture(bNodeTree *ntree)
   }
 
   LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
-    if (node->flag & NODE_ACTIVE_TEXTURE) {
+    if (node->flag & sub_activity) {
       activetexnode = node;
       /* if active we can return immediately */
       if (node->flag & NODE_ACTIVE) {
         return node;
       }
     }
-    else if (!inactivenode && node->typeinfo->nclass == NODE_CLASS_TEXTURE) {
+    else if (!inactivenode && nodeSupportsActiveFlag(node, sub_activity)) {
       inactivenode = node;
     }
     else if (node->type == NODE_GROUP) {
@@ -195,7 +208,7 @@ bNode *nodeGetActiveTexture(bNodeTree *ntree)
 
   /* first, check active group for textures */
   if (activegroup) {
-    bNode *tnode = nodeGetActiveTexture((bNodeTree *)activegroup->id);
+    bNode *tnode = node_get_active((bNodeTree *)activegroup->id, sub_activity);
     /* active node takes priority, so ignore any other possible nodes here */
     if (tnode) {
       return tnode;
@@ -210,8 +223,8 @@ bNode *nodeGetActiveTexture(bNodeTree *ntree)
     /* node active texture node in this tree, look inside groups */
     LISTBASE_FOREACH (bNode *, node, &ntree->nodes) {
       if (node->type == NODE_GROUP) {
-        bNode *tnode = nodeGetActiveTexture((bNodeTree *)node->id);
-        if (tnode && ((tnode->flag & NODE_ACTIVE_TEXTURE) || !inactivenode)) {
+        bNode *tnode = node_get_active((bNodeTree *)node->id, sub_activity);
+        if (tnode && ((tnode->flag & sub_activity) || !inactivenode)) {
           return tnode;
         }
       }
@@ -221,6 +234,16 @@ bNode *nodeGetActiveTexture(bNodeTree *ntree)
   return inactivenode;
 }
 
+bNode *nodeGetActiveTexture(bNodeTree *ntree)
+{
+  return node_get_active(ntree, NODE_ACTIVE_TEXTURE);
+}
+
+bNode *nodeGetActivePaintCanvas(bNodeTree *ntree)
+{
+  return node_get_active(ntree, NODE_ACTIVE_PAINT_CANVAS);
+}
+
 void ntreeExecGPUNodes(bNodeTreeExec *exec, GPUMaterial *mat, bNode *output_node)
 {
   bNodeExec *nodeexec;