Geometry Nodes: support randomly picking instances from collection

This uses the "id" attribute to randomly pick instances from a collection
for each point.

There is one issue. When the collection is updated (e.g. when an object is
added to it), the nodes modifier is not automatically updated. It seems
like we don't have the infrastructure to support this dependency yet.
The same issue exists in the Boolean modifier and with collision collections.
This should be solved separately soonish.

When "Whole Collection" is disabled, one direct child of the input collection
is instanced at each point. A direct child can be an object or a collection.

Currently, all objects are picked approximately equally often. In the future,
we will provide more control over which point gets which instance.

Differential Revision: https://developer.blender.org/D9884

Ref T82372.

6 files changed, 141 insertions, 50 deletions

diff --git a/source/blender/editors/space_node/drawnode.c b/source/blender/editors/space_node/drawnode.c
index b6744719cad..652c70155ab 100644
--- a/source/blender/editors/space_node/drawnode.c
+++ b/source/blender/editors/space_node/drawnode.c
@@ -3196,6 +3196,9 @@ static void node_geometry_buts_point_instance(uiLayout *layout,
                                               PointerRNA *ptr)
 {
   uiItemR(layout, ptr, "instance_type", DEFAULT_FLAGS | UI_ITEM_R_EXPAND, NULL, ICON_NONE);
+  if (RNA_enum_get(ptr, "instance_type") == GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION) {
+    uiItemR(layout, ptr, "use_whole_collection", DEFAULT_FLAGS, NULL, ICON_NONE);
+  }
 }
 
 static void node_geometry_buts_attribute_fill(uiLayout *layout,
diff --git a/source/blender/makesrna/intern/rna_nodetree.c b/source/blender/makesrna/intern/rna_nodetree.c
index 25bfc70ef8f..cbd677582a9 100644
--- a/source/blender/makesrna/intern/rna_nodetree.c
+++ b/source/blender/makesrna/intern/rna_nodetree.c
@@ -8502,6 +8502,11 @@ static void def_geo_point_instance(StructRNA *srna)
   RNA_def_property_enum_default(prop, GEO_NODE_POINT_INSTANCE_TYPE_OBJECT);
   RNA_def_property_ui_text(prop, "Instance Type", "");
   RNA_def_property_update(prop, NC_NODE | NA_EDITED, "rna_Node_socket_update");
+
+  prop = RNA_def_property(srna, "use_whole_collection", PROP_BOOLEAN, PROP_NONE);
+  RNA_def_property_boolean_negative_sdna(prop, NULL, "custom2", 1);
+  RNA_def_property_ui_text(prop, "Whole Collection", "Instance entire collection on each point");
+  RNA_def_property_update(prop, 0, "rna_Node_socket_update");
 }
 
 static void def_geo_attribute_mix(StructRNA *srna)
diff --git a/source/blender/modifiers/intern/MOD_nodes.cc b/source/blender/modifiers/intern/MOD_nodes.cc
index bbf0c7601f5..9730f734a17 100644
--- a/source/blender/modifiers/intern/MOD_nodes.cc
+++ b/source/blender/modifiers/intern/MOD_nodes.cc
@@ -154,6 +154,7 @@ static void updateDepsgraph(ModifierData *md, const ModifierUpdateDepsgraphConte
   }
 
   /* TODO: Add relations for IDs in settings. */
+  /* TODO: Add dependency for collection changes. */
 }
 
 static void foreachIDLink(ModifierData *md, Object *ob, IDWalkFunc walk, void *userData)
diff --git a/source/blender/nodes/geometry/node_geometry_util.hh b/source/blender/nodes/geometry/node_geometry_util.hh
index 3102690d2ed..7c12611a898 100644
--- a/source/blender/nodes/geometry/node_geometry_util.hh
+++ b/source/blender/nodes/geometry/node_geometry_util.hh
@@ -50,4 +50,7 @@ void poisson_disk_point_elimination(Vector<float3> const *input_points,
                                     float maximum_distance,
                                     float3 boundbox);
 
+Array<uint32_t> get_geometry_element_ids_as_uints(const GeometryComponent &component,
+                                                  const AttributeDomain domain);
+
 }  // namespace blender::nodes
