1 files changed, 35 insertions, 0 deletions
diff --git a/source/blender/nodes/geometry/node_geometry_tree.cc b/source/blender/nodes/geometry/node_geometry_tree.cc
index d4a9805f311..69f9f7fb4ed 100644
--- a/source/blender/nodes/geometry/node_geometry_tree.cc
+++ b/source/blender/nodes/geometry/node_geometry_tree.cc
@@ -20,16 +20,49 @@
 
 #include "NOD_geometry.h"
 
+#include "BKE_context.h"
 #include "BKE_node.h"
+#include "BKE_object.h"
 
 #include "BLT_translation.h"
 
+#include "DNA_modifier_types.h"
 #include "DNA_node_types.h"
+#include "DNA_space_types.h"
 
 #include "RNA_access.h"
 
 bNodeTreeType *ntreeType_Geometry;
 
+static void geometry_node_tree_get_from_context(const bContext *C,
+                                                bNodeTreeType *UNUSED(treetype),
+                                                bNodeTree **r_ntree,
+                                                ID **r_id,
+                                                ID **r_from)
+{
+  ViewLayer *view_layer = CTX_data_view_layer(C);
+  Object *ob = OBACT(view_layer);
+
+  if (ob == nullptr) {
+    return;
+  }
+
+  const ModifierData *md = BKE_object_active_modifier(ob);
+
+  if (md == nullptr) {
+    return;
+  }
+
+  if (md->type == eModifierType_Nodes) {
+    NodesModifierData *nmd = (NodesModifierData *)md;
+    if (nmd->node_group != nullptr) {
+      *r_from = &ob->id;
+      *r_id = &ob->id;
+      *r_ntree = nmd->node_group;
+    }
+  }
+}
+
 void register_node_tree_type_geo(void)
 {
   bNodeTreeType *tt = ntreeType_Geometry = static_cast<bNodeTreeType *>(
@@ -41,5 +74,7 @@ void register_node_tree_type_geo(void)
   strcpy(tt->ui_description, N_("Geometry nodes"));
   tt->rna_ext.srna = &RNA_GeometryNodeTree;
 
+  tt->get_from_context = geometry_node_tree_get_from_context;
+
   ntreeTypeAdd(tt);
 }