astar.h

#pragma once
#include "graph.h"
#include "hashbag.h"
#include "parlay/internal/get_time.h"
using namespace std;
using namespace parlay;

constexpr size_t LOCAL_QUEUE_SIZE = 4096;
constexpr size_t DEG_THLD = 20;
constexpr size_t NUM_SAMPLES = 1000;

enum Algorithm { rho_stepping = 0, delta_stepping, bellman_ford };

class ASTAR {
 protected:
  const Graph &G;
  const sequence<double> &coords;
  const bool bidirection;
  size_t n;
  size_t dimension;

  bool sparse;
  const int sd_scale;
  size_t frontier_size;
  hashbag<NodeId> bag;
  sequence<EdgeTy> dist;
  sequence<NodeId> frontier;
  sequence<atomic<bool>> in_frontier;
  sequence<atomic<bool>> in_next_frontier;
  //sequence<double> heuristic_seq;
  NodeId source;
  NodeId destination;
  EdgeTy gth;

  inline NodeId Node(NodeId u) { return u << bidirection; }
  inline NodeId get_vertex(NodeId u) { return u >> bidirection; }
  inline NodeId get_source(NodeId u) { return bidirection && (u & 1); }

#ifdef OSM
  inline double toRadians(const double &degree) {
    double one_deg = (M_PI) / 180;
    return (one_deg * degree);
  }

  inline double distance(double lat1, double long1, double lat2, double long2) {
    lat1 = toRadians(lat1);
    long1 = toRadians(long1);
    lat2 = toRadians(lat2);
    long2 = toRadians(long2);
    double dlong = long2 - long1;
    double dlat = lat2 - lat1;
    double ans =
        pow(sin(dlat / 2), 2) + cos(lat1) * cos(lat2) * pow(sin(dlong / 2), 2);
    ans = 2 * asin(sqrt(ans));
    double R = 6371;
    ans = ans * R * 1000;
    return ans;
  }

  inline EdgeTy heuristic_func(NodeId u, NodeId v) {
    return distance(coords[u * 2 + 1], coords[u * 2], coords[v * 2 + 1],
                    coords[v * 2]);
  }
#else
  inline EdgeTy heuristic_func(NodeId u, NodeId v) {
    double diff = 0;
    for (size_t i = 0; i < dimension; ++i) {
      diff += (coords[u * dimension + i] - coords[v * dimension + i]) *
              (coords[u * dimension + i] - coords[v * dimension + i]);
    }
    return sqrt(diff);
  }
#endif

  inline EdgeTy heuristic(NodeId u, NodeId v) {
    NodeId s = source;
    NodeId t = destination;
    if (v == t) {
      return heuristic_func(u, t);
      //if (heuristic_seq[u] == DIST_MAX) {
        //heuristic_seq[u] = heuristic_func(u, t);
      //}
      //return heuristic_seq[u];
    } else {
      return heuristic_func(u, s);
      //if (heuristic_seq[u + G.n] == DIST_MAX) {
        //heuristic_seq[u + G.n] = heuristic_func(u, s);
      //}
      //return heuristic_seq[u + G.n];
    }
  }

  inline EdgeTy estimate_distance(NodeId u) {
    if (bidirection) {
      NodeId s = source;
      NodeId t = destination;
      NodeId v = get_vertex(u);
      if (get_source(u) == false) {
        return dist[u] + (heuristic(v, t) - heuristic(v, s)) / 2;
      } else {
        return dist[u] + (heuristic(v, s) - heuristic(v, t)) / 2;
      }
    } else {
      return dist[u] + heuristic(u, destination);
    }
  }