diff --git a/source/blender/nodes/geometry/nodes/node_geo_attribute_randomize.cc b/source/blender/nodes/geometry/nodes/node_geo_attribute_randomize.cc
index 2c3acfc9735..b8b53a17ecc 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_attribute_randomize.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_attribute_randomize.cc
@@ -74,54 +74,57 @@ static float noise_from_index(const int seed, const int hash)
   return BLI_hash_int_01(combined_hash);
 }
 
-static void randomize_attribute(BooleanWriteAttribute &attribute, Span<int> hashes, const int seed)
+static void randomize_attribute(BooleanWriteAttribute &attribute,
+                                Span<uint32_t> hashes,
+                                const int seed)
 {
   MutableSpan<bool> attribute_span = attribute.get_span();
   for (const int i : IndexRange(attribute.size())) {
-    const bool value = noise_from_index(seed, hashes[i]) > 0.5f;
+    const bool value = noise_from_index(seed, (int)hashes[i]) > 0.5f;
     attribute_span[i] = value;
   }
   attribute.apply_span();
 }
 
 static void randomize_attribute(
-    FloatWriteAttribute &attribute, float min, float max, Span<int> hashes, const int seed)
+    FloatWriteAttribute &attribute, float min, float max, Span<uint32_t> hashes, const int seed)
 {
   MutableSpan<float> attribute_span = attribute.get_span();
   for (const int i : IndexRange(attribute.size())) {
-    const float value = noise_from_index(seed, hashes[i]) * (max - min) + min;
+    const float value = noise_from_index(seed, (int)hashes[i]) * (max - min) + min;
     attribute_span[i] = value;
   }
   attribute.apply_span();
 }
 
 static void randomize_attribute(
-    Float3WriteAttribute &attribute, float3 min, float3 max, Span<int> hashes, const int seed)
+    Float3WriteAttribute &attribute, float3 min, float3 max, Span<uint32_t> hashes, const int seed)
 {
   MutableSpan<float3> attribute_span = attribute.get_span();
   for (const int i : IndexRange(attribute.size())) {
-    const float x = noise_from_index_and_mutator(seed, hashes[i], 47);
-    const float y = noise_from_index_and_mutator(seed, hashes[i], 8);
-    const float z = noise_from_index_and_mutator(seed, hashes[i], 64);
+    const float x = noise_from_index_and_mutator(seed, (int)hashes[i], 47);
+    const float y = noise_from_index_and_mutator(seed, (int)hashes[i], 8);
+    const float z = noise_from_index_and_mutator(seed, (int)hashes[i], 64);
     const float3 value = float3(x, y, z) * (max - min) + min;
     attribute_span[i] = value;
   }
   attribute.apply_span();
 }
 
