geometry_nodes_lazy_function.cc

Go to the documentation of this file.

/* SPDX-FileCopyrightText: 2023 Blender Authors
 *
 * SPDX-License-Identifier: GPL-2.0-or-later */

 
#include "NOD_geometry_exec.hh"
#include "NOD_geometry_nodes_lazy_function.hh"
#include "NOD_multi_function.hh"
#include "NOD_node_declaration.hh"
 
#include "BLI_array_utils.hh"
#include "BLI_bit_group_vector.hh"
#include "BLI_bit_span_ops.hh"
#include "BLI_cpp_types.hh"
#include "BLI_dot_export.hh"
#include "BLI_hash.h"
#include "BLI_hash_md5.hh"
#include "BLI_lazy_threading.hh"
#include "BLI_map.hh"
 
#include "DNA_ID.h"
 
#include "BKE_anonymous_attribute_make.hh"
#include "BKE_compute_contexts.hh"
#include "BKE_curves.hh"
#include "BKE_geometry_nodes_gizmos_transforms.hh"
#include "BKE_geometry_set.hh"
#include "BKE_grease_pencil.hh"
#include "BKE_node_socket_value.hh"
#include "BKE_node_tree_anonymous_attributes.hh"
#include "BKE_node_tree_zones.hh"
#include "BKE_type_conversions.hh"
 
#include "FN_lazy_function_execute.hh"
#include "FN_lazy_function_graph_executor.hh"
 
#include "DEG_depsgraph_query.hh"
 
#include "GEO_extract_elements.hh"
#include "GEO_join_geometries.hh"
 
#include <fmt/format.h>
#include <sstream>
 
namespace blender::nodes {
 
namespace aai = bke::anonymous_attribute_inferencing;
using bke::bNodeTreeZone;
using bke::bNodeTreeZones;
using bke::SocketValueVariant;
 
static const CPPType *get_socket_cpp_type(const bke::bNodeSocketType &typeinfo)
{
  const CPPType *type = typeinfo.geometry_nodes_cpp_type;
  if (type == nullptr) {
    return nullptr;
  }
  BLI_assert(type->has_special_member_functions());
  return type;
}
 
static const CPPType *get_socket_cpp_type(const bNodeSocket &socket)
{
  return get_socket_cpp_type(*socket.typeinfo);
}
 
static const CPPType *get_vector_type(const CPPType &type)
{
  const VectorCPPType *vector_type = VectorCPPType::get_from_value(type);
  if (vector_type == nullptr) {
    return nullptr;
  }
  return &vector_type->self;
}

static void lazy_function_interface_from_node(const bNode &node,
                                              Vector<lf::Input> &r_inputs,
                                              Vector<lf::Output> &r_outputs,
                                              MutableSpan<int> r_lf_index_by_bsocket)
{
  const bool is_muted = node.is_muted();
  const lf::ValueUsage input_usage = lf::ValueUsage::Used;
  for (const bNodeSocket *socket : node.input_sockets()) {
    if (!socket->is_available()) {
      continue;
    }
    const CPPType *type = get_socket_cpp_type(*socket);
    if (type == nullptr) {
      continue;
    }
    if (socket->is_multi_input() && !is_muted) {
      type = get_vector_type(*type);
    }
    r_lf_index_by_bsocket[socket->index_in_tree()] = r_inputs.append_and_get_index_as(
        socket->name, *type, input_usage);
  }
  for (const bNodeSocket *socket : node.output_sockets()) {
    if (!socket->is_available()) {
      continue;
    }
    const CPPType *type = get_socket_cpp_type(*socket);
    if (type == nullptr) {
      continue;
    }
    r_lf_index_by_bsocket[socket->index_in_tree()] = r_outputs.append_and_get_index_as(
        socket->name, *type);
  }
}

class LazyFunctionForGeometryNode : public LazyFunction {
 private:
  const bNode &node_;
  const GeometryNodesLazyFunctionGraphInfo &own_lf_graph_info_;
  Vector<bool> is_attribute_output_bsocket_;
 
 public:
  LazyFunctionForGeometryNode(const bNode &node,
                              GeometryNodesLazyFunctionGraphInfo &own_lf_graph_info)
      : node_(node),
        own_lf_graph_info_(own_lf_graph_info),
        is_attribute_output_bsocket_(node.output_sockets().size(), false)
  {
    BLI_assert(node.typeinfo->geometry_node_execute != nullptr);
    debug_name_ = node.name;
    lazy_function_interface_from_node(
        node, inputs_, outputs_, own_lf_graph_info.mapping.lf_index_by_bsocket);
 
    const NodeDeclaration &node_decl = *node.declaration();
    const aal::RelationsInNode *relations = node_decl.anonymous_attribute_relations();
    if (relations == nullptr) {
      return;
    }
    if (!relations->available_relations.is_empty()) {
      /* Inputs are only used when an output is used that is not just outputting an anonymous
       * attribute field. */
      for (lf::Input &input : inputs_) {
        input.usage = lf::ValueUsage::Maybe;
      }
      for (const aal::AvailableRelation &relation : relations->available_relations) {
        is_attribute_output_bsocket_[relation.field_output] = true;
      }
    }
    Vector<const bNodeSocket *> handled_field_outputs;
    for (const aal::AvailableRelation &relation : relations->available_relations) {
      const bNodeSocket &output_bsocket = node.output_socket(relation.field_output);
      if (output_bsocket.is_available() && !handled_field_outputs.contains(&output_bsocket)) {
        handled_field_outputs.append(&output_bsocket);
        const int lf_index = inputs_.append_and_get_index_as("Output Used", CPPType::get<bool>());
        own_lf_graph_info.mapping
            .lf_input_index_for_output_bsocket_usage[output_bsocket.index_in_all_outputs()] =
            lf_index;
      }
    }
 
    Vector<const bNodeSocket *> handled_geometry_outputs;
    for (const aal::PropagateRelation &relation : relations->propagate_relations) {
      const bNodeSocket &output_bsocket = node.output_socket(relation.to_geometry_output);
      if (output_bsocket.is_available() && !handled_geometry_outputs.contains(&output_bsocket)) {
        handled_geometry_outputs.append(&output_bsocket);
        const int lf_index = inputs_.append_and_get_index_as(
            "Propagate to Output", CPPType::get<bke::AnonymousAttributeSet>());
        own_lf_graph_info.mapping.lf_input_index_for_attribute_propagation_to_output
            [output_bsocket.index_in_all_outputs()] = lf_index;
      }
    }
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const ScopedNodeTimer node_timer{context, node_};
 
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
 
    bool used_non_attribute_output_exists = false;
    for (const int output_bsocket_index : node_.output_sockets().index_range()) {
      const bNodeSocket &output_bsocket = node_.output_socket(output_bsocket_index);
      const int lf_index =
          own_lf_graph_info_.mapping.lf_index_by_bsocket[output_bsocket.index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      const lf::ValueUsage output_usage = params.get_output_usage(lf_index);
      if (output_usage == lf::ValueUsage::Unused) {
        continue;
      }
      if (is_attribute_output_bsocket_[output_bsocket_index]) {
        if (params.output_was_set(lf_index)) {
          continue;
        }
        this->output_anonymous_attribute_field(params, *user_data, lf_index, output_bsocket);
      }
      else {
        if (output_usage == lf::ValueUsage::Used) {
          used_non_attribute_output_exists = true;
        }
      }
    }
 
    if (!used_non_attribute_output_exists) {
      /* Only attribute outputs are used currently, no need to evaluate the full node and its
       * inputs. */
      return;
    }
 
    bool missing_input = false;
    for (const int lf_index : inputs_.index_range()) {
      if (params.try_get_input_data_ptr_or_request(lf_index) == nullptr) {
        missing_input = true;
      }
    }
    if (missing_input) {
      /* Wait until all inputs are available. */
      return;
    }
 
    auto get_anonymous_attribute_name = [&](const int i) {
      return this->anonymous_attribute_name_for_output(*user_data, i);
    };
 
    GeoNodeExecParams geo_params{
        node_,
        params,
        context,
        own_lf_graph_info_.mapping.lf_input_index_for_output_bsocket_usage,
        own_lf_graph_info_.mapping.lf_input_index_for_attribute_propagation_to_output,
        get_anonymous_attribute_name};
 
    node_.typeinfo->geometry_node_execute(geo_params);
  }
 
  std::string input_name(const int index) const override
  {
    for (const bNodeSocket *bsocket : node_.output_sockets()) {
      {
        const int lf_index =
            own_lf_graph_info_.mapping
                .lf_input_index_for_output_bsocket_usage[bsocket->index_in_all_outputs()];
        if (index == lf_index) {
          return StringRef("Use Output '") + bsocket->name + "'";
        }
      }
      {
        const int lf_index =
            own_lf_graph_info_.mapping.lf_input_index_for_attribute_propagation_to_output
                [bsocket->index_in_all_outputs()];
        if (index == lf_index) {
          return StringRef("Propagate to '") + bsocket->name + "'";
        }
      }
    }
    return inputs_[index].debug_name;
  }
 
  std::string output_name(const int index) const override
  {
    return outputs_[index].debug_name;
  }
 
  void output_anonymous_attribute_field(lf::Params &params,
                                        const GeoNodesLFUserData &user_data,
                                        const int lf_index,
                                        const bNodeSocket &socket) const
  {
    std::string attribute_name = this->anonymous_attribute_name_for_output(user_data,
                                                                           socket.index());
    std::string socket_inspection_name = make_anonymous_attribute_socket_inspection_string(socket);
    auto attribute_field = std::make_shared<AttributeFieldInput>(
        std::move(attribute_name),
        *socket.typeinfo->base_cpp_type,
        std::move(socket_inspection_name));
 
    void *r_value = params.get_output_data_ptr(lf_index);
    new (r_value) SocketValueVariant(GField(std::move(attribute_field)));
    params.output_set(lf_index);
  }
 
  std::string anonymous_attribute_name_for_output(const GeoNodesLFUserData &user_data,
                                                  const int output_index) const
  {
    return bke::hash_to_anonymous_attribute_name(user_data.call_data->self_object()->id.name,
                                                 user_data.compute_context->hash(),
                                                 node_.identifier,
                                                 node_.output_socket(output_index).identifier);
  }
};

class LazyFunctionForMultiInput : public LazyFunction {
 private:
  const CPPType *base_type_;
 
 public:
  Vector<const bNodeLink *> links;
 
  LazyFunctionForMultiInput(const bNodeSocket &socket)
  {
    debug_name_ = "Multi Input";
    base_type_ = get_socket_cpp_type(socket);
    BLI_assert(base_type_ != nullptr);
    BLI_assert(socket.is_multi_input());
    for (const bNodeLink *link : socket.directly_linked_links()) {
      if (link->is_muted() || !link->fromsock->is_available() ||
          link->fromnode->is_dangling_reroute())
      {
        continue;
      }
      inputs_.append({"Input", *base_type_});
      this->links.append(link);
    }
    const CPPType *vector_type = get_vector_type(*base_type_);
    BLI_assert(vector_type != nullptr);
    outputs_.append({"Output", *vector_type});
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    /* Currently we only have multi-inputs for geometry and value sockets. This could be
     * generalized in the future. */
    base_type_->to_static_type_tag<GeometrySet, SocketValueVariant>([&](auto type_tag) {
      using T = typename decltype(type_tag)::type;
      if constexpr (std::is_void_v<T>) {
        /* This type is not supported in this node for now. */
        BLI_assert_unreachable();
      }
      else {
        void *output_ptr = params.get_output_data_ptr(0);
        Vector<T> &values = *new (output_ptr) Vector<T>();
        for (const int i : inputs_.index_range()) {
          values.append(params.extract_input<T>(i));
        }
        params.output_set(0);
      }
    });
  }
};

class LazyFunctionForRerouteNode : public LazyFunction {
 public:
  LazyFunctionForRerouteNode(const CPPType &type)
  {
    debug_name_ = "Reroute";
    inputs_.append({"Input", type});
    outputs_.append({"Output", type});
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    void *input_value = params.try_get_input_data_ptr(0);
    void *output_value = params.get_output_data_ptr(0);
    BLI_assert(input_value != nullptr);
    BLI_assert(output_value != nullptr);
    const CPPType &type = *inputs_[0].type;
    type.move_construct(input_value, output_value);
    params.output_set(0);
  }
};

class LazyFunctionForUndefinedNode : public LazyFunction {
  const bNode &node_;
 
 public:
  LazyFunctionForUndefinedNode(const bNode &node, MutableSpan<int> r_lf_index_by_bsocket)
      : node_(node)
  {
    debug_name_ = "Undefined";
    Vector<lf::Input> dummy_inputs;
    lazy_function_interface_from_node(node, dummy_inputs, outputs_, r_lf_index_by_bsocket);
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    set_default_remaining_node_outputs(params, node_);
  }
};
 
static void set_default_value_for_output_socket(lf::Params &params,
                                                const int lf_index,
                                                const bNodeSocket &bsocket)
{
  const CPPType &cpp_type = *bsocket.typeinfo->geometry_nodes_cpp_type;
  void *output_value = params.get_output_data_ptr(lf_index);
  if (bsocket.typeinfo->geometry_nodes_default_cpp_value) {
    cpp_type.copy_construct(bsocket.typeinfo->geometry_nodes_default_cpp_value, output_value);
  }
  else {
    cpp_type.value_initialize(output_value);
  }
  params.output_set(lf_index);
}
 
void set_default_remaining_node_outputs(lf::Params &params, const bNode &node)
{
  const bNodeTree &ntree = node.owner_tree();
  const Span<int> lf_index_by_bsocket =
      ntree.runtime->geometry_nodes_lazy_function_graph_info->mapping.lf_index_by_bsocket;
  for (const bNodeSocket *bsocket : node.output_sockets()) {
    const int lf_index = lf_index_by_bsocket[bsocket->index_in_tree()];
    if (lf_index == -1) {
      continue;
    }
    if (params.output_was_set(lf_index)) {
      continue;
    }
    set_default_value_for_output_socket(params, lf_index, *bsocket);
  }
}
 
std::string make_anonymous_attribute_socket_inspection_string(const bNodeSocket &socket)
{
  return make_anonymous_attribute_socket_inspection_string(socket.owner_node().label_or_name(),
                                                           socket.name);
}
std::string make_anonymous_attribute_socket_inspection_string(StringRef node_name,
                                                              StringRef socket_name)
{
  return fmt::format(TIP_("\"{}\" from {}"), socket_name, node_name);
}
 
static void execute_multi_function_on_value_variant__single(
    const MultiFunction &fn,
    const Span<SocketValueVariant *> input_values,
    const Span<SocketValueVariant *> output_values)
{
  /* In this case, the multi-function is evaluated directly. */
  const IndexMask mask(1);
  mf::ParamsBuilder params{fn, &mask};
  mf::ContextBuilder context;
 
  for (const int i : input_values.index_range()) {
    SocketValueVariant &input_variant = *input_values[i];
    input_variant.convert_to_single();
    const void *value = input_variant.get_single_ptr_raw();
    const mf::ParamType param_type = fn.param_type(params.next_param_index());
    const CPPType &cpp_type = param_type.data_type().single_type();
    params.add_readonly_single_input(GPointer{cpp_type, value});
  }
  for (const int i : output_values.index_range()) {
    if (output_values[i] == nullptr) {
      params.add_ignored_single_output("");
      continue;
    }
    SocketValueVariant &output_variant = *output_values[i];
    const mf::ParamType param_type = fn.param_type(params.next_param_index());
    const CPPType &cpp_type = param_type.data_type().single_type();
    const eNodeSocketDatatype socket_type =
        bke::geo_nodes_base_cpp_type_to_socket_type(cpp_type).value();
    void *value = output_variant.allocate_single(socket_type);
    params.add_uninitialized_single_output(GMutableSpan{cpp_type, value, 1});
  }
  fn.call(mask, params, context);
}
 
static void execute_multi_function_on_value_variant__field(
    const MultiFunction &fn,
    const std::shared_ptr<MultiFunction> &owned_fn,
    const Span<SocketValueVariant *> input_values,
    const Span<SocketValueVariant *> output_values)
{
  /* Convert all inputs into fields, so that they can be used as input in the new field. */
  Vector<GField> input_fields;
  for (const int i : input_values.index_range()) {
    input_fields.append(input_values[i]->extract<GField>());
  }
 
  /* Construct the new field node. */
  std::shared_ptr<fn::FieldOperation> operation;
  if (owned_fn) {
    operation = fn::FieldOperation::Create(owned_fn, std::move(input_fields));
  }
  else {
    operation = fn::FieldOperation::Create(fn, std::move(input_fields));
  }
 
  /* Store the new fields in the output. */
  for (const int i : output_values.index_range()) {
    if (output_values[i] == nullptr) {
      continue;
    }
    output_values[i]->set(GField{operation, i});
  }
}

static void execute_multi_function_on_value_variant(const MultiFunction &fn,
                                                    const std::shared_ptr<MultiFunction> &owned_fn,
                                                    const Span<SocketValueVariant *> input_values,
                                                    const Span<SocketValueVariant *> output_values)
{
  /* Check input types which determine how the function is evaluated. */
  bool any_input_is_field = false;
  for (const int i : input_values.index_range()) {
    const SocketValueVariant &value = *input_values[i];
    if (value.is_context_dependent_field()) {
      any_input_is_field = true;
    }
  }
 
  if (any_input_is_field) {
    execute_multi_function_on_value_variant__field(fn, owned_fn, input_values, output_values);
  }
  else {
    execute_multi_function_on_value_variant__single(fn, input_values, output_values);
  }
}

class LazyFunctionForMutedNode : public LazyFunction {
 private:
  const bNode &node_;
  Span<int> lf_index_by_bsocket_;
  Array<const bNodeSocket *> input_by_output_index_;
 
 public:
  LazyFunctionForMutedNode(const bNode &node, MutableSpan<int> r_lf_index_by_bsocket)
      : node_(node), lf_index_by_bsocket_(r_lf_index_by_bsocket)
  {
    debug_name_ = "Muted";
    lazy_function_interface_from_node(node, inputs_, outputs_, r_lf_index_by_bsocket);
    for (lf::Input &fn_input : inputs_) {
      fn_input.usage = lf::ValueUsage::Maybe;
    }
 
    for (lf::Input &fn_input : inputs_) {
      fn_input.usage = lf::ValueUsage::Unused;
    }
 
    input_by_output_index_.reinitialize(node.output_sockets().size());
    input_by_output_index_.fill(nullptr);
    for (const bNodeLink &internal_link : node.internal_links()) {
      const int input_i = r_lf_index_by_bsocket[internal_link.fromsock->index_in_tree()];
      const int output_i = r_lf_index_by_bsocket[internal_link.tosock->index_in_tree()];
      if (ELEM(-1, input_i, output_i)) {
        continue;
      }
      input_by_output_index_[internal_link.tosock->index()] = internal_link.fromsock;
      inputs_[input_i].usage = lf::ValueUsage::Maybe;
    }
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    for (const bNodeSocket *output_bsocket : node_.output_sockets()) {
      const int lf_output_index = lf_index_by_bsocket_[output_bsocket->index_in_tree()];
      if (lf_output_index == -1) {
        continue;
      }
      if (params.output_was_set(lf_output_index)) {
        continue;
      }
      if (params.get_output_usage(lf_output_index) != lf::ValueUsage::Used) {
        continue;
      }
      const bNodeSocket *input_bsocket = input_by_output_index_[output_bsocket->index()];
      if (input_bsocket == nullptr) {
        set_default_value_for_output_socket(params, lf_output_index, *output_bsocket);
        continue;
      }
      const int lf_input_index = lf_index_by_bsocket_[input_bsocket->index_in_tree()];
      const void *input_value = params.try_get_input_data_ptr_or_request(lf_input_index);
      if (input_value == nullptr) {
        /* Wait for value to be available. */
        continue;
      }
      void *output_value = params.get_output_data_ptr(lf_output_index);
      if (input_bsocket->type == output_bsocket->type) {
        inputs_[lf_input_index].type->copy_construct(input_value, output_value);
        params.output_set(lf_output_index);
        continue;
      }
      const bke::DataTypeConversions &conversions = bke::get_implicit_type_conversions();
      if (conversions.is_convertible(*input_bsocket->typeinfo->base_cpp_type,
                                     *output_bsocket->typeinfo->base_cpp_type))
      {
        const MultiFunction &multi_fn = *conversions.get_conversion_multi_function(
            mf::DataType::ForSingle(*input_bsocket->typeinfo->base_cpp_type),
            mf::DataType::ForSingle(*output_bsocket->typeinfo->base_cpp_type));
        SocketValueVariant input_variant = *static_cast<const SocketValueVariant *>(input_value);
        SocketValueVariant *output_variant = new (output_value) SocketValueVariant();
        execute_multi_function_on_value_variant(multi_fn, {}, {&input_variant}, {output_variant});
        params.output_set(lf_output_index);
        continue;
      }
      set_default_value_for_output_socket(params, lf_output_index, *output_bsocket);
    }
  }
};

class LazyFunctionForMultiFunctionConversion : public LazyFunction {
 private:
  const MultiFunction &fn_;
 
 public:
  LazyFunctionForMultiFunctionConversion(const MultiFunction &fn) : fn_(fn)
  {
    debug_name_ = "Convert";
    inputs_.append_as("From", CPPType::get<SocketValueVariant>());
    outputs_.append_as("To", CPPType::get<SocketValueVariant>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    SocketValueVariant *from_value = params.try_get_input_data_ptr<SocketValueVariant>(0);
    SocketValueVariant *to_value = new (params.get_output_data_ptr(0)) SocketValueVariant();
    BLI_assert(from_value != nullptr);
    BLI_assert(to_value != nullptr);
 
    execute_multi_function_on_value_variant(fn_, {}, {from_value}, {to_value});
 
    params.output_set(0);
  }
};

class LazyFunctionForMultiFunctionNode : public LazyFunction {
 private:
  const NodeMultiFunctions::Item fn_item_;
 
 public:
  LazyFunctionForMultiFunctionNode(const bNode &node,
                                   NodeMultiFunctions::Item fn_item,
                                   MutableSpan<int> r_lf_index_by_bsocket)
      : fn_item_(std::move(fn_item))
  {
    BLI_assert(fn_item_.fn != nullptr);
    debug_name_ = node.name;
    lazy_function_interface_from_node(node, inputs_, outputs_, r_lf_index_by_bsocket);
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    Vector<SocketValueVariant *> input_values(inputs_.size());
    Vector<SocketValueVariant *> output_values(outputs_.size());
    for (const int i : inputs_.index_range()) {
      input_values[i] = params.try_get_input_data_ptr<SocketValueVariant>(i);
    }
    for (const int i : outputs_.index_range()) {
      if (params.get_output_usage(i) != lf::ValueUsage::Unused) {
        output_values[i] = new (params.get_output_data_ptr(i)) SocketValueVariant();
      }
      else {
        output_values[i] = nullptr;
      }
    }
    execute_multi_function_on_value_variant(
        *fn_item_.fn, fn_item_.owned_fn, input_values, output_values);
    for (const int i : outputs_.index_range()) {
      if (params.get_output_usage(i) != lf::ValueUsage::Unused) {
        params.output_set(i);
      }
    }
  }
};

class LazyFunctionForImplicitInput : public LazyFunction {
 private:
  std::function<void(void *)> init_fn_;
 
 public:
  LazyFunctionForImplicitInput(const CPPType &type, std::function<void(void *)> init_fn)
      : init_fn_(std::move(init_fn))
  {
    debug_name_ = "Input";
    outputs_.append({"Output", type});
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    void *value = params.get_output_data_ptr(0);
    init_fn_(value);
    params.output_set(0);
  }
};

class LazyFunctionForViewerNode : public LazyFunction {
 private:
  const bNode &bnode_;
  bool use_field_input_ = true;
 
 public:
  LazyFunctionForViewerNode(const bNode &bnode, MutableSpan<int> r_lf_index_by_bsocket)
      : bnode_(bnode)
  {
    debug_name_ = "Viewer";
    lazy_function_interface_from_node(bnode, inputs_, outputs_, r_lf_index_by_bsocket);
 
    /* Remove field input if it is not used. */
    for (const bNodeSocket *bsocket : bnode.input_sockets().drop_front(1)) {
      if (!bsocket->is_available()) {
        continue;
      }
      const Span<const bNodeLink *> links = bsocket->directly_linked_links();
      if (links.is_empty() || links.first()->fromnode->is_dangling_reroute()) {
        use_field_input_ = false;
        inputs_.pop_last();
        r_lf_index_by_bsocket[bsocket->index_in_tree()] = -1;
      }
    }
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
    geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(user_data);
    if (tree_logger == nullptr) {
      return;
    }
 
    GeometrySet geometry = params.extract_input<GeometrySet>(0);
    const NodeGeometryViewer *storage = static_cast<NodeGeometryViewer *>(bnode_.storage);
 
    if (use_field_input_) {
      SocketValueVariant *value_variant = params.try_get_input_data_ptr<SocketValueVariant>(1);
      BLI_assert(value_variant != nullptr);
      GField field = value_variant->extract<GField>();
      const AttrDomain domain = AttrDomain(storage->domain);
      const StringRefNull viewer_attribute_name = ".viewer";
      if (domain == AttrDomain::Instance) {
        if (geometry.has_instances()) {
          GeometryComponent &component = geometry.get_component_for_write(
              bke::GeometryComponent::Type::Instance);
          bke::try_capture_field_on_geometry(
              component, viewer_attribute_name, AttrDomain::Instance, field);
        }
      }
      else {
        geometry.modify_geometry_sets([&](GeometrySet &geometry) {
          for (const bke::GeometryComponent::Type type :
               {bke::GeometryComponent::Type::Mesh,
                bke::GeometryComponent::Type::PointCloud,
                bke::GeometryComponent::Type::Curve,
                bke::GeometryComponent::Type::GreasePencil})
          {
            if (geometry.has(type)) {
              GeometryComponent &component = geometry.get_component_for_write(type);
              AttrDomain used_domain = domain;
              if (used_domain == AttrDomain::Auto) {
                if (const std::optional<AttrDomain> detected_domain = bke::try_detect_field_domain(
                        component, field))
                {
                  used_domain = *detected_domain;
                }
                else {
                  used_domain = AttrDomain::Point;
                }
              }
              bke::try_capture_field_on_geometry(
                  component, viewer_attribute_name, used_domain, field);
            }
          }
        });
      }
    }
 
    tree_logger->log_viewer_node(bnode_, std::move(geometry));
  }
};

class LazyFunctionForViewerInputUsage : public LazyFunction {
 private:
  const lf::FunctionNode &lf_viewer_node_;
 
 public:
  LazyFunctionForViewerInputUsage(const lf::FunctionNode &lf_viewer_node)
      : lf_viewer_node_(lf_viewer_node)
  {
    debug_name_ = "Viewer Input Usage";
    outputs_.append_as("Viewer is Used", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
    if (!user_data->call_data->side_effect_nodes) {
      params.set_output<bool>(0, false);
      return;
    }
    const ComputeContextHash &context_hash = user_data->compute_context->hash();
    const Span<const lf::FunctionNode *> nodes_with_side_effects =
        user_data->call_data->side_effect_nodes->nodes_by_context.lookup(context_hash);
 
    const bool viewer_is_used = nodes_with_side_effects.contains(&lf_viewer_node_);
    params.set_output(0, viewer_is_used);
  }
};

static bool gizmo_is_used(const GeoNodesLFUserData &user_data,
                          const lf::FunctionNode &lf_gizmo_node)
{
  if (!user_data.call_data->side_effect_nodes) {
    return false;
  }
  const Span<const lf::FunctionNode *> nodes_with_side_effects =
      user_data.call_data->side_effect_nodes->nodes_by_context.lookup(
          user_data.compute_context->hash());
  const bool is_used = nodes_with_side_effects.contains(&lf_gizmo_node);
  return is_used;
}

class LazyFunctionForGizmoNode : public LazyFunction {
 private:
  const bNode &bnode_;
 
