//===- RootOrdering.cpp - Optimal root ordering ---------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// An implementation of Edmonds' optimal branching algorithm. This is a
// directed analogue of the minimum spanning tree problem for a given root.
//
//===----------------------------------------------------------------------===//

#include "RootOrdering.h"

#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SmallVector.h"
#include <queue>
#include <utility>

using namespace mlir;
using namespace mlir::pdl_to_pdl_interp;

/// Returns the cycle implied by the specified parent relation, starting at the
/// given node.
static SmallVector<Value> getCycle(const DenseMap<Value, Value> &parents,
                                   Value rep) {
  SmallVector<Value> cycle;
  Value node = rep;
  do {
    cycle.push_back(node);
    node = parents.lookup(node);
    assert(node && "got an empty value in the cycle");
  } while (node != rep);
  return cycle;
}

/// Contracts the specified cycle in the given graph in-place.
/// The parentsCost map specifies, for each node in the cycle, the lowest cost
/// among the edges entering that node. Then, the nodes in the cycle C are
/// replaced with a single node v_C (the first node in the cycle). All edges
/// (u, v) entering the cycle, v \in C, are replaced with a single edge
/// (u, v_C) with an appropriately chosen cost, and the selected node v is
/// marked in the output map actualTarget[u]. All edges (u, v) leaving the
/// cycle, u \in C, are replaced with a single edge (v_C, v), and the selected
/// node u is marked in the ouptut map actualSource[v].
static void contract(RootOrderingGraph &graph, ArrayRef<Value> cycle,
                     const DenseMap<Value, unsigned> &parentCosts,
                     DenseMap<Value, Value> &actualSource,
                     DenseMap<Value, Value> &actualTarget) {
  Value rep = cycle.front();
  DenseSet<Value> cycleSet(cycle.begin(), cycle.end());

  // Now, contract the cycle, marking the actual sources and targets.
  DenseMap<Value, RootOrderingCost> repCosts;
  for (auto outer = graph.begin(), e = graph.end(); outer != e; ++outer) {
    Value target = outer->first;
    if (cycleSet.contains(target)) {
      // Target in the cycle => edges incoming to the cycle or within the cycle.
      unsigned parentCost = parentCosts.lookup(target);
      for (const auto &inner : outer->second) {
        Value source = inner.first;
        // Ignore edges within the cycle.
        if (cycleSet.contains(source))
          continue;

        // Edge incoming to the cycle.
        std::pair<unsigned, unsigned> cost = inner.second.cost;
        assert(parentCost <= cost.first && "invalid parent cost");

        // Subtract the cost of the parent within the cycle from the cost of
        // the edge incoming to the cycle. This update ensures that the cost
        // of the minimum-weight spanning arborescence of the entire graph is
        // the cost of arborescence for the contracted graph plus the cost of
        // the cycle, no matter which edge in the cycle we choose to drop.
        cost.first -= parentCost;
        auto it = repCosts.find(source);
        if (it == repCosts.end() || it->second.cost > cost) {
          actualTarget[source] = target;
          // Do not bother populating the connector (the connector is only
          // relevant for the final traversal, not for the optimal branching).
          repCosts[source].cost = cost;
        }
      }
      // Erase the node in the cycle.
      graph.erase(outer);
    } else {
      // Target not in cycle => edges going away from or unrelated to the cycle.
      DenseMap<Value, RootOrderingCost> &costs = outer->second;
      Value bestSource;
      std::pair<unsigned, unsigned> bestCost;
      auto inner = costs.begin(), inner_e = costs.end();
      while (inner != inner_e) {
        Value source = inner->first;
        if (cycleSet.contains(source)) {
          // Going-away edge => get its cost and erase it.
          if (!bestSource || bestCost > inner->second.cost) {
            bestSource = source;
            bestCost = inner->second.cost;
          }
          costs.erase(inner++);
        } else {
          ++inner;
        }
      }

      // There were going-away edges, contract them.
      if (bestSource) {
        costs[rep].cost = bestCost;
        actualSource[target] = bestSource;
      }
    }
  }

  // Store the edges to the representative.
  graph[rep] = std::move(repCosts);
}

