AdjacencyList.h

// Copyright (C) 2019-2025 Garth N. Wells
//
// This file is part of DOLFINx (https://www.fenicsproject.org)
//
// SPDX-License-Identifier:    LGPL-3.0-or-later
 
#pragma once
 
#include <cassert>
#include <concepts>
#include <cstdint>
#include <format>
#include <iterator>
#include <numeric>
#include <optional>
#include <span>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>
 
namespace dolfinx::graph
{
template <typename LinkData, typename NodeData = std::nullptr_t>
class AdjacencyList
{
public:
  using link_type = LinkData;
  using node_data_type = NodeData;

  explicit AdjacencyList(const std::int32_t n) : _array(n), _offsets(n + 1)
  {
    std::iota(_array.begin(), _array.end(), 0);
    std::iota(_offsets.begin(), _offsets.end(), 0);
  }

  template <typename U, typename V>
    requires std::is_convertible_v<std::remove_cvref_t<U>,
                                   std::vector<LinkData>>
                 and std::is_convertible_v<std::remove_cvref_t<V>,
                                           std::vector<std::int32_t>>
  AdjacencyList(U&& data, V&& offsets)
      : _array(std::forward<U>(data)), _offsets(std::forward<V>(offsets))
  {
    if (_offsets.back() != (std::int32_t)_array.size())
    {
      throw std::runtime_error("Last offset must equal the size of the data "
                               "array.");
    }
  }

  template <typename U, typename V, typename W>
    requires std::is_convertible_v<std::remove_cvref_t<U>,
                                   std::vector<LinkData>>
                 and std::is_convertible_v<std::remove_cvref_t<V>,
                                           std::vector<std::int32_t>>
                 and std::is_convertible_v<std::remove_cvref_t<W>,
                                           std::vector<NodeData>>
  AdjacencyList(U&& data, V&& offsets, W&& node_data)
      : _array(std::forward<U>(data)), _offsets(std::forward<V>(offsets)),
        _node_data(std::forward<W>(node_data))
  {
    if (!_node_data.has_value() or _node_data->size() != _offsets.size() - 1)
    {
      throw std::runtime_error("Node data size must equal the number of "
                               "nodes.");
    }
    _array.reserve(_offsets.back());
    if (_offsets.back() != (std::int32_t)_array.size())
    {
      throw std::runtime_error("Last offset must equal the size of the data "
                               "array.");
    }
  }

  template <typename X>
  explicit AdjacencyList(const std::vector<X>& data)
  {
    // Initialize offsets and compute total size
    _offsets.reserve(data.size() + 1);
    _offsets.push_back(0);
    for (auto& row : data)
      _offsets.push_back(_offsets.back() + row.size());
 
    _array.reserve(_offsets.back());
    for (auto& e : data)
      _array.insert(_array.end(), e.begin(), e.end());
  }

  AdjacencyList(const AdjacencyList& list) = default;

  AdjacencyList(AdjacencyList&& list) = default;

  ~AdjacencyList() = default;

  AdjacencyList& operator=(const AdjacencyList& list) = default;

  AdjacencyList& operator=(AdjacencyList&& list) = default;

  bool operator==(const AdjacencyList& list) const
  {
    return this->_array == list._array and this->_offsets == list._offsets;
  }

  std::int32_t num_nodes() const { return _offsets.size() - 1; }

  int num_links(std::size_t node) const
  {
    assert((node + 1) < _offsets.size());
    return _offsets[node + 1] - _offsets[node];
  }

  std::span<LinkData> links(std::size_t node)
  {
    auto it = std::next(_offsets.begin(), node);
    return std::span<LinkData>(std::next(_array.begin(), *it),
                               std::next(_array.begin(), *(it + 1)));
  }

  std::span<const LinkData> links(std::size_t node) const
  {
    auto it = std::next(_offsets.begin(), node);
    return std::span<const LinkData>(std::next(_array.begin(), *it),
                                     std::next(_array.begin(), *(it + 1)));
  }

  const std::vector<LinkData>& array() const { return _array; }

  std::vector<LinkData>& array() { return _array; }

  const std::vector<std::int32_t>& offsets() const { return _offsets; }

  std::vector<std::int32_t>& offsets() { return _offsets; }

  const std::optional<std::vector<NodeData>>& node_data() const
  {
    return _node_data;
  }

  std::optional<std::vector<NodeData>>& node_data() { return _node_data; }

  std::string str() const
  {
    std::string s
        = std::format("<AdjacencyList> with {} nodes\n", this->num_nodes());
    for (std::size_t e = 0; e < _offsets.size() - 1; ++e)
    {
      std::format_to(std::back_inserter(s), "  {}: [", e);
      for (auto link : this->links(e))
        std::format_to(std::back_inserter(s), "{} ", link);
      s += "]\n";
    }
    return s;
  }
 
private:
  // Connections (links/edges) for all entities stored as a contiguous
  // array
  std::vector<LinkData> _array;
 
  // Position of first connection for each entity (using local index)
  std::vector<std::int32_t> _offsets;
 
  // Node data, where _node_data[i] is the data associated with node `i`
  std::optional<std::vector<NodeData>> _node_data = std::nullopt;
};

template <typename T, typename U>
AdjacencyList(T, U) -> AdjacencyList<typename T::value_type, std::nullptr_t>;

template <typename T, typename U, typename W>
AdjacencyList(T, U, W)
    -> AdjacencyList<typename T::value_type, typename W::value_type>;

template <typename V = std::nullptr_t, typename U>
  requires requires {
    typename std::decay_t<U>::value_type;
    requires std::convertible_to<
        U, std::vector<typename std::decay_t<U>::value_type>>;
  }
AdjacencyList<typename std::decay_t<U>::value_type, V>
regular_adjacency_list(U&& data, int degree)
{
  if (degree == 0 and !data.empty())
  {
    throw std::runtime_error("Degree is zero but data is not empty for "
                             "constant degree AdjacencyList");
  }
 
  if (degree > 0 and data.size() % degree != 0)
  {
    throw std::runtime_error(
        "Incompatible data size and degree for constant degree AdjacencyList");
  }
 
  std::int32_t num_nodes = degree == 0 ? data.size() : data.size() / degree;
  std::vector<std::int32_t> offsets(num_nodes + 1, 0);
  for (std::size_t i = 1; i < offsets.size(); ++i)
    offsets[i] = offsets[i - 1] + degree;
  return AdjacencyList<typename std::decay_t<U>::value_type, V>(
      std::forward<U>(data), std::move(offsets));
}
 
} // namespace dolfinx::graph