 public:
  const lf::FunctionNode *self_node = nullptr;
  Vector<const bNodeLink *> gizmo_links;
 
  LazyFunctionForGizmoNode(const bNode &bnode, MutableSpan<int> r_lf_index_by_bsocket)
      : bnode_(bnode)
  {
    debug_name_ = bnode.name;
    const bNodeSocket &gizmo_socket = bnode.input_socket(0);
    /* Create inputs for every input of the multi-input socket to make sure that they can be
     * logged. */
    for (const bNodeLink *link : gizmo_socket.directly_linked_links()) {
      if (!link->is_used()) {
        continue;
      }
      if (link->fromnode->is_dangling_reroute()) {
        continue;
      }
      inputs_.append_and_get_index_as(gizmo_socket.identifier,
                                      *gizmo_socket.typeinfo->geometry_nodes_cpp_type,
                                      lf::ValueUsage::Maybe);
      gizmo_links.append(link);
    }
    for (const bNodeSocket *socket : bnode.input_sockets().drop_front(1)) {
      r_lf_index_by_bsocket[socket->index_in_tree()] = inputs_.append_and_get_index_as(
          socket->identifier, *socket->typeinfo->geometry_nodes_cpp_type, lf::ValueUsage::Maybe);
    }
    r_lf_index_by_bsocket[bnode.output_socket(0).index_in_tree()] =
        outputs_.append_and_get_index_as("Transform", CPPType::get<GeometrySet>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    if (!gizmo_is_used(user_data, *this->self_node)) {
      set_default_remaining_node_outputs(params, bnode_);
      return;
    }
    if (!params.output_was_set(0)) {
      GeometrySet geometry;
      GeometryComponentEditData &edit_data =
          geometry.get_component_for_write<GeometryComponentEditData>();
      edit_data.gizmo_edit_hints_ = std::make_unique<bke::GizmoEditHints>();
      edit_data.gizmo_edit_hints_->gizmo_transforms.add(
          {user_data.compute_context->hash(), bnode_.identifier}, float4x4::identity());
      params.set_output(0, std::move(geometry));
    }
 
    /* Request all inputs so that their values can be logged. */
    for (const int i : inputs_.index_range()) {
      params.try_get_input_data_ptr_or_request(i);
    }
 
    const auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
    if (geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(user_data))
    {
      tree_logger->evaluated_gizmo_nodes.append(*tree_logger->allocator, {bnode_.identifier});
    }
  }
};
 
class LazyFunctionForGizmoInputsUsage : public LazyFunction {
 private:
  const lf::FunctionNode *lf_gizmo_node_ = nullptr;
 
 public:
  LazyFunctionForGizmoInputsUsage(const bNode &gizmo_node, const lf::FunctionNode &lf_gizmo_node)
      : lf_gizmo_node_(&lf_gizmo_node)
  {
    debug_name_ = gizmo_node.name;
    outputs_.append_as("Need Inputs", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const bool is_used = gizmo_is_used(user_data, *lf_gizmo_node_);
    params.set_output(0, is_used);
  }
};
 
class LazyFunctionForSimulationInputsUsage : public LazyFunction {
 private:
  const bNode *output_bnode_;
 
 public:
  LazyFunctionForSimulationInputsUsage(const bNode &output_bnode) : output_bnode_(&output_bnode)
  {
    debug_name_ = "Simulation Inputs Usage";
    outputs_.append_as("Need Input Inputs", CPPType::get<bool>());
    outputs_.append_as("Need Output Inputs", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const GeoNodesCallData &call_data = *user_data.call_data;
    if (!call_data.simulation_params) {
      this->set_default_outputs(params);
      return;
    }
    const std::optional<FoundNestedNodeID> found_id = find_nested_node_id(
        user_data, output_bnode_->identifier);
    if (!found_id) {
      this->set_default_outputs(params);
      return;
    }
    if (found_id->is_in_loop) {
      this->set_default_outputs(params);
      return;
    }
    SimulationZoneBehavior *zone_behavior = call_data.simulation_params->get(found_id->id);
    if (!zone_behavior) {
      this->set_default_outputs(params);
      return;
    }
 
    bool solve_contains_side_effect = false;
    if (call_data.side_effect_nodes) {
      const Span<const lf::FunctionNode *> side_effect_nodes =
          call_data.side_effect_nodes->nodes_by_context.lookup(user_data.compute_context->hash());
      solve_contains_side_effect = !side_effect_nodes.is_empty();
    }
 
    params.set_output(0, std::holds_alternative<sim_input::PassThrough>(zone_behavior->input));
    params.set_output(
        1,
        solve_contains_side_effect ||
            std::holds_alternative<sim_output::StoreNewState>(zone_behavior->output));
  }
 
  void set_default_outputs(lf::Params &params) const
  {
    params.set_output(0, false);
    params.set_output(1, false);
  }
};
 
class LazyFunctionForBakeInputsUsage : public LazyFunction {
 private:
  const bNode *bnode_;
 
 public:
  LazyFunctionForBakeInputsUsage(const bNode &bnode) : bnode_(&bnode)
  {
    debug_name_ = "Bake Inputs Usage";
    outputs_.append_as("Used", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    if (!user_data.call_data->bake_params) {
      this->set_default_outputs(params);
      return;
    }
    const std::optional<FoundNestedNodeID> found_id = find_nested_node_id(user_data,
                                                                          bnode_->identifier);
    if (!found_id) {
      this->set_default_outputs(params);
      return;
    }
    if (found_id->is_in_loop || found_id->is_in_simulation) {
      this->set_default_outputs(params);
      return;
    }
    BakeNodeBehavior *behavior = user_data.call_data->bake_params->get(found_id->id);
    if (!behavior) {
      this->set_default_outputs(params);
      return;
    }
    const bool need_inputs = std::holds_alternative<sim_output::PassThrough>(behavior->behavior) ||
                             std::holds_alternative<sim_output::StoreNewState>(behavior->behavior);
    params.set_output(0, need_inputs);
  }
 
  void set_default_outputs(lf::Params &params) const
  {
    params.set_output(0, false);
  }
};
 
static bool should_log_socket_values_for_context(const GeoNodesLFUserData &user_data,
                                                 const ComputeContextHash hash)
{
  if (const Set<ComputeContextHash> *contexts = user_data.call_data->socket_log_contexts) {
    return contexts->contains(hash);
  }
  else if (user_data.call_data->operator_data) {
    return false;
  }
  return true;
}

class LazyFunctionForGroupNode : public LazyFunction {
 private:
  const bNode &group_node_;
  const LazyFunction &group_lazy_function_;
  bool has_many_nodes_ = false;
 
  struct Storage {
    void *group_storage = nullptr;
    /* To avoid computing the hash more than once. */
    std::optional<ComputeContextHash> context_hash_cache;
  };
 
 public:
  LazyFunctionForGroupNode(const bNode &group_node,
                           const GeometryNodesLazyFunctionGraphInfo &group_lf_graph_info,
                           GeometryNodesLazyFunctionGraphInfo &own_lf_graph_info)
      : group_node_(group_node), group_lazy_function_(*group_lf_graph_info.function.function)
  {
    debug_name_ = group_node.name;
    allow_missing_requested_inputs_ = true;
 
    /* This wrapper has the same interface as the actual underlying node group. */
    inputs_ = group_lf_graph_info.function.function->inputs();
    outputs_ = group_lf_graph_info.function.function->outputs();
 
    has_many_nodes_ = group_lf_graph_info.num_inline_nodes_approximate > 1000;
 
    /* Add a boolean input for every output bsocket that indicates whether that socket is used. */
    for (const int i : group_node.output_sockets().index_range()) {
      own_lf_graph_info.mapping.lf_input_index_for_output_bsocket_usage
          [group_node.output_socket(i).index_in_all_outputs()] =
          group_lf_graph_info.function.inputs.output_usages[i];
    }
 
    /* Add an attribute set input for every output geometry socket that can propagate attributes
     * from inputs. */
    for (const int i : group_lf_graph_info.function.inputs.attributes_to_propagate.geometry_outputs
                           .index_range())
    {
      const int lf_index = group_lf_graph_info.function.inputs.attributes_to_propagate.range[i];
      const int output_index =
          group_lf_graph_info.function.inputs.attributes_to_propagate.geometry_outputs[i];
      const bNodeSocket &output_bsocket = group_node_.output_socket(output_index);
      own_lf_graph_info.mapping.lf_input_index_for_attribute_propagation_to_output
          [output_bsocket.index_in_all_outputs()] = lf_index;
    }
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const ScopedNodeTimer node_timer{context, group_node_};
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
 
    if (has_many_nodes_) {
      /* If the called node group has many nodes, it's likely that executing it takes a while even
       * if every individual node is very small. */
      lazy_threading::send_hint();
    }
 
    Storage *storage = static_cast<Storage *>(context.storage);
 
    /* The compute context changes when entering a node group. */
    bke::GroupNodeComputeContext compute_context{
        user_data->compute_context, group_node_.identifier, storage->context_hash_cache};
    storage->context_hash_cache = compute_context.hash();
 
    GeoNodesLFUserData group_user_data = *user_data;
    group_user_data.compute_context = &compute_context;
    group_user_data.log_socket_values = should_log_socket_values_for_context(
        *user_data, compute_context.hash());
 
    GeoNodesLFLocalUserData group_local_user_data{group_user_data};
    lf::Context group_context{storage->group_storage, &group_user_data, &group_local_user_data};
 
    ScopedComputeContextTimer timer(group_context);
    group_lazy_function_.execute(params, group_context);
  }
 
  void *init_storage(LinearAllocator<> &allocator) const override
  {
    Storage *s = allocator.construct<Storage>().release();
    s->group_storage = group_lazy_function_.init_storage(allocator);
    return s;
  }
 
  void destruct_storage(void *storage) const override
  {
    Storage *s = static_cast<Storage *>(storage);
    group_lazy_function_.destruct_storage(s->group_storage);
    std::destroy_at(s);
  }
 
  std::string name() const override
  {
    return fmt::format(TIP_("Group '{}' ({})"), group_node_.id->name + 2, group_node_.name);
  }
 
  std::string input_name(const int i) const override
  {
    return group_lazy_function_.input_name(i);
  }
 
  std::string output_name(const int i) const override
  {
    return group_lazy_function_.output_name(i);
  }
};
 
static GMutablePointer get_socket_default_value(LinearAllocator<> &allocator,
                                                const bNodeSocket &bsocket)
{
  const bke::bNodeSocketType &typeinfo = *bsocket.typeinfo;
  const CPPType *type = get_socket_cpp_type(typeinfo);
  if (type == nullptr) {
    return {};
  }
  void *buffer = allocator.allocate(type->size(), type->alignment());
  typeinfo.get_geometry_nodes_cpp_value(bsocket.default_value, buffer);
  return {type, buffer};
}

class LazyFunctionForLogicalOr : public lf::LazyFunction {
 public:
  LazyFunctionForLogicalOr(const int inputs_num)
  {
    debug_name_ = "Logical Or";
    for ([[maybe_unused]] const int i : IndexRange(inputs_num)) {
      inputs_.append_as("Input", CPPType::get<bool>(), lf::ValueUsage::Maybe);
    }
    outputs_.append_as("Output", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    Vector<int, 16> unavailable_inputs;
    /* First check all inputs for available values without requesting more inputs. If any of the
     * available inputs is true already, the others don't have to be requested anymore. */
    for (const int i : inputs_.index_range()) {
      if (const bool *value = params.try_get_input_data_ptr<bool>(i)) {
        if (*value) {
          params.set_output(0, true);
          return;
        }
      }
      else {
        unavailable_inputs.append(i);
      }
    }
    if (unavailable_inputs.is_empty()) {
      params.set_output(0, false);
      return;
    }
    /* Request the next unavailable input. Note that a value might be available now even if it was
     * not available before, because it might have been computed in the mean-time. */
    for (const int i : unavailable_inputs) {
      if (const bool *value = params.try_get_input_data_ptr_or_request<bool>(i)) {
        if (*value) {
          params.set_output(0, true);
          return;
        }
      }
      else {
        /* The input has been requested and it's not available yet, so wait until it is ready. */
        return;
      }
    }
    /* All inputs were available now and all of them were false, so the final output is false. */
    params.set_output(0, false);
  }
};

class LazyFunctionForSwitchSocketUsage : public lf::LazyFunction {
 public:
  LazyFunctionForSwitchSocketUsage()
  {
    debug_name_ = "Switch Socket Usage";
    inputs_.append_as("Condition", CPPType::get<SocketValueVariant>());
    outputs_.append_as("False", CPPType::get<bool>());
    outputs_.append_as("True", CPPType::get<bool>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    const SocketValueVariant &condition_variant = params.get_input<SocketValueVariant>(0);
    if (condition_variant.is_context_dependent_field()) {
      params.set_output(0, true);
      params.set_output(1, true);
    }
    else {
      const bool value = condition_variant.get<bool>();
      params.set_output(0, !value);
      params.set_output(1, value);
    }
  }
};

class LazyFunctionForIndexSwitchSocketUsage : public lf::LazyFunction {
 public:
  LazyFunctionForIndexSwitchSocketUsage(const bNode &bnode)
  {
    debug_name_ = "Index Switch Socket Usage";
    inputs_.append_as("Index", CPPType::get<SocketValueVariant>());
    for (const bNodeSocket *socket : bnode.input_sockets().drop_front(1)) {
      outputs_.append_as(socket->identifier, CPPType::get<bool>());
    }
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    const SocketValueVariant &index_variant = params.get_input<SocketValueVariant>(0);
    if (index_variant.is_context_dependent_field()) {
      for (const int i : outputs_.index_range()) {
        params.set_output(i, true);
      }
    }
    else {
      const int value = index_variant.get<int>();
      for (const int i : outputs_.index_range()) {
        params.set_output(i, i == value);
      }
    }
  }
};

class LazyFunctionForAnonymousAttributeSetExtract : public lf::LazyFunction {
 public:
  LazyFunctionForAnonymousAttributeSetExtract()
  {
    debug_name_ = "Extract Attribute Set";
    inputs_.append_as("Use", CPPType::get<bool>());
    inputs_.append_as("Field", CPPType::get<SocketValueVariant>(), lf::ValueUsage::Maybe);
    outputs_.append_as("Attributes", CPPType::get<bke::AnonymousAttributeSet>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    const bool use = params.get_input<bool>(0);
    if (!use) {
      params.set_output<bke::AnonymousAttributeSet>(0, {});
      return;
    }
    const SocketValueVariant *value_variant =
        params.try_get_input_data_ptr_or_request<SocketValueVariant>(1);
    if (value_variant == nullptr) {
      /* Wait until the field is computed. */
      return;
    }
 
    bke::AnonymousAttributeSet attributes;
    if (value_variant->is_context_dependent_field()) {
      const GField &field = value_variant->get<GField>();
      field.node().for_each_field_input_recursive([&](const FieldInput &field_input) {
        if (const auto *attr_field_input = dynamic_cast<const AttributeFieldInput *>(&field_input))
        {
          const StringRef name = attr_field_input->attribute_name();
          if (bke::attribute_name_is_anonymous(name)) {
            if (!attributes.names) {
              attributes.names = std::make_shared<Set<std::string>>();
            }
            attributes.names->add_as(name);
          }
        }
      });
    }
    params.set_output(0, std::move(attributes));
  }
};

class LazyFunctionForAnonymousAttributeSetJoin : public lf::LazyFunction {
  const int amount_;
 
 public:
  LazyFunctionForAnonymousAttributeSetJoin(const int amount) : amount_(amount)
  {
    debug_name_ = "Join Attribute Sets";
    for ([[maybe_unused]] const int i : IndexRange(amount)) {
      inputs_.append_as("Use", CPPType::get<bool>());
      inputs_.append_as(
          "Attribute Set", CPPType::get<bke::AnonymousAttributeSet>(), lf::ValueUsage::Maybe);
    }
    outputs_.append_as("Attribute Set", CPPType::get<bke::AnonymousAttributeSet>());
  }
 
  void execute_impl(lf::Params &params, const lf::Context & /*context*/) const override
  {
    Vector<bke::AnonymousAttributeSet *> sets;
    bool set_is_missing = false;
    for (const int i : IndexRange(amount_)) {
      if (params.get_input<bool>(this->get_use_input(i))) {
        if (bke::AnonymousAttributeSet *set =
                params.try_get_input_data_ptr_or_request<bke::AnonymousAttributeSet>(
                    this->get_attribute_set_input(i)))
        {
          sets.append(set);
        }
        else {
          set_is_missing = true;
        }
      }
    }
    if (set_is_missing) {
      return;
    }
    bke::AnonymousAttributeSet joined_set;
    if (sets.is_empty()) {
      /* Nothing to do. */
    }
    else if (sets.size() == 1) {
      joined_set.names = std::move(sets[0]->names);
    }
    else {
      joined_set.names = std::make_shared<Set<std::string>>();
      for (const bke::AnonymousAttributeSet *set : sets) {
        if (set->names) {
          for (const std::string &name : *set->names) {
            joined_set.names->add(name);
          }
        }
      }
    }
    params.set_output(0, std::move(joined_set));
  }
 
  int get_use_input(const int i) const
  {
    return 2 * i;
  }
 
  int get_attribute_set_input(const int i) const
  {
    return 2 * i + 1;
  }

  static const LazyFunctionForAnonymousAttributeSetJoin &get_cached(const int amount,
                                                                    ResourceScope &scope)
  {
    constexpr int cache_amount = 16;
    static std::array<LazyFunctionForAnonymousAttributeSetJoin, cache_amount> cached_functions =
        get_cache(std::make_index_sequence<cache_amount>{});
    if (amount < cached_functions.size()) {
      return cached_functions[amount];
    }
 
    return scope.construct<LazyFunctionForAnonymousAttributeSetJoin>(amount);
  }
 
 private:
  template<size_t... I>
  static std::array<LazyFunctionForAnonymousAttributeSetJoin, sizeof...(I)> get_cache(
      std::index_sequence<I...> /*indices*/)
  {
    return {LazyFunctionForAnonymousAttributeSetJoin(I)...};
  }
};
 
class LazyFunctionForSimulationZone : public LazyFunction {
 private:
  const bNode &sim_output_bnode_;
  const LazyFunction &fn_;
 
 public:
  LazyFunctionForSimulationZone(const bNode &sim_output_bnode, const LazyFunction &fn)
      : sim_output_bnode_(sim_output_bnode), fn_(fn)
  {
    debug_name_ = "Simulation Zone";
    inputs_ = fn.inputs();
    outputs_ = fn.outputs();
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    ScopedNodeTimer node_timer{context, sim_output_bnode_};
    GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
 
    bke::SimulationZoneComputeContext compute_context{user_data.compute_context,
                                                      sim_output_bnode_};
 
    GeoNodesLFUserData zone_user_data = user_data;
    zone_user_data.compute_context = &compute_context;
    zone_user_data.log_socket_values = should_log_socket_values_for_context(
        user_data, compute_context.hash());
 
    GeoNodesLFLocalUserData zone_local_user_data{zone_user_data};
    lf::Context zone_context{context.storage, &zone_user_data, &zone_local_user_data};
    fn_.execute(params, zone_context);
  }
 
  void *init_storage(LinearAllocator<> &allocator) const override
  {
    return fn_.init_storage(allocator);
  }
 
  void destruct_storage(void *storage) const override
  {
    fn_.destruct_storage(storage);
  }
 
  std::string input_name(const int i) const override
  {
    return fn_.input_name(i);
  }
 
  std::string output_name(const int i) const override
  {
    return fn_.output_name(i);
  }
};
 
using JoinAttributeSetsCache = Map<Vector<lf::OutputSocket *>, lf::OutputSocket *>;
 
struct BuildGraphParams {
  lf::Graph &lf_graph;
  MultiValueMap<const bNodeSocket *, lf::InputSocket *> lf_inputs_by_bsocket;
  Map<const bNodeSocket *, lf::OutputSocket *> lf_output_by_bsocket;
  Map<const bNodeSocket *, lf::OutputSocket *> usage_by_bsocket;
  Map<const bNodeSocket *, lf::InputSocket *> lf_attribute_set_input_by_output_geometry_bsocket;
  Map<const bNodeLink *, lf::InputSocket *> lf_input_by_multi_input_link;
  Map<const bNodeLink *, lf::InputSocket *> lf_input_by_border_link;
  Set<lf::InputSocket *> socket_usage_inputs;
  MultiValueMap<int, lf::InputSocket *> lf_attribute_set_input_by_field_source_index;
  MultiValueMap<int, lf::InputSocket *> lf_attribute_set_input_by_caller_propagation_index;
  Map<Vector<lf::OutputSocket *>, lf::OutputSocket *> socket_usages_combination_cache;
 
  BuildGraphParams(lf::Graph &lf_graph) : lf_graph(lf_graph) {}
};
 
struct ZoneFunctionIndices {
  struct {
    Vector<int> main;
    Vector<int> border_links;
    Vector<int> output_usages;
    Map<int, int> attributes_by_field_source_index;
    Map<int, int> attributes_by_caller_propagation_index;
  } inputs;
  struct {
    Vector<int> main;
    Vector<int> border_link_usages;
    Vector<int> input_usages;
  } outputs;
};
 
struct ZoneBuildInfo {
  const LazyFunction *lazy_function = nullptr;

  ZoneFunctionIndices indices;
};

struct ZoneBodyFunction {
  const LazyFunction *function = nullptr;
  ZoneFunctionIndices indices;
};
 
static bool ignore_zone_bsocket(const bNodeSocket &bsocket)
{
  if (!bsocket.is_available()) {
    return true;
  }
  if (!bsocket.typeinfo->geometry_nodes_cpp_type) {
    /* These are typically extend sockets. */
    return true;
  }
  return false;
}
 
static void initialize_zone_wrapper(const bNodeTreeZone &zone,
                                    ZoneBuildInfo &zone_info,
                                    const ZoneBodyFunction &body_fn,
                                    Vector<lf::Input> &r_inputs,
                                    Vector<lf::Output> &r_outputs)
{
  for (const bNodeSocket *socket : zone.input_node->input_sockets()) {
    if (ignore_zone_bsocket(*socket)) {
      continue;
    }
    zone_info.indices.inputs.main.append(r_inputs.append_and_get_index_as(
        socket->name, *socket->typeinfo->geometry_nodes_cpp_type, lf::ValueUsage::Maybe));
  }
 
  for (const bNodeLink *link : zone.border_links) {
    zone_info.indices.inputs.border_links.append(
        r_inputs.append_and_get_index_as(link->fromsock->name,
                                         *link->tosock->typeinfo->geometry_nodes_cpp_type,
                                         lf::ValueUsage::Maybe));
  }
 
  for (const bNodeSocket *socket : zone.output_node->output_sockets()) {
    if (ignore_zone_bsocket(*socket)) {
      continue;
    }
    const CPPType *cpp_type = socket->typeinfo->geometry_nodes_cpp_type;
    zone_info.indices.inputs.output_usages.append(
        r_inputs.append_and_get_index_as("Usage", CPPType::get<bool>(), lf::ValueUsage::Maybe));
    zone_info.indices.outputs.main.append(
        r_outputs.append_and_get_index_as(socket->name, *cpp_type));
  }
 
  for ([[maybe_unused]] const bNodeSocket *socket : zone.input_node->input_sockets()) {
    if (ignore_zone_bsocket(*socket)) {
      continue;
    }
    zone_info.indices.outputs.input_usages.append(
        r_outputs.append_and_get_index_as("Usage", CPPType::get<bool>()));
  }
 
  for ([[maybe_unused]] const bNodeLink *link : zone.border_links) {
    zone_info.indices.outputs.border_link_usages.append(
        r_outputs.append_and_get_index_as("Border Link Usage", CPPType::get<bool>()));
  }
 
  for (const auto item : body_fn.indices.inputs.attributes_by_field_source_index.items()) {
    zone_info.indices.inputs.attributes_by_field_source_index.add_new(
        item.key,
        r_inputs.append_and_get_index_as(
            "Attribute Set", CPPType::get<bke::AnonymousAttributeSet>(), lf::ValueUsage::Maybe));
  }
  for (const auto item : body_fn.indices.inputs.attributes_by_caller_propagation_index.items()) {
    zone_info.indices.inputs.attributes_by_caller_propagation_index.add_new(
        item.key,
        r_inputs.append_and_get_index_as(
            "Attribute Set", CPPType::get<bke::AnonymousAttributeSet>(), lf::ValueUsage::Maybe));
  }
}
 
static std::string zone_wrapper_input_name(const ZoneBuildInfo &zone_info,
                                           const bNodeTreeZone &zone,
                                           const Span<lf::Input> inputs,
                                           const int lf_socket_i)
{
  if (zone_info.indices.inputs.output_usages.contains(lf_socket_i)) {
    const int output_usage_i = lf_socket_i - zone_info.indices.inputs.output_usages.first();
    int current_valid_i = 0;
    for (const bNodeSocket *bsocket : zone.output_node->output_sockets()) {
      if (ignore_zone_bsocket(*bsocket)) {
        continue;
      }
      if (current_valid_i == output_usage_i) {
        return "Usage: " + StringRef(bsocket->name);
      }
      current_valid_i++;
    }
  }
  return inputs[lf_socket_i].debug_name;
}
 
static std::string zone_wrapper_output_name(const ZoneBuildInfo &zone_info,
                                            const bNodeTreeZone &zone,
                                            const Span<lf::Output> outputs,
                                            const int lf_socket_i)
{
  if (zone_info.indices.outputs.input_usages.contains(lf_socket_i)) {
    const int input_usage_i = lf_socket_i - zone_info.indices.outputs.input_usages.first();
    int current_valid_i = 0;
    for (const bNodeSocket *bsocket : zone.input_node->input_sockets()) {
      if (ignore_zone_bsocket(*bsocket)) {
        continue;
      }
      if (current_valid_i == input_usage_i) {
        return "Usage: " + StringRef(bsocket->name);
      }
      current_valid_i++;
    }
  }
  return outputs[lf_socket_i].debug_name;
}

class RepeatBodyNodeExecuteWrapper : public lf::GraphExecutorNodeExecuteWrapper {
 public:
  const bNode *repeat_output_bnode_ = nullptr;
  VectorSet<lf::FunctionNode *> *lf_body_nodes_ = nullptr;
 
  void execute_node(const lf::FunctionNode &node,
                    lf::Params &params,
                    const lf::Context &context) const
  {
    GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const int iteration = lf_body_nodes_->index_of_try(const_cast<lf::FunctionNode *>(&node));
    const LazyFunction &fn = node.function();
    if (iteration == -1) {
      /* The node is not a loop body node, just execute it normally. */
      fn.execute(params, context);
      return;
    }
 
    /* Setup context for the loop body evaluation. */
    bke::RepeatZoneComputeContext body_compute_context{
        user_data.compute_context, *repeat_output_bnode_, iteration};
    GeoNodesLFUserData body_user_data = user_data;
    body_user_data.compute_context = &body_compute_context;
    body_user_data.log_socket_values = should_log_socket_values_for_context(
        user_data, body_compute_context.hash());
 
    GeoNodesLFLocalUserData body_local_user_data{body_user_data};
    lf::Context body_context{context.storage, &body_user_data, &body_local_user_data};
    fn.execute(params, body_context);
  }
};

class RepeatZoneSideEffectProvider : public lf::GraphExecutorSideEffectProvider {
 public:
  const bNode *repeat_output_bnode_ = nullptr;
  Span<lf::FunctionNode *> lf_body_nodes_;
 