OptimalBranching::OptimalBranching(RootOrderingGraph graph, Value root)
    : graph(std::move(graph)), root(root) {}

unsigned OptimalBranching::solve() {
  // Initialize the parents and total cost.
  parents.clear();
  parents[root] = Value();
  unsigned totalCost = 0;

  // A map that stores the cost of the optimal local choice for each node
  // in a directed cycle. This map is cleared every time we seed the search.
  DenseMap<Value, unsigned> parentCosts;
  parentCosts.reserve(graph.size());

  // Determine if the optimal local choice results in an acyclic graph. This is
  // done by computing the optimal local choice and traversing up the computed
  // parents. On success, `parents` will contain the parent of each node.
  for (const auto &outer : graph) {
    Value node = outer.first;
    if (parents.count(node)) // already visited
      continue;

    // Follow the trail of best sources until we reach an already visited node.
    // The code will assert if we cannot reach an already visited node, i.e.,
    // the graph is not strongly connected.
    parentCosts.clear();
    do {
      auto it = graph.find(node);
      assert(it != graph.end() && "the graph is not strongly connected");

      Value &bestSource = parents[node];
      unsigned &bestCost = parentCosts[node];
      for (const auto &inner : it->second) {
        const RootOrderingCost &cost = inner.second;
        if (!bestSource /* initial */ || bestCost > cost.cost.first) {
          bestSource = inner.first;
          bestCost = cost.cost.first;
        }
      }
      assert(bestSource && "the graph is not strongly connected");
      node = bestSource;
      totalCost += bestCost;
    } while (!parents.count(node));

    // If we reached a non-root node, we have a cycle.
    if (parentCosts.count(node)) {
      // Determine the cycle starting at the representative node.
      SmallVector<Value> cycle = getCycle(parents, node);

      // The following maps disambiguate the source / target of the edges
      // going out of / into the cycle.
      DenseMap<Value, Value> actualSource, actualTarget;

      // Contract the cycle and recurse.
      contract(graph, cycle, parentCosts, actualSource, actualTarget);
      totalCost = solve();

      // Redirect the going-away edges.
      for (auto &p : parents)
        if (p.second == node)
          // The parent is the node representating the cycle; replace it
          // with the actual (best) source in the cycle.
          p.second = actualSource.lookup(p.first);

      // Redirect the unique incoming edge and copy the cycle.
      Value parent = parents.lookup(node);
      Value entry = actualTarget.lookup(parent);
      cycle.push_back(node); // complete the cycle
      for (size_t i = 0, e = cycle.size() - 1; i < e; ++i) {
        totalCost += parentCosts.lookup(cycle[i]);
        if (cycle[i] == entry)
          parents[cycle[i]] = parent; // break the cycle
        else
          parents[cycle[i]] = cycle[i + 1];
      }

      // `parents` has a complete solution.
      break;
    }
  }

  return totalCost;
}

OptimalBranching::EdgeList
OptimalBranching::preOrderTraversal(ArrayRef<Value> nodes) const {
  // Invert the parent mapping.
  DenseMap<Value, std::vector<Value>> children;
  for (Value node : nodes) {
    if (node != root) {
      Value parent = parents.lookup(node);
      assert(parent && "invalid parent");
      children[parent].push_back(node);
    }
  }

  // The result which simultaneously acts as a queue.
  EdgeList result;
  result.reserve(nodes.size());
  result.emplace_back(root, Value());

  // Perform a BFS, pushing into the queue.
  for (size_t i = 0; i < result.size(); ++i) {
    Value node = result[i].first;
    for (Value child : children[node])
      result.emplace_back(child, node);
  }

  return result;
}