Function ListSortByAssociatedKeysNative

extensions/lists_functions.cc:283–324

template <typename ValueType>
absl::StatusOr<Value> ListSortByAssociatedKeysNative(
    const ListValue& list, const ListValue& keys,
    const google::protobuf::DescriptorPool* absl_nonnull descriptor_pool,
    google::protobuf::MessageFactory* absl_nonnull message_factory,
    google::protobuf::Arena* absl_nonnull arena) {
  CEL_ASSIGN_OR_RETURN(size_t size, list.Size());
  // If the list is empty or has a single element, we can return it as is.
  if (size < 2) {
    return list;
  }
  std::vector<ValueType> keys_vec;
  absl::Status status = keys.ForEach(
      [&keys_vec](const Value& value) -> absl::StatusOr<bool> {
        if (auto typed_value = value.As<ValueType>(); typed_value.has_value()) {
          keys_vec.push_back(*typed_value);
        } else {
          return absl::InvalidArgumentError(
              "sort(): list elements must have the same type");
        }
        return true;
      },
      descriptor_pool, message_factory, arena);
  if (!status.ok()) {
    return ErrorValue(status);
  }
  ABSL_ASSERT(keys_vec.size() == size);  // Already checked by the caller.
  std::vector<int64_t> sorted_indices(keys_vec.size());
  std::iota(sorted_indices.begin(), sorted_indices.end(), 0);
  std::sort(
      sorted_indices.begin(), sorted_indices.end(),
      [&](int64_t a, int64_t b) -> bool { return keys_vec[a] < keys_vec[b]; });

  // Now sorted_indices contains the indices of the keys in sorted order.
  // We can use it to build the sorted list.
  auto builder = NewListValueBuilder(arena);
  for (const auto& index : sorted_indices) {
    CEL_ASSIGN_OR_RETURN(
        Value value, list.Get(index, descriptor_pool, message_factory, arena));
    CEL_RETURN_IF_ERROR(builder->Add(value));
  }
  return std::move(*builder).Build();
}

// Internal function used for the implementation of sort() and sortBy().
//