  Vector<const lf::FunctionNode *> get_nodes_with_side_effects(
      const lf::Context &context) const override
  {
    GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const GeoNodesCallData &call_data = *user_data.call_data;
    if (!call_data.side_effect_nodes) {
      return {};
    }
    const ComputeContextHash &context_hash = user_data.compute_context->hash();
    const Span<int> iterations_with_side_effects =
        call_data.side_effect_nodes->iterations_by_iteration_zone.lookup(
            {context_hash, repeat_output_bnode_->identifier});
 
    Vector<const lf::FunctionNode *> lf_nodes;
    for (const int i : iterations_with_side_effects) {
      if (i >= 0 && i < lf_body_nodes_.size()) {
        lf_nodes.append(lf_body_nodes_[i]);
      }
    }
    return lf_nodes;
  }
};
 
struct RepeatEvalStorage {
  LinearAllocator<> allocator;
  VectorSet<lf::FunctionNode *> lf_body_nodes;
  lf::Graph graph;
  std::optional<LazyFunctionForLogicalOr> or_function;
  std::optional<RepeatZoneSideEffectProvider> side_effect_provider;
  std::optional<RepeatBodyNodeExecuteWrapper> body_execute_wrapper;
  std::optional<lf::GraphExecutor> graph_executor;
  Array<SocketValueVariant> index_values;
  void *graph_executor_storage = nullptr;
  bool multi_threading_enabled = false;
  Vector<int> input_index_map;
  Vector<int> output_index_map;
};
 
class LazyFunctionForRepeatZone : public LazyFunction {
 private:
  const bNodeTree &btree_;
  const bNodeTreeZone &zone_;
  const bNode &repeat_output_bnode_;
  const ZoneBuildInfo &zone_info_;
  const ZoneBodyFunction &body_fn_;
 
 public:
  LazyFunctionForRepeatZone(const bNodeTree &btree,
                            const bNodeTreeZone &zone,
                            ZoneBuildInfo &zone_info,
                            const ZoneBodyFunction &body_fn)
      : btree_(btree),
        zone_(zone),
        repeat_output_bnode_(*zone.output_node),
        zone_info_(zone_info),
        body_fn_(body_fn)
  {
    debug_name_ = "Repeat Zone";
 
    initialize_zone_wrapper(zone, zone_info, body_fn, inputs_, outputs_);
    /* Iterations input is always used. */
    inputs_[zone_info.indices.inputs.main[0]].usage = lf::ValueUsage::Used;
  }
 
  void *init_storage(LinearAllocator<> &allocator) const override
  {
    return allocator.construct<RepeatEvalStorage>().release();
  }
 
  void destruct_storage(void *storage) const override
  {
    RepeatEvalStorage *s = static_cast<RepeatEvalStorage *>(storage);
    if (s->graph_executor_storage) {
      s->graph_executor->destruct_storage(s->graph_executor_storage);
    }
    std::destroy_at(s);
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const ScopedNodeTimer node_timer{context, repeat_output_bnode_};
 
    auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
 
    const NodeGeometryRepeatOutput &node_storage = *static_cast<const NodeGeometryRepeatOutput *>(
        repeat_output_bnode_.storage);
    RepeatEvalStorage &eval_storage = *static_cast<RepeatEvalStorage *>(context.storage);
 
    const int iterations_usage_index = zone_info_.indices.outputs.input_usages[0];
    if (!params.output_was_set(iterations_usage_index)) {
      /* The iterations input is always used. */
      params.set_output(iterations_usage_index, true);
    }
 
    if (!eval_storage.graph_executor) {
      /* Create the execution graph in the first evaluation. */
      this->initialize_execution_graph(
          params, eval_storage, node_storage, user_data, local_user_data);
    }
 
    /* Execute the graph for the repeat zone. */
    lf::RemappedParams eval_graph_params{*eval_storage.graph_executor,
                                         params,
                                         eval_storage.input_index_map,
                                         eval_storage.output_index_map,
                                         eval_storage.multi_threading_enabled};
    lf::Context eval_graph_context{
        eval_storage.graph_executor_storage, context.user_data, context.local_user_data};
    eval_storage.graph_executor->execute(eval_graph_params, eval_graph_context);
  }

  void initialize_execution_graph(lf::Params &params,
                                  RepeatEvalStorage &eval_storage,
                                  const NodeGeometryRepeatOutput &node_storage,
                                  GeoNodesLFUserData &user_data,
                                  GeoNodesLFLocalUserData &local_user_data) const
  {
    const int num_repeat_items = node_storage.items_num;
    const int num_border_links = body_fn_.indices.inputs.border_links.size();
 
    /* Number of iterations to evaluate. */
    const int iterations = std::max<int>(
        0, params.get_input<SocketValueVariant>(zone_info_.indices.inputs.main[0]).get<int>());
 
    /* Show a warning when the inspection index is out of range. */
    if (node_storage.inspection_index > 0) {
      if (node_storage.inspection_index >= iterations) {
        if (geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(
                user_data))
        {
          tree_logger->node_warnings.append(
              *tree_logger->allocator,
              {repeat_output_bnode_.identifier,
               {NodeWarningType::Info, N_("Inspection index is out of range")}});
        }
      }
    }
 
    /* Take iterations input into account. */
    const int main_inputs_offset = 1;
    const int body_inputs_offset = 1;
 
    lf::Graph &lf_graph = eval_storage.graph;
 
    Vector<lf::GraphInputSocket *> lf_inputs;
    Vector<lf::GraphOutputSocket *> lf_outputs;
 
    for (const int i : inputs_.index_range()) {
      const lf::Input &input = inputs_[i];
      lf_inputs.append(&lf_graph.add_input(*input.type, this->input_name(i)));
    }
    for (const int i : outputs_.index_range()) {
      const lf::Output &output = outputs_[i];
      lf_outputs.append(&lf_graph.add_output(*output.type, this->output_name(i)));
    }
 
    /* Create body nodes. */
    VectorSet<lf::FunctionNode *> &lf_body_nodes = eval_storage.lf_body_nodes;
    for ([[maybe_unused]] const int i : IndexRange(iterations)) {
      lf::FunctionNode &lf_node = lf_graph.add_function(*body_fn_.function);
      lf_body_nodes.add_new(&lf_node);
    }
 
    /* Create nodes for combining border link usages. A border link is used when any of the loop
     * bodies uses the border link, so an "or" node is necessary. */
    Array<lf::FunctionNode *> lf_border_link_usage_or_nodes(num_border_links);
    eval_storage.or_function.emplace(iterations);
    for (const int i : IndexRange(num_border_links)) {
      lf::FunctionNode &lf_node = lf_graph.add_function(*eval_storage.or_function);
      lf_border_link_usage_or_nodes[i] = &lf_node;
    }
 
    const bool use_index_values = zone_.input_node->output_socket(0).is_directly_linked();
 
    if (use_index_values) {
      eval_storage.index_values.reinitialize(iterations);
      threading::parallel_for(IndexRange(iterations), 1024, [&](const IndexRange range) {
        for (const int i : range) {
          eval_storage.index_values[i].set(i);
        }
      });
    }
 
    /* Handle body nodes one by one. */
    static const SocketValueVariant static_unused_index{-1};
    for (const int iter_i : lf_body_nodes.index_range()) {
      lf::FunctionNode &lf_node = *lf_body_nodes[iter_i];
      const SocketValueVariant *index_value = use_index_values ?
                                                  &eval_storage.index_values[iter_i] :
                                                  &static_unused_index;
      lf_node.input(body_fn_.indices.inputs.main[0]).set_default_value(index_value);
      for (const int i : IndexRange(num_border_links)) {
        lf_graph.add_link(*lf_inputs[zone_info_.indices.inputs.border_links[i]],
                          lf_node.input(body_fn_.indices.inputs.border_links[i]));
        lf_graph.add_link(lf_node.output(body_fn_.indices.outputs.border_link_usages[i]),
                          lf_border_link_usage_or_nodes[i]->input(iter_i));
      }
      for (const auto item : body_fn_.indices.inputs.attributes_by_field_source_index.items()) {
        lf_graph.add_link(
            *lf_inputs[zone_info_.indices.inputs.attributes_by_field_source_index.lookup(
                item.key)],
            lf_node.input(item.value));
      }
      for (const auto item :
           body_fn_.indices.inputs.attributes_by_caller_propagation_index.items())
      {
        lf_graph.add_link(
            *lf_inputs[zone_info_.indices.inputs.attributes_by_caller_propagation_index.lookup(
                item.key)],
            lf_node.input(item.value));
      }
    }
 
    static bool static_true = true;
 
    /* Handle body nodes pair-wise. */
    for (const int iter_i : lf_body_nodes.index_range().drop_back(1)) {
      lf::FunctionNode &lf_node = *lf_body_nodes[iter_i];
      lf::FunctionNode &lf_next_node = *lf_body_nodes[iter_i + 1];
      for (const int i : IndexRange(num_repeat_items)) {
        lf_graph.add_link(
            lf_node.output(body_fn_.indices.outputs.main[i]),
            lf_next_node.input(body_fn_.indices.inputs.main[i + body_inputs_offset]));
        /* TODO: Add back-link after being able to check for cyclic dependencies. */
        // lf_graph.add_link(lf_next_node.output(body_fn_.indices.outputs.input_usages[i]),
        //                   lf_node.input(body_fn_.indices.inputs.output_usages[i]));
        lf_node.input(body_fn_.indices.inputs.output_usages[i]).set_default_value(&static_true);
      }
    }
 
    /* Handle border link usage outputs. */
    for (const int i : IndexRange(num_border_links)) {
      lf_graph.add_link(lf_border_link_usage_or_nodes[i]->output(0),
                        *lf_outputs[zone_info_.indices.outputs.border_link_usages[i]]);
    }
 
    if (iterations > 0) {
      {
        /* Link first body node to input/output nodes. */
        lf::FunctionNode &lf_first_body_node = *lf_body_nodes[0];
        for (const int i : IndexRange(num_repeat_items)) {
          lf_graph.add_link(
              *lf_inputs[zone_info_.indices.inputs.main[i + main_inputs_offset]],
              lf_first_body_node.input(body_fn_.indices.inputs.main[i + body_inputs_offset]));
          lf_graph.add_link(
              lf_first_body_node.output(
                  body_fn_.indices.outputs.input_usages[i + body_inputs_offset]),
              *lf_outputs[zone_info_.indices.outputs.input_usages[i + main_inputs_offset]]);
        }
      }
      {
        /* Link last body node to input/output nodes. */
        lf::FunctionNode &lf_last_body_node = *lf_body_nodes.as_span().last();
        for (const int i : IndexRange(num_repeat_items)) {
          lf_graph.add_link(lf_last_body_node.output(body_fn_.indices.outputs.main[i]),
                            *lf_outputs[zone_info_.indices.outputs.main[i]]);
          lf_graph.add_link(*lf_inputs[zone_info_.indices.inputs.output_usages[i]],
                            lf_last_body_node.input(body_fn_.indices.inputs.output_usages[i]));
        }
      }
    }
    else {
      /* There are no iterations, just link the input directly to the output. */
      for (const int i : IndexRange(num_repeat_items)) {
        lf_graph.add_link(*lf_inputs[zone_info_.indices.inputs.main[i + main_inputs_offset]],
                          *lf_outputs[zone_info_.indices.outputs.main[i]]);
        lf_graph.add_link(
            *lf_inputs[zone_info_.indices.inputs.output_usages[i]],
            *lf_outputs[zone_info_.indices.outputs.input_usages[i + main_inputs_offset]]);
      }
      for (const int i : IndexRange(num_border_links)) {
        static bool static_false = false;
        lf_outputs[zone_info_.indices.outputs.border_link_usages[i]]->set_default_value(
            &static_false);
      }
    }
 
    lf_outputs[zone_info_.indices.outputs.input_usages[0]]->set_default_value(&static_true);
 
    /* The graph is ready, update the node indices which are required by the executor. */
    lf_graph.update_node_indices();
 
    // std::cout << "\n\n" << lf_graph.to_dot() << "\n\n";
 
    /* Create a mapping from parameter indices inside of this graph to parameters of the repeat
     * zone. The main complexity below stems from the fact that the iterations input is handled
     * outside of this graph. */
    eval_storage.output_index_map.reinitialize(outputs_.size() - 1);
    eval_storage.input_index_map.resize(inputs_.size() - 1);
    array_utils::fill_index_range<int>(eval_storage.input_index_map, 1);
 
    Vector<const lf::GraphInputSocket *> lf_graph_inputs = lf_inputs.as_span().drop_front(1);
 
    const int iteration_usage_index = zone_info_.indices.outputs.input_usages[0];
    array_utils::fill_index_range<int>(
        eval_storage.output_index_map.as_mutable_span().take_front(iteration_usage_index));
    array_utils::fill_index_range<int>(
        eval_storage.output_index_map.as_mutable_span().drop_front(iteration_usage_index),
        iteration_usage_index + 1);
 
    Vector<const lf::GraphOutputSocket *> lf_graph_outputs = lf_outputs.as_span().take_front(
        iteration_usage_index);
    lf_graph_outputs.extend(lf_outputs.as_span().drop_front(iteration_usage_index + 1));
 
    eval_storage.body_execute_wrapper.emplace();
    eval_storage.body_execute_wrapper->repeat_output_bnode_ = &repeat_output_bnode_;
    eval_storage.body_execute_wrapper->lf_body_nodes_ = &lf_body_nodes;
    eval_storage.side_effect_provider.emplace();
    eval_storage.side_effect_provider->repeat_output_bnode_ = &repeat_output_bnode_;
    eval_storage.side_effect_provider->lf_body_nodes_ = lf_body_nodes;
 
    eval_storage.graph_executor.emplace(lf_graph,
                                        std::move(lf_graph_inputs),
                                        std::move(lf_graph_outputs),
                                        nullptr,
                                        &*eval_storage.side_effect_provider,
                                        &*eval_storage.body_execute_wrapper);
    eval_storage.graph_executor_storage = eval_storage.graph_executor->init_storage(
        eval_storage.allocator);
 
    /* Log graph for debugging purposes. */
    bNodeTree &btree_orig = *reinterpret_cast<bNodeTree *>(
        DEG_get_original_id(const_cast<ID *>(&btree_.id)));
    if (btree_orig.runtime->logged_zone_graphs) {
      std::lock_guard lock{btree_orig.runtime->logged_zone_graphs->mutex};
      btree_orig.runtime->logged_zone_graphs->graph_by_zone_id.lookup_or_add_cb(
          repeat_output_bnode_.identifier, [&]() { return lf_graph.to_dot(); });
    }
  }
 
  std::string input_name(const int i) const override
  {
    return zone_wrapper_input_name(zone_info_, zone_, inputs_, i);
  }
 
  std::string output_name(const int i) const override
  {
    return zone_wrapper_output_name(zone_info_, zone_, outputs_, i);
  }
};
 
class LazyFunctionForForeachGeometryElementZone;
struct ForeachGeometryElementEvalStorage;
 
struct ForeachElementComponentID {
  GeometryComponent::Type component_type;
  AttrDomain domain;
  std::optional<int> layer_index;
};

struct ForeachElementComponent {
  ForeachElementComponentID id;
  std::optional<bke::GeometryFieldContext> field_context;
  std::optional<FieldEvaluator> field_evaluator;
  Array<SocketValueVariant> index_values;
  Array<Array<SocketValueVariant>> item_input_values;
  std::optional<Array<GeometrySet>> element_geometries;
  IndexRange body_nodes_range;
 
  void emplace_field_context(const GeometrySet &geometry)
  {
    if (this->id.component_type == GeometryComponent::Type::GreasePencil &&
        ELEM(this->id.domain, AttrDomain::Point, AttrDomain::Curve))
    {
      this->field_context.emplace(
          *geometry.get_grease_pencil(), this->id.domain, *this->id.layer_index);
    }
    else {
      this->field_context.emplace(*geometry.get_component(this->id.component_type),
                                  this->id.domain);
    }
  }
 
  AttributeAccessor input_attributes() const
  {
    return *this->field_context->attributes();
  }
 
  MutableAttributeAccessor attributes_for_write(GeometrySet &geometry) const
  {
    if (this->id.component_type == GeometryComponent::Type::GreasePencil &&
        ELEM(this->id.domain, AttrDomain::Point, AttrDomain::Curve))
    {
      BLI_assert(this->id.layer_index.has_value());
      GreasePencil *grease_pencil = geometry.get_grease_pencil_for_write();
      const bke::greasepencil::Layer &layer = grease_pencil->layer(*this->id.layer_index);
      bke::greasepencil::Drawing *drawing = grease_pencil->get_eval_drawing(layer);
      BLI_assert(drawing);
      return drawing->strokes_for_write().attributes_for_write();
    }
    GeometryComponent &component = geometry.get_component_for_write(this->id.component_type);
    return *component.attributes_for_write();
  }
};

struct LazyFunctionForReduceForeachGeometryElement : public LazyFunction {
  const LazyFunctionForForeachGeometryElementZone &parent_;
  ForeachGeometryElementEvalStorage &eval_storage_;
 
 public:
  LazyFunctionForReduceForeachGeometryElement(
      const LazyFunctionForForeachGeometryElementZone &parent,
      ForeachGeometryElementEvalStorage &eval_storage);
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override;
 
  void handle_main_items_and_geometry(lf::Params &params, const lf::Context &context) const;
  void handle_generation_items(lf::Params &params, const lf::Context &context) const;
  int handle_invalid_generation_items(lf::Params &params) const;
  void handle_generation_item_groups(lf::Params &params,
                                     const lf::Context &context,
                                     int first_valid_item_i) const;
  void handle_generation_items_group(lf::Params &params,
                                     const lf::Context &context,
                                     int geometry_item_i,
                                     IndexRange generation_items_range) const;
  bool handle_generation_items_group_lazyness(lf::Params &params,
                                              const lf::Context &context,
                                              int geometry_item_i,
                                              IndexRange generation_items_range) const;
};

class ForeachGeometryElementNodeExecuteWrapper : public lf::GraphExecutorNodeExecuteWrapper {
 public:
  const bNode *output_bnode_ = nullptr;
  VectorSet<lf::FunctionNode *> *lf_body_nodes_ = nullptr;
 
  void execute_node(const lf::FunctionNode &node,
                    lf::Params &params,
                    const lf::Context &context) const
  {
    GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const int index = lf_body_nodes_->index_of_try(const_cast<lf::FunctionNode *>(&node));
    const LazyFunction &fn = node.function();
    if (index == -1) {
      /* The node is not a loop body node, just execute it normally. */
      fn.execute(params, context);
      return;
    }
 
    /* Setup context for the loop body evaluation. */
    bke::ForeachGeometryElementZoneComputeContext body_compute_context{
        user_data.compute_context, *output_bnode_, index};
    GeoNodesLFUserData body_user_data = user_data;
    body_user_data.compute_context = &body_compute_context;
    body_user_data.log_socket_values = should_log_socket_values_for_context(
        user_data, body_compute_context.hash());
 
    GeoNodesLFLocalUserData body_local_user_data{body_user_data};
    lf::Context body_context{context.storage, &body_user_data, &body_local_user_data};
    fn.execute(params, body_context);
  }
};

class ForeachGeometryElementZoneSideEffectProvider : public lf::GraphExecutorSideEffectProvider {
 public:
  const bNode *output_bnode_ = nullptr;
  Span<lf::FunctionNode *> lf_body_nodes_;
 
  Vector<const lf::FunctionNode *> get_nodes_with_side_effects(
      const lf::Context &context) const override
  {
    GeoNodesLFUserData &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const GeoNodesCallData &call_data = *user_data.call_data;
    if (!call_data.side_effect_nodes) {
      return {};
    }
    const ComputeContextHash &context_hash = user_data.compute_context->hash();
    const Span<int> iterations_with_side_effects =
        call_data.side_effect_nodes->iterations_by_iteration_zone.lookup(
            {context_hash, output_bnode_->identifier});
 
    Vector<const lf::FunctionNode *> lf_nodes;
    for (const int i : iterations_with_side_effects) {
      if (i >= 0 && i < lf_body_nodes_.size()) {
        lf_nodes.append(lf_body_nodes_[i]);
      }
    }
    return lf_nodes;
  }
};

struct ForeachGeometryElementEvalStorage {
  LinearAllocator<> allocator;

  lf::Graph graph;
  std::optional<ForeachGeometryElementZoneSideEffectProvider> side_effect_provider;
  std::optional<ForeachGeometryElementNodeExecuteWrapper> body_execute_wrapper;
  std::optional<lf::GraphExecutor> graph_executor;
  void *graph_executor_storage = nullptr;

  std::optional<LazyFunctionForLogicalOr> or_function;
  std::optional<LazyFunctionForReduceForeachGeometryElement> reduce_function;

  VectorSet<lf::FunctionNode *> lf_body_nodes;

  GeometrySet main_geometry;
  Array<ForeachElementComponent> components;
  int total_iterations_num = 0;
};
 
class LazyFunctionForForeachGeometryElementZone : public LazyFunction {
 private:
  const bNodeTree &btree_;
  const bNodeTreeZone &zone_;
  const bNode &output_bnode_;
  const ZoneBuildInfo &zone_info_;
  const ZoneBodyFunction &body_fn_;
 
  struct ItemIndices {
    /* `outer` refers to sockets on the outside of the zone, and `inner` to the sockets on the
     * inside. The `lf` and `bsocket` indices are similar, but the `lf` indices skip unavailable
     * and extend sockets. */
    IndexRange lf_outer;
    IndexRange lf_inner;
    IndexRange bsocket_outer;
    IndexRange bsocket_inner;
  };

  struct {
    ItemIndices inputs;
    ItemIndices main;
    ItemIndices generation;
  } indices_;
 
  friend LazyFunctionForReduceForeachGeometryElement;
 
 public:
  LazyFunctionForForeachGeometryElementZone(const bNodeTree &btree,
                                            const bNodeTreeZone &zone,
                                            ZoneBuildInfo &zone_info,
                                            const ZoneBodyFunction &body_fn)
      : btree_(btree),
        zone_(zone),
        output_bnode_(*zone.output_node),
        zone_info_(zone_info),
        body_fn_(body_fn)
  {
    debug_name_ = "Foreach Geometry Element";
 
    initialize_zone_wrapper(zone, zone_info, body_fn, inputs_, outputs_);
    /* All main inputs are always used for now. */
    for (const int i : zone_info.indices.inputs.main) {
      inputs_[i].usage = lf::ValueUsage::Used;
    }
 
    const auto &node_storage = *static_cast<const NodeGeometryForeachGeometryElementOutput *>(
        output_bnode_.storage);
    const AttrDomain iteration_domain = AttrDomain(node_storage.domain);
    BLI_assert(zone_.input_node->output_socket(1).is_available() ==
               (iteration_domain != AttrDomain::Corner));
 
    const int input_items_num = node_storage.input_items.items_num;
    const int main_items_num = node_storage.main_items.items_num;
    const int generation_items_num = node_storage.generation_items.items_num;
 
    indices_.inputs.lf_outer = IndexRange::from_begin_size(2, input_items_num);
    indices_.inputs.lf_inner = IndexRange::from_begin_size(
        iteration_domain == AttrDomain::Corner ? 1 : 2, input_items_num);
    indices_.inputs.bsocket_outer = indices_.inputs.lf_outer;
    indices_.inputs.bsocket_inner = indices_.inputs.lf_inner;
 
    indices_.main.lf_outer = IndexRange::from_begin_size(1, main_items_num);
    indices_.main.lf_inner = IndexRange::from_begin_size(0, main_items_num);
    indices_.main.bsocket_outer = indices_.main.lf_outer;
    indices_.main.bsocket_inner = indices_.main.lf_inner;
 
    indices_.generation.lf_outer = IndexRange::from_begin_size(1 + main_items_num,
                                                               generation_items_num);
    indices_.generation.lf_inner = IndexRange::from_begin_size(main_items_num,
                                                               generation_items_num);
    indices_.generation.bsocket_outer = IndexRange::from_begin_size(2 + main_items_num,
                                                                    generation_items_num);
    indices_.generation.bsocket_inner = IndexRange::from_begin_size(1 + main_items_num,
                                                                    generation_items_num);
  }
 
  void *init_storage(LinearAllocator<> &allocator) const override
  {
    return allocator.construct<ForeachGeometryElementEvalStorage>().release();
  }
 
  void destruct_storage(void *storage) const override
  {
    auto *s = static_cast<ForeachGeometryElementEvalStorage *>(storage);
    if (s->graph_executor_storage) {
      s->graph_executor->destruct_storage(s->graph_executor_storage);
    }
    std::destroy_at(s);
  }
 
  void execute_impl(lf::Params &params, const lf::Context &context) const override
  {
    const ScopedNodeTimer node_timer{context, output_bnode_};
 
    auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
 
    const auto &node_storage = *static_cast<const NodeGeometryForeachGeometryElementOutput *>(
        output_bnode_.storage);
    auto &eval_storage = *static_cast<ForeachGeometryElementEvalStorage *>(context.storage);
    geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(user_data);
 
    if (!eval_storage.graph_executor) {
      /* Create the execution graph in the first evaluation. */
      this->initialize_execution_graph(params, eval_storage, node_storage);
 
      if (tree_logger) {
        if (eval_storage.total_iterations_num == 0) {
          if (!eval_storage.main_geometry.is_empty()) {
            tree_logger->node_warnings.append(
                *tree_logger->allocator,
                {zone_.input_node->identifier,
                 {NodeWarningType::Info,
                  N_("Input geometry has no elements in the iteration domain.")}});
          }
        }
      }
    }
 
    lf::Context eval_graph_context{
        eval_storage.graph_executor_storage, context.user_data, context.local_user_data};
 
    eval_storage.graph_executor->execute(params, eval_graph_context);
  }
 
