#include "config.h"
#include <ccan/bitmap/bitmap.h>
#include <ccan/tal/tal.h>
#include <plugins/askrene/child/algorithm.h>
#include <plugins/askrene/child/priorityqueue.h>
static const s64 INFINITE = INT64_MAX;
#define MAX(x, y) (((x) > (y)) ? (x) : (y))
#define MIN(x, y) (((x) < (y)) ? (x) : (y))
bool BFS_path(const tal_t *ctx, const struct graph *graph,
const struct node source, const struct node destination,
const s64 *capacity, const s64 cap_threshold, struct arc *prev)
{
const tal_t *this_ctx = tal(ctx, tal_t);
bool target_found = false;
assert(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
assert(source.idx < max_num_nodes);
assert(capacity);
assert(prev);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(prev) == max_num_nodes);
for (size_t i = 0; i < max_num_nodes; i++)
prev[i].idx = INVALID_INDEX;
u32 *queue = tal_arr(this_ctx, u32, max_num_nodes);
size_t queue_start = 0, queue_end = 0;
queue[queue_end++] = source.idx;
while (queue_start < queue_end) {
struct node cur = {.idx = queue[queue_start++]};
if (cur.idx == destination.idx) {
target_found = true;
break;
}
for (struct arc arc = node_adjacency_begin(graph, cur);
!node_adjacency_end(arc);
arc = node_adjacency_next(graph, arc)) {
if (capacity[arc.idx] < cap_threshold)
continue;
const struct node next = arc_head(graph, arc);
if (prev[next.idx].idx != INVALID_INDEX ||
next.idx == source.idx)
continue;
prev[next.idx] = arc;
assert(queue_end < max_num_nodes);
queue[queue_end++] = next.idx;
}
}
tal_free(this_ctx);
return target_found;
}
bool dijkstra_path(const tal_t *ctx, const struct graph *graph,
const struct node source, const struct node destination,
bool prune, const s64 *capacity, const s64 cap_threshold,
const s64 *cost, const s64 *potential, struct arc *prev,
s64 *distance)
{
bool target_found = false;
assert(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
const tal_t *this_ctx = tal(ctx, tal_t);
assert(source.idx<max_num_nodes);
assert(cost);
assert(capacity);
assert(prev);
assert(distance);
assert(destination.idx < max_num_nodes || !prune);
assert(tal_count(cost) == max_num_arcs);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(prev) == max_num_nodes);
assert(tal_count(distance) == max_num_nodes);
bitmap *visited = tal_arrz(this_ctx, bitmap,
BITMAP_NWORDS(max_num_nodes));
for (size_t i = 0; i < max_num_nodes; ++i)
prev[i].idx = INVALID_INDEX;
struct priorityqueue *q;
q = priorityqueue_new(this_ctx, max_num_nodes);
const s64 *const dijkstra_distance = priorityqueue_value(q);
priorityqueue_init(q);
priorityqueue_update(q, source.idx, 0);
while (!priorityqueue_empty(q)) {
const u32 cur = priorityqueue_top(q);
priorityqueue_pop(q);
if (bitmap_test_bit(visited, cur))
continue;
bitmap_set_bit(visited, cur);
if (cur == destination.idx) {
target_found = true;
if (prune)
break;
}
for (struct arc arc =
node_adjacency_begin(graph, node_obj(cur));
!node_adjacency_end(arc);
arc = node_adjacency_next(graph, arc)) {
if (capacity[arc.idx] < cap_threshold)
continue;
const struct node next = arc_head(graph, arc);
const s64 cij = cost[arc.idx] - potential[cur] +
potential[next.idx];
assert(cij >= 0);
if (dijkstra_distance[next.idx] <=
dijkstra_distance[cur] + cij)
continue;
priorityqueue_update(q, next.idx,
dijkstra_distance[cur] + cij);
prev[next.idx] = arc;
}
}
for (size_t i = 0; i < max_num_nodes; i++)
distance[i] = dijkstra_distance[i];
tal_free(this_ctx);
return target_found;
}
static s64 get_augmenting_flow(const struct graph *graph,
const struct node source,
const struct node target, const s64 *capacity,
const struct arc *prev)
{
const size_t max_num_nodes = graph_max_num_nodes(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
assert(max_num_nodes == tal_count(prev));
assert(max_num_arcs == tal_count(capacity));
int path_length = 0;
s64 flow = INFINITE;
struct node cur = target;
while (cur.idx != source.idx) {
assert(cur.idx < max_num_nodes);
const struct arc arc = prev[cur.idx];
assert(arc.idx < max_num_arcs);
flow = MIN(flow, capacity[arc.idx]);
cur = arc_tail(graph, arc);
path_length++;
if(path_length >= max_num_nodes){
flow = -1;
break;
}
}
assert(flow < INFINITE && flow > 0);
return flow;
}
static void sendflow(const struct graph *graph, const struct arc arc,
const s64 flow, s64 *arc_capacity, s64 *node_balance)
{
const struct arc dual = arc_dual(graph, arc);
arc_capacity[arc.idx] -= flow;
arc_capacity[dual.idx] += flow;
if (node_balance) {
const struct node src = arc_tail(graph, arc),
dst = arc_tail(graph, dual);
node_balance[src.idx] -= flow;
node_balance[dst.idx] += flow;
}
}
static void augment_flow(const struct graph *graph,
const struct node source,
const struct node target,
const struct arc *prev,
s64 *excess,
s64 *capacity,
s64 flow)
{
const size_t max_num_nodes = graph_max_num_nodes(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
assert(max_num_nodes == tal_count(prev));
assert(max_num_arcs == tal_count(capacity));
struct node cur = target;
int path_length = 0;
while (cur.idx != source.idx) {
assert(cur.idx < max_num_nodes);
const struct arc arc = prev[cur.idx];
sendflow(graph, arc, flow, capacity, excess);
cur = arc_tail(graph, arc);
path_length++;
if (path_length >= max_num_nodes)
break;
}
assert(path_length < max_num_nodes);
}
bool simple_feasibleflow(const tal_t *ctx,
const struct graph *graph,
const struct node source,
const struct node destination,
s64 *capacity,
s64 amount)
{
const tal_t *this_ctx = tal(ctx, tal_t);
assert(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
assert(amount > 0);
assert(source.idx < max_num_nodes);
assert(destination.idx < max_num_nodes);
assert(capacity);
assert(tal_count(capacity) == max_num_arcs);
struct arc *prev = tal_arr(this_ctx, struct arc, max_num_nodes);
if (!prev)
goto finish;
while (amount > 0) {
if (!BFS_path(this_ctx, graph, source, destination, capacity, 1,
prev))
goto finish;
s64 delta = get_augmenting_flow(graph, source, destination,
capacity, prev);
delta = MIN(amount, delta);
assert(delta > 0 && delta <= amount);
augment_flow(graph, source, destination, prev, NULL, capacity,
delta);
amount -= delta;
}
finish:
tal_free(this_ctx);
return amount == 0;
}
s64 node_balance(const struct graph *graph,
const struct node node,
const s64 *capacity)
{
s64 balance = 0;
for (struct arc arc = node_adjacency_begin(graph, node);
!node_adjacency_end(arc); arc = node_adjacency_next(graph, arc)) {
struct arc dual = arc_dual(graph, arc);
if (arc_is_dual(graph, arc))
balance += capacity[arc.idx];
else
balance -= capacity[dual.idx];
}
return balance;
}
static s64 reduced_cost(const struct graph *graph, const struct arc arc,
const s64 *cost, const s64 *potential)
{
struct node src = arc_tail(graph, arc);
struct node dst = arc_head(graph, arc);
return cost[arc.idx] - potential[src.idx] + potential[dst.idx];
}
static struct node dijkstra_nearest_sink(const tal_t *ctx,
const struct graph *graph,
const struct node source,
const s64 *node_balance,
const s64 *capacity,
const s64 cap_threshold,
const s64 *cost,
const s64 *potential,
struct arc *prev,
s64 *distance)
{
struct node target = {.idx = INVALID_INDEX};
const tal_t *this_ctx = tal(ctx, tal_t);
assert(graph);
assert(node_balance);
assert(capacity);
assert(cost);
assert(potential);
assert(prev);
assert(distance);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
assert(source.idx < max_num_nodes);
assert(tal_count(node_balance) == max_num_nodes);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(cost) == max_num_arcs);
assert(tal_count(potential) == max_num_nodes);
assert(tal_count(prev) == max_num_nodes);
assert(tal_count(distance) == max_num_nodes);
for (size_t i = 0; i < max_num_arcs; i++) {
if (capacity[i] < cap_threshold)
continue;
struct arc arc = {.idx = i};
struct node tail = arc_tail(graph, arc);
struct node head = arc_head(graph, arc);
s64 red_cost =
cost[i] - potential[tail.idx] + potential[head.idx];
if (red_cost < 0)
goto finish;
}
for (size_t i = 0; i < max_num_nodes; ++i)
prev[i].idx = INVALID_INDEX;
#ifdef ASKRENE_UNITTEST
bitmap *visited =
tal_arrz(this_ctx, bitmap, BITMAP_NWORDS(max_num_nodes));
#endif
struct priorityqueue *q;
q = priorityqueue_new(this_ctx, max_num_nodes);
const s64 *const dijkstra_distance = priorityqueue_value(q);
priorityqueue_init(q);
priorityqueue_update(q, source.idx, 0);
while (!priorityqueue_empty(q)) {
const u32 idx = priorityqueue_top(q);
const struct node cur = {.idx = idx};
priorityqueue_pop(q);
#ifdef ASKRENE_UNITTEST
assert(!bitmap_test_bit(visited, cur.idx));
bitmap_set_bit(visited, cur.idx);
#endif
if (node_balance[cur.idx] < 0) {
target = cur;
break;
}
for (struct arc arc = node_adjacency_begin(graph, cur);
!node_adjacency_end(arc);
arc = node_adjacency_next(graph, arc)) {
if (capacity[arc.idx] < cap_threshold)
continue;
const struct node next = arc_head(graph, arc);
const s64 cij = cost[arc.idx] - potential[cur.idx] +
potential[next.idx];
assert(cij >= 0);
if (dijkstra_distance[next.idx] <=
dijkstra_distance[cur.idx] + cij)
continue;
priorityqueue_update(q, next.idx,
dijkstra_distance[cur.idx] + cij);
prev[next.idx] = arc;
}
}
for (size_t i = 0; i < max_num_nodes; i++)
distance[i] = dijkstra_distance[i];
finish:
tal_free(this_ctx);
return target;
}
bool mcf_refinement(const tal_t *ctx,
const struct graph *graph,
s64 *excess,
s64 *capacity,
const s64 *cost,
s64 *potential)
{
bool solved = false;
const tal_t *this_ctx = tal(ctx, tal_t);
assert(graph);
assert(excess);
assert(capacity);
assert(cost);
assert(potential);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
assert(tal_count(excess) == max_num_nodes);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(cost) == max_num_arcs);
assert(tal_count(potential) == max_num_nodes);
s64 total_excess = 0;
for (u32 i = 0; i < max_num_nodes; i++)
total_excess += excess[i];
if (total_excess)
goto finish;
for (u32 arc_id = 0; arc_id < max_num_arcs; arc_id++) {
struct arc arc = {.idx = arc_id};
if(!arc_enabled(graph, arc))
continue;
const s64 r = capacity[arc.idx];
if (reduced_cost(graph, arc, cost, potential) < 0 && r > 0) {
sendflow(graph, arc, r, capacity, excess);
}
}
struct arc *prev = tal_arr(this_ctx, struct arc, max_num_nodes);
s64 *distance = tal_arrz(this_ctx, s64, max_num_nodes);
if (!prev || !distance)
goto finish;
for (u32 node_id = 0; node_id < max_num_nodes; node_id++) {
struct node src = {.idx = node_id};
while (excess[src.idx] > 0) {
struct node dst = dijkstra_nearest_sink(
this_ctx, graph, src, excess, capacity, 1, cost,
potential, prev, distance);
if (dst.idx >= max_num_nodes)
goto finish;
s64 delta = get_augmenting_flow(graph, src, dst,
capacity, prev);
delta = MIN(excess[src.idx], delta);
delta = MIN(-excess[dst.idx], delta);
assert(delta > 0);
augment_flow(graph, src, dst, prev, excess, capacity,
delta);
for (u32 n = 0; n < max_num_nodes; n++) {
potential[n] -=
MIN(distance[dst.idx], distance[n]);
}
}
}
#ifdef ASKRENE_UNITTEST
for (u32 i = 0; i < max_num_nodes; i++) {
assert(excess[i] == 0);
}
for (u32 i = 0; i < max_num_arcs; i++) {
struct arc arc = {.idx = i};
if(!arc_enabled(graph, arc))
continue;
const s64 cap = capacity[arc.idx];
const s64 rc = reduced_cost(graph, arc, cost, potential);
assert(cap >= 0);
assert(!(rc < 0) || cap == 0);
}
#endif
solved = true;
finish:
tal_free(this_ctx);
return solved;
}
bool simple_mcf(const tal_t *ctx, const struct graph *graph,
const struct node source, const struct node destination,
s64 *capacity, s64 amount, const s64 *cost)
{
const tal_t *this_ctx = tal(ctx, tal_t);
assert(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
const size_t max_num_nodes = graph_max_num_nodes(graph);
assert(amount > 0);
assert(source.idx < max_num_nodes);
assert(destination.idx < max_num_nodes);
assert(capacity);
assert(cost);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(cost) == max_num_arcs);
s64 *potential = tal_arrz(this_ctx, s64, max_num_nodes);
s64 *excess = tal_arrz(this_ctx, s64, max_num_nodes);
excess[source.idx] = amount;
excess[destination.idx] = -amount;
if (!mcf_refinement(this_ctx, graph, excess, capacity, cost, potential))
goto fail;
tal_free(this_ctx);
return true;
fail:
tal_free(this_ctx);
return false;
}
s64 flow_cost(const struct graph *graph, const s64 *capacity, const s64 *cost)
{
assert(graph);
const size_t max_num_arcs = graph_max_num_arcs(graph);
s64 total_cost = 0;
assert(capacity);
assert(cost);
assert(tal_count(capacity) == max_num_arcs);
assert(tal_count(cost) == max_num_arcs);
for (u32 i = 0; i < max_num_arcs; i++) {
struct arc arc = {.idx = i};
struct arc dual = arc_dual(graph, arc);
if (arc_is_dual(graph, arc))
continue;
total_cost += capacity[dual.idx] * cost[arc.idx];
}
return total_cost;
}