  bool get_processing_status() {
    if (gth != DIST_MAX) {
      return true;
    }
    NodeId processing_set = 0;
    if (sparse) {
      auto state = delayed_seq<NodeId>(frontier_size, [&](size_t i) {
        return static_cast<NodeId>(1 << get_source(frontier[i]));
      });
      processing_set = reduce(state, parlay::bit_or<NodeId>());
    } else {
      auto state = delayed_seq<NodeId>(n, [&](size_t i) {
        if (!in_frontier[i]) {
          return static_cast<NodeId>(0);
        }
        return static_cast<NodeId>(1 << get_source(i));
      });
      processing_set = reduce(state, parlay::bit_or<NodeId>());
    }
    return processing_set == 3;
  }

  void update_global_threshold(NodeId v) {
    if (bidirection) {
      if (dist[v | 1] != DIST_MAX && dist[v & ~1] != DIST_MAX) {
        write_min(&gth, dist[v | 1] + dist[v & ~1]);
      }
    } else if (v == destination) {
      write_min(&gth, dist[v]);
    }
  }

  bool prune(NodeId v) {
    if (bidirection) {
      return estimate_distance(v) * 2 >= gth;
    } else {
      return estimate_distance(v) >= gth;
    }
  }

  void add_to_frontier(NodeId v) {
    if (prune(v)) {
      return;
    }
    if (sparse) {
      if (!in_frontier[v] &&
          compare_and_swap(&in_next_frontier[v], false, true)) {
        bag.insert(v);
      }
    } else {  // dense
      if (!in_frontier[v] && !in_next_frontier[v]) {
        in_next_frontier[v] = true;
      }
    }
  }

  size_t estimate_size() {
    static uint32_t seed = 10086;
    size_t hits = 0;
    for (size_t i = 0; i < NUM_SAMPLES; i++) {
      NodeId u = hash32(seed) % n;
      if (in_frontier[u]) {
        hits++;
      }
      seed++;
    }
    return hits * n / NUM_SAMPLES;
  }

  size_t sparse_relax() {
    bool super_sparse = true;

    EdgeTy th = get_threshold();

    parallel_for(0, frontier_size, [&](size_t i) {
      NodeId f = frontier[i];
      in_frontier[f] = false;
      if (estimate_distance(f) > th) {
        add_to_frontier(f);
      } else {
        size_t _n = G.offset[get_vertex(f) + 1] - G.offset[get_vertex(f)];
        if (super_sparse && _n < LOCAL_QUEUE_SIZE) {
          NodeId local_queue[LOCAL_QUEUE_SIZE];
          size_t front = 0, rear = 0;
          local_queue[rear++] = f;
          while (front < rear && rear < LOCAL_QUEUE_SIZE) {
            NodeId u = local_queue[front++];
            size_t deg = G.offset[get_vertex(u) + 1] - G.offset[get_vertex(u)];
            if (deg >= LOCAL_QUEUE_SIZE || estimate_distance(u) > th) {
              add_to_frontier(u);
              continue;
            }
            if (G.symmetrized) {
              EdgeTy temp_dist = dist[u];
              for (EdgeId es = G.offset[get_vertex(u)];
                   es < G.offset[get_vertex(u) + 1]; es++) {
                NodeId v = G.edge[es].v;
                v = Node(v) + get_source(u);
                EdgeTy w = G.edge[es].w;
                if (dist[v] != DIST_MAX) {
                  temp_dist = min(temp_dist, dist[v] + w);
                }
              }
              if (write_min(&dist[u], temp_dist)) {
                update_global_threshold(u);
              }
            }
            for (EdgeId es = G.offset[get_vertex(u)];
                 es < G.offset[get_vertex(u) + 1]; es++) {
              NodeId v = G.edge[es].v;
              v = Node(v) + get_source(u);
              EdgeTy w = G.edge[es].w;
              if (write_min(&dist[v], dist[u] + w)) {
                update_global_threshold(v);
                if (rear < LOCAL_QUEUE_SIZE) {
                  if (prune(v)) {
                    continue;
                  }
                  local_queue[rear++] = v;
                } else {
                  add_to_frontier(v);
                }
              }
            }
          }
          for (size_t j = front; j < rear; j++) {
            add_to_frontier(local_queue[j]);
          }
        } else {
          blocked_for(G.offset[get_vertex(f)], G.offset[get_vertex(f) + 1],
                      BLOCK_SIZE, [&](size_t, size_t _s, size_t _e) {
                        if (G.symmetrized) {
                          EdgeTy temp_dist = dist[f];
                          for (EdgeId es = _s; es < _e; es++) {
                            NodeId v = G.edge[es].v;
                            v = Node(v) + get_source(f);
                            EdgeTy w = G.edge[es].w;
                            if (dist[v] != DIST_MAX) {
                              temp_dist = min(temp_dist, dist[v] + w);
                            }
                          }
                          if (write_min(&dist[f], temp_dist)) {
                            update_global_threshold(f);
                            add_to_frontier(f);
                          }
                        }
                        for (EdgeId es = _s; es < _e; es++) {
                          NodeId v = G.edge[es].v;
                          v = Node(v) + get_source(f);
                          EdgeTy w = G.edge[es].w;
                          if (write_min(&dist[v], dist[f] + w)) {
                            update_global_threshold(v);
                            add_to_frontier(v);
                          }
                        }
                      });
        }
      }
    });
    swap(in_frontier, in_next_frontier);
    return bag.pack_into(make_slice(frontier));
  }