  void initialize_execution_graph(
      lf::Params &params,
      ForeachGeometryElementEvalStorage &eval_storage,
      const NodeGeometryForeachGeometryElementOutput &node_storage) const
  {
    eval_storage.main_geometry = params.extract_input<GeometrySet>(
        zone_info_.indices.inputs.main[0]);
 
    /* Find all the things we need to iterate over in the input geometry. */
    this->prepare_components(params, eval_storage, node_storage);
 
    /* Add interface sockets for the zone graph. Those are the same as for the entire zone, even
     * though some of the inputs are not strictly needed anymore. It's easier to avoid another
     * level of index remapping though. */
    lf::Graph &lf_graph = eval_storage.graph;
    Vector<lf::GraphInputSocket *> graph_inputs;
    Vector<lf::GraphOutputSocket *> graph_outputs;
    for (const int i : inputs_.index_range()) {
      const lf::Input &input = inputs_[i];
      graph_inputs.append(&lf_graph.add_input(*input.type, this->input_name(i)));
    }
    for (const int i : outputs_.index_range()) {
      const lf::Output &output = outputs_[i];
      graph_outputs.append(&lf_graph.add_output(*output.type, this->output_name(i)));
    }
 
    /* Add all the nodes and links to the graph. */
    this->build_graph_contents(eval_storage, node_storage, graph_inputs, graph_outputs);
 
    eval_storage.side_effect_provider.emplace();
    eval_storage.side_effect_provider->output_bnode_ = &output_bnode_;
    eval_storage.side_effect_provider->lf_body_nodes_ = eval_storage.lf_body_nodes;
 
    eval_storage.body_execute_wrapper.emplace();
    eval_storage.body_execute_wrapper->output_bnode_ = &output_bnode_;
    eval_storage.body_execute_wrapper->lf_body_nodes_ = &eval_storage.lf_body_nodes;
 
    lf_graph.update_node_indices();
    eval_storage.graph_executor.emplace(lf_graph,
                                        graph_inputs.as_span(),
                                        graph_outputs.as_span(),
                                        nullptr,
                                        &*eval_storage.side_effect_provider,
                                        &*eval_storage.body_execute_wrapper);
    eval_storage.graph_executor_storage = eval_storage.graph_executor->init_storage(
        eval_storage.allocator);
 
    /* Log graph for debugging purposes. */
    bNodeTree &btree_orig = *reinterpret_cast<bNodeTree *>(
        DEG_get_original_id(const_cast<ID *>(&btree_.id)));
    if (btree_orig.runtime->logged_zone_graphs) {
      std::lock_guard lock{btree_orig.runtime->logged_zone_graphs->mutex};
      btree_orig.runtime->logged_zone_graphs->graph_by_zone_id.lookup_or_add_cb(
          output_bnode_.identifier, [&]() { return lf_graph.to_dot(); });
    }
  }
 
  void prepare_components(lf::Params &params,
                          ForeachGeometryElementEvalStorage &eval_storage,
                          const NodeGeometryForeachGeometryElementOutput &node_storage) const
  {
    const AttrDomain iteration_domain = AttrDomain(node_storage.domain);
 
    /* TODO: Get propagation info from input, but that's not necessary for correctness for now. */
    AttributeFilter attribute_filter;
 
    const bNodeSocket &element_geometry_bsocket = zone_.input_node->output_socket(1);
    const bool create_element_geometries = element_geometry_bsocket.is_available() &&
                                           element_geometry_bsocket.is_directly_linked();
 
    /* Gather components to process. */
    Vector<ForeachElementComponentID> component_ids;
    for (const GeometryComponent *src_component : eval_storage.main_geometry.get_components()) {
      const GeometryComponent::Type component_type = src_component->type();
      if (src_component->type() == GeometryComponent::Type::GreasePencil &&
          ELEM(iteration_domain, AttrDomain::Point, AttrDomain::Curve))
      {
        const GreasePencil &grease_pencil = *eval_storage.main_geometry.get_grease_pencil();
        for (const int layer_i : grease_pencil.layers().index_range()) {
          const bke::greasepencil::Drawing *drawing = grease_pencil.get_eval_drawing(
              grease_pencil.layer(layer_i));
          if (drawing == nullptr) {
            continue;
          }
          const bke::CurvesGeometry &curves = drawing->strokes();
          if (curves.curves_num() == 0) {
            continue;
          }
          component_ids.append({component_type, iteration_domain, layer_i});
        }
      }
      else {
        const int domain_size = src_component->attribute_domain_size(iteration_domain);
        if (domain_size > 0) {
          component_ids.append({component_type, iteration_domain});
        }
      }
    }
 
    const Field<bool> selection_field = params
                                            .extract_input<SocketValueVariant>(
                                                zone_info_.indices.inputs.main[1])
                                            .extract<Field<bool>>();
 
    /* Evaluate the selection and field inputs for all components. */
    int body_nodes_offset = 0;
    eval_storage.components.reinitialize(component_ids.size());
    for (const int component_i : component_ids.index_range()) {
      const ForeachElementComponentID id = component_ids[component_i];
      ForeachElementComponent &component_info = eval_storage.components[component_i];
      component_info.id = id;
      component_info.emplace_field_context(eval_storage.main_geometry);
 
      const int domain_size = component_info.input_attributes().domain_size(id.domain);
      BLI_assert(domain_size > 0);
 
      /* Prepare field evaluation for the zone inputs. */
      component_info.field_evaluator.emplace(*component_info.field_context, domain_size);
      component_info.field_evaluator->set_selection(selection_field);
      for (const int item_i : IndexRange(node_storage.input_items.items_num)) {
        const GField item_field =
            params
                .get_input<SocketValueVariant>(
                    zone_info_.indices.inputs.main[indices_.inputs.lf_outer[item_i]])
                .get<GField>();
        component_info.field_evaluator->add(item_field);
      }
 
      /* Evaluate all fields passed to the zone input. */
      component_info.field_evaluator->evaluate();
 
      /* The mask contains all the indices that should be iterated over in the component. */
      const IndexMask mask = component_info.field_evaluator->get_evaluated_selection_as_mask();
      component_info.body_nodes_range = IndexRange::from_begin_size(body_nodes_offset,
                                                                    mask.size());
      body_nodes_offset += mask.size();
 
      /* Prepare indices that are passed into each iteration. */
      component_info.index_values.reinitialize(mask.size());
      mask.foreach_index(
          [&](const int i, const int pos) { component_info.index_values[pos].set(i); });
 
      if (create_element_geometries) {
        component_info.element_geometries = this->try_extract_element_geometries(
            eval_storage.main_geometry, id, mask, attribute_filter);
      }
 
      /* Prepare remaining inputs that come from the field evaluation.*/
      component_info.item_input_values.reinitialize(node_storage.input_items.items_num);
      for (const int item_i : IndexRange(node_storage.input_items.items_num)) {
        const NodeForeachGeometryElementInputItem &item = node_storage.input_items.items[item_i];
        const eNodeSocketDatatype socket_type = eNodeSocketDatatype(item.socket_type);
        component_info.item_input_values[item_i].reinitialize(mask.size());
        const GVArray &values = component_info.field_evaluator->get_evaluated(item_i);
        mask.foreach_index(GrainSize(1024), [&](const int i, const int pos) {
          SocketValueVariant &value_variant = component_info.item_input_values[item_i][pos];
          void *buffer = value_variant.allocate_single(socket_type);
          values.get_to_uninitialized(i, buffer);
        });
      }
    }
 
    eval_storage.total_iterations_num = body_nodes_offset;
  }
 
  std::optional<Array<GeometrySet>> try_extract_element_geometries(
      const GeometrySet &main_geometry,
      const ForeachElementComponentID &id,
      const IndexMask &mask,
      const AttributeFilter &attribute_filter) const
  {
    switch (id.component_type) {
      case GeometryComponent::Type::Mesh: {
        const Mesh &main_mesh = *main_geometry.get_mesh();
        Array<Mesh *> meshes;
        switch (id.domain) {
          case AttrDomain::Point: {
            meshes = geometry::extract_mesh_vertices(main_mesh, mask, attribute_filter);
            break;
          }
          case AttrDomain::Edge: {
            meshes = geometry::extract_mesh_edges(main_mesh, mask, attribute_filter);
            break;
          }
          case AttrDomain::Face: {
            meshes = geometry::extract_mesh_faces(main_mesh, mask, attribute_filter);
            break;
          }
          default: {
            return std::nullopt;
          }
        }
        Array<GeometrySet> element_geometries(meshes.size());
        for (const int i : meshes.index_range()) {
          element_geometries[i].replace_mesh(meshes[i]);
        }
        return element_geometries;
      }
      case GeometryComponent::Type::PointCloud: {
        if (id.domain != AttrDomain::Point) {
          return std::nullopt;
        }
        const PointCloud &main_pointcloud = *main_geometry.get_pointcloud();
        Array<PointCloud *> pointclouds = geometry::extract_pointcloud_points(
            main_pointcloud, mask, attribute_filter);
        Array<GeometrySet> element_geometries(pointclouds.size());
        for (const int i : pointclouds.index_range()) {
          element_geometries[i].replace_pointcloud(pointclouds[i]);
        }
        return element_geometries;
      }
      case GeometryComponent::Type::Curve: {
        const Curves &main_curves = *main_geometry.get_curves();
        Array<Curves *> element_curves;
        switch (id.domain) {
          case AttrDomain::Point: {
            element_curves = geometry::extract_curves_points(main_curves, mask, attribute_filter);
            break;
          }
          case AttrDomain::Curve: {
            element_curves = geometry::extract_curves(main_curves, mask, attribute_filter);
            break;
          }
          default:
            return std::nullopt;
        }
        Array<GeometrySet> element_geometries(element_curves.size());
        for (const int i : element_curves.index_range()) {
          element_geometries[i].replace_curves(element_curves[i]);
        }
        return element_geometries;
      }
      case GeometryComponent::Type::Instance: {
        if (id.domain != AttrDomain::Instance) {
          return std::nullopt;
        }
        const bke::Instances &main_instances = *main_geometry.get_instances();
        Array<bke::Instances *> element_instances = geometry::extract_instances(
            main_instances, mask, attribute_filter);
        Array<GeometrySet> element_geometries(element_instances.size());
        for (const int i : element_instances.index_range()) {
          element_geometries[i].replace_instances(element_instances[i]);
        }
        return element_geometries;
      }
      case GeometryComponent::Type::GreasePencil: {
        const GreasePencil &main_grease_pencil = *main_geometry.get_grease_pencil();
        Array<GreasePencil *> element_grease_pencils;
        switch (id.domain) {
          case AttrDomain::Layer: {
            element_grease_pencils = geometry::extract_greasepencil_layers(
                main_grease_pencil, mask, attribute_filter);
            break;
          }
          case AttrDomain::Point: {
            element_grease_pencils = geometry::extract_greasepencil_layer_points(
                main_grease_pencil, *id.layer_index, mask, attribute_filter);
            break;
          }
          case AttrDomain::Curve: {
            element_grease_pencils = geometry::extract_greasepencil_layer_curves(
                main_grease_pencil, *id.layer_index, mask, attribute_filter);
            break;
          }
          default:
            return std::nullopt;
        }
        Array<GeometrySet> element_geometries(element_grease_pencils.size());
        for (const int i : element_geometries.index_range()) {
          element_geometries[i].replace_grease_pencil(element_grease_pencils[i]);
        }
        return element_geometries;
      }
      default:
        break;
    }
    return std::nullopt;
  }
 
  void build_graph_contents(ForeachGeometryElementEvalStorage &eval_storage,
                            const NodeGeometryForeachGeometryElementOutput &node_storage,
                            Span<lf::GraphInputSocket *> graph_inputs,
                            Span<lf::GraphOutputSocket *> graph_outputs) const
  {
    lf::Graph &lf_graph = eval_storage.graph;
 
    /* Create body nodes. */
    VectorSet<lf::FunctionNode *> &lf_body_nodes = eval_storage.lf_body_nodes;
    for ([[maybe_unused]] const int i : IndexRange(eval_storage.total_iterations_num)) {
      lf::FunctionNode &lf_node = lf_graph.add_function(*body_fn_.function);
      lf_body_nodes.add_new(&lf_node);
    }
 
    /* Link up output usages to body nodes. */
    for (const int zone_output_i : body_fn_.indices.inputs.output_usages.index_range()) {
      /* +1 because of geometry output. */
      lf::GraphInputSocket &lf_graph_input =
          *graph_inputs[zone_info_.indices.inputs.output_usages[1 + zone_output_i]];
      for (const int i : lf_body_nodes.index_range()) {
        lf::FunctionNode &lf_node = *lf_body_nodes[i];
        lf_graph.add_link(lf_graph_input,
                          lf_node.input(body_fn_.indices.inputs.output_usages[zone_output_i]));
      }
    }
 
    const bNodeSocket &element_geometry_bsocket = zone_.input_node->output_socket(1);
 
    static const GeometrySet empty_geometry;
    for (const ForeachElementComponent &component_info : eval_storage.components) {
      for (const int i : component_info.body_nodes_range.index_range()) {
        const int body_i = component_info.body_nodes_range[i];
        lf::FunctionNode &lf_body_node = *lf_body_nodes[body_i];
        /* Set index input for loop body. */
        lf_body_node.input(body_fn_.indices.inputs.main[0])
            .set_default_value(&component_info.index_values[i]);
        /* Set geometry element input for loop body. */
        if (element_geometry_bsocket.is_available()) {
          const GeometrySet *element_geometry = component_info.element_geometries.has_value() ?
                                                    &(*component_info.element_geometries)[i] :
                                                    &empty_geometry;
          lf_body_node.input(body_fn_.indices.inputs.main[1]).set_default_value(element_geometry);
        }
        /* Set main input values for loop body. */
        for (const int item_i : IndexRange(node_storage.input_items.items_num)) {
          lf_body_node.input(body_fn_.indices.inputs.main[indices_.inputs.lf_inner[item_i]])
              .set_default_value(&component_info.item_input_values[item_i][i]);
        }
        /* Link up border-link inputs to the loop body. */
        for (const int border_link_i : zone_info_.indices.inputs.border_links.index_range()) {
          lf_graph.add_link(
              *graph_inputs[zone_info_.indices.inputs.border_links[border_link_i]],
              lf_body_node.input(body_fn_.indices.inputs.border_links[border_link_i]));
        }
        /* Link up attribute propagation information. */
        for (const auto item : body_fn_.indices.inputs.attributes_by_field_source_index.items()) {
          lf_graph.add_link(
              *graph_inputs[zone_info_.indices.inputs.attributes_by_field_source_index.lookup(
                  item.key)],
              lf_body_node.input(item.value));
        }
        for (const auto item :
             body_fn_.indices.inputs.attributes_by_caller_propagation_index.items())
        {
          lf_graph.add_link(
              *graph_inputs[zone_info_.indices.inputs.attributes_by_caller_propagation_index
                                .lookup(item.key)],
              lf_body_node.input(item.value));
        }
      }
    }
 
    /* Add the reduce function that has all outputs from the zone bodies as input. */
    eval_storage.reduce_function.emplace(*this, eval_storage);
    lf::FunctionNode &lf_reduce = lf_graph.add_function(*eval_storage.reduce_function);
 
    /* Link up body outputs to reduce function. */
    const int body_main_outputs_num = node_storage.main_items.items_num +
                                      node_storage.generation_items.items_num;
    BLI_assert(body_main_outputs_num == body_fn_.indices.outputs.main.size());
    for (const int i : IndexRange(eval_storage.total_iterations_num)) {
      lf::FunctionNode &lf_body_node = *lf_body_nodes[i];
      for (const int item_i : IndexRange(node_storage.main_items.items_num)) {
        lf_graph.add_link(lf_body_node.output(body_fn_.indices.outputs.main[item_i]),
                          lf_reduce.input(i * body_main_outputs_num + item_i));
      }
      for (const int item_i : IndexRange(node_storage.generation_items.items_num)) {
        const int body_output_i = item_i + node_storage.main_items.items_num;
        lf_graph.add_link(lf_body_node.output(body_fn_.indices.outputs.main[body_output_i]),
                          lf_reduce.input(i * body_main_outputs_num + body_output_i));
      }
    }
 
    /* Link up reduce function outputs to final zone outputs. */
    lf_graph.add_link(lf_reduce.output(0), *graph_outputs[zone_info_.indices.outputs.main[0]]);
    for (const int item_i : IndexRange(node_storage.main_items.items_num)) {
      const int output_i = indices_.main.lf_outer[item_i];
      lf_graph.add_link(lf_reduce.output(output_i),
                        *graph_outputs[zone_info_.indices.outputs.main[output_i]]);
    }
    for (const int item_i : IndexRange(node_storage.generation_items.items_num)) {
      const int output_i = indices_.generation.lf_outer[item_i];
      lf_graph.add_link(lf_reduce.output(output_i),
                        *graph_outputs[zone_info_.indices.outputs.main[output_i]]);
    }
 
    /* All zone inputs are used for now. */
    static bool static_true{true};
    for (const int i : zone_info_.indices.outputs.input_usages) {
      graph_outputs[i]->set_default_value(&static_true);
    }
 
    /* Handle usage outputs for border-links. A border-link is used if it's used by any of the
     * iterations. */
    eval_storage.or_function.emplace(eval_storage.total_iterations_num);
    for (const int border_link_i : zone_.border_links.index_range()) {
      lf::FunctionNode &lf_or = lf_graph.add_function(*eval_storage.or_function);
      for (const int i : lf_body_nodes.index_range()) {
        lf::FunctionNode &lf_body_node = *lf_body_nodes[i];
        lf_graph.add_link(
            lf_body_node.output(body_fn_.indices.outputs.border_link_usages[border_link_i]),
            lf_or.input(i));
      }
      lf_graph.add_link(
          lf_or.output(0),
          *graph_outputs[zone_info_.indices.outputs.border_link_usages[border_link_i]]);
    }
  }
 
  std::string input_name(const int i) const override
  {
    return zone_wrapper_input_name(zone_info_, zone_, inputs_, i);
  }
 
  std::string output_name(const int i) const override
  {
    return zone_wrapper_output_name(zone_info_, zone_, outputs_, i);
  }
};
 
LazyFunctionForReduceForeachGeometryElement::LazyFunctionForReduceForeachGeometryElement(
    const LazyFunctionForForeachGeometryElementZone &parent,
    ForeachGeometryElementEvalStorage &eval_storage)
    : parent_(parent), eval_storage_(eval_storage)
{
  debug_name_ = "Reduce";
 
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent.output_bnode_.storage);
 
  inputs_.reserve(eval_storage.total_iterations_num *
                  (node_storage.main_items.items_num + node_storage.generation_items.items_num));
 
  for ([[maybe_unused]] const int i : eval_storage.lf_body_nodes.index_range()) {
    /* Add parameters for main items. */
    for (const int item_i : IndexRange(node_storage.main_items.items_num)) {
      const NodeForeachGeometryElementMainItem &item = node_storage.main_items.items[item_i];
      const bNodeSocket &socket = parent.output_bnode_.input_socket(
          parent_.indices_.main.bsocket_inner[item_i]);
      inputs_.append_as(
          item.name, *socket.typeinfo->geometry_nodes_cpp_type, lf::ValueUsage::Used);
    }
    /* Add parameters for generation items. */
    for (const int item_i : IndexRange(node_storage.generation_items.items_num)) {
      const NodeForeachGeometryElementGenerationItem &item =
          node_storage.generation_items.items[item_i];
      const bNodeSocket &socket = parent.output_bnode_.input_socket(
          parent_.indices_.generation.bsocket_inner[item_i]);
      inputs_.append_as(
          item.name, *socket.typeinfo->geometry_nodes_cpp_type, lf::ValueUsage::Maybe);
    }
  }
 
  /* Add output for main geometry. */
  outputs_.append_as("Geometry", CPPType::get<GeometrySet>());
  /* Add outputs for main items. */
  for (const int item_i : IndexRange(node_storage.main_items.items_num)) {
    const NodeForeachGeometryElementMainItem &item = node_storage.main_items.items[item_i];
    const bNodeSocket &socket = parent.output_bnode_.output_socket(
        parent_.indices_.main.bsocket_outer[item_i]);
    outputs_.append_as(item.name, *socket.typeinfo->geometry_nodes_cpp_type);
  }
  /* Add outputs for generation items. */
  for (const int item_i : IndexRange(node_storage.generation_items.items_num)) {
    const NodeForeachGeometryElementGenerationItem &item =
        node_storage.generation_items.items[item_i];
    const bNodeSocket &socket = parent.output_bnode_.output_socket(
        parent_.indices_.generation.bsocket_outer[item_i]);
    outputs_.append_as(item.name, *socket.typeinfo->geometry_nodes_cpp_type);
  }
}

static std::optional<AttrDomain> get_foreach_attribute_propagation_target_domain(
    const GeometryComponent::Type component_type)
{
  switch (component_type) {
    case GeometryComponent::Type::Mesh:
    case GeometryComponent::Type::PointCloud:
      return AttrDomain::Point;
    case GeometryComponent::Type::Curve:
      return AttrDomain::Curve;
    case GeometryComponent::Type::Instance:
      return AttrDomain::Instance;
    case GeometryComponent::Type::GreasePencil:
      return AttrDomain::Layer;
    default:
      break;
  }
  return std::nullopt;
}
 
void LazyFunctionForReduceForeachGeometryElement::execute_impl(lf::Params &params,
                                                               const lf::Context &context) const
{
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
 
  this->handle_main_items_and_geometry(params, context);
  if (node_storage.generation_items.items_num == 0) {
    return;
  }
  this->handle_generation_items(params, context);
}
 
void LazyFunctionForReduceForeachGeometryElement::handle_main_items_and_geometry(
    lf::Params &params, const lf::Context &context) const
{
  auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
  const int body_main_outputs_num = node_storage.main_items.items_num +
                                    node_storage.generation_items.items_num;
 
  const int main_geometry_output = 0;
  if (params.output_was_set(main_geometry_output)) {
    /* Done already. */
    return;
  }
 
  GeometrySet output_geometry = eval_storage_.main_geometry;
 
  for (const int item_i : IndexRange(node_storage.main_items.items_num)) {
    const NodeForeachGeometryElementMainItem &item = node_storage.main_items.items[item_i];
    const eNodeSocketDatatype socket_type = eNodeSocketDatatype(item.socket_type);
    const CPPType *base_cpp_type = bke::socket_type_to_geo_nodes_base_cpp_type(socket_type);
    if (!base_cpp_type) {
      continue;
    }
    const eCustomDataType cd_type = bke::cpp_type_to_custom_data_type(*base_cpp_type);
 
    /* Compute output attribute name for this item. */
    const std::string attribute_name = bke::hash_to_anonymous_attribute_name(
        user_data.call_data->self_object()->id.name,
        user_data.compute_context->hash(),
        parent_.output_bnode_.identifier,
        item.identifier);
 
    /* Create a new output attribute for the current item on each iteration component. */
    for (const ForeachElementComponent &component_info : eval_storage_.components) {
      MutableAttributeAccessor attributes = component_info.attributes_for_write(output_geometry);
      const int domain_size = attributes.domain_size(component_info.id.domain);
      const IndexMask mask = component_info.field_evaluator->get_evaluated_selection_as_mask();
 
      /* Actually create the attribute. */
      GSpanAttributeWriter attribute = attributes.lookup_or_add_for_write_only_span(
          attribute_name, component_info.id.domain, cd_type);
 
      /* Fill the elements of the attribute that we didn't iterate over because they were not
       * selected. */
      IndexMaskMemory memory;
      const IndexMask inverted_mask = mask.complement(IndexRange(domain_size), memory);
      base_cpp_type->value_initialize_indices(attribute.span.data(), inverted_mask);
 
      /* Copy the values from each iteration into the attribute. */
      mask.foreach_index([&](const int i, const int pos) {
        const int lf_param_index = pos * body_main_outputs_num + item_i;
        SocketValueVariant &value_variant = params.get_input<SocketValueVariant>(lf_param_index);
        value_variant.convert_to_single();
        const void *value = value_variant.get_single_ptr_raw();
        base_cpp_type->copy_construct(value, attribute.span[i]);
      });
 
      attribute.finish();
    }
 
    /* Output the field for the anonymous attribute. */
    auto attribute_field = std::make_shared<AttributeFieldInput>(
        attribute_name,
        *base_cpp_type,
        make_anonymous_attribute_socket_inspection_string(
            parent_.output_bnode_.output_socket(parent_.indices_.main.bsocket_outer[item_i])));
    SocketValueVariant attribute_value_variant{GField(std::move(attribute_field))};
    params.set_output(1 + item_i, std::move(attribute_value_variant));
  }
 
  /* Output the original geometry with potentially additional attributes. */
  params.set_output(main_geometry_output, std::move(output_geometry));
}
 
void LazyFunctionForReduceForeachGeometryElement::handle_generation_items(
    lf::Params &params, const lf::Context &context) const
{
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
 
  const int first_valid_item_i = this->handle_invalid_generation_items(params);
  if (first_valid_item_i == node_storage.generation_items.items_num) {
    return;
  }
  this->handle_generation_item_groups(params, context, first_valid_item_i);
}
 
int LazyFunctionForReduceForeachGeometryElement::handle_invalid_generation_items(
    lf::Params &params) const
{
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
 
  int item_i = 0;
  /* Handle invalid generation items that come before a geometry. */
  for (; item_i < node_storage.generation_items.items_num; item_i++) {
    const NodeForeachGeometryElementGenerationItem &item =
        node_storage.generation_items.items[item_i];
    const eNodeSocketDatatype socket_type = eNodeSocketDatatype(item.socket_type);
    if (socket_type == SOCK_GEOMETRY) {
      break;
    }
    const int lf_socket_i = parent_.indices_.generation.lf_outer[item_i];
    if (!params.output_was_set(lf_socket_i)) {
      const int bsocket_i = parent_.indices_.generation.bsocket_outer[item_i];
      set_default_value_for_output_socket(
          params, lf_socket_i, parent_.zone_.output_node->output_socket(bsocket_i));
    }
  }
  return item_i;
}
 
void LazyFunctionForReduceForeachGeometryElement::handle_generation_item_groups(
    lf::Params &params, const lf::Context &context, const int first_valid_item_i) const
{
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
  int previous_geometry_item_i = first_valid_item_i;
  /* Iterate over all groups. A group starts with a geometry socket followed by an arbitrary number
   * of non-geometry sockets. */
  for (const int item_i :
       IndexRange::from_begin_end(first_valid_item_i + 1, node_storage.generation_items.items_num))
  {
    const NodeForeachGeometryElementGenerationItem &item =
        node_storage.generation_items.items[item_i];
    const eNodeSocketDatatype socket_type = eNodeSocketDatatype(item.socket_type);
    if (socket_type == SOCK_GEOMETRY) {
      this->handle_generation_items_group(
          params,
          context,
          previous_geometry_item_i,
          IndexRange::from_begin_end(previous_geometry_item_i + 1, item_i));
      previous_geometry_item_i = item_i;
    }
  }
  this->handle_generation_items_group(
      params,
      context,
      previous_geometry_item_i,
      IndexRange::from_begin_end(previous_geometry_item_i + 1,
                                 node_storage.generation_items.items_num));
}
 
void LazyFunctionForReduceForeachGeometryElement::handle_generation_items_group(
    lf::Params &params,
    const lf::Context &context,
    const int geometry_item_i,
    const IndexRange generation_items_range) const
{
  auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
  const int body_main_outputs_num = node_storage.main_items.items_num +
                                    node_storage.generation_items.items_num;
 
  /* Handle the case when the output is not needed or the inputs have not been computed yet. */
  if (!this->handle_generation_items_group_lazyness(
          params, context, geometry_item_i, generation_items_range))
  {
    return;
  }
 
  /* TODO: Get propagation info from input, but that's not necessary for correctness for now. */
  AttributeFilter attribute_filter;
 
  const int bodies_num = eval_storage_.lf_body_nodes.size();
  Array<GeometrySet> geometries(bodies_num + 1);
 
  /* Create attribute names for the outputs. */
  Array<std::string> attribute_names(generation_items_range.size());
  for (const int i : generation_items_range.index_range()) {
    const int item_i = generation_items_range[i];
    const NodeForeachGeometryElementGenerationItem &item =
        node_storage.generation_items.items[item_i];
    attribute_names[i] = bke::hash_to_anonymous_attribute_name(
        user_data.call_data->self_object()->id.name,
        user_data.compute_context->hash(),
        parent_.output_bnode_.identifier,
        item.identifier);
  }
 
  for (const ForeachElementComponent &component_info : eval_storage_.components) {
    const AttributeAccessor src_attributes = component_info.input_attributes();
 
    /* These are the attributes we need to propagate from the original input geometry. */
    struct NameWithType {
      StringRef name;
      eCustomDataType type;
    };
    Vector<NameWithType> attributes_to_propagate;
    src_attributes.foreach_attribute([&](const bke::AttributeIter &iter) {
      if (iter.data_type == CD_PROP_STRING) {
        return;
      }
      if (attribute_filter.allow_skip(iter.name)) {
        return;
      }
      attributes_to_propagate.append({iter.name, iter.data_type});
    });
    Map<StringRef, GVArray> cached_adapted_src_attributes;
 
    const IndexMask mask = component_info.field_evaluator->get_evaluated_selection_as_mask();
 
    /* Add attributes for each field on the geometry created by each iteration. */
    mask.foreach_index([&](const int element_i, const int local_body_i) {
      const int body_i = component_info.body_nodes_range[local_body_i];
      const int geometry_param_i = body_i * body_main_outputs_num +
                                   parent_.indices_.generation.lf_inner[geometry_item_i];
      GeometrySet &geometry = geometries[body_i];
      geometry = params.extract_input<GeometrySet>(geometry_param_i);
 
      for (const GeometryComponent::Type dst_component_type :
           {GeometryComponent::Type::Mesh,
            GeometryComponent::Type::PointCloud,
            GeometryComponent::Type::Curve,
            GeometryComponent::Type::GreasePencil,
            GeometryComponent::Type::Instance})
      {
        if (!geometry.has(dst_component_type)) {
          continue;
        }
        GeometryComponent &dst_component = geometry.get_component_for_write(dst_component_type);
        MutableAttributeAccessor dst_attributes = *dst_component.attributes_for_write();
 
        /* Determine the domain that we propagate the input attribute to. Technically, this is only
         * a single value for the entire geometry, but we can't optimize for that yet. */
        const std::optional<AttrDomain> propagation_domain =
            get_foreach_attribute_propagation_target_domain(dst_component_type);
        if (!propagation_domain) {
          continue;
        }
 
        /* Propagate attributes from the input geometry. */
        for (const NameWithType &name_with_type : attributes_to_propagate) {
          const StringRef name = name_with_type.name;
          const eCustomDataType cd_type = name_with_type.type;
          if (src_attributes.is_builtin(name) && !dst_attributes.is_builtin(name)) {
            continue;
          }
          if (dst_attributes.contains(name)) {
            /* Attributes created in the zone shouldn't be overridden. */
            continue;
          }
          /* Get the source attribute adapted to the iteration domain. */
          const GVArray &src_attribute = cached_adapted_src_attributes.lookup_or_add_cb(
              name, [&]() {
                GAttributeReader attribute = src_attributes.lookup(name);
                return src_attributes.adapt_domain(
                    *attribute, attribute.domain, component_info.id.domain);
              });
          if (!src_attribute) {
            continue;
          }
          const CPPType &type = src_attribute.type();
          BUFFER_FOR_CPP_TYPE_VALUE(type, element_value);
          src_attribute.get_to_uninitialized(element_i, element_value);
 
          /* Actually create the attribute. */
          GSpanAttributeWriter dst_attribute = dst_attributes.lookup_or_add_for_write_only_span(
              name, *propagation_domain, cd_type);
          type.fill_assign_n(element_value, dst_attribute.span.data(), dst_attribute.span.size());
          dst_attribute.finish();
 
          type.destruct(element_value);
        }
      }
 
      /* Create an attribute for each field that corresponds to the current geometry. */
      for (const int local_item_i : generation_items_range.index_range()) {
        const int item_i = generation_items_range[local_item_i];
        const NodeForeachGeometryElementGenerationItem &item =
            node_storage.generation_items.items[item_i];
        const AttrDomain capture_domain = AttrDomain(item.domain);
        const int field_param_i = body_i * body_main_outputs_num +
                                  parent_.indices_.generation.lf_inner[item_i];
        GField field = params.get_input<SocketValueVariant>(field_param_i).get<GField>();
 
        if (capture_domain == AttrDomain::Instance) {
          if (geometry.has_instances()) {
            bke::try_capture_field_on_geometry(
                geometry.get_component_for_write(GeometryComponent::Type::Instance),
                attribute_names[local_item_i],
                capture_domain,
                field);
          }
        }
        else {
          geometry.modify_geometry_sets([&](GeometrySet &sub_geometry) {
            for (const GeometryComponent::Type component_type :
                 {GeometryComponent::Type::Mesh,
                  GeometryComponent::Type::PointCloud,
                  GeometryComponent::Type::Curve,
                  GeometryComponent::Type::GreasePencil})
            {
              if (sub_geometry.has(component_type)) {
                bke::try_capture_field_on_geometry(
                    sub_geometry.get_component_for_write(component_type),
                    attribute_names[local_item_i],
                    capture_domain,
                    field);
              }
            }
          });
        }
      }
    });
  }
 
  /* The last geometry contains the edit data from the main geometry. */
  GeometrySet &edit_data_geometry = geometries.last();
  edit_data_geometry = eval_storage_.main_geometry;
  edit_data_geometry.keep_only({GeometryComponent::Type::Edit});
 
  /* Join the geometries from all iterations into a single one. */
  GeometrySet joined_geometry = geometry::join_geometries(geometries, attribute_filter);
 
  /* Output the joined geometry. */
  params.set_output(parent_.indices_.generation.lf_outer[geometry_item_i],
                    std::move(joined_geometry));
 
  /* Output the anonymous attribute fields. */
  for (const int local_item_i : generation_items_range.index_range()) {
    const int item_i = generation_items_range[local_item_i];
    const NodeForeachGeometryElementGenerationItem &item =
        node_storage.generation_items.items[item_i];
    const eNodeSocketDatatype socket_type = eNodeSocketDatatype(item.socket_type);
    const CPPType &base_cpp_type = *bke::socket_type_to_geo_nodes_base_cpp_type(socket_type);
    const StringRef attribute_name = attribute_names[local_item_i];
    auto attribute_field = std::make_shared<AttributeFieldInput>(
        attribute_name,
        base_cpp_type,
        make_anonymous_attribute_socket_inspection_string(
            parent_.output_bnode_.output_socket(2 + node_storage.main_items.items_num + item_i)));
    SocketValueVariant attribute_value_variant{GField(std::move(attribute_field))};
    params.set_output(parent_.indices_.generation.lf_outer[item_i],
                      std::move(attribute_value_variant));
  }
}
 
bool LazyFunctionForReduceForeachGeometryElement::handle_generation_items_group_lazyness(
    lf::Params &params,
    const lf::Context & /*context*/,
    const int geometry_item_i,
    const IndexRange generation_items_range) const
{
  const auto &node_storage = *static_cast<NodeGeometryForeachGeometryElementOutput *>(
      parent_.output_bnode_.storage);
  const int body_main_outputs_num = node_storage.main_items.items_num +
                                    node_storage.generation_items.items_num;
 
  const int geometry_output_param = parent_.indices_.generation.lf_outer[geometry_item_i];
 
  if (params.output_was_set(geometry_output_param)) {
    /* Done already. */
    return false;
  }
  const lf::ValueUsage geometry_output_usage = params.get_output_usage(geometry_output_param);
  if (geometry_output_usage == lf::ValueUsage::Unused) {
    /* Output dummy values. */
    const int start_bsocket_i = parent_.indices_.generation.bsocket_outer[geometry_item_i];
    for (const int i : IndexRange(1 + generation_items_range.size())) {
      const bNodeSocket &bsocket = parent_.output_bnode_.output_socket(start_bsocket_i + i);
      set_default_value_for_output_socket(params, geometry_output_param + i, bsocket);
    }
    return false;
  }
  bool any_output_used = false;
  for (const int i : IndexRange(1 + generation_items_range.size())) {
    const lf::ValueUsage usage = params.get_output_usage(geometry_output_param + i);
    if (usage == lf::ValueUsage::Used) {
      any_output_used = true;
      break;
    }
  }
  if (!any_output_used) {
    /* Only execute below if we are sure that the output is actually needed. */
    return false;
  }
  const int bodies_num = eval_storage_.lf_body_nodes.size();
 
  /* Check if all inputs are available, and request them if not. */
  bool has_missing_input = false;
  for (const int body_i : IndexRange(bodies_num)) {
    const int offset = body_i * body_main_outputs_num +
                       parent_.indices_.generation.lf_inner[geometry_item_i];
    for (const int i : IndexRange(1 + generation_items_range.size())) {
      const bool is_available = params.try_get_input_data_ptr_or_request(offset + i) != nullptr;
      if (!is_available) {
        has_missing_input = true;
      }
    }
  }
  if (has_missing_input) {
    /* Come back when all inputs are available. */
    return false;
  }
  return true;
}

class GeometryNodesLazyFunctionLogger : public lf::GraphExecutor::Logger {
 private:
  const GeometryNodesLazyFunctionGraphInfo &lf_graph_info_;
 
