DependencyResolver.h

// Copyright 2024, UChicago Argonne, LLC
// All Rights Reserved
// Software Name: NEML2 -- the New Engineering material Model Library, version 2
// By: Argonne National Laboratory
// OPEN SOURCE LICENSE (MIT)
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.
 
#pragma once
 
#include <vector>
#include <map>
#include <algorithm>
 
#include "neml2/base/DependencyDefinition.h"
 
namespace neml2
{
template <typename Node, typename ItemType>
class DependencyResolver
{
public:
  struct Item
  {
    Item(Node * const node, const ItemType & item)
      : parent(node),
        value(item)
    {
    }
 
    Node * const parent;
 
    const ItemType value;
 
    bool operator==(const Item & other) const
    {
      return parent == other.parent && value == other.value;
    }
 
    bool operator!=(const Item & other) const
    {
      return parent != other.parent || value != other.value;
    }
 
    bool operator<(const Item & other) const
    {
      return parent != other.parent ? (parent < other.parent) : (value < other.value);
    }
  };
 
  DependencyResolver() = default;
 
  void add_node(DependencyDefinition<ItemType> *);
 
  void add_additional_outbound_item(const ItemType & item);
 
  void set_priority(DependencyDefinition<ItemType> *, size_t);
 
  void resolve();
 
  const std::vector<Node *> & resolution() const { return _resolution; }
 
  const std::map<Item, std::set<Item>> & item_providers() const { return _item_provider_graph; }
 
  const std::map<Item, std::set<Item>> & item_consumers() const { return _item_consumer_graph; }
 
  const std::map<Node *, std::set<Node *>> & node_providers() const { return _node_provider_graph; }
 
  const std::map<Node *, std::set<Node *>> & node_consumers() const { return _node_consumer_graph; }
 
  const std::set<Node *> & end_nodes() const { return _end_nodes; }
 
  const std::set<Item> & outbound_items() const { return _out_items; }
 
  const std::set<Node *> & start_nodes() const { return _start_nodes; }
 
  const std::set<Item> & inbound_items() const { return _in_items; }
 
  bool & unique_item_provider() { return _unique_item_provider; }
 
  bool & unique_item_consumer() { return _unique_item_consumer; }
 
private:
  void build_graph();
 
  void resolve(Node *);
 
  bool _unique_item_provider = true;
 
  bool _unique_item_consumer = false;
 
  std::set<Node *> _nodes;
 
  std::set<Item> _consumed_items;
 
  std::set<Item> _provided_items;
 
  std::map<Item, std::set<Item>> _item_provider_graph;
 
  std::map<Item, std::set<Item>> _item_consumer_graph;
 
  std::map<Node *, std::set<Node *>> _node_provider_graph;
 
  std::map<Node *, std::set<Node *>> _node_consumer_graph;
 
  std::set<Node *> _end_nodes;
 
  std::set<Node *> _start_nodes;
 
  std::set<Item> _out_items;
 
  std::set<Item> _in_items;
 
  std::vector<Node *> _resolution;
 
  std::map<Node *, int> _status;
 