  size_t dense_relax() {
    while (estimate_size() >= n / sd_scale) {
      EdgeTy th = get_threshold();
      parallel_for(0, n, [&](NodeId u) {
        if (in_frontier[u]) {
          in_frontier[u] = false;
          if (estimate_distance(u) > th) {
            in_next_frontier[u] = true;
          } else {
            blocked_for(G.offset[get_vertex(u)], G.offset[get_vertex(u) + 1],
                        BLOCK_SIZE, [&](size_t, size_t _s, size_t _e) {
                          if (G.symmetrized) {
                            EdgeTy temp_dist = dist[u];
                            for (size_t es = _s; es < _e; es++) {
                              NodeId v = G.edge[es].v;
                              v = Node(v) + get_source(u);
                              EdgeTy w = G.edge[es].w;
                              if (dist[v] != DIST_MAX) {
                                temp_dist = min(temp_dist, dist[v] + w);
                              }
                            }
                            if (write_min(&dist[u], temp_dist)) {
                              update_global_threshold(u);
                              add_to_frontier(u);
                            }
                          }
                          for (size_t es = _s; es < _e; es++) {
                            NodeId v = G.edge[es].v;
                            v = Node(v) + get_source(u);
                            EdgeTy w = G.edge[es].w;
                            if (write_min(&dist[v], dist[u] + w)) {
                              update_global_threshold(v);
                              add_to_frontier(v);
                            }
                          }
                        });
          }
        }
      });
      swap(in_frontier, in_next_frontier);
    }
    return count(in_frontier, true);
  }

  void sparse2dense() {
    // parallel_for(0, frontier_size, [&](size_t i) {
    // NodeId u = frontier[i];
    // assert(in_frontier[u] == true);
    //  in_frontier[u] = true;
    //});
  }

  void dense2sparse() {
    auto identity = delayed_seq<NodeId>(n, [&](NodeId i) { return i; });
    pack_into_uninitialized(identity, in_frontier, frontier);
  }

  virtual void init() = 0;
  virtual EdgeTy get_threshold() = 0;

 public:
  ASTAR() = delete;
  ASTAR(const Graph &_G, const sequence<double> &_coords, bool _bidirection)
      : G(_G),
        coords(_coords),
        bidirection(_bidirection),
        n(bidirection ? 2 * G.n : G.n),
        sd_scale(G.m / G.n),
        bag(n) {
    dist = sequence<EdgeTy>::uninitialized(n);
    frontier = sequence<NodeId>::uninitialized(n);
    in_frontier = sequence<atomic<bool>>::uninitialized(n);
    in_next_frontier = sequence<atomic<bool>>::uninitialized(n);
    //heuristic_seq = sequence<double>::uninitialized(n);
    if (!coords.empty()) {
      dimension = coords.size() / G.n;
    }
  }