 public:
  GeometryNodesLazyFunctionLogger(const GeometryNodesLazyFunctionGraphInfo &lf_graph_info)
      : lf_graph_info_(lf_graph_info)
  {
  }
 
  void log_socket_value(const lf::Socket &lf_socket,
                        const GPointer value,
                        const lf::Context &context) const override
  {
    auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    if (!user_data.log_socket_values) {
      return;
    }
    auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
    geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(user_data);
    if (tree_logger == nullptr) {
      return;
    }
 
    const Span<const bNodeSocket *> bsockets =
        lf_graph_info_.mapping.bsockets_by_lf_socket_map.lookup(&lf_socket);
    if (bsockets.is_empty()) {
      return;
    }
 
    for (const bNodeSocket *bsocket : bsockets) {
      /* Avoid logging to some sockets when the same value will also be logged to a linked socket.
       * This reduces the number of logged values without losing information. */
      if (bsocket->is_input() && bsocket->is_directly_linked()) {
        continue;
      }
      const bNode &bnode = bsocket->owner_node();
      if (bnode.is_reroute()) {
        continue;
      }
      tree_logger->log_value(bsocket->owner_node(), *bsocket, value);
    }
  }
 
  static inline std::mutex dump_error_context_mutex;
 
  void dump_when_outputs_are_missing(const lf::FunctionNode &node,
                                     Span<const lf::OutputSocket *> missing_sockets,
                                     const lf::Context &context) const override
  {
    std::lock_guard lock{dump_error_context_mutex};
 
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
    user_data->compute_context->print_stack(std::cout, node.name());
    std::cout << "Missing outputs:\n";
    for (const lf::OutputSocket *socket : missing_sockets) {
      std::cout << "  " << socket->name() << "\n";
    }
  }
 
  void dump_when_input_is_set_twice(const lf::InputSocket &target_socket,
                                    const lf::OutputSocket &from_socket,
                                    const lf::Context &context) const override
  {
    std::lock_guard lock{dump_error_context_mutex};
 
    std::stringstream ss;
    ss << from_socket.node().name() << ":" << from_socket.name() << " -> "
       << target_socket.node().name() << ":" << target_socket.name();
 
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
    user_data->compute_context->print_stack(std::cout, ss.str());
  }
 
  void log_before_node_execute(const lf::FunctionNode &node,
                               const lf::Params & /*params*/,
                               const lf::Context &context) const override
  {
    /* Enable this to see the threads that invoked a node. */
    if constexpr (false) {
      this->add_thread_id_debug_message(node, context);
    }
  }
 
  void add_thread_id_debug_message(const lf::FunctionNode &node, const lf::Context &context) const
  {
    static std::atomic<int> thread_id_source = 0;
    static thread_local const int thread_id = thread_id_source.fetch_add(1);
    static thread_local const std::string thread_id_str = "Thread: " + std::to_string(thread_id);
 
    const auto &user_data = *static_cast<GeoNodesLFUserData *>(context.user_data);
    const auto &local_user_data = *static_cast<GeoNodesLFLocalUserData *>(context.local_user_data);
    geo_eval_log::GeoTreeLogger *tree_logger = local_user_data.try_get_tree_logger(user_data);
    if (tree_logger == nullptr) {
      return;
    }
 
    /* Find corresponding node based on the socket mapping. */
    auto check_sockets = [&](const Span<const lf::Socket *> lf_sockets) {
      for (const lf::Socket *lf_socket : lf_sockets) {
        const Span<const bNodeSocket *> bsockets =
            lf_graph_info_.mapping.bsockets_by_lf_socket_map.lookup(lf_socket);
        if (!bsockets.is_empty()) {
          const bNodeSocket &bsocket = *bsockets[0];
          const bNode &bnode = bsocket.owner_node();
          tree_logger->debug_messages.append(*tree_logger->allocator,
                                             {bnode.identifier, thread_id_str});
          return true;
        }
      }
      return false;
    };
 
    if (check_sockets(node.inputs().cast<const lf::Socket *>())) {
      return;
    }
    check_sockets(node.outputs().cast<const lf::Socket *>());
  }
};

class GeometryNodesLazyFunctionSideEffectProvider : public lf::GraphExecutor::SideEffectProvider {
 private:
  Span<const lf::FunctionNode *> local_side_effect_nodes_;
 
 public:
  GeometryNodesLazyFunctionSideEffectProvider(
      Span<const lf::FunctionNode *> local_side_effect_nodes = {})
      : local_side_effect_nodes_(local_side_effect_nodes)
  {
  }
 
  Vector<const lf::FunctionNode *> get_nodes_with_side_effects(
      const lf::Context &context) const override
  {
    GeoNodesLFUserData *user_data = dynamic_cast<GeoNodesLFUserData *>(context.user_data);
    BLI_assert(user_data != nullptr);
    const GeoNodesCallData &call_data = *user_data->call_data;
    if (!call_data.side_effect_nodes) {
      return {};
    }
    const ComputeContextHash &context_hash = user_data->compute_context->hash();
    Vector<const lf::FunctionNode *> side_effect_nodes =
        call_data.side_effect_nodes->nodes_by_context.lookup(context_hash);
    side_effect_nodes.extend(local_side_effect_nodes_);
    return side_effect_nodes;
  }
};

struct GeometryNodesLazyFunctionBuilder {
 private:
  const bNodeTree &btree_;
  const aai::AnonymousAttributeInferencingResult &attribute_inferencing_;
  ResourceScope &scope_;
  NodeMultiFunctions &node_multi_functions_;
  GeometryNodesLazyFunctionGraphInfo *lf_graph_info_;
  GeometryNodeLazyFunctionGraphMapping *mapping_;
  const bke::DataTypeConversions *conversions_;

  Map<const bNode *, lf::Node *> simulation_inputs_usage_nodes_;
 
  const bNodeTreeZones *tree_zones_;
  MutableSpan<ZoneBuildInfo> zone_build_infos_;
 
  std::optional<BuildGraphParams> root_graph_build_params_;

  Vector<const lf::GraphInputSocket *> group_input_sockets_;
  Vector<const lf::GraphOutputSocket *> standard_group_output_sockets_;
  Vector<lf::GraphInputSocket *> group_output_used_sockets_;
  Vector<const lf::GraphOutputSocket *> group_input_usage_sockets_;
  Map<int, const lf::GraphInputSocket *> attribute_set_by_geometry_output_;
 
  friend class UsedSocketVisualizeOptions;
 
 public:
  GeometryNodesLazyFunctionBuilder(const bNodeTree &btree,
                                   GeometryNodesLazyFunctionGraphInfo &lf_graph_info)
      : btree_(btree),
        attribute_inferencing_(*btree.runtime->anonymous_attribute_inferencing),
        scope_(lf_graph_info.scope),
        node_multi_functions_(lf_graph_info.scope.construct<NodeMultiFunctions>(btree)),
        lf_graph_info_(&lf_graph_info)
  {
  }
 
  void build()
  {
    btree_.ensure_topology_cache();
    btree_.ensure_interface_cache();
 
    mapping_ = &lf_graph_info_->mapping;
    conversions_ = &bke::get_implicit_type_conversions();
    tree_zones_ = btree_.zones();
 
    this->initialize_mapping_arrays();
    this->build_zone_functions();
    this->build_root_graph();
    this->build_geometry_nodes_group_function();
  }
 
 private:
  void initialize_mapping_arrays()
  {
    mapping_->lf_input_index_for_output_bsocket_usage.reinitialize(
        btree_.all_output_sockets().size());
    mapping_->lf_input_index_for_output_bsocket_usage.fill(-1);
    mapping_->lf_input_index_for_attribute_propagation_to_output.reinitialize(
        btree_.all_output_sockets().size());
    mapping_->lf_input_index_for_attribute_propagation_to_output.fill(-1);
    mapping_->lf_index_by_bsocket.reinitialize(btree_.all_sockets().size());
    mapping_->lf_index_by_bsocket.fill(-1);
  }

  void build_zone_functions()
  {
    zone_build_infos_ = scope_.construct<Array<ZoneBuildInfo>>(tree_zones_->zones.size());
 
    const Array<int> zone_build_order = this->compute_zone_build_order();
 
    for (const int zone_i : zone_build_order) {
      const bNodeTreeZone &zone = *tree_zones_->zones[zone_i];
      switch (zone.output_node->type) {
        case GEO_NODE_SIMULATION_OUTPUT: {
          this->build_simulation_zone_function(zone);
          break;
        }
        case GEO_NODE_REPEAT_OUTPUT: {
          this->build_repeat_zone_function(zone);
          break;
        }
        case GEO_NODE_FOREACH_GEOMETRY_ELEMENT_OUTPUT:
          this->build_foreach_geometry_element_zone_function(zone);
          break;
        default: {
          BLI_assert_unreachable();
          break;
        }
      }
    }
  }
 
  Array<int> compute_zone_build_order()
  {
    /* Build nested zones first. */
    Array<int> zone_build_order(tree_zones_->zones.size());
    array_utils::fill_index_range<int>(zone_build_order);
    std::sort(
        zone_build_order.begin(), zone_build_order.end(), [&](const int zone_a, const int zone_b) {
          return tree_zones_->zones[zone_a]->depth > tree_zones_->zones[zone_b]->depth;
        });
    return zone_build_order;
  }

  void build_simulation_zone_function(const bNodeTreeZone &zone)
  {
    const int zone_i = zone.index;
    ZoneBuildInfo &zone_info = zone_build_infos_[zone_i];
    lf::Graph &lf_graph = scope_.construct<lf::Graph>();
    const auto &sim_output_storage = *static_cast<const NodeGeometrySimulationOutput *>(
        zone.output_node->storage);
 
    Vector<lf::GraphInputSocket *> lf_zone_inputs;
    Vector<lf::GraphOutputSocket *> lf_zone_outputs;
 
    if (zone.input_node != nullptr) {
      for (const bNodeSocket *bsocket : zone.input_node->input_sockets().drop_back(1)) {
        zone_info.indices.inputs.main.append(lf_zone_inputs.append_and_get_index(
            &lf_graph.add_input(*bsocket->typeinfo->geometry_nodes_cpp_type, bsocket->name)));
        zone_info.indices.outputs.input_usages.append(lf_zone_outputs.append_and_get_index(
            &lf_graph.add_output(CPPType::get<bool>(), "Usage: " + StringRef(bsocket->name))));
      }
    }
 
    this->build_zone_border_links_inputs(
        zone, lf_graph, lf_zone_inputs, zone_info.indices.inputs.border_links);
    this->build_zone_border_link_input_usages(
        zone, lf_graph, lf_zone_outputs, zone_info.indices.outputs.border_link_usages);
 
    for (const bNodeSocket *bsocket : zone.output_node->output_sockets().drop_back(1)) {
      zone_info.indices.outputs.main.append(lf_zone_outputs.append_and_get_index(
          &lf_graph.add_output(*bsocket->typeinfo->geometry_nodes_cpp_type, bsocket->name)));
      zone_info.indices.inputs.output_usages.append(lf_zone_inputs.append_and_get_index(
          &lf_graph.add_input(CPPType::get<bool>(), "Usage: " + StringRef(bsocket->name))));
    }
 
    lf::Node &lf_simulation_usage_node = [&]() -> lf::Node & {
      auto &lazy_function = scope_.construct<LazyFunctionForSimulationInputsUsage>(
          *zone.output_node);
      lf::Node &lf_node = lf_graph.add_function(lazy_function);
 
      for (const int i : zone_info.indices.outputs.input_usages) {
        lf_graph.add_link(lf_node.output(0), *lf_zone_outputs[i]);
      }
 
      return lf_node;
    }();
 
    BuildGraphParams graph_params{lf_graph};
 
    lf::FunctionNode *lf_simulation_input = nullptr;
    if (zone.input_node) {
      lf_simulation_input = this->insert_simulation_input_node(
          btree_, *zone.input_node, graph_params);
    }
    lf::FunctionNode &lf_simulation_output = this->insert_simulation_output_node(*zone.output_node,
                                                                                 graph_params);
 
    for (const bNodeSocket *bsocket : zone.output_node->input_sockets().drop_back(1)) {
      graph_params.usage_by_bsocket.add(bsocket, &lf_simulation_usage_node.output(1));
    }
 
    /* Link simulation input node directly to simulation output node for skip behavior. */
    for (const int i : IndexRange(sim_output_storage.items_num)) {
      lf::InputSocket &lf_to = lf_simulation_output.input(i + 1);
      if (lf_simulation_input) {
        lf::OutputSocket &lf_from = lf_simulation_input->output(i + 1);
        lf_graph.add_link(lf_from, lf_to);
      }
      else {
        lf_to.set_default_value(lf_to.type().default_value());
      }
    }
 
    this->insert_nodes_and_zones(zone.child_nodes, zone.child_zones, graph_params);
 
    if (zone.input_node) {
      this->build_output_socket_usages(*zone.input_node, graph_params);
    }
    for (const auto item : graph_params.lf_output_by_bsocket.items()) {
      this->insert_links_from_socket(*item.key, *item.value, graph_params);
    }
 
    this->link_border_link_inputs_and_usages(zone,
                                             lf_zone_inputs,
                                             zone_info.indices.inputs.border_links,
                                             lf_zone_outputs,
                                             zone_info.indices.outputs.border_link_usages,
                                             graph_params);
 
    for (const int i : zone_info.indices.inputs.main) {
      lf_graph.add_link(*lf_zone_inputs[i], lf_simulation_input->input(i));
    }
 
    for (const int i : zone_info.indices.outputs.main.index_range()) {
      lf_graph.add_link(lf_simulation_output.output(i),
                        *lf_zone_outputs[zone_info.indices.outputs.main[i]]);
    }
 
    this->add_default_inputs(graph_params);
 
    Map<int, lf::OutputSocket *> lf_attribute_set_by_field_source_index;
    Map<int, lf::OutputSocket *> lf_attribute_set_by_caller_propagation_index;
    this->build_attribute_set_inputs_for_zone(graph_params,
                                              lf_attribute_set_by_field_source_index,
                                              lf_attribute_set_by_caller_propagation_index);
    for (const auto item : lf_attribute_set_by_field_source_index.items()) {
      lf::OutputSocket &lf_attribute_set_socket = *item.value;
      if (lf_attribute_set_socket.node().is_interface()) {
        zone_info.indices.inputs.attributes_by_field_source_index.add_new(
            item.key, lf_zone_inputs.append_and_get_index(&lf_attribute_set_socket));
      }
    }
    for (const auto item : lf_attribute_set_by_caller_propagation_index.items()) {
      lf::OutputSocket &lf_attribute_set_socket = *item.value;
      if (lf_attribute_set_socket.node().is_interface()) {
        zone_info.indices.inputs.attributes_by_caller_propagation_index.add_new(
            item.key, lf_zone_inputs.append_and_get_index(&lf_attribute_set_socket));
      }
    }
    this->link_attribute_set_inputs(lf_graph,
                                    graph_params,
                                    lf_attribute_set_by_field_source_index,
                                    lf_attribute_set_by_caller_propagation_index);
    this->fix_link_cycles(lf_graph, graph_params.socket_usage_inputs);
 
    lf_graph.update_node_indices();
 
    auto &logger = scope_.construct<GeometryNodesLazyFunctionLogger>(*lf_graph_info_);
    auto &side_effect_provider = scope_.construct<GeometryNodesLazyFunctionSideEffectProvider>();
 
    const auto &lf_graph_fn = scope_.construct<lf::GraphExecutor>(lf_graph,
                                                                  lf_zone_inputs.as_span(),
                                                                  lf_zone_outputs.as_span(),
                                                                  &logger,
                                                                  &side_effect_provider,
                                                                  nullptr);
    const auto &zone_function = scope_.construct<LazyFunctionForSimulationZone>(*zone.output_node,
                                                                                lf_graph_fn);
    zone_info.lazy_function = &zone_function;
 
    lf_graph_info_->debug_zone_body_graphs.add(zone.output_node->identifier, &lf_graph);
    // std::cout << "\n\n" << lf_graph.to_dot() << "\n\n";
  }

  void build_repeat_zone_function(const bNodeTreeZone &zone)
  {
    ZoneBuildInfo &zone_info = zone_build_infos_[zone.index];
    /* Build a function for the loop body. */
    ZoneBodyFunction &body_fn = this->build_zone_body_function(zone, "Repeat Body");
    /* Wrap the loop body by another function that implements the repeat behavior. */
    auto &zone_fn = scope_.construct<LazyFunctionForRepeatZone>(btree_, zone, zone_info, body_fn);
    zone_info.lazy_function = &zone_fn;
  }
 
  void build_foreach_geometry_element_zone_function(const bNodeTreeZone &zone)
  {
    ZoneBuildInfo &zone_info = zone_build_infos_[zone.index];
    /* Build a function for the loop body. */
    ZoneBodyFunction &body_fn = this->build_zone_body_function(zone, "Foreach Body");
    /* Wrap the loop body in another function that implements the foreach behavior. */
    auto &zone_fn = scope_.construct<LazyFunctionForForeachGeometryElementZone>(
        btree_, zone, zone_info, body_fn);
    zone_info.lazy_function = &zone_fn;
  }

  ZoneBodyFunction &build_zone_body_function(const bNodeTreeZone &zone, const StringRef name)
  {
    lf::Graph &lf_body_graph = scope_.construct<lf::Graph>(name);
 
    BuildGraphParams graph_params{lf_body_graph};
 
    Vector<lf::GraphInputSocket *> lf_body_inputs;
    Vector<lf::GraphOutputSocket *> lf_body_outputs;
    ZoneBodyFunction &body_fn = scope_.construct<ZoneBodyFunction>();
 
    for (const bNodeSocket *bsocket : zone.input_node->output_sockets()) {
      if (ignore_zone_bsocket(*bsocket)) {
        continue;
      }
      lf::GraphInputSocket &lf_input = lf_body_graph.add_input(
          *bsocket->typeinfo->geometry_nodes_cpp_type, bsocket->name);
      lf::GraphOutputSocket &lf_input_usage = lf_body_graph.add_output(
          CPPType::get<bool>(), "Usage: " + StringRef(bsocket->name));
      body_fn.indices.inputs.main.append(lf_body_inputs.append_and_get_index(&lf_input));
      body_fn.indices.outputs.input_usages.append(
          lf_body_outputs.append_and_get_index(&lf_input_usage));
      graph_params.lf_output_by_bsocket.add_new(bsocket, &lf_input);
    }
 
    this->build_zone_border_links_inputs(
        zone, lf_body_graph, lf_body_inputs, body_fn.indices.inputs.border_links);
    this->build_zone_border_link_input_usages(
        zone, lf_body_graph, lf_body_outputs, body_fn.indices.outputs.border_link_usages);
 
    for (const bNodeSocket *bsocket : zone.output_node->input_sockets()) {
      if (ignore_zone_bsocket(*bsocket)) {
        continue;
      }
      lf::GraphOutputSocket &lf_output = lf_body_graph.add_output(
          *bsocket->typeinfo->geometry_nodes_cpp_type, bsocket->name);
      lf::GraphInputSocket &lf_output_usage = lf_body_graph.add_input(
          CPPType::get<bool>(), "Usage: " + StringRef(bsocket->name));
      graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_output);
      graph_params.usage_by_bsocket.add(bsocket, &lf_output_usage);
      body_fn.indices.outputs.main.append(lf_body_outputs.append_and_get_index(&lf_output));
      body_fn.indices.inputs.output_usages.append(
          lf_body_inputs.append_and_get_index(&lf_output_usage));
    }
 