-static Array<int> get_element_hashes(GeometryComponent &component,
-                                     const AttributeDomain domain,
-                                     const int attribute_size)
+Array<uint32_t> get_geometry_element_ids_as_uints(const GeometryComponent &component,
+                                                  const AttributeDomain domain)
 {
+  const int domain_size = component.attribute_domain_size(domain);
+
   /* Hash the reserved name attribute "id" as a (hopefully) stable seed for each point. */
   ReadAttributePtr hash_attribute = component.attribute_try_get_for_read("id", domain);
-  Array<int> hashes(attribute_size);
+  Array<uint32_t> hashes(domain_size);
   if (hash_attribute) {
     BLI_assert(hashes.size() == hash_attribute->size());
     const CPPType &cpp_type = hash_attribute->cpp_type();
     fn::GSpan items = hash_attribute->get_span();
     for (const int i : hashes.index_range()) {
-      hashes[i] = (int)cpp_type.hash(items[i]);
+      hashes[i] = cpp_type.hash(items[i]);
     }
   }
   else {
@@ -129,7 +132,7 @@ static Array<int> get_element_hashes(GeometryComponent &component,
     RandomNumberGenerator rng;
     rng.seed(0);
     for (const int i : hashes.index_range()) {
-      hashes[i] = rng.get_int32();
+      hashes[i] = rng.get_uint32();
     }
   }
 
@@ -154,7 +157,7 @@ static void randomize_attribute(GeometryComponent &component,
     return;
   }
 
-  Array<int> hashes = get_element_hashes(component, domain, attribute->size());
+  Array<uint32_t> hashes = get_geometry_element_ids_as_uints(component, domain);
 
   switch (data_type) {
     case CD_PROP_FLOAT: {
diff --git a/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc b/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
index e030bc3eec6..4274ded2024 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_point_instance.cc
@@ -16,16 +16,21 @@
 
 #include "BKE_mesh.h"
 #include "BKE_persistent_data_handle.hh"
+
+#include "DNA_collection_types.h"
 #include "DNA_mesh_types.h"
 #include "DNA_meshdata_types.h"
 #include "DNA_pointcloud_types.h"
 
+#include "BLI_hash.h"
+
 #include "node_geometry_util.hh"
 
 static bNodeSocketTemplate geo_node_point_instance_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
     {SOCK_OBJECT, N_("Object")},
     {SOCK_COLLECTION, N_("Collection")},
+    {SOCK_INT, N_("Seed"), 0, 0, 0, 0, -10000, 10000},
     {-1, ""},
 };
 
@@ -40,70 +45,141 @@ static void geo_node_point_instance_update(bNodeTree *UNUSED(tree), bNode *node)
 {
   bNodeSocket *object_socket = (bNodeSocket *)BLI_findlink(&node->inputs, 1);
   bNodeSocket *collection_socket = object_socket->next;
+  bNodeSocket *seed_socket = collection_socket->next;
 
   GeometryNodePointInstanceType type = (GeometryNodePointInstanceType)node->custom1;
+  const bool use_whole_collection = node->custom2 == 0;
 
   nodeSetSocketAvailability(object_socket, type == GEO_NODE_POINT_INSTANCE_TYPE_OBJECT);
   nodeSetSocketAvailability(collection_socket, type == GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION);
+  nodeSetSocketAvailability(
+      seed_socket, type == GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION && !use_whole_collection);
+}
+
+static void get_instanced_data__object(const GeoNodeExecParams &params,
+                                       MutableSpan<std::optional<InstancedData>> r_instances_data)
+{
+  bke::PersistentObjectHandle object_handle = params.get_input<bke::PersistentObjectHandle>(
+      "Object");
+  Object *object = params.handle_map().lookup(object_handle);
+  if (object == params.self_object()) {
+    object = nullptr;
+  }
+  if (object != nullptr) {
+    InstancedData instance;
+    instance.type = INSTANCE_DATA_TYPE_OBJECT;
+    instance.data.object = object;
+    r_instances_data.fill(instance);
+  }
+}
+
+static void get_instanced_data__collection(
+    const GeoNodeExecParams &params,
+    const GeometryComponent &component,
+    MutableSpan<std::optional<InstancedData>> r_instances_data)
+{
+  const bNode &node = params.node();
+  bke::PersistentCollectionHandle collection_handle =
+      params.get_input<bke::PersistentCollectionHandle>("Collection");
+  Collection *collection = params.handle_map().lookup(collection_handle);
+  if (collection != nullptr) {
+    const bool use_whole_collection = node.custom2 == 0;
+    if (use_whole_collection) {
+      InstancedData instance;
+      instance.type = INSTANCE_DATA_TYPE_COLLECTION;
+      instance.data.collection = collection;
+      r_instances_data.fill(instance);
+    }
+    else {
+      Vector<InstancedData> possible_instances;
+      /* Direct child objects are instanced as objects. */
+      LISTBASE_FOREACH (CollectionObject *, cob, &collection->gobject) {
+        Object *object = cob->ob;
+        InstancedData instance;
+        instance.type = INSTANCE_DATA_TYPE_OBJECT;
+        instance.data.object = object;
+        possible_instances.append(instance);
+      }
+      /* Direct child collections are instanced as collections. */
+      LISTBASE_FOREACH (CollectionChild *, child, &collection->children) {
+        Collection *child_collection = child->collection;
+        InstancedData instance;
+        instance.type = INSTANCE_DATA_TYPE_COLLECTION;
+        instance.data.collection = child_collection;
+        possible_instances.append(instance);
+      }
+
+      if (!possible_instances.is_empty()) {
+        const int seed = params.get_input<int>("Seed");
+        Array<uint32_t> ids = get_geometry_element_ids_as_uints(component, ATTR_DOMAIN_POINT);
+        for (const int i : r_instances_data.index_range()) {
+          const int index = BLI_hash_int_2d(ids[i], seed) % possible_instances.size();
+          r_instances_data[i] = possible_instances[index];
+        }
+      }
+    }
+  }
+}
+
+static Array<std::optional<InstancedData>> get_instanced_data(const GeoNodeExecParams &params,
+                                                              const GeometryComponent &component,
+                                                              const int amount)
+{
+  const bNode &node = params.node();
+  const GeometryNodePointInstanceType type = (GeometryNodePointInstanceType)node.custom1;
+
+  Array<std::optional<InstancedData>> instances_data(amount);
+
+  switch (type) {
+    case GEO_NODE_POINT_INSTANCE_TYPE_OBJECT: {
+      get_instanced_data__object(params, instances_data);
+      break;
+    }
+    case GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION: {
+      get_instanced_data__collection(params, component, instances_data);
+      break;
+    }
+  }
+  return instances_data;
 }
 
 static void add_instances_from_geometry_component(InstancesComponent &instances,
                                                   const GeometryComponent &src_geometry,
-                                                  Object *object,
-                                                  Collection *collection)
+                                                  const GeoNodeExecParams &params)
 {
+  const AttributeDomain domain = ATTR_DOMAIN_POINT;
+
+  const int domain_size = src_geometry.attribute_domain_size(domain);
+  Array<std::optional<InstancedData>> instances_data = get_instanced_data(
+      params, src_geometry, domain_size);
+
   Float3ReadAttribute positions = src_geometry.attribute_get_for_read<float3>(
-      "position", ATTR_DOMAIN_POINT, {0, 0, 0});
+      "position", domain, {0, 0, 0});
   Float3ReadAttribute rotations = src_geometry.attribute_get_for_read<float3>(
-      "rotation", ATTR_DOMAIN_POINT, {0, 0, 0});
+      "rotation", domain, {0, 0, 0});
   Float3ReadAttribute scales = src_geometry.attribute_get_for_read<float3>(
-      "scale", ATTR_DOMAIN_POINT, {1, 1, 1});
+      "scale", domain, {1, 1, 1});
 