  std::map<Node *, size_t> _priority;
};
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::add_node(DependencyDefinition<ItemType> * def)
{
  auto node = dynamic_cast<Node *>(def);
  _nodes.emplace(node);
 
  for (const auto & item : node->consumed_items())
    _consumed_items.emplace(node, item);
 
  for (const auto & item : node->provided_items())
    _provided_items.emplace(node, item);
}
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::add_additional_outbound_item(const ItemType & item)
{
  _consumed_items.emplace(nullptr, item);
}
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::set_priority(DependencyDefinition<ItemType> * def,
                                                 size_t priority)
{
  auto node = dynamic_cast<Node *>(def);
  _priority[node] = priority;
}
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::build_graph()
{
  // Clear the previous graph
  _item_provider_graph.clear();
  _item_consumer_graph.clear();
  _node_provider_graph.clear();
  _node_consumer_graph.clear();
  _start_nodes.clear();
  _end_nodes.clear();
  _in_items.clear();
  _out_items.clear();
 
  // Build the adjacency matrix for item providers and node providers
  for (const auto & itemi : _consumed_items)
  {
    std::vector<Item> providers;
 
    for (const auto & itemj : _provided_items)
    {
      // Match consumer with provider
      if (itemi.value != itemj.value)
        continue;
 
      // No self dependency
      if (itemi.parent == itemj.parent)
        continue;
 
      // Enforce priority
      if (_priority[itemi.parent] > _priority[itemj.parent])
        continue;
 
      providers.push_back(itemj);
    }
 
    // If the user asks for unique providers, we should error if multiple providers have been
    // found. Otherwise, just put the first provider into the graph.
    if (!providers.empty())
    {
      if (_unique_item_provider)
        neml_assert(
            providers.size() == 1, "Multiple providers have been found for item ", itemi.value);
      _item_provider_graph[itemi].insert(providers[0]);
      if (itemi.parent)
        _node_provider_graph[itemi.parent].insert(providers[0].parent);
    }
  }
 
  // Build the adjacency matrix for item consumers
  for (const auto & itemi : _provided_items)
  {
    std::vector<Item> consumers;
 
    for (const auto & itemj : _consumed_items)
    {
      // Skip additional outbound item
      if (!itemj.parent)
        continue;
 
      // Match provider with consumer
      if (itemi.value != itemj.value)
        continue;
 
      // No self dependency
      if (itemi.parent == itemj.parent)
        continue;
 
      // Enforce priority
      if (_priority[itemi.parent] < _priority[itemj.parent])
        continue;
 
      consumers.push_back(itemj);
    }
 
    // If the user asks for unique consumers, we should error if multiple consumers have been
    // found. Otherwise, just put the first consumer into the graph.
    if (!consumers.empty())
    {
      if (_unique_item_consumer)
        neml_assert(
            consumers.size() == 1, "Multiple consumers have been found for item ", itemi.value);
      _item_consumer_graph[itemi].insert(consumers[0]);
      _node_consumer_graph[itemi.parent].insert(consumers[0].parent);
    }
  }
 
  // Find start nodes
  for (const auto & node : _nodes)
    if (_node_provider_graph.count(node) == 0)
      _start_nodes.insert(node);
 
  // Find end nodes
  for (const auto & node : _nodes)
    if (_node_consumer_graph.count(node) == 0)
      _end_nodes.insert(node);
 
  // Find inbound items
  for (const auto & item : _consumed_items)
    if (_item_provider_graph.count(item) == 0)
      _in_items.insert(item);
 
  // Find outbound items
  for (const auto & item : _provided_items)
    if (_item_consumer_graph.count(item) == 0)
      _out_items.insert(item);
 
  // Additional outbound items
  for (const auto & item : _consumed_items)
    if (!item.parent)
    {
      neml_assert(_item_provider_graph.count(item),
                  "Unable to find provider of the additional outbound item ",
                  item.value);
      for (const auto & provider : _item_provider_graph[item])
      {
        _out_items.insert(provider);
        _end_nodes.insert(provider.parent);
      }
    }
}
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::resolve()
{
  build_graph();
 
  _status.clear();
  _resolution.clear();
  for (const auto & node : _end_nodes)
    if (!_status[node])
      resolve(node);
 
  // Make sure each node appears in the resolution once and only once
  for (const auto & node : _nodes)
  {
    auto count = std::count(_resolution.begin(), _resolution.end(), node);
    neml_assert(count > 0,
                "Each node must appear in the dependency resolution. Node ",
                node->name(),
                " is missing. This is an internal error -- consider filing a bug report.");
    neml_assert(count == 1,
                "Each node must appear in the dependency resolution once and only once. Node ",
                node->name(),
                " appeared ",
                count,
                " times. This indicates cyclic dependency.");
  }
}
 
template <typename Node, typename ItemType>
void
DependencyResolver<Node, ItemType>::resolve(Node * node)
{
  // Mark the current node as visiting (so that we know there is circular dependency when a
  // "visiting" node is visited again).
  _status[node] += 1;
 
  // Recurse for all the dependent nodes
  if (_node_provider_graph.count(node))
    for (const auto & dep : _node_provider_graph[node])
    {
      // The dependent node must either be "not visited" or "visited".
      // If the dependent node is "being visited", there must be cyclic dependency.
      neml_assert(_status[dep] != 1,
                  "While resolving dependency, two nodes '",
                  node->name(),
                  "' and '",
                  dep->name(),
                  "' have (possibly indirect) cyclic dependency. The cyclic dependency can be "
                  "resolved by explicitly setting the node priorities.");
 
      if (!_status[dep])
        resolve(dep);
    }
 
  // At this point, all the dependent nodes must have been pushed into the resolution. It is
  // therefore safe to push the current node into the resolution.
  _resolution.push_back(node);
 
  // Finished visiting this node
  _status[node] += 1;
}
} // namespace neml2