    this->insert_nodes_and_zones(zone.child_nodes, zone.child_zones, graph_params);
 
    this->build_output_socket_usages(*zone.input_node, graph_params);
 
    {
      int valid_socket_i = 0;
      for (const bNodeSocket *bsocket : zone.input_node->output_sockets()) {
        if (ignore_zone_bsocket(*bsocket)) {
          continue;
        }
        lf::OutputSocket *lf_usage = graph_params.usage_by_bsocket.lookup_default(bsocket,
                                                                                  nullptr);
        lf::GraphOutputSocket &lf_usage_output =
            *lf_body_outputs[body_fn.indices.outputs.input_usages[valid_socket_i]];
        if (lf_usage) {
          lf_body_graph.add_link(*lf_usage, lf_usage_output);
        }
        else {
          static const bool static_false = false;
          lf_usage_output.set_default_value(&static_false);
        }
        valid_socket_i++;
      }
    }
 
    for (const auto item : graph_params.lf_output_by_bsocket.items()) {
      this->insert_links_from_socket(*item.key, *item.value, graph_params);
    }
 
    this->link_border_link_inputs_and_usages(zone,
                                             lf_body_inputs,
                                             body_fn.indices.inputs.border_links,
                                             lf_body_outputs,
                                             body_fn.indices.outputs.border_link_usages,
                                             graph_params);
 
    this->add_default_inputs(graph_params);
 
    Map<int, lf::OutputSocket *> lf_attribute_set_by_field_source_index;
    Map<int, lf::OutputSocket *> lf_attribute_set_by_caller_propagation_index;
 
    this->build_attribute_set_inputs_for_zone(graph_params,
                                              lf_attribute_set_by_field_source_index,
                                              lf_attribute_set_by_caller_propagation_index);
    for (const auto item : lf_attribute_set_by_field_source_index.items()) {
      lf::OutputSocket &lf_attribute_set_socket = *item.value;
      if (lf_attribute_set_socket.node().is_interface()) {
        body_fn.indices.inputs.attributes_by_field_source_index.add_new(
            item.key, lf_body_inputs.append_and_get_index(&lf_attribute_set_socket));
      }
    }
    for (const auto item : lf_attribute_set_by_caller_propagation_index.items()) {
      lf::OutputSocket &lf_attribute_set_socket = *item.value;
      if (lf_attribute_set_socket.node().is_interface()) {
        body_fn.indices.inputs.attributes_by_caller_propagation_index.add_new(
            item.key, lf_body_inputs.append_and_get_index(&lf_attribute_set_socket));
      }
    }
    this->link_attribute_set_inputs(lf_body_graph,
                                    graph_params,
                                    lf_attribute_set_by_field_source_index,
                                    lf_attribute_set_by_caller_propagation_index);
    this->fix_link_cycles(lf_body_graph, graph_params.socket_usage_inputs);
 
    lf_body_graph.update_node_indices();
 
    auto &logger = scope_.construct<GeometryNodesLazyFunctionLogger>(*lf_graph_info_);
    auto &side_effect_provider = scope_.construct<GeometryNodesLazyFunctionSideEffectProvider>();
 
    body_fn.function = &scope_.construct<lf::GraphExecutor>(lf_body_graph,
                                                            lf_body_inputs.as_span(),
                                                            lf_body_outputs.as_span(),
                                                            &logger,
                                                            &side_effect_provider,
                                                            nullptr);
 
    lf_graph_info_->debug_zone_body_graphs.add(zone.output_node->identifier, &lf_body_graph);
 
    // std::cout << "\n\n" << lf_body_graph.to_dot() << "\n\n";
 
    return body_fn;
  }
 
  void build_zone_border_links_inputs(const bNodeTreeZone &zone,
                                      lf::Graph &lf_graph,
                                      Vector<lf::GraphInputSocket *> &r_lf_graph_inputs,
                                      Vector<int> &r_indices)
  {
    for (const bNodeLink *border_link : zone.border_links) {
      r_indices.append(r_lf_graph_inputs.append_and_get_index(
          &lf_graph.add_input(*border_link->tosock->typeinfo->geometry_nodes_cpp_type,
                              StringRef("Link from ") + border_link->fromsock->name)));
    }
  }
 
  void build_zone_border_link_input_usages(const bNodeTreeZone &zone,
                                           lf::Graph &lf_graph,
                                           Vector<lf::GraphOutputSocket *> &r_lf_graph_outputs,
                                           Vector<int> &r_indices)
  {
    for (const bNodeLink *border_link : zone.border_links) {
      r_indices.append(r_lf_graph_outputs.append_and_get_index(&lf_graph.add_output(
          CPPType::get<bool>(), StringRef("Usage: Link from ") + border_link->fromsock->name)));
    }
  }
 
  void build_attribute_set_inputs_for_zone(
      BuildGraphParams &graph_params,
      Map<int, lf::OutputSocket *> &lf_attribute_set_by_field_source_index,
      Map<int, lf::OutputSocket *> &lf_attribute_set_by_caller_propagation_index)
  {
    const Vector<int> all_required_field_sources = this->find_all_required_field_source_indices(
        graph_params.lf_attribute_set_input_by_output_geometry_bsocket,
        graph_params.lf_attribute_set_input_by_field_source_index);
    const Vector<int> all_required_caller_propagation_indices =
        this->find_all_required_caller_propagation_indices(
            graph_params.lf_attribute_set_input_by_output_geometry_bsocket,
            graph_params.lf_attribute_set_input_by_caller_propagation_index);
 
    Map<int, int> input_by_field_source_index;
 
    for (const int field_source_index : all_required_field_sources) {
      const aai::FieldSource &field_source =
          attribute_inferencing_.all_field_sources[field_source_index];
      if ([[maybe_unused]] const auto *input_field_source = std::get_if<aai::InputFieldSource>(
              &field_source.data))
      {
        input_by_field_source_index.add_new(field_source_index,
                                            input_by_field_source_index.size());
      }
      else {
        const auto &socket_field_source = std::get<aai::SocketFieldSource>(field_source.data);
        const bNodeSocket &bsocket = *socket_field_source.socket;
        if (lf::OutputSocket *lf_field_socket = graph_params.lf_output_by_bsocket.lookup_default(
                &bsocket, nullptr))
        {
          lf::OutputSocket *lf_usage_socket = graph_params.usage_by_bsocket.lookup_default(
              &bsocket, nullptr);
          lf::OutputSocket &lf_attribute_set_socket = this->get_extracted_attributes(
              *lf_field_socket,
              lf_usage_socket,
              graph_params.lf_graph,
              graph_params.socket_usage_inputs);
          lf_attribute_set_by_field_source_index.add(field_source_index, &lf_attribute_set_socket);
        }
        else {
          input_by_field_source_index.add_new(field_source_index,
                                              input_by_field_source_index.size());
        }
      }
    }
 
    {
      Vector<lf::GraphInputSocket *> attribute_set_inputs;
      const int num = input_by_field_source_index.size() +
                      all_required_caller_propagation_indices.size();
      for ([[maybe_unused]] const int i : IndexRange(num)) {
        attribute_set_inputs.append(&graph_params.lf_graph.add_input(
            CPPType::get<bke::AnonymousAttributeSet>(), "Attribute Set"));
      }
 
      for (const auto item : input_by_field_source_index.items()) {
        const int field_source_index = item.key;
        const int attribute_set_index = item.value;
        lf::GraphInputSocket &lf_attribute_set_socket = *attribute_set_inputs[attribute_set_index];
        lf_attribute_set_by_field_source_index.add(field_source_index, &lf_attribute_set_socket);
      }
      for (const int i : all_required_caller_propagation_indices.index_range()) {
        const int caller_propagation_index = all_required_caller_propagation_indices[i];
        lf::GraphInputSocket &lf_attribute_set_socket =
            *attribute_set_inputs[input_by_field_source_index.size() + i];
        lf_attribute_set_by_caller_propagation_index.add_new(caller_propagation_index,
                                                             &lf_attribute_set_socket);
      }
    }
  }

  void build_root_graph()
  {
    lf::Graph &lf_graph = lf_graph_info_->graph;
 
    this->build_main_group_inputs(lf_graph);
    if (btree_.group_output_node() == nullptr) {
      this->build_fallback_group_outputs(lf_graph);
    }
 
    for (const bNodeTreeInterfaceSocket *interface_input : btree_.interface_inputs()) {
      lf::GraphOutputSocket &lf_socket = lf_graph.add_output(
          CPPType::get<bool>(),
          StringRef("Usage: ") + (interface_input->name ? interface_input->name : ""));
      group_input_usage_sockets_.append(&lf_socket);
    }
 
    Vector<lf::GraphInputSocket *> lf_output_usages;
    for (const bNodeTreeInterfaceSocket *interface_output : btree_.interface_outputs()) {
      lf::GraphInputSocket &lf_socket = lf_graph.add_input(
          CPPType::get<bool>(),
          StringRef("Usage: ") + (interface_output->name ? interface_output->name : ""));
      group_output_used_sockets_.append(&lf_socket);
      lf_output_usages.append(&lf_socket);
    }
 
    BuildGraphParams &graph_params = root_graph_build_params_.emplace(lf_graph);
    if (const bNode *group_output_bnode = btree_.group_output_node()) {
      for (const bNodeSocket *bsocket : group_output_bnode->input_sockets().drop_back(1)) {
        graph_params.usage_by_bsocket.add(bsocket, lf_output_usages[bsocket->index()]);
      }
    }
 
    this->insert_nodes_and_zones(
        tree_zones_->nodes_outside_zones, tree_zones_->root_zones, graph_params);
 
    for (const auto item : graph_params.lf_output_by_bsocket.items()) {
      this->insert_links_from_socket(*item.key, *item.value, graph_params);
    }
    this->build_group_input_usages(graph_params);
    this->add_default_inputs(graph_params);
 
    this->build_attribute_propagation_input_node(lf_graph);
 
    Map<int, lf::OutputSocket *> lf_attribute_set_by_field_source_index;
    Map<int, lf::OutputSocket *> lf_attribute_set_by_caller_propagation_index;
    this->build_attribute_set_inputs_outside_of_zones(
        graph_params,
        lf_attribute_set_by_field_source_index,
        lf_attribute_set_by_caller_propagation_index);
    this->link_attribute_set_inputs(lf_graph,
                                    graph_params,
                                    lf_attribute_set_by_field_source_index,
                                    lf_attribute_set_by_caller_propagation_index);
 
    this->fix_link_cycles(lf_graph, graph_params.socket_usage_inputs);
 
    // std::cout << "\n\n" << lf_graph.to_dot() << "\n\n";
 
    lf_graph.update_node_indices();
    lf_graph_info_->num_inline_nodes_approximate += lf_graph.nodes().size();
  }

  void build_geometry_nodes_group_function()
  {
    GeometryNodesGroupFunction &function = lf_graph_info_->function;
 
    Vector<const lf::GraphInputSocket *> lf_graph_inputs;
    Vector<const lf::GraphOutputSocket *> lf_graph_outputs;
 
    lf_graph_inputs.extend(group_input_sockets_);
    function.inputs.main = lf_graph_inputs.index_range().take_back(group_input_sockets_.size());
 
    lf_graph_inputs.extend(group_output_used_sockets_);
    function.inputs.output_usages = lf_graph_inputs.index_range().take_back(
        group_output_used_sockets_.size());
 
    for (auto [output_index, lf_socket] : attribute_set_by_geometry_output_.items()) {
      lf_graph_inputs.append(lf_socket);
      function.inputs.attributes_to_propagate.geometry_outputs.append(output_index);
    }
    function.inputs.attributes_to_propagate.range = lf_graph_inputs.index_range().take_back(
        attribute_set_by_geometry_output_.size());
 
    lf_graph_outputs.extend(standard_group_output_sockets_);
    function.outputs.main = lf_graph_outputs.index_range().take_back(
        standard_group_output_sockets_.size());
 
    lf_graph_outputs.extend(group_input_usage_sockets_);
    function.outputs.input_usages = lf_graph_outputs.index_range().take_back(
        group_input_usage_sockets_.size());
 
    Vector<const lf::FunctionNode *> &local_side_effect_nodes =
        scope_.construct<Vector<const lf::FunctionNode *>>();
    for (const bNode *bnode : btree_.nodes_by_type("GeometryNodeWarning")) {
      if (bnode->output_socket(0).is_directly_linked()) {
        /* The warning node is not a side-effect node. Instead, the user explicitly used the output
         * socket to specify when the warning node should be used. */
        continue;
      }
      if (tree_zones_->get_zone_by_node(bnode->identifier)) {
        /* "Global" warning nodes that are evaluated whenever the node group is evaluated must not
         * be in a zone. */
        continue;
      }
      /* Add warning node as side-effect node so that it is always evaluated if the node group is
       * evaluated. */
      const lf::Socket *lf_socket = root_graph_build_params_->lf_inputs_by_bsocket.lookup(
          &bnode->input_socket(0))[0];
      const lf::FunctionNode &lf_node = static_cast<const lf::FunctionNode &>(lf_socket->node());
      local_side_effect_nodes.append(&lf_node);
    }
 
    function.function = &scope_.construct<lf::GraphExecutor>(
        lf_graph_info_->graph,
        std::move(lf_graph_inputs),
        std::move(lf_graph_outputs),
        &scope_.construct<GeometryNodesLazyFunctionLogger>(*lf_graph_info_),
        &scope_.construct<GeometryNodesLazyFunctionSideEffectProvider>(local_side_effect_nodes),
        nullptr);
  }
 
  void build_attribute_set_inputs_outside_of_zones(
      BuildGraphParams &graph_params,
      Map<int, lf::OutputSocket *> &lf_attribute_set_by_field_source_index,
      Map<int, lf::OutputSocket *> &lf_attribute_set_by_caller_propagation_index)
  {
    const Vector<int> all_required_field_sources = this->find_all_required_field_source_indices(
        graph_params.lf_attribute_set_input_by_output_geometry_bsocket,
        graph_params.lf_attribute_set_input_by_field_source_index);
 
    for (const int field_source_index : all_required_field_sources) {
      const aai::FieldSource &field_source =
          attribute_inferencing_.all_field_sources[field_source_index];
      lf::OutputSocket *lf_attribute_set_socket;
      if (const auto *input_field_source = std::get_if<aai::InputFieldSource>(&field_source.data))
      {
        const int input_index = input_field_source->input_index;
        lf::OutputSocket &lf_field_socket = const_cast<lf::OutputSocket &>(
            *group_input_sockets_[input_index]);
        lf::OutputSocket *lf_usage_socket = const_cast<lf::OutputSocket *>(
            group_input_usage_sockets_[input_index]->origin());
        lf_attribute_set_socket = &this->get_extracted_attributes(
            lf_field_socket,
            lf_usage_socket,
            graph_params.lf_graph,
            graph_params.socket_usage_inputs);
      }
      else {
        const auto &socket_field_source = std::get<aai::SocketFieldSource>(field_source.data);
        const bNodeSocket &bsocket = *socket_field_source.socket;
        lf::OutputSocket &lf_field_socket = *graph_params.lf_output_by_bsocket.lookup(&bsocket);
        lf::OutputSocket *lf_usage_socket = graph_params.usage_by_bsocket.lookup_default(&bsocket,
                                                                                         nullptr);
        lf_attribute_set_socket = &this->get_extracted_attributes(
            lf_field_socket,
            lf_usage_socket,
            graph_params.lf_graph,
            graph_params.socket_usage_inputs);
      }
      lf_attribute_set_by_field_source_index.add_new(field_source_index, lf_attribute_set_socket);
    }
 
    for (const int caller_propagation_index :
         attribute_inferencing_.propagated_output_geometry_indices.index_range())
    {
      const int group_output_index =
          attribute_inferencing_.propagated_output_geometry_indices[caller_propagation_index];
      lf::OutputSocket &lf_attribute_set_socket = const_cast<lf::OutputSocket &>(
          *attribute_set_by_geometry_output_.lookup(group_output_index));
      lf_attribute_set_by_caller_propagation_index.add(caller_propagation_index,
                                                       &lf_attribute_set_socket);
    }
  }
 
  Vector<int> find_all_required_field_source_indices(
      const Map<const bNodeSocket *, lf::InputSocket *>
          &lf_attribute_set_input_by_output_geometry_bsocket,
      const MultiValueMap<int, lf::InputSocket *> &lf_attribute_set_input_by_field_source_index)
  {
    BitVector<> all_required_field_sources(attribute_inferencing_.all_field_sources.size(), false);
    for (const bNodeSocket *geometry_output_bsocket :
         lf_attribute_set_input_by_output_geometry_bsocket.keys())
    {
      all_required_field_sources |=
          attribute_inferencing_
              .required_fields_by_geometry_socket[geometry_output_bsocket->index_in_tree()];
    }
    for (const int field_source_index : lf_attribute_set_input_by_field_source_index.keys()) {
      all_required_field_sources[field_source_index].set();
    }
 
    Vector<int> indices;
    bits::foreach_1_index(all_required_field_sources, [&](const int i) { indices.append(i); });
    return indices;
  }
 
  Vector<int> find_all_required_caller_propagation_indices(
      const Map<const bNodeSocket *, lf::InputSocket *>
          &lf_attribute_set_input_by_output_geometry_bsocket,
      const MultiValueMap<int, lf::InputSocket *>
          &lf_attribute_set_input_by_caller_propagation_index)
  {
    BitVector<> all_required_caller_propagation_indices(
        attribute_inferencing_.propagated_output_geometry_indices.size(), false);
    for (const bNodeSocket *geometry_output_bs :
         lf_attribute_set_input_by_output_geometry_bsocket.keys())
    {
      all_required_caller_propagation_indices |=
          attribute_inferencing_
              .propagate_to_output_by_geometry_socket[geometry_output_bs->index_in_tree()];
    }
    for (const int caller_propagation_index :
         lf_attribute_set_input_by_caller_propagation_index.keys())
    {
      all_required_caller_propagation_indices[caller_propagation_index].set();
    }
 
    Vector<int> indices;
    bits::foreach_1_index(all_required_caller_propagation_indices,
                          [&](const int i) { indices.append(i); });
    return indices;
  }
 
  void link_attribute_set_inputs(
      lf::Graph &lf_graph,
      BuildGraphParams &graph_params,
      const Map<int, lf::OutputSocket *> &lf_attribute_set_by_field_source_index,
      const Map<int, lf::OutputSocket *> &lf_attribute_set_by_caller_propagation_index)
  {
    JoinAttributeSetsCache join_attribute_sets_cache;
 
    for (const MapItem<const bNodeSocket *, lf::InputSocket *> item :
         graph_params.lf_attribute_set_input_by_output_geometry_bsocket.items())
    {
      const bNodeSocket &geometry_output_bsocket = *item.key;
      lf::InputSocket &lf_attribute_set_input = *item.value;
 
      Vector<lf::OutputSocket *> lf_attribute_set_sockets;
 
      const BoundedBitSpan required_fields =
          attribute_inferencing_
              .required_fields_by_geometry_socket[geometry_output_bsocket.index_in_tree()];
      bits::foreach_1_index(required_fields, [&](const int field_source_index) {
        const auto &field_source = attribute_inferencing_.all_field_sources[field_source_index];
        if (const auto *socket_field_source = std::get_if<aai::SocketFieldSource>(
                &field_source.data))
        {
          if (&socket_field_source->socket->owner_node() == &geometry_output_bsocket.owner_node())
          {
            return;
          }
        }
        lf_attribute_set_sockets.append(
            lf_attribute_set_by_field_source_index.lookup(field_source_index));
      });
 
      const BoundedBitSpan required_caller_propagations =
          attribute_inferencing_
              .propagate_to_output_by_geometry_socket[geometry_output_bsocket.index_in_tree()];
      bits::foreach_1_index(required_caller_propagations, [&](const int caller_propagation_index) {
        lf_attribute_set_sockets.append(
            lf_attribute_set_by_caller_propagation_index.lookup(caller_propagation_index));
      });
 
      if (lf::OutputSocket *lf_attribute_set = this->join_attribute_sets(
              lf_attribute_set_sockets,
              join_attribute_sets_cache,
              lf_graph,
              graph_params.socket_usage_inputs))
      {
        lf_graph.add_link(*lf_attribute_set, lf_attribute_set_input);
      }
      else {
        static const bke::AnonymousAttributeSet empty_set;
        lf_attribute_set_input.set_default_value(&empty_set);
      }
    }
 
    for (const auto item : graph_params.lf_attribute_set_input_by_field_source_index.items()) {
      const int field_source_index = item.key;
      lf::OutputSocket &lf_attribute_set_socket = *lf_attribute_set_by_field_source_index.lookup(
          field_source_index);
      for (lf::InputSocket *lf_attribute_set_input : item.value) {
        lf_graph.add_link(lf_attribute_set_socket, *lf_attribute_set_input);
      }
    }
    for (const auto item : graph_params.lf_attribute_set_input_by_caller_propagation_index.items())
    {
      const int caller_propagation_index = item.key;
      lf::OutputSocket &lf_attribute_set_socket =
          *lf_attribute_set_by_caller_propagation_index.lookup(caller_propagation_index);
      for (lf::InputSocket *lf_attribute_set_input : item.value) {
        lf_graph.add_link(lf_attribute_set_socket, *lf_attribute_set_input);
      }
    }
  }
 
  void insert_nodes_and_zones(const Span<const bNode *> bnodes,
                              const Span<const bNodeTreeZone *> zones,
                              BuildGraphParams &graph_params)
  {
    Vector<const bNode *> nodes_to_insert = bnodes;
    Map<const bNode *, const bNodeTreeZone *> zone_by_output;
    for (const bNodeTreeZone *zone : zones) {
      nodes_to_insert.append(zone->output_node);
      zone_by_output.add(zone->output_node, zone);
    }
    /* Insert nodes from right to left so that usage sockets can be build in the same pass. */
    std::sort(nodes_to_insert.begin(), nodes_to_insert.end(), [](const bNode *a, const bNode *b) {
      return a->runtime->toposort_right_to_left_index < b->runtime->toposort_right_to_left_index;
    });
 
    for (const bNode *bnode : nodes_to_insert) {
      this->build_output_socket_usages(*bnode, graph_params);
      if (const bNodeTreeZone *zone = zone_by_output.lookup_default(bnode, nullptr)) {
        this->insert_child_zone_node(*zone, graph_params);
      }
      else {
        this->insert_node_in_graph(*bnode, graph_params);
      }
    }
  }
 
  void link_border_link_inputs_and_usages(const bNodeTreeZone &zone,
                                          const Span<lf::GraphInputSocket *> lf_inputs,
                                          const Span<int> lf_border_link_input_indices,
                                          const Span<lf::GraphOutputSocket *> lf_usages,
                                          const Span<int> lf_border_link_usage_indices,
                                          BuildGraphParams &graph_params)
  {
    lf::Graph &lf_graph = graph_params.lf_graph;
    for (const int border_link_i : zone.border_links.index_range()) {
      const bNodeLink &border_link = *zone.border_links[border_link_i];
      lf::GraphInputSocket &lf_from = *lf_inputs[lf_border_link_input_indices[border_link_i]];
      const Vector<lf::InputSocket *> lf_link_targets = this->find_link_targets(border_link,
                                                                                graph_params);
      for (lf::InputSocket *lf_to : lf_link_targets) {
        lf_graph.add_link(lf_from, *lf_to);
      }
      lf::GraphOutputSocket &lf_usage_output =
          *lf_usages[lf_border_link_usage_indices[border_link_i]];
      if (lf::OutputSocket *lf_usage = graph_params.usage_by_bsocket.lookup_default(
              border_link.tosock, nullptr))
      {
        lf_graph.add_link(*lf_usage, lf_usage_output);
      }
      else {
        static const bool static_false = false;
        lf_usage_output.set_default_value(&static_false);
      }
    }
  }
 
  lf::OutputSocket &get_extracted_attributes(lf::OutputSocket &lf_field_socket,
                                             lf::OutputSocket *lf_usage_socket,
                                             lf::Graph &lf_graph,
                                             Set<lf::InputSocket *> &socket_usage_inputs)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForAnonymousAttributeSetExtract>();
    lf::Node &lf_node = lf_graph.add_function(lazy_function);
    lf::InputSocket &lf_use_input = lf_node.input(0);
    lf::InputSocket &lf_field_input = lf_node.input(1);
    socket_usage_inputs.add_new(&lf_use_input);
    if (lf_usage_socket) {
      lf_graph.add_link(*lf_usage_socket, lf_use_input);
    }
    else {
      static const bool static_false = false;
      lf_use_input.set_default_value(&static_false);
    }
    lf_graph.add_link(lf_field_socket, lf_field_input);
    return lf_node.output(0);
  }

  lf::OutputSocket *join_attribute_sets(const Span<lf::OutputSocket *> lf_attribute_set_sockets,
                                        JoinAttributeSetsCache &cache,
                                        lf::Graph &lf_graph,
                                        Set<lf::InputSocket *> &socket_usage_inputs)
  {
    if (lf_attribute_set_sockets.is_empty()) {
      return nullptr;
    }
    if (lf_attribute_set_sockets.size() == 1) {
      return lf_attribute_set_sockets[0];
    }
 
    Vector<lf::OutputSocket *, 16> key = lf_attribute_set_sockets;
    std::sort(key.begin(), key.end());
    return cache.lookup_or_add_cb(key, [&]() {
      const auto &lazy_function = LazyFunctionForAnonymousAttributeSetJoin::get_cached(
          lf_attribute_set_sockets.size(), scope_);
      lf::Node &lf_node = lf_graph.add_function(lazy_function);
      for (const int i : lf_attribute_set_sockets.index_range()) {
        lf::OutputSocket &lf_attribute_set_socket = *lf_attribute_set_sockets[i];
        lf::InputSocket &lf_use_input = lf_node.input(lazy_function.get_use_input(i));
 
        /* Some attribute sets could potentially be set unused in the future based on more dynamic
         * analysis of the node tree. */
        static const bool static_true = true;
        lf_use_input.set_default_value(&static_true);
 
        socket_usage_inputs.add(&lf_use_input);
        lf::InputSocket &lf_attribute_set_input = lf_node.input(
            lazy_function.get_attribute_set_input(i));
        lf_graph.add_link(lf_attribute_set_socket, lf_attribute_set_input);
      }
      return &lf_node.output(0);
    });
  }
 
