source/blender/blenlib/BLI_bit_array_vector.hh

@@ -66,6 +66,11 @@ class BitArrayVector {
     return aligned_group_size_ == 0 ? 0 : data_.size() / aligned_group_size_;
   }
 
+  int64_t group_size() const
+  {
+    return group_size_;
+  }
+
   IndexRange index_range() const
   {
     return IndexRange{this->size()};

source/blender/blenlib/BLI_bit_span.hh

@@ -326,6 +326,21 @@ class MutableBoundedBitSpan : public MutableBitSpan {
   }
 };
 
+/**
+ * This is overloaded in BLI_bit_vector.hh. The purpose is to make passing #BitVector into bit span
+ * operations more simpler and efficient (interpreting it as `BoundedBitSpan` instead of just
+ * `BitSpan`).
+ */
+template<typename T> inline T to_best_bit_span(const T &data)
+{
+  static_assert(is_same_any_v<std::decay_t<T>,
+                              BitSpan,
+                              MutableBitSpan,
+                              BoundedBitSpan,
+                              MutableBoundedBitSpan>);
+  return data;
+}
+
 template<typename... Args>
 constexpr bool all_bounded_spans =
     (is_same_any_v<std::decay_t<Args>, BoundedBitSpan, MutableBoundedBitSpan> && ...);

source/blender/blenlib/BLI_bit_span_ops.hh