  EdgeTy astar(NodeId s, NodeId t) {
    if (!G.weighted) {
      fprintf(stderr, "Error: Input graph is unweighted\n");
      exit(EXIT_FAILURE);
    }
    gth = DIST_MAX;
    source = s;
    destination = t;

    init();
    // parallel_for(0, n, [&](NodeId i) {
    //   dist[i] = DIST_MAX;
    //   in_frontier[i] = in_next_frontier[i] = false;
    // });

    frontier_size = 0;
    if (bidirection) {
      NodeId s_d = Node(s), t_d = Node(t) + 1;
      dist[s_d] = dist[t_d] = 0;
      in_frontier[s_d] = in_frontier[t_d] = true;
      frontier[frontier_size++] = s_d;
      frontier[frontier_size++] = t_d;
    } else {
      dist[s] = 0;
      in_frontier[s] = true;
      frontier[frontier_size++] = s;
    }
    sparse = true;

    while (frontier_size) {
      if (bidirection) {
        if (!get_processing_status()) {
          break;
        }
      }
      if (sparse) {
        frontier_size = sparse_relax();
      } else {
        frontier_size = dense_relax();
      }
      bool next_sparse = (frontier_size < n / sd_scale) ? true : false;
      if (sparse && !next_sparse) {
        sparse2dense();
      } else if (!sparse && next_sparse) {
        dense2sparse();
      }
      sparse = next_sparse;
    }
    return gth;
  }

  void reset() {
    parallel_for(0, n, [&](NodeId i) {
      dist[i] = DIST_MAX;
      //heuristic_seq[i] = DIST_MAX;
      in_frontier[i] = in_next_frontier[i] = false;
    });
  }
};

class Rho_Stepping : public ASTAR {
  size_t rho;
  uint32_t seed;

 public:
  Rho_Stepping(const Graph &_G, const sequence<double> &_coords,
               bool _bidirection, size_t _rho = 1 << 20)
      : ASTAR(_G, _coords, _bidirection), rho(_rho) {
    seed = 0;
  }
  void init() override {}
  EdgeTy get_threshold() override {
    if (frontier_size <= rho) {
      if (sparse) {
        auto _dist = delayed_seq<EdgeTy>(frontier_size, [&](size_t i) {
          return estimate_distance(frontier[i]);
        });
        return *max_element(_dist);
      } else {
        return DIST_MAX;
      }
    }
    EdgeTy sample_dist[NUM_SAMPLES + 1];
    for (size_t i = 0; i <= NUM_SAMPLES; i++) {
      if (sparse) {
        NodeId v = frontier[hash32(seed + i) % frontier_size];
        sample_dist[i] = estimate_distance(v);
      } else {
        NodeId v = hash32(seed + i) % n;
        if (in_frontier[v]) {
          sample_dist[i] = estimate_distance(v);
        } else {
          sample_dist[i] = DIST_MAX;
        }
      }
    }
    seed += NUM_SAMPLES + 1;
    size_t id = 1.0 * rho / frontier_size * NUM_SAMPLES;
    sort(sample_dist, sample_dist + NUM_SAMPLES + 1);
    return sample_dist[id];
  }
};

class Delta_Stepping : public ASTAR {
  EdgeTy delta;
  EdgeTy thres;

 public:
  Delta_Stepping(const Graph &_G, const sequence<double> &_coords,
                 bool _bidirection, EdgeTy _delta = 1 << 15)
      : ASTAR(_G, _coords, _bidirection), delta(_delta) {}
  void init() override { thres = 0; }
  EdgeTy get_threshold() override {
    thres += delta;
    return thres;
  }
};

class Bellman_Ford : public ASTAR {
 public:
  Bellman_Ford(const Graph &_G, const sequence<double> &_coords,
               bool _bidirection)
      : ASTAR(_G, _coords, _bidirection) {}
  void init() override {}
  EdgeTy get_threshold() override { return DIST_MAX; }
};