  void insert_child_zone_node(const bNodeTreeZone &child_zone, BuildGraphParams &graph_params)
  {
    const int child_zone_i = child_zone.index;
    ZoneBuildInfo &child_zone_info = zone_build_infos_[child_zone_i];
    lf::FunctionNode &child_zone_node = graph_params.lf_graph.add_function(
        *child_zone_info.lazy_function);
    mapping_->zone_node_map.add_new(&child_zone, &child_zone_node);
 
    {
      int valid_socket_i = 0;
      for (const bNodeSocket *bsocket : child_zone.input_node->input_sockets()) {
        if (ignore_zone_bsocket(*bsocket)) {
          continue;
        }
        lf::InputSocket &lf_input_socket = child_zone_node.input(
            child_zone_info.indices.inputs.main[valid_socket_i]);
        lf::OutputSocket &lf_usage_socket = child_zone_node.output(
            child_zone_info.indices.outputs.input_usages[valid_socket_i]);
        mapping_->bsockets_by_lf_socket_map.add(&lf_input_socket, bsocket);
        graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_input_socket);
        graph_params.usage_by_bsocket.add(bsocket, &lf_usage_socket);
        valid_socket_i++;
      }
    }
    {
      int valid_socket_i = 0;
      for (const bNodeSocket *bsocket : child_zone.output_node->output_sockets()) {
        if (ignore_zone_bsocket(*bsocket)) {
          continue;
        }
        lf::OutputSocket &lf_output_socket = child_zone_node.output(
            child_zone_info.indices.outputs.main[valid_socket_i]);
        lf::InputSocket &lf_usage_input = child_zone_node.input(
            child_zone_info.indices.inputs.output_usages[valid_socket_i]);
        mapping_->bsockets_by_lf_socket_map.add(&lf_output_socket, bsocket);
        graph_params.lf_output_by_bsocket.add(bsocket, &lf_output_socket);
        graph_params.socket_usage_inputs.add(&lf_usage_input);
        if (lf::OutputSocket *lf_usage = graph_params.usage_by_bsocket.lookup_default(bsocket,
                                                                                      nullptr))
        {
          graph_params.lf_graph.add_link(*lf_usage, lf_usage_input);
        }
        else {
          static const bool static_false = false;
          lf_usage_input.set_default_value(&static_false);
        }
        valid_socket_i++;
      }
    }
 
    const Span<const bNodeLink *> child_border_links = child_zone.border_links;
    for (const int child_border_link_i : child_border_links.index_range()) {
      lf::InputSocket &child_border_link_input = child_zone_node.input(
          child_zone_info.indices.inputs.border_links[child_border_link_i]);
      const bNodeLink &link = *child_border_links[child_border_link_i];
      graph_params.lf_input_by_border_link.add(&link, &child_border_link_input);
      lf::OutputSocket &lf_usage = child_zone_node.output(
          child_zone_info.indices.outputs.border_link_usages[child_border_link_i]);
      graph_params.lf_inputs_by_bsocket.add(link.tosock, &child_border_link_input);
      graph_params.usage_by_bsocket.add(link.tosock, &lf_usage);
    }
 
    for (const auto item : child_zone_info.indices.inputs.attributes_by_field_source_index.items())
    {
      const int field_source_index = item.key;
      const int child_zone_input_index = item.value;
      lf::InputSocket &lf_attribute_set_input = child_zone_node.input(child_zone_input_index);
      graph_params.lf_attribute_set_input_by_field_source_index.add(field_source_index,
                                                                    &lf_attribute_set_input);
    }
    for (const auto item :
         child_zone_info.indices.inputs.attributes_by_caller_propagation_index.items())
    {
      const int caller_propagation_index = item.key;
      const int child_zone_input_index = item.value;
      lf::InputSocket &lf_attribute_set_input = child_zone_node.input(child_zone_input_index);
      BLI_assert(lf_attribute_set_input.type().is<bke::AnonymousAttributeSet>());
      graph_params.lf_attribute_set_input_by_caller_propagation_index.add(caller_propagation_index,
                                                                          &lf_attribute_set_input);
    }
  }
 
  void build_main_group_inputs(lf::Graph &lf_graph)
  {
    const Span<const bNodeTreeInterfaceSocket *> interface_inputs = btree_.interface_inputs();
    for (const bNodeTreeInterfaceSocket *interface_input : interface_inputs) {
      const bke::bNodeSocketType *typeinfo = interface_input->socket_typeinfo();
      lf::GraphInputSocket &lf_socket = lf_graph.add_input(
          *typeinfo->geometry_nodes_cpp_type, interface_input->name ? interface_input->name : "");
      group_input_sockets_.append(&lf_socket);
    }
  }

  void build_fallback_group_outputs(lf::Graph &lf_graph)
  {
    for (const bNodeTreeInterfaceSocket *interface_output : btree_.interface_outputs()) {
      const bke::bNodeSocketType *typeinfo = interface_output->socket_typeinfo();
      const CPPType &type = *typeinfo->geometry_nodes_cpp_type;
      lf::GraphOutputSocket &lf_socket = lf_graph.add_output(
          type, interface_output->name ? interface_output->name : "");
      const void *default_value = typeinfo->geometry_nodes_default_cpp_value;
      if (default_value == nullptr) {
        default_value = type.default_value();
      }
      lf_socket.set_default_value(default_value);
      standard_group_output_sockets_.append(&lf_socket);
    }
  }
 