@@ -6,6 +6,7 @@
 
 namespace blender::bits {
 
+namespace detail {
 /**
  * Evaluates the expression on one or more bit spans and stores the result in the first.
  *
@@ -92,7 +93,7 @@ inline bool any_set_expr(ExprFn &&expr, const FirstBitSpanT &first_arg, const Bi
     /* Fallback or arbitrary bit spans. This could be implemented more efficiently but adds more
      * complexity and is not necessary yet. */
     for (const int64_t i : IndexRange(size)) {
-      const BitInt result = fn(BitInt(first_arg[i].test()), BitInt(args[i].test())...);
+      const BitInt result = expr(BitInt(first_arg[i].test()), BitInt(args[i].test())...);
       if (result != 0) {
         return true;
       }
@@ -137,7 +138,7 @@ inline void foreach_1_index_expr(ExprFn &&expr,
     if (const int64_t final_bits = first_arg.final_bits_num()) {
       BitInt tmp = expr(first_data[full_ints_num] >> first_arg.offset(),
                         (*args.data()[full_ints_num] >> args.offset())...) &
-                   mask_first_n_bits(size);
+                   mask_first_n_bits(final_bits);
       const int64_t offset = full_ints_num << BitToIntIndexShift;
       while (tmp != 0) {
         const int index = bitscan_forward_uint64(tmp);
@@ -158,6 +159,32 @@ inline void foreach_1_index_expr(ExprFn &&expr,
   }
 }
 
+}  // namespace detail
+
+template<typename ExprFn, typename FirstBitSpanT, typename... BitSpanT>
+inline void mix_into_first_expr(ExprFn &&expr,
+                                const FirstBitSpanT &first_arg,
+                                const BitSpanT &...args)
+{
+  detail::mix_into_first_expr(expr, to_best_bit_span(first_arg), to_best_bit_span(args)...);
+}
+
+template<typename ExprFn, typename FirstBitSpanT, typename... BitSpanT>
+inline bool any_set_expr(ExprFn &&expr, const FirstBitSpanT &first_arg, const BitSpanT &...args)
+{
+  return detail::any_set_expr(expr, to_best_bit_span(first_arg), to_best_bit_span(args)...);
+}
+
+template<typename ExprFn, typename HandleFn, typename FirstBitSpanT, typename... BitSpanT>
+inline void foreach_1_index_expr(ExprFn &&expr,
+                                 HandleFn &&handle,
+                                 const FirstBitSpanT &first_arg,
+                                 const BitSpanT &...args)
+{
+  detail::foreach_1_index_expr(
+      expr, handle, to_best_bit_span(first_arg), to_best_bit_span(args)...);
+}
+
 template<typename FirstBitSpanT, typename... BitSpanT>
 inline void inplace_or(FirstBitSpanT &first_arg, const BitSpanT &...args)
 {

source/blender/blenlib/BLI_bit_vector.hh

@@ -170,12 +170,12 @@ class BitVector {
     return move_assign_container(*this, std::move(other));
   }
 
-  operator BitSpan() const
+  operator BoundedBitSpan() const
   {
     return {data_, IndexRange(size_in_bits_)};
   }
 
-  operator MutableBitSpan()
+  operator MutableBoundedBitSpan()
   {
     return {data_, IndexRange(size_in_bits_)};
   }
@@ -367,6 +367,18 @@ class BitVector {
   }
 };
 
+template<int64_t InlineBufferCapacity, typename Allocator>
+inline BoundedBitSpan to_best_bit_span(const BitVector<InlineBufferCapacity, Allocator> &data)
+{
+  return data;
+}
+
+template<int64_t InlineBufferCapacity, typename Allocator>
+inline MutableBoundedBitSpan to_best_bit_span(BitVector<InlineBufferCapacity, Allocator> &data)
+{
+  return data;
+}
+
 }  // namespace blender::bits
 
 namespace blender {

source/blender/nodes/intern/geometry_nodes_lazy_function.cc

@@ -18,6 +18,8 @@
 #include "NOD_multi_function.hh"
 #include "NOD_node_declaration.hh"
 
+#include "BLI_bit_array_vector.hh"
+#include "BLI_bit_span_ops.hh"
 #include "BLI_cpp_types.hh"
 #include "BLI_dot_export.hh"
 #include "BLI_hash.h"
@@ -2298,15 +2300,21 @@ struct GeometryNodesLazyFunctionGraphBuilder {
     this->build_attribute_references(
         relations_by_node, attribute_reference_keys, attribute_reference_infos);
 
-    MultiValueMap<const bNodeSocket *, int> referenced_by_field_socket;
-    MultiValueMap<const bNodeSocket *, int> propagated_to_geometry_socket;
+    const int sockets_num = btree_.all_sockets().size();
+    const int attribute_references_num = attribute_reference_keys.size();
+
+    /* The code below uses #BitArrayVector to store a set of attribute references per socket. Each
+     * socket has a bit span where each bit corresponds to one attribute reference. */
+    BitArrayVector<> referenced_by_field_socket(sockets_num, attribute_references_num, false);
+    BitArrayVector<> propagated_to_geometry_socket(sockets_num, attribute_references_num, false);
     this->gather_referenced_and_potentially_propagated_data(relations_by_node,
                                                             attribute_reference_keys,
                                                             attribute_reference_infos,
                                                             referenced_by_field_socket,
                                                             propagated_to_geometry_socket);
 
-    MultiValueMap<const bNodeSocket *, int> required_propagated_to_geometry_socket;
+    BitArrayVector<> required_propagated_to_geometry_socket(
+        sockets_num, attribute_references_num, false);
     this->gather_required_propagated_data(relations_by_node,
                                           attribute_reference_keys,
                                           referenced_by_field_socket,
@@ -2407,10 +2415,10 @@ struct GeometryNodesLazyFunctionGraphBuilder {
       const Span<const aal::RelationsInNode *> relations_by_node,
       const Span<AttributeReferenceKey> attribute_reference_keys,
       const Span<AttributeReferenceInfo> attribute_reference_infos,
-      MultiValueMap<const bNodeSocket *, int> &r_referenced_by_field_socket,
-      MultiValueMap<const bNodeSocket *, int> &r_propagated_to_geometry_socket)
+      BitArrayVector<> &r_referenced_by_field_socket,
+      BitArrayVector<> &r_propagated_to_geometry_socket)
   {
-    /* Initialize maps. */
+    /* Insert initial referenced/propagated attributes. */
     for (const int key_index : attribute_reference_keys.index_range()) {
       const AttributeReferenceKey &key = attribute_reference_keys[key_index];
       const AttributeReferenceInfo &info = attribute_reference_infos[key_index];
@@ -2418,7 +2426,7 @@ struct GeometryNodesLazyFunctionGraphBuilder {
         case AttributeReferenceKeyType::InputField: {
           for (const bNode *bnode : btree_.group_input_nodes()) {
             const bNodeSocket &bsocket = bnode->output_socket(key.index);
-            r_referenced_by_field_socket.add(&bsocket, key_index);
+            r_referenced_by_field_socket[bsocket.index_in_tree()][key_index].set();
           }
           break;
         }
@@ -2426,34 +2434,28 @@ struct GeometryNodesLazyFunctionGraphBuilder {
           break;
         }
         case AttributeReferenceKeyType::Socket: {
-          r_referenced_by_field_socket.add(key.bsocket, key_index);
+          r_referenced_by_field_socket[key.bsocket->index_in_tree()][key_index].set();
           break;
         }
       }
       for (const bNodeSocket *geometry_bsocket : info.initial_geometry_sockets) {
-        r_propagated_to_geometry_socket.add(geometry_bsocket, key_index);
+        r_propagated_to_geometry_socket[geometry_bsocket->index_in_tree()][key_index].set();
       }
     }
     /* Propagate attribute usages from left to right. */
     for (const bNode *bnode : btree_.toposort_left_to_right()) {
       for (const bNodeSocket *bsocket : bnode->input_sockets()) {
         if (bsocket->is_available()) {
-          Vector<int> referenced_keys;
-          Vector<int> propagated_keys;
+          const int dst_index = bsocket->index_in_tree();
+          MutableBoundedBitSpan referenced_dst = r_referenced_by_field_socket[dst_index];
+          MutableBoundedBitSpan propagated_dst = r_propagated_to_geometry_socket[dst_index];
           for (const bNodeLink *blink : bsocket->directly_linked_links()) {
             if (blink->is_used()) {
-              referenced_keys.extend_non_duplicates(
-                  r_referenced_by_field_socket.lookup(blink->fromsock));
-              propagated_keys.extend_non_duplicates(
-                  r_propagated_to_geometry_socket.lookup(blink->fromsock));
+              const int src_index = blink->fromsock->index_in_tree();
+              referenced_dst |= r_referenced_by_field_socket[src_index];
+              propagated_dst |= r_propagated_to_geometry_socket[src_index];
             }
           }
-          if (!referenced_keys.is_empty()) {
-            r_referenced_by_field_socket.add_multiple(bsocket, referenced_keys);
-          }
-          if (!propagated_keys.is_empty()) {
-            r_propagated_to_geometry_socket.add_multiple(bsocket, propagated_keys);
-          }
         }
       }
       const aal::RelationsInNode &relations = *relations_by_node[bnode->index()];
@@ -2463,8 +2465,8 @@ struct GeometryNodesLazyFunctionGraphBuilder {
         if (!input_bsocket.is_available() || !output_bsocket.is_available()) {
           continue;
         }
-        r_referenced_by_field_socket.add_multiple(
-            &output_bsocket, Vector<int>(r_referenced_by_field_socket.lookup(&input_bsocket)));
+        r_referenced_by_field_socket[output_bsocket.index_in_tree()] |=
+            r_referenced_by_field_socket[input_bsocket.index_in_tree()];
       }
       for (const aal::PropagateRelation &relation : relations.propagate_relations) {
         const bNodeSocket &input_bsocket = bnode->input_socket(relation.from_geometry_input);
@@ -2472,8 +2474,8 @@ struct GeometryNodesLazyFunctionGraphBuilder {
         if (!input_bsocket.is_available() || !output_bsocket.is_available()) {
           continue;
         }
-        r_propagated_to_geometry_socket.add_multiple(
-            &output_bsocket, Vector<int>(r_propagated_to_geometry_socket.lookup(&input_bsocket)));
+        r_propagated_to_geometry_socket[output_bsocket.index_in_tree()] |=
+            r_propagated_to_geometry_socket[input_bsocket.index_in_tree()];
       }
     }
   }
@@ -2484,12 +2486,14 @@ struct GeometryNodesLazyFunctionGraphBuilder {
   void gather_required_propagated_data(
       const Span<const aal::RelationsInNode *> relations_by_node,
       const VectorSet<AttributeReferenceKey> &attribute_reference_keys,
-      const MultiValueMap<const bNodeSocket *, int> &referenced_by_field_socket,
-      const MultiValueMap<const bNodeSocket *, int> &propagated_to_geometry_socket,
-      MultiValueMap<const bNodeSocket *, int> &r_required_propagated_to_geometry_socket)
+      const BitArrayVector<> &referenced_by_field_socket,
+      const BitArrayVector<> &propagated_to_geometry_socket,
+      BitArrayVector<> &r_required_propagated_to_geometry_socket)
   {
     const aal::RelationsInNode &tree_relations = *btree_.runtime->anonymous_attribute_relations;
-    MultiValueMap<const bNodeSocket *, int> required_by_geometry_socket;
+    const int sockets_num = btree_.all_sockets().size();
+    const int attribute_references_num = referenced_by_field_socket.group_size();
+    BitArrayVector<> required_by_geometry_socket(sockets_num, attribute_references_num, false);
 
     /* Initialize required attributes at group output. */
     if (const bNode *group_output_bnode = btree_.group_output_node()) {
@@ -2498,74 +2502,66 @@ struct GeometryNodesLazyFunctionGraphBuilder {
         key.type = AttributeReferenceKeyType::OutputGeometry;
         key.index = relation.to_geometry_output;
         const int key_index = attribute_reference_keys.index_of(key);
-        required_by_geometry_socket.add(
-            &group_output_bnode->input_socket(relation.to_geometry_output), key_index);
+        required_by_geometry_socket[group_output_bnode->input_socket(relation.to_geometry_output)
+                                        .index_in_tree()][key_index]
+            .set();
       }
       for (const aal::AvailableRelation &relation : tree_relations.available_relations) {
         const bNodeSocket &geometry_bsocket = group_output_bnode->input_socket(
             relation.geometry_output);
         const bNodeSocket &field_bsocket = group_output_bnode->input_socket(relation.field_output);
-        required_by_geometry_socket.add_multiple(
-            &geometry_bsocket, referenced_by_field_socket.lookup(&field_bsocket));
+        required_by_geometry_socket[geometry_bsocket.index_in_tree()] |=
+            referenced_by_field_socket[field_bsocket.index_in_tree()];
       }
     }
 
     /* Propagate attribute usages from right to left. */
+    BitVector<> required_attributes(attribute_references_num);
     for (const bNode *bnode : btree_.toposort_right_to_left()) {
       const aal::RelationsInNode &relations = *relations_by_node[bnode->index()];
       for (const bNodeSocket *bsocket : bnode->output_sockets()) {
         if (!bsocket->is_available()) {
           continue;
         }
-        Vector<int> required_attributes;
+        required_attributes.fill(false);
         for (const bNodeLink *blink : bsocket->directly_linked_links()) {
           if (blink->is_used()) {
             const bNodeSocket &to_socket = *blink->tosock;
-            required_attributes.extend_non_duplicates(
-                required_by_geometry_socket.lookup(&to_socket));
+            required_attributes |= required_by_geometry_socket[to_socket.index_in_tree()];
           }
         }
-        const Span<int> available_attributes = propagated_to_geometry_socket.lookup(bsocket);
-        for (const int key_index : required_attributes) {
-          if (available_attributes.contains(key_index)) {
-            required_by_geometry_socket.add(bsocket, key_index);
-
-            const AttributeReferenceKey &key = attribute_reference_keys[key_index];
-            if (key.type != AttributeReferenceKeyType::Socket ||
-                &key.bsocket->owner_node() != bnode) {
-              r_required_propagated_to_geometry_socket.add(bsocket, key_index);
-            }
+        required_attributes &= propagated_to_geometry_socket[bsocket->index_in_tree()];
+        required_by_geometry_socket[bsocket->index_in_tree()] |= required_attributes;
+        bits::foreach_1_index(required_attributes, [&](const int key_index) {
+          const AttributeReferenceKey &key = attribute_reference_keys[key_index];
+          if (key.type != AttributeReferenceKeyType::Socket ||
+              &key.bsocket->owner_node() != bnode) {
+            r_required_propagated_to_geometry_socket[bsocket->index_in_tree()][key_index].set();
           }
-        }
+        });
       }
 
       for (const bNodeSocket *bsocket : bnode->input_sockets()) {
         if (!bsocket->is_available()) {
           continue;
         }
-        Vector<int> required_attributes;
+        required_attributes.fill(false);
         for (const aal::PropagateRelation &relation : relations.propagate_relations) {
           if (relation.from_geometry_input == bsocket->index()) {
             const bNodeSocket &output_bsocket = bnode->output_socket(relation.to_geometry_output);
-            required_attributes.extend_non_duplicates(
-                required_by_geometry_socket.lookup(&output_bsocket));
+            required_attributes |= required_by_geometry_socket[output_bsocket.index_in_tree()];
           }
         }
         for (const aal::EvalRelation &relation : relations.eval_relations) {
           if (relation.geometry_input == bsocket->index()) {
             const bNodeSocket &field_bsocket = bnode->input_socket(relation.field_input);
             if (field_bsocket.is_available()) {
-              required_attributes.extend_non_duplicates(
-                  referenced_by_field_socket.lookup(&field_bsocket));
+              required_attributes |= referenced_by_field_socket[field_bsocket.index_in_tree()];
             }
           }
         }
-        const Span<int> available_attributes = propagated_to_geometry_socket.lookup(bsocket);
-        for (const int key_index : required_attributes) {
-          if (available_attributes.contains(key_index)) {
-            required_by_geometry_socket.add(bsocket, key_index);
-          }
-        }
+        required_attributes &= propagated_to_geometry_socket[bsocket->index_in_tree()];
+        required_by_geometry_socket[bsocket->index_in_tree()] |= required_attributes;
       }
     }
   }
@@ -2577,20 +2573,19 @@ struct GeometryNodesLazyFunctionGraphBuilder {
   void build_attribute_sets_to_propagate(
       const Span<AttributeReferenceKey> attribute_reference_keys,
       const Span<AttributeReferenceInfo> attribute_reference_infos,
-      const MultiValueMap<const bNodeSocket *, int> &required_propagated_to_geometry_socket)
+      const BitArrayVector<> &required_propagated_to_geometry_socket)
   {
     JoinAttibuteSetsCache join_attribute_sets_cache;
 
     for (const auto [geometry_output_bsocket, lf_attribute_set_input] :
          attribute_set_propagation_map_.items()) {
-      const Span<int> required = required_propagated_to_geometry_socket.lookup(
-          geometry_output_bsocket);
+      const BoundedBitSpan required =
+          required_propagated_to_geometry_socket[geometry_output_bsocket->index_in_tree()];
 
       Vector<lf::OutputSocket *> attribute_set_sockets;
       Vector<lf::OutputSocket *> used_sockets;
 
-      for (const int i : required.index_range()) {
-        const int key_index = required[i];
+      bits::foreach_1_index(required, [&](const int key_index) {
         const AttributeReferenceKey &key = attribute_reference_keys[key_index];
         const AttributeReferenceInfo &info = attribute_reference_infos[key_index];
         lf::OutputSocket *lf_socket_usage = nullptr;
@@ -2615,7 +2610,7 @@ struct GeometryNodesLazyFunctionGraphBuilder {
           attribute_set_sockets.append(info.lf_attribute_set_socket);
           used_sockets.append(lf_socket_usage);
         }
-      }
+      });
       if (lf::OutputSocket *joined_attribute_set = this->join_attribute_sets(
               attribute_set_sockets, used_sockets, join_attribute_sets_cache)) {
         lf_graph_->add_link(*joined_attribute_set, *lf_attribute_set_input);