-  for (const int i : IndexRange(positions.size())) {
-    if (object != nullptr) {
-      instances.add_instance(object, positions[i], rotations[i], scales[i]);
-    }
-    if (collection != nullptr) {
-      instances.add_instance(collection, positions[i], rotations[i], scales[i]);
+  for (const int i : IndexRange(domain_size)) {
+    if (instances_data[i].has_value()) {
+      instances.add_instance(*instances_data[i], positions[i], rotations[i], scales[i]);
     }
   }
 }
 
 static void geo_node_point_instance_exec(GeoNodeExecParams params)
 {
-  GeometryNodePointInstanceType type = (GeometryNodePointInstanceType)params.node().custom1;
   GeometrySet geometry_set = params.extract_input<GeometrySet>("Geometry");
   GeometrySet geometry_set_out;
 
-  Object *object = nullptr;
-  Collection *collection = nullptr;
-
-  if (type == GEO_NODE_POINT_INSTANCE_TYPE_OBJECT) {
-    bke::PersistentObjectHandle object_handle = params.extract_input<bke::PersistentObjectHandle>(
-        "Object");
-    object = params.handle_map().lookup(object_handle);
-    /* Avoid accidental recursion of instances. */
-    if (object == params.self_object()) {
-      object = nullptr;
-    }
-  }
-  else if (type == GEO_NODE_POINT_INSTANCE_TYPE_COLLECTION) {
-    bke::PersistentCollectionHandle collection_handle =
-        params.extract_input<bke::PersistentCollectionHandle>("Collection");
-    collection = params.handle_map().lookup(collection_handle);
-  }
-
   InstancesComponent &instances = geometry_set_out.get_component_for_write<InstancesComponent>();
   if (geometry_set.has<MeshComponent>()) {
     add_instances_from_geometry_component(
-        instances, *geometry_set.get_component_for_read<MeshComponent>(), object, collection);
+        instances, *geometry_set.get_component_for_read<MeshComponent>(), params);
   }
   if (geometry_set.has<PointCloudComponent>()) {
     add_instances_from_geometry_component(
-        instances,
-        *geometry_set.get_component_for_read<PointCloudComponent>(),
-        object,
-        collection);
+        instances, *geometry_set.get_component_for_read<PointCloudComponent>(), params);
   }
 
   params.set_output("Geometry", std::move(geometry_set_out));