  void insert_node_in_graph(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const bke::bNodeType *node_type = bnode.typeinfo;
    if (node_type == nullptr) {
      return;
    }
    if (bnode.is_muted()) {
      this->build_muted_node(bnode, graph_params);
      return;
    }
    switch (node_type->type) {
      case NODE_FRAME: {
        /* Ignored. */
        break;
      }
      case NODE_REROUTE: {
        this->build_reroute_node(bnode, graph_params);
        break;
      }
      case NODE_GROUP_INPUT: {
        this->handle_group_input_node(bnode, graph_params);
        break;
      }
      case NODE_GROUP_OUTPUT: {
        this->build_group_output_node(bnode, graph_params);
        break;
      }
      case NODE_CUSTOM_GROUP:
      case NODE_GROUP: {
        this->build_group_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_VIEWER: {
        this->build_viewer_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_SWITCH: {
        this->build_switch_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_INDEX_SWITCH: {
        this->build_index_switch_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_WARNING: {
        this->build_warning_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_GIZMO_LINEAR:
      case GEO_NODE_GIZMO_DIAL:
      case GEO_NODE_GIZMO_TRANSFORM: {
        this->build_gizmo_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_BAKE: {
        this->build_bake_node(bnode, graph_params);
        break;
      }
      case GEO_NODE_MENU_SWITCH: {
        this->build_menu_switch_node(bnode, graph_params);
        break;
      }
      default: {
        if (node_type->geometry_node_execute) {
          this->build_geometry_node(bnode, graph_params);
          break;
        }
        const NodeMultiFunctions::Item &fn_item = node_multi_functions_.try_get(bnode);
        if (fn_item.fn != nullptr) {
          this->build_multi_function_node(bnode, fn_item, graph_params);
          break;
        }
        if (node_type == &bke::NodeTypeUndefined) {
          this->build_undefined_node(bnode, graph_params);
          break;
        }
        /* Nodes that don't match any of the criteria above are just ignored. */
        break;
      }
    }
  }
 
  void build_muted_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForMutedNode>(
        bnode, mapping_->lf_index_by_bsocket);
    lf::Node &lf_node = graph_params.lf_graph.add_function(lazy_function);
    for (const bNodeSocket *bsocket : bnode.input_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::InputSocket &lf_socket = lf_node.input(lf_index);
      graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::OutputSocket &lf_socket = lf_node.output(lf_index);
      graph_params.lf_output_by_bsocket.add_new(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
 
    this->build_muted_node_usages(bnode, graph_params);
  }

  void build_muted_node_usages(const bNode &bnode, BuildGraphParams &graph_params)
  {
    /* Find all outputs that use a specific input. */
    MultiValueMap<const bNodeSocket *, const bNodeSocket *> outputs_by_input;
    for (const bNodeLink &blink : bnode.internal_links()) {
      outputs_by_input.add(blink.fromsock, blink.tosock);
    }
    for (const auto item : outputs_by_input.items()) {
      const bNodeSocket &input_bsocket = *item.key;
      const Span<const bNodeSocket *> output_bsockets = item.value;
 
      /* The input is used if any of the internally linked outputs is used. */
      Vector<lf::OutputSocket *> lf_socket_usages;
      for (const bNodeSocket *output_bsocket : output_bsockets) {
        if (lf::OutputSocket *lf_socket = graph_params.usage_by_bsocket.lookup_default(
                output_bsocket, nullptr))
        {
          lf_socket_usages.append(lf_socket);
        }
      }
      graph_params.usage_by_bsocket.add(&input_bsocket,
                                        this->or_socket_usages(lf_socket_usages, graph_params));
    }
  }
 
  void build_reroute_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const bNodeSocket &input_bsocket = bnode.input_socket(0);
    const bNodeSocket &output_bsocket = bnode.output_socket(0);
    const CPPType *type = get_socket_cpp_type(input_bsocket);
    if (type == nullptr) {
      return;
    }
 
    auto &lazy_function = scope_.construct<LazyFunctionForRerouteNode>(*type);
    lf::Node &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    lf::InputSocket &lf_input = lf_node.input(0);
    lf::OutputSocket &lf_output = lf_node.output(0);
    graph_params.lf_inputs_by_bsocket.add(&input_bsocket, &lf_input);
    graph_params.lf_output_by_bsocket.add_new(&output_bsocket, &lf_output);
    mapping_->bsockets_by_lf_socket_map.add(&lf_input, &input_bsocket);
    mapping_->bsockets_by_lf_socket_map.add(&lf_output, &output_bsocket);
 
    if (lf::OutputSocket *lf_usage = graph_params.usage_by_bsocket.lookup_default(
            &bnode.output_socket(0), nullptr))
    {
      graph_params.usage_by_bsocket.add(&bnode.input_socket(0), lf_usage);
    }
  }
 
  void handle_group_input_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    for (const int i : btree_.interface_inputs().index_range()) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::GraphInputSocket &lf_socket = *const_cast<lf::GraphInputSocket *>(
          group_input_sockets_[i]);
      graph_params.lf_output_by_bsocket.add_new(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
  }
 
  void build_group_output_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    Vector<lf::GraphOutputSocket *> lf_graph_outputs;
 
    for (const int i : btree_.interface_outputs().index_range()) {
      const bNodeTreeInterfaceSocket &interface_output = *btree_.interface_outputs()[i];
      const bNodeSocket &bsocket = bnode.input_socket(i);
      const bke::bNodeSocketType *typeinfo = interface_output.socket_typeinfo();
      const CPPType &type = *typeinfo->geometry_nodes_cpp_type;
      lf::GraphOutputSocket &lf_socket = graph_params.lf_graph.add_output(
          type, interface_output.name ? interface_output.name : "");
      lf_graph_outputs.append(&lf_socket);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
 
    if (&bnode == btree_.group_output_node()) {
      standard_group_output_sockets_ = lf_graph_outputs.as_span();
    }
  }
 
  void build_group_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const bNodeTree *group_btree = reinterpret_cast<bNodeTree *>(bnode.id);
    if (group_btree == nullptr) {
      return;
    }
    const GeometryNodesLazyFunctionGraphInfo *group_lf_graph_info =
        ensure_geometry_nodes_lazy_function_graph(*group_btree);
    if (group_lf_graph_info == nullptr) {
      return;
    }
 
    auto &lazy_function = scope_.construct<LazyFunctionForGroupNode>(
        bnode, *group_lf_graph_info, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const int i : bnode.input_sockets().index_range()) {
      const bNodeSocket &bsocket = bnode.input_socket(i);
      BLI_assert(!bsocket.is_multi_input());
      lf::InputSocket &lf_socket = lf_node.input(group_lf_graph_info->function.inputs.main[i]);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range()) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::OutputSocket &lf_socket = lf_node.output(group_lf_graph_info->function.outputs.main[i]);
      graph_params.lf_output_by_bsocket.add_new(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    mapping_->group_node_map.add(&bnode, &lf_node);
    lf_graph_info_->num_inline_nodes_approximate +=
        group_lf_graph_info->num_inline_nodes_approximate;
    static const bool static_false = false;
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      {
        const int lf_input_index =
            mapping_->lf_input_index_for_output_bsocket_usage[bsocket->index_in_all_outputs()];
        if (lf_input_index != -1) {
          lf::InputSocket &lf_input = lf_node.input(lf_input_index);
          lf_input.set_default_value(&static_false);
          graph_params.socket_usage_inputs.add(&lf_input);
        }
      }
      {
        /* Keep track of attribute set inputs that need to be populated later. */
        const int lf_input_index = mapping_->lf_input_index_for_attribute_propagation_to_output
                                       [bsocket->index_in_all_outputs()];
        if (lf_input_index != -1) {
          lf::InputSocket &lf_input = lf_node.input(lf_input_index);
          graph_params.lf_attribute_set_input_by_output_geometry_bsocket.add(bsocket, &lf_input);
        }
      }
    }
 
    this->build_group_node_socket_usage(bnode, lf_node, graph_params, *group_lf_graph_info);
  }
 
  void build_group_node_socket_usage(const bNode &bnode,
                                     lf::FunctionNode &lf_group_node,
                                     BuildGraphParams &graph_params,
                                     const GeometryNodesLazyFunctionGraphInfo &group_lf_graph_info)
  {
    for (const bNodeSocket *input_bsocket : bnode.input_sockets()) {
      const int input_index = input_bsocket->index();
      const InputUsageHint &input_usage_hint =
          group_lf_graph_info.mapping.group_input_usage_hints[input_index];
      switch (input_usage_hint.type) {
        case InputUsageHintType::Never: {
          /* Nothing to do. */
          break;
        }
        case InputUsageHintType::DependsOnOutput: {
          Vector<lf::OutputSocket *> output_usages;
          for (const int i : input_usage_hint.output_dependencies) {
            if (lf::OutputSocket *lf_socket = graph_params.usage_by_bsocket.lookup_default(
                    &bnode.output_socket(i), nullptr))
            {
              output_usages.append(lf_socket);
            }
          }
          graph_params.usage_by_bsocket.add(input_bsocket,
                                            this->or_socket_usages(output_usages, graph_params));
          break;
        }
        case InputUsageHintType::DynamicSocket: {
          graph_params.usage_by_bsocket.add(
              input_bsocket,
              &lf_group_node.output(
                  group_lf_graph_info.function.outputs.input_usages[input_index]));
          break;
        }
      }
    }
 
    for (const bNodeSocket *output_bsocket : bnode.output_sockets()) {
      const int lf_input_index =
          mapping_
              ->lf_input_index_for_output_bsocket_usage[output_bsocket->index_in_all_outputs()];
      BLI_assert(lf_input_index >= 0);
      lf::InputSocket &lf_socket = lf_group_node.input(lf_input_index);
      if (lf::OutputSocket *lf_output_is_used = graph_params.usage_by_bsocket.lookup_default(
              output_bsocket, nullptr))
      {
        graph_params.lf_graph.add_link(*lf_output_is_used, lf_socket);
      }
      else {
        static const bool static_false = false;
        lf_socket.set_default_value(&static_false);
      }
    }
  }
 
  void build_geometry_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForGeometryNode>(bnode, *lf_graph_info_);
    lf::Node &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const bNodeSocket *bsocket : bnode.input_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::InputSocket &lf_socket = lf_node.input(lf_index);
 
      if (bsocket->is_multi_input()) {
        auto &multi_input_lazy_function = scope_.construct<LazyFunctionForMultiInput>(*bsocket);
        lf::Node &lf_multi_input_node = graph_params.lf_graph.add_function(
            multi_input_lazy_function);
        graph_params.lf_graph.add_link(lf_multi_input_node.output(0), lf_socket);
        for (const int i : multi_input_lazy_function.links.index_range()) {
          lf::InputSocket &lf_multi_input_socket = lf_multi_input_node.input(i);
          const bNodeLink *link = multi_input_lazy_function.links[i];
          graph_params.lf_input_by_multi_input_link.add(link, &lf_multi_input_socket);
          mapping_->bsockets_by_lf_socket_map.add(&lf_multi_input_socket, bsocket);
          const void *default_value = lf_multi_input_socket.type().default_value();
          lf_multi_input_socket.set_default_value(default_value);
        }
      }
      else {
        graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_socket);
        mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
      }
    }
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::OutputSocket &lf_socket = lf_node.output(lf_index);
      graph_params.lf_output_by_bsocket.add_new(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
 
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      {
        const int lf_input_index =
            mapping_->lf_input_index_for_output_bsocket_usage[bsocket->index_in_all_outputs()];
        if (lf_input_index != -1) {
          lf::InputSocket &lf_input_socket = lf_node.input(lf_input_index);
          if (lf::OutputSocket *lf_usage = graph_params.usage_by_bsocket.lookup_default(bsocket,
                                                                                        nullptr))
          {
            graph_params.lf_graph.add_link(*lf_usage, lf_input_socket);
          }
          else {
            static const bool static_false = false;
            lf_input_socket.set_default_value(&static_false);
          }
          graph_params.socket_usage_inputs.add_new(&lf_node.input(lf_input_index));
        }
      }
      {
        /* Keep track of attribute set inputs that need to be populated later. */
        const int lf_input_index = mapping_->lf_input_index_for_attribute_propagation_to_output
                                       [bsocket->index_in_all_outputs()];
        if (lf_input_index != -1) {
          graph_params.lf_attribute_set_input_by_output_geometry_bsocket.add(
              bsocket, &lf_node.input(lf_input_index));
        }
      }
    }
 
    this->build_standard_node_input_socket_usage(bnode, graph_params);
  }
 
  void build_standard_node_input_socket_usage(const bNode &bnode, BuildGraphParams &graph_params)
  {
    if (bnode.input_sockets().is_empty()) {
      return;
    }
 
    Vector<lf::OutputSocket *> output_usages;
    for (const bNodeSocket *output_socket : bnode.output_sockets()) {
      if (!output_socket->is_available()) {
        continue;
      }
      if (lf::OutputSocket *is_used_socket = graph_params.usage_by_bsocket.lookup_default(
              output_socket, nullptr))
      {
        output_usages.append_non_duplicates(is_used_socket);
      }
    }
 
    /* Assume every input is used when any output is used. */
    lf::OutputSocket *lf_usage = this->or_socket_usages(output_usages, graph_params);
    if (lf_usage == nullptr) {
      return;
    }
 
    for (const bNodeSocket *input_socket : bnode.input_sockets()) {
      if (input_socket->is_available()) {
        graph_params.usage_by_bsocket.add(input_socket, lf_usage);
      }
    }
  }
 
  void build_multi_function_node(const bNode &bnode,
                                 const NodeMultiFunctions::Item &fn_item,
                                 BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForMultiFunctionNode>(
        bnode, fn_item, mapping_->lf_index_by_bsocket);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const bNodeSocket *bsocket : bnode.input_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      BLI_assert(!bsocket->is_multi_input());
      lf::InputSocket &lf_socket = lf_node.input(lf_index);
      graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::OutputSocket &lf_socket = lf_node.output(lf_index);
      graph_params.lf_output_by_bsocket.add(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
 
    this->build_standard_node_input_socket_usage(bnode, graph_params);
  }
 
  void build_viewer_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForViewerNode>(
        bnode, mapping_->lf_index_by_bsocket);
    lf::FunctionNode &lf_viewer_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const bNodeSocket *bsocket : bnode.input_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::InputSocket &lf_socket = lf_viewer_node.input(lf_index);
      graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
 
    mapping_->possible_side_effect_node_map.add(&bnode, &lf_viewer_node);
 
    {
      auto &usage_lazy_function = scope_.construct<LazyFunctionForViewerInputUsage>(
          lf_viewer_node);
      lf::FunctionNode &lf_usage_node = graph_params.lf_graph.add_function(usage_lazy_function);
 
      for (const bNodeSocket *bsocket : bnode.input_sockets()) {
        if (bsocket->is_available()) {
          graph_params.usage_by_bsocket.add(bsocket, &lf_usage_node.output(0));
        }
      }
    }
  }
 
  void build_gizmo_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForGizmoNode>(
        bnode, mapping_->lf_index_by_bsocket);
    lf::FunctionNode &lf_gizmo_node = graph_params.lf_graph.add_function(lazy_function);
    lazy_function.self_node = &lf_gizmo_node;
 
    for (const int i : lazy_function.gizmo_links.index_range()) {
      const bNodeLink &link = *lazy_function.gizmo_links[i];
      lf::InputSocket &lf_socket = lf_gizmo_node.input(i);
      graph_params.lf_input_by_multi_input_link.add(&link, &lf_socket);
    }
    for (const int i : bnode.input_sockets().drop_front(1).index_range()) {
      lf::InputSocket &lf_socket = lf_gizmo_node.input(i + lazy_function.gizmo_links.size());
      const bNodeSocket &bsocket = bnode.input_socket(i + 1);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range()) {
      lf::OutputSocket &lf_socket = lf_gizmo_node.output(i);
      const bNodeSocket &bsocket = bnode.output_socket(i);
      graph_params.lf_output_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
 
    this->build_gizmo_node_socket_usage(bnode, graph_params, lf_gizmo_node);
 
    mapping_->possible_side_effect_node_map.add(&bnode, &lf_gizmo_node);
  }
 
  void build_gizmo_node_socket_usage(const bNode &bnode,
                                     BuildGraphParams &graph_params,
                                     const lf::FunctionNode &lf_gizmo_node)
  {
    const auto &usage_fn = scope_.construct<LazyFunctionForGizmoInputsUsage>(bnode, lf_gizmo_node);
    lf::FunctionNode &lf_usage_node = graph_params.lf_graph.add_function(usage_fn);
    for (const bNodeSocket *bsocket : bnode.input_sockets()) {
      graph_params.usage_by_bsocket.add(bsocket, &lf_usage_node.output(0));
    }
  }
 
  lf::FunctionNode *insert_simulation_input_node(const bNodeTree &node_tree,
                                                 const bNode &bnode,
                                                 BuildGraphParams &graph_params)
  {
    const NodeGeometrySimulationInput *storage = static_cast<const NodeGeometrySimulationInput *>(
        bnode.storage);
    if (node_tree.node_by_id(storage->output_node_id) == nullptr) {
      return nullptr;
    }
 
    std::unique_ptr<LazyFunction> lazy_function = get_simulation_input_lazy_function(
        node_tree, bnode, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    for (const int i : bnode.input_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.input_socket(i);
      lf::InputSocket &lf_socket = lf_node.input(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::OutputSocket &lf_socket = lf_node.output(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_output_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    return &lf_node;
  }
 
  lf::FunctionNode &insert_simulation_output_node(const bNode &bnode,
                                                  BuildGraphParams &graph_params)
  {
    std::unique_ptr<LazyFunction> lazy_function = get_simulation_output_lazy_function(
        bnode, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    for (const int i : bnode.input_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.input_socket(i);
      lf::InputSocket &lf_socket = lf_node.input(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::OutputSocket &lf_socket = lf_node.output(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_output_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
 
    mapping_->possible_side_effect_node_map.add(&bnode, &lf_node);
 
    return lf_node;
  }
 
  void build_bake_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    std::unique_ptr<LazyFunction> lazy_function = get_bake_lazy_function(bnode, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    for (const int i : bnode.input_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.input_socket(i);
      lf::InputSocket &lf_socket = lf_node.input(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range().drop_back(1)) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::OutputSocket &lf_socket = lf_node.output(
          mapping_->lf_index_by_bsocket[bsocket.index_in_tree()]);
      graph_params.lf_output_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
 
    mapping_->possible_side_effect_node_map.add(&bnode, &lf_node);
 
    this->build_bake_node_socket_usage(bnode, graph_params);
  }
 
  void build_bake_node_socket_usage(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const LazyFunction &usage_fn = scope_.construct<LazyFunctionForBakeInputsUsage>(bnode);
    lf::FunctionNode &lf_usage_node = graph_params.lf_graph.add_function(usage_fn);
    const int items_num = bnode.input_sockets().size() - 1;
    for (const int i : IndexRange(items_num)) {
      graph_params.usage_by_bsocket.add(&bnode.input_socket(i), &lf_usage_node.output(0));
    }
  }
 
  void build_switch_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    std::unique_ptr<LazyFunction> lazy_function = get_switch_node_lazy_function(bnode);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    for (const int i : bnode.input_sockets().index_range()) {
      graph_params.lf_inputs_by_bsocket.add(&bnode.input_socket(i), &lf_node.input(i));
      mapping_->bsockets_by_lf_socket_map.add(&lf_node.input(i), &bnode.input_socket(i));
    }
 
    graph_params.lf_output_by_bsocket.add(&bnode.output_socket(0), &lf_node.output(0));
    mapping_->bsockets_by_lf_socket_map.add(&lf_node.output(0), &bnode.output_socket(0));
 
    this->build_switch_node_socket_usage(bnode, graph_params);
  }
 
  void build_switch_node_socket_usage(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const bNodeSocket &switch_input_bsocket = bnode.input_socket(0);
    const bNodeSocket &false_input_bsocket = bnode.input_socket(1);
    const bNodeSocket &true_input_bsocket = bnode.input_socket(2);
    const bNodeSocket &output_bsocket = bnode.output_socket(0);
    lf::OutputSocket *output_is_used_socket = graph_params.usage_by_bsocket.lookup_default(
        &output_bsocket, nullptr);
    if (output_is_used_socket == nullptr) {
      return;
    }
    graph_params.usage_by_bsocket.add(&switch_input_bsocket, output_is_used_socket);
    if (switch_input_bsocket.is_directly_linked()) {
      /* The condition input is dynamic, so the usage of the other inputs is as well. */
      static const LazyFunctionForSwitchSocketUsage switch_socket_usage_fn;
      lf::Node &lf_node = graph_params.lf_graph.add_function(switch_socket_usage_fn);
      graph_params.lf_inputs_by_bsocket.add(&switch_input_bsocket, &lf_node.input(0));
      graph_params.usage_by_bsocket.add(&false_input_bsocket, &lf_node.output(0));
      graph_params.usage_by_bsocket.add(&true_input_bsocket, &lf_node.output(1));
    }
    else {
      if (switch_input_bsocket.default_value_typed<bNodeSocketValueBoolean>()->value) {
        graph_params.usage_by_bsocket.add(&true_input_bsocket, output_is_used_socket);
      }
      else {
        graph_params.usage_by_bsocket.add(&false_input_bsocket, output_is_used_socket);
      }
    }
  }
 
  void build_index_switch_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    std::unique_ptr<LazyFunction> lazy_function = get_index_switch_node_lazy_function(
        bnode, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    for (const int i : bnode.input_sockets().drop_back(1).index_range()) {
      graph_params.lf_inputs_by_bsocket.add(&bnode.input_socket(i), &lf_node.input(i));
      mapping_->bsockets_by_lf_socket_map.add(&lf_node.input(i), &bnode.input_socket(i));
    }
 
    graph_params.lf_output_by_bsocket.add(&bnode.output_socket(0), &lf_node.output(0));
    mapping_->bsockets_by_lf_socket_map.add(&lf_node.output(0), &bnode.output_socket(0));
 
    this->build_index_switch_node_socket_usage(bnode, graph_params);
  }
 
  void build_index_switch_node_socket_usage(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const bNodeSocket &index_socket = bnode.input_socket(0);
    const int items_num = bnode.input_sockets().size() - 1;
 
    lf::OutputSocket *output_is_used = graph_params.usage_by_bsocket.lookup_default(
        &bnode.output_socket(0), nullptr);
    if (output_is_used == nullptr) {
      return;
    }
    graph_params.usage_by_bsocket.add(&index_socket, output_is_used);
    if (index_socket.is_directly_linked()) {
      /* The condition input is dynamic, so the usage of the other inputs is as well. */
      auto usage_fn = std::make_unique<LazyFunctionForIndexSwitchSocketUsage>(bnode);
      lf::Node &lf_node = graph_params.lf_graph.add_function(*usage_fn);
      scope_.add(std::move(usage_fn));
 
      graph_params.lf_inputs_by_bsocket.add(&index_socket, &lf_node.input(0));
      for (const int i : IndexRange(items_num)) {
        graph_params.usage_by_bsocket.add(&bnode.input_socket(i + 1), &lf_node.output(i));
      }
    }
    else {
      const int index = index_socket.default_value_typed<bNodeSocketValueInt>()->value;
      if (IndexRange(items_num).contains(index)) {
        graph_params.usage_by_bsocket.add(&bnode.input_socket(index + 1), output_is_used);
      }
    }
  }
 
  void build_warning_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto lazy_function_ptr = get_warning_node_lazy_function(bnode);
    LazyFunction &lazy_function = *lazy_function_ptr;
    scope_.add(std::move(lazy_function_ptr));
 
    lf::Node &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const int i : bnode.input_sockets().index_range()) {
      const bNodeSocket &bsocket = bnode.input_socket(i);
      lf::InputSocket &lf_socket = lf_node.input(i);
      graph_params.lf_inputs_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
    for (const int i : bnode.output_sockets().index_range()) {
      const bNodeSocket &bsocket = bnode.output_socket(i);
      lf::OutputSocket &lf_socket = lf_node.output(i);
      graph_params.lf_output_by_bsocket.add(&bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, &bsocket);
    }
 
    const bNodeSocket &output_bsocket = bnode.output_socket(0);
 
    lf::OutputSocket *lf_usage = nullptr;
    if (output_bsocket.is_directly_linked()) {
      /* The warning node is only used if the output socket is used. */
      lf_usage = graph_params.usage_by_bsocket.lookup_default(&output_bsocket, nullptr);
    }
    else {
      /* The warning node is used if any of the output sockets is used. */
      lf_usage = this->or_socket_usages(group_output_used_sockets_, graph_params);
    }
    if (lf_usage) {
      for (const bNodeSocket *socket : bnode.input_sockets()) {
        graph_params.usage_by_bsocket.add(socket, lf_usage);
      }
    }
  }
 
  void build_menu_switch_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    std::unique_ptr<LazyFunction> lazy_function = get_menu_switch_node_lazy_function(
        bnode, *lf_graph_info_);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
    scope_.add(std::move(lazy_function));
 
    int input_index = 0;
    for (const bNodeSocket *bsocket : bnode.input_sockets().drop_back(1)) {
      if (bsocket->is_available()) {
        lf::InputSocket &lf_socket = lf_node.input(input_index);
        graph_params.lf_inputs_by_bsocket.add(bsocket, &lf_socket);
        mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
        input_index++;
      }
    }
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      if (bsocket->is_available()) {
        lf::OutputSocket &lf_socket = lf_node.output(0);
        graph_params.lf_output_by_bsocket.add(bsocket, &lf_socket);
        mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
        break;
      }
    }
 
    this->build_menu_switch_node_socket_usage(bnode, graph_params);
  }
 
  void build_menu_switch_node_socket_usage(const bNode &bnode, BuildGraphParams &graph_params)
  {
    const NodeMenuSwitch &storage = *static_cast<NodeMenuSwitch *>(bnode.storage);
    const NodeEnumDefinition &enum_def = storage.enum_definition;
 
    const bNodeSocket *switch_input_bsocket = bnode.input_sockets()[0];
    Vector<const bNodeSocket *> input_bsockets(enum_def.items_num);
    for (const int i : IndexRange(enum_def.items_num)) {
      input_bsockets[i] = bnode.input_sockets()[i + 1];
    }
    const bNodeSocket *output_bsocket = bnode.output_sockets()[0];
 
    lf::OutputSocket *output_is_used_socket = graph_params.usage_by_bsocket.lookup_default(
        output_bsocket, nullptr);
    if (output_is_used_socket == nullptr) {
      return;
    }
    graph_params.usage_by_bsocket.add(switch_input_bsocket, output_is_used_socket);
    if (switch_input_bsocket->is_directly_linked()) {
      /* The condition input is dynamic, so the usage of the other inputs is as well. */
      std::unique_ptr<LazyFunction> lazy_function =
          get_menu_switch_node_socket_usage_lazy_function(bnode);
      lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(*lazy_function);
      scope_.add(std::move(lazy_function));
 
      graph_params.lf_inputs_by_bsocket.add(switch_input_bsocket, &lf_node.input(0));
      for (const int i : IndexRange(enum_def.items_num)) {
        graph_params.usage_by_bsocket.add(input_bsockets[i], &lf_node.output(i));
      }
    }
    else {
      const int condition =
          switch_input_bsocket->default_value_typed<bNodeSocketValueMenu>()->value;
      for (const int i : IndexRange(enum_def.items_num)) {
        const NodeEnumItem &enum_item = enum_def.items()[i];
        if (enum_item.identifier == condition) {
          graph_params.usage_by_bsocket.add(input_bsockets[i], output_is_used_socket);
          break;
        }
      }
    }
  }
 
  void build_undefined_node(const bNode &bnode, BuildGraphParams &graph_params)
  {
    auto &lazy_function = scope_.construct<LazyFunctionForUndefinedNode>(
        bnode, mapping_->lf_index_by_bsocket);
    lf::FunctionNode &lf_node = graph_params.lf_graph.add_function(lazy_function);
 
    for (const bNodeSocket *bsocket : bnode.output_sockets()) {
      const int lf_index = mapping_->lf_index_by_bsocket[bsocket->index_in_tree()];
      if (lf_index == -1) {
        continue;
      }
      lf::OutputSocket &lf_socket = lf_node.output(lf_index);
      graph_params.lf_output_by_bsocket.add(bsocket, &lf_socket);
      mapping_->bsockets_by_lf_socket_map.add(&lf_socket, bsocket);
    }
  }
 
  struct TypeWithLinks {
    const bke::bNodeSocketType *typeinfo;
    Vector<const bNodeLink *> links;
  };
 
  void insert_links_from_socket(const bNodeSocket &from_bsocket,
                                lf::OutputSocket &from_lf_socket,
                                BuildGraphParams &graph_params)
  {
    if (from_bsocket.owner_node().is_dangling_reroute()) {
      return;
    }
 
    const bke::bNodeSocketType &from_typeinfo = *from_bsocket.typeinfo;
 
    /* Group available target sockets by type so that they can be handled together. */
    const Vector<TypeWithLinks> types_with_links = this->group_link_targets_by_type(from_bsocket);
 
    for (const TypeWithLinks &type_with_links : types_with_links) {
      if (type_with_links.typeinfo == nullptr) {
        continue;
      }
      if (type_with_links.typeinfo->geometry_nodes_cpp_type == nullptr) {
        continue;
      }
      const bke::bNodeSocketType &to_typeinfo = *type_with_links.typeinfo;
      const CPPType &to_type = *to_typeinfo.geometry_nodes_cpp_type;
      const Span<const bNodeLink *> links = type_with_links.links;
 
      lf::OutputSocket *converted_from_lf_socket = this->insert_type_conversion_if_necessary(
          from_lf_socket, from_typeinfo, to_typeinfo, graph_params.lf_graph);
 
      for (const bNodeLink *link : links) {
        const Vector<lf::InputSocket *> lf_link_targets = this->find_link_targets(*link,
                                                                                  graph_params);
        if (converted_from_lf_socket == nullptr) {
          const void *default_value = to_type.default_value();
          for (lf::InputSocket *to_lf_socket : lf_link_targets) {
            to_lf_socket->set_default_value(default_value);
          }
        }
        else {
          for (lf::InputSocket *to_lf_socket : lf_link_targets) {
            graph_params.lf_graph.add_link(*converted_from_lf_socket, *to_lf_socket);
          }
        }
      }
    }
  }
 
  Vector<TypeWithLinks> group_link_targets_by_type(const bNodeSocket &from_bsocket)
  {
    const Span<const bNodeLink *> links_from_bsocket = from_bsocket.directly_linked_links();
    Vector<TypeWithLinks> types_with_links;
    for (const bNodeLink *link : links_from_bsocket) {
      if (link->is_muted()) {
        continue;
      }
      if (!link->is_available()) {
        continue;
      }
      const bNodeSocket &to_bsocket = *link->tosock;
      bool inserted = false;
      for (TypeWithLinks &types_with_links : types_with_links) {
        if (types_with_links.typeinfo == to_bsocket.typeinfo) {
          types_with_links.links.append(link);
          inserted = true;
          break;
        }
      }
      if (inserted) {
        continue;
      }
      types_with_links.append({to_bsocket.typeinfo, {link}});
    }
    return types_with_links;
  }
 
  Vector<lf::InputSocket *> find_link_targets(const bNodeLink &link,
                                              const BuildGraphParams &graph_params)
  {
    if (lf::InputSocket *lf_input_socket = graph_params.lf_input_by_border_link.lookup_default(
            &link, nullptr))
    {
      return {lf_input_socket};
    }
 
    const bNodeSocket &to_bsocket = *link.tosock;
    if (to_bsocket.is_multi_input()) {
      /* TODO: Cache this index on the link. */
      int link_index = 0;
      for (const bNodeLink *multi_input_link : to_bsocket.directly_linked_links()) {
        if (multi_input_link == &link) {
          break;
        }
        if (multi_input_link->is_muted() || !multi_input_link->fromsock->is_available() ||
            multi_input_link->fromnode->is_dangling_reroute())
        {
          continue;
        }
        link_index++;
      }
      if (to_bsocket.owner_node().is_muted()) {
        if (link_index == 0) {
          return Vector<lf::InputSocket *>(graph_params.lf_inputs_by_bsocket.lookup(&to_bsocket));
        }
      }
      else {
        lf::InputSocket *lf_multi_input_socket =
            graph_params.lf_input_by_multi_input_link.lookup_default(&link, nullptr);
        if (!lf_multi_input_socket) {
          return {};
        }
        return {lf_multi_input_socket};
      }
    }
    else {
      return Vector<lf::InputSocket *>(graph_params.lf_inputs_by_bsocket.lookup(&to_bsocket));
    }
    return {};
  }
 
  lf::OutputSocket *insert_type_conversion_if_necessary(lf::OutputSocket &from_socket,
                                                        const bke::bNodeSocketType &from_typeinfo,
                                                        const bke::bNodeSocketType &to_typeinfo,
                                                        lf::Graph &lf_graph)
  {
    if (from_typeinfo.type == to_typeinfo.type) {
      return &from_socket;
    }
    if (from_typeinfo.base_cpp_type && to_typeinfo.base_cpp_type) {
      if (conversions_->is_convertible(*from_typeinfo.base_cpp_type, *to_typeinfo.base_cpp_type)) {
        const MultiFunction &multi_fn = *conversions_->get_conversion_multi_function(
            mf::DataType::ForSingle(*from_typeinfo.base_cpp_type),
            mf::DataType::ForSingle(*to_typeinfo.base_cpp_type));
        auto &fn = scope_.construct<LazyFunctionForMultiFunctionConversion>(multi_fn);
        lf::Node &conversion_node = lf_graph.add_function(fn);
        lf_graph.add_link(from_socket, conversion_node.input(0));
        return &conversion_node.output(0);
      }
    }
    return nullptr;
  }
 
  void add_default_inputs(BuildGraphParams &graph_params)
  {
    for (auto item : graph_params.lf_inputs_by_bsocket.items()) {
      const bNodeSocket &bsocket = *item.key;
      const Span<lf::InputSocket *> lf_sockets = item.value;
      for (lf::InputSocket *lf_socket : lf_sockets) {
        if (lf_socket->origin() != nullptr) {
          /* Is linked already. */
          continue;
        }
        this->add_default_input(bsocket, *lf_socket, graph_params);
      }
    }
  }
 
  void add_default_input(const bNodeSocket &input_bsocket,
                         lf::InputSocket &input_lf_socket,
                         BuildGraphParams &graph_params)
  {
    if (this->try_add_implicit_input(input_bsocket, input_lf_socket, graph_params)) {
      return;
    }
    GMutablePointer value = get_socket_default_value(scope_.linear_allocator(), input_bsocket);
    if (value.get() == nullptr) {
      /* Not possible to add a default value. */
      return;
    }
    input_lf_socket.set_default_value(value.get());
    if (!value.type()->is_trivially_destructible()) {
      scope_.add_destruct_call([value]() mutable { value.destruct(); });
    }
  }
 
  bool try_add_implicit_input(const bNodeSocket &input_bsocket,
                              lf::InputSocket &input_lf_socket,
                              BuildGraphParams &graph_params)
  {
    const bNode &bnode = input_bsocket.owner_node();
    const SocketDeclaration *socket_decl = input_bsocket.runtime->declaration;
    if (socket_decl == nullptr) {
      return false;
    }
    if (socket_decl->input_field_type != InputSocketFieldType::Implicit) {
      return false;
    }
    const ImplicitInputValueFn *implicit_input_fn = socket_decl->implicit_input_fn.get();
    if (implicit_input_fn == nullptr) {
      return false;
    }
    std::function<void(void *)> init_fn = [&bnode, implicit_input_fn](void *r_value) {
      (*implicit_input_fn)(bnode, r_value);
    };
    const CPPType &type = input_lf_socket.type();
    auto &lazy_function = scope_.construct<LazyFunctionForImplicitInput>(type, std::move(init_fn));
    lf::Node &lf_node = graph_params.lf_graph.add_function(lazy_function);
    graph_params.lf_graph.add_link(lf_node.output(0), input_lf_socket);
    return true;
  }

  void build_attribute_propagation_input_node(lf::Graph &lf_graph)
  {
    const aal::RelationsInNode &tree_relations =
        btree_.runtime->anonymous_attribute_inferencing->tree_relations;
    Vector<int> output_indices;
    for (const aal::PropagateRelation &relation : tree_relations.propagate_relations) {
      output_indices.append_non_duplicates(relation.to_geometry_output);
    }
 
    for (const int i : output_indices.index_range()) {
      const int output_index = output_indices[i];
      const char *name = btree_.interface_outputs()[output_index]->name;
      lf::GraphInputSocket &lf_socket = lf_graph.add_input(
          CPPType::get<bke::AnonymousAttributeSet>(),
          StringRef("Propagate: ") + (name ? name : ""));
      attribute_set_by_geometry_output_.add(output_index, &lf_socket);
    }
  }

  lf::OutputSocket *or_socket_usages(MutableSpan<lf::OutputSocket *> usages,
                                     BuildGraphParams &graph_params)
  {
    if (usages.is_empty()) {
      return nullptr;
    }
    if (usages.size() == 1) {
      return usages[0];
    }
 
    std::sort(usages.begin(), usages.end());
    return graph_params.socket_usages_combination_cache.lookup_or_add_cb_as(usages, [&]() {
      auto &logical_or_fn = scope_.construct<LazyFunctionForLogicalOr>(usages.size());
      lf::Node &logical_or_node = graph_params.lf_graph.add_function(logical_or_fn);
 
      for (const int i : usages.index_range()) {
        graph_params.lf_graph.add_link(*usages[i], logical_or_node.input(i));
      }
      return &logical_or_node.output(0);
    });
  }
 
  void build_output_socket_usages(const bNode &bnode, BuildGraphParams &graph_params)
  {
    /* Output sockets are used when any of their linked inputs are used. */
    for (const bNodeSocket *socket : bnode.output_sockets()) {
      if (!socket->is_available()) {
        continue;
      }
      /* Determine when linked target sockets are used. */
      Vector<lf::OutputSocket *> target_usages;
      for (const bNodeLink *link : socket->directly_linked_links()) {
        if (!link->is_used()) {
          continue;
        }
        const bNodeSocket &target_socket = *link->tosock;
        if (lf::OutputSocket *is_used_socket = graph_params.usage_by_bsocket.lookup_default(
                &target_socket, nullptr))
        {
          target_usages.append_non_duplicates(is_used_socket);
        }
      }
      /* Combine target socket usages into the usage of the current socket. */
      graph_params.usage_by_bsocket.add(socket,
                                        this->or_socket_usages(target_usages, graph_params));
    }
  }
 
  void build_group_input_usages(BuildGraphParams &graph_params)
  {
    const Span<const bNode *> group_input_nodes = btree_.group_input_nodes();
    for (const int i : btree_.interface_inputs().index_range()) {
      Vector<lf::OutputSocket *> target_usages;
      for (const bNode *group_input_node : group_input_nodes) {
        if (lf::OutputSocket *lf_socket = graph_params.usage_by_bsocket.lookup_default(
                &group_input_node->output_socket(i), nullptr))
        {
          target_usages.append_non_duplicates(lf_socket);
        }
      }
 
      lf::OutputSocket *lf_socket = this->or_socket_usages(target_usages, graph_params);
      lf::InputSocket *lf_group_output = const_cast<lf::InputSocket *>(
          group_input_usage_sockets_[i]);
      InputUsageHint input_usage_hint;
      if (lf_socket == nullptr) {
        static const bool static_false = false;
        lf_group_output->set_default_value(&static_false);
        input_usage_hint.type = InputUsageHintType::Never;
      }
      else {
        graph_params.lf_graph.add_link(*lf_socket, *lf_group_output);
        if (lf_socket->node().is_interface()) {
          /* Can support slightly more complex cases where it depends on more than one output in
           * the future. */
          input_usage_hint.type = InputUsageHintType::DependsOnOutput;
          input_usage_hint.output_dependencies = {
              group_output_used_sockets_.first_index_of(lf_socket)};
        }
        else {
          input_usage_hint.type = InputUsageHintType::DynamicSocket;
        }
      }
      lf_graph_info_->mapping.group_input_usage_hints.append(std::move(input_usage_hint));
    }
  }

  void fix_link_cycles(lf::Graph &lf_graph, const Set<lf::InputSocket *> &socket_usage_inputs)
  {
    lf_graph.update_socket_indices();
    const int sockets_num = lf_graph.socket_num();
 
    struct SocketState {
      bool done = false;
      bool in_stack = false;
    };
 
    Array<SocketState> socket_states(sockets_num);
 
    Vector<lf::Socket *> lf_sockets_to_check;
    for (lf::Node *lf_node : lf_graph.nodes()) {
      if (lf_node->is_function()) {
        for (lf::OutputSocket *lf_socket : lf_node->outputs()) {
          if (lf_socket->targets().is_empty()) {
            lf_sockets_to_check.append(lf_socket);
          }
        }
      }
      if (lf_node->outputs().is_empty()) {
        for (lf::InputSocket *lf_socket : lf_node->inputs()) {
          lf_sockets_to_check.append(lf_socket);
        }
      }
    }
    Vector<lf::Socket *> lf_socket_stack;
    while (!lf_sockets_to_check.is_empty()) {
      lf::Socket *lf_inout_socket = lf_sockets_to_check.last();
      lf::Node &lf_node = lf_inout_socket->node();
      SocketState &state = socket_states[lf_inout_socket->index_in_graph()];
 
      if (!state.in_stack) {
        lf_socket_stack.append(lf_inout_socket);
        state.in_stack = true;
      }
 
      Vector<lf::Socket *, 16> lf_origin_sockets;
      if (lf_inout_socket->is_input()) {
        lf::InputSocket &lf_input_socket = lf_inout_socket->as_input();
        if (lf::OutputSocket *lf_origin_socket = lf_input_socket.origin()) {
          lf_origin_sockets.append(lf_origin_socket);
        }
      }
      else {
        lf::OutputSocket &lf_output_socket = lf_inout_socket->as_output();
        if (lf_node.is_function()) {
          lf::FunctionNode &lf_function_node = static_cast<lf::FunctionNode &>(lf_node);
          const lf::LazyFunction &fn = lf_function_node.function();
          fn.possible_output_dependencies(
              lf_output_socket.index(), [&](const Span<int> input_indices) {
                for (const int input_index : input_indices) {
                  lf_origin_sockets.append(&lf_node.input(input_index));
                }
              });
        }
      }
 
      bool pushed_socket = false;
      bool detected_cycle = false;
      for (lf::Socket *lf_origin_socket : lf_origin_sockets) {
        if (socket_states[lf_origin_socket->index_in_graph()].in_stack) {
          /* A cycle has been detected. The cycle is broken by removing a link and replacing it
           * with a constant "true" input. This can only affect inputs which determine whether a
           * specific value is used. Therefore, setting it to a constant true can result in more
           * computation later, but does not change correctness.
           *
           * After the cycle is broken, the cycle-detection is "rolled back" to the socket where
           * the first socket of the cycle was found. This is necessary in case another cycle
           * goes through this socket. */
 
          detected_cycle = true;
          const int index_in_socket_stack = lf_socket_stack.first_index_of(lf_origin_socket);
          const int index_in_sockets_to_check = lf_sockets_to_check.first_index_of(
              lf_origin_socket);
          const Span<lf::Socket *> cycle = lf_socket_stack.as_span().drop_front(
              index_in_socket_stack);
 
          bool broke_cycle = false;
          for (lf::Socket *lf_cycle_socket : cycle) {
            if (lf_cycle_socket->is_input() &&
                socket_usage_inputs.contains(&lf_cycle_socket->as_input()))
            {
              lf::InputSocket &lf_cycle_input_socket = lf_cycle_socket->as_input();
              lf_graph.clear_origin(lf_cycle_input_socket);
              static const bool static_true = true;
              lf_cycle_input_socket.set_default_value(&static_true);
              broke_cycle = true;
            }
            /* This is actually removed from the stack when it is resized below. */
            SocketState &lf_cycle_socket_state = socket_states[lf_cycle_socket->index_in_graph()];
            lf_cycle_socket_state.in_stack = false;
          }
          if (!broke_cycle) {
            BLI_assert_unreachable();
          }
          /* Roll back algorithm by removing the sockets that corresponded to the cycle from the
           * stacks. */
          lf_socket_stack.resize(index_in_socket_stack);
          /* The +1 is there so that the socket itself is not removed. */
          lf_sockets_to_check.resize(index_in_sockets_to_check + 1);
          break;
        }
        else if (!socket_states[lf_origin_socket->index_in_graph()].done) {
          lf_sockets_to_check.append(lf_origin_socket);
          pushed_socket = true;
        }
      }
      if (detected_cycle) {
        continue;
      }
      if (pushed_socket) {
        continue;
      }
 
      state.done = true;
      state.in_stack = false;
      lf_sockets_to_check.pop_last();
      lf_socket_stack.pop_last();
    }
  }
};
 
const GeometryNodesLazyFunctionGraphInfo *ensure_geometry_nodes_lazy_function_graph(
    const bNodeTree &btree)
{
  btree.ensure_topology_cache();
  btree.ensure_interface_cache();
  if (btree.has_available_link_cycle()) {
    return nullptr;
  }
  const bNodeTreeZones *tree_zones = btree.zones();
  if (tree_zones == nullptr) {
    return nullptr;
  }
  for (const std::unique_ptr<bNodeTreeZone> &zone : tree_zones->zones) {
    if (zone->input_node == nullptr || zone->output_node == nullptr) {
      /* Simulations and repeats need input and output nodes. */
      return nullptr;
    }
  }
  if (const ID *id_orig = DEG_get_original_id(const_cast<ID *>(&btree.id))) {
    if (id_orig->tag & ID_TAG_MISSING) {
      return nullptr;
    }
  }
  for (const bNodeTreeInterfaceSocket *interface_bsocket : btree.interface_inputs()) {
    const bke::bNodeSocketType *typeinfo = interface_bsocket->socket_typeinfo();
    if (typeinfo->geometry_nodes_cpp_type == nullptr) {
      return nullptr;
    }
  }
  for (const bNodeTreeInterfaceSocket *interface_bsocket : btree.interface_outputs()) {
    const bke::bNodeSocketType *typeinfo = interface_bsocket->socket_typeinfo();
    if (typeinfo->geometry_nodes_cpp_type == nullptr) {
      return nullptr;
    }
  }
 
  std::unique_ptr<GeometryNodesLazyFunctionGraphInfo> &lf_graph_info_ptr =
      btree.runtime->geometry_nodes_lazy_function_graph_info;
 
  if (lf_graph_info_ptr) {
    return lf_graph_info_ptr.get();
  }
  std::lock_guard lock{btree.runtime->geometry_nodes_lazy_function_graph_info_mutex};
  if (lf_graph_info_ptr) {
    return lf_graph_info_ptr.get();
  }
 
  auto lf_graph_info = std::make_unique<GeometryNodesLazyFunctionGraphInfo>();
  GeometryNodesLazyFunctionBuilder builder{btree, *lf_graph_info};
  builder.build();
 
  lf_graph_info_ptr = std::move(lf_graph_info);
  return lf_graph_info_ptr.get();
}
 
destruct_ptr<lf::LocalUserData> GeoNodesLFUserData::get_local(LinearAllocator<> &allocator)
{
  return allocator.construct<GeoNodesLFLocalUserData>(*this);
}
 
void GeoNodesLFLocalUserData::ensure_tree_logger(const GeoNodesLFUserData &user_data) const
{
  if (geo_eval_log::GeoModifierLog *log = user_data.call_data->eval_log) {
    tree_logger_.emplace(&log->get_local_tree_logger(*user_data.compute_context));
    return;
  }
  this->tree_logger_.emplace(nullptr);
}
 
std::optional<FoundNestedNodeID> find_nested_node_id(const GeoNodesLFUserData &user_data,
                                                     const int node_id)
{
  FoundNestedNodeID found;
  Vector<int> node_ids;
  for (const ComputeContext *context = user_data.compute_context; context != nullptr;
       context = context->parent())
  {
    if (const auto *node_context = dynamic_cast<const bke::GroupNodeComputeContext *>(context)) {
      node_ids.append(node_context->node_id());
    }
    else if (dynamic_cast<const bke::RepeatZoneComputeContext *>(context) != nullptr) {
      found.is_in_loop = true;
    }
    else if (dynamic_cast<const bke::SimulationZoneComputeContext *>(context) != nullptr) {
      found.is_in_simulation = true;
    }
    else if (dynamic_cast<const bke::ForeachGeometryElementZoneComputeContext *>(context) !=
             nullptr)
    {
      found.is_in_loop = true;
    }
  }
  std::reverse(node_ids.begin(), node_ids.end());
  node_ids.append(node_id);
  const bNestedNodeRef *nested_node_ref =
      user_data.call_data->root_ntree->nested_node_ref_from_node_id_path(node_ids);
  if (nested_node_ref == nullptr) {
    return std::nullopt;
  }
  found.id = nested_node_ref->id;
  return found;
}
 
GeoNodesOperatorDepsgraphs::~GeoNodesOperatorDepsgraphs()
{
  if (Depsgraph *graph = this->extra) {
    DEG_graph_free(graph);
  }
}
 
static const ID *get_only_evaluated_id(const Depsgraph &depsgraph, const ID &id_orig)
{
  const ID *id = DEG_get_evaluated_id(&depsgraph, const_cast<ID *>(&id_orig));
  if (id == &id_orig) {
    return nullptr;
  }
  return id;
}
 
const ID *GeoNodesOperatorDepsgraphs::get_evaluated_id(const ID &id_orig) const
{
  if (const Depsgraph *graph = this->active) {
    if (const ID *id = get_only_evaluated_id(*graph, id_orig)) {
      return id;
    }
  }
  if (const Depsgraph *graph = this->extra) {
    if (const ID *id = get_only_evaluated_id(*graph, id_orig)) {
      return id;
    }
  }
  return nullptr;
}
 
const Object *GeoNodesCallData::self_object() const
{
  if (this->modifier_data) {
    return this->modifier_data->self_object;
  }
  if (this->operator_data) {
    return DEG_get_evaluated_object(this->operator_data->depsgraphs->active,
                                    const_cast<Object *>(this->operator_data->self_object_orig));
  }
  return nullptr;
}
 
}  // namespace blender::nodes