From 3e070c1f4cd4d444adb43731b455d508c5e964cd Mon Sep 17 00:00:00 2001
From: Gonzalo Brito Gadeschi <gonzalob@nvidia.com>
Date: Mon, 12 Aug 2024 10:45:43 -0700
Subject: [PATCH] refactor more allocator apis

---
 src/all_pairs.h | 2 ++
 src/alloc.h     | 4 ++++
 src/bvh.h       | 6 +++---
 src/system.h    | 5 ++---
 4 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/src/all_pairs.h b/src/all_pairs.h
index 68d7cfa..fac5a9a 100644
--- a/src/all_pairs.h
+++ b/src/all_pairs.h
@@ -38,6 +38,8 @@ void all_pairs_collapsed_force(System<T, N>& system) {
       return;
     }
 
+    // TODO: we should exploit the symmetry of the force pairs to do
+    // (N^2)/2 computations here by taking this "a" and doing "s.a[j] -= a / mj * mi".
     auto pi = s.x[i];
     auto pj = s.x[j];
     auto a  = s.c * s.m[j] * (pj - pi) / dist3(pi, pj);
diff --git a/src/alloc.h b/src/alloc.h
index 1469047..9ab4602 100644
--- a/src/alloc.h
+++ b/src/alloc.h
@@ -2,6 +2,7 @@
 #include <cstdint>
 #include <cstdlib>
 #include <memory>
+#include <vector>
 
 template <typename T>
 void alloc(T*& ptr, std::size_t n) {
@@ -50,3 +51,6 @@ template <typename T, typename U>
 constexpr bool operator!=(allocator<T> const &, allocator<U> const &) noexcept {
   return false;
 }
+
+template <typename T>
+using vector = std::vector<T, allocator<T>>;
diff --git a/src/bvh.h b/src/bvh.h
index 1f87632..2ee8186 100644
--- a/src/bvh.h
+++ b/src/bvh.h
@@ -35,7 +35,7 @@ void hilbert_sort(System<T, N>& system, aabb<T, N> bbox) {
 
   // Compute the Hilbert index for each body in the Cartesian grid
   auto bids = system.body_indices();
-  static std::vector<uint64_t> hilbert_ids(system.size);
+  static vector<uint64_t> hilbert_ids(system.size);
   std::for_each(par_unseq, bids.begin(), bids.end(),
                 [hids = hilbert_ids.data(), x = system.x.data(), mins = bbox.xmin, grid_cell_size](auto idx) {
                   // Bucket the body into a Cartesian grid cell:
@@ -59,7 +59,7 @@ void hilbert_sort(System<T, N>& system, aabb<T, N> bbox) {
   // TODO: sort an array of keys and then apply a permutation in O(N) time and O(1) storage
   // (instead of O(N) time and O(N) storage).
 
-  static std::vector<std::pair<uint64_t, std::size_t>> hilbert_index_map(system.size);
+  static vector<std::pair<uint64_t, std::size_t>> hilbert_index_map(system.size);
   std::for_each_n(par_unseq, counting_iterator<std::size_t>(0), system.size,
                   [hmap = hilbert_index_map.data(), hids = hilbert_ids.data()](std::size_t idx) {
                     hmap[idx] = std::make_pair(hids[idx], idx);
@@ -70,7 +70,7 @@ void hilbert_sort(System<T, N>& system, aabb<T, N> bbox) {
 
   // create temp copy of system so that we don't get race conditions when
   // rearranging values in the next step
-  static std::vector<std::tuple<vec<T, N>, T, vec<T, N>, vec<T, N>, vec<T, N>>> tmp_system(system.size);
+  static vector<std::tuple<vec<T, N>, T, vec<T, N>, vec<T, N>, vec<T, N>>> tmp_system(system.size);
   std::for_each_n(par_unseq, counting_iterator<std::size_t>(0), system.size,
                   [tmp_sys = tmp_system.data(), x = system.x.data(), m = system.m.data(), v = system.v.data(),
                    a = system.a.data(), ao = system.ao.data()](std::size_t idx) {
diff --git a/src/system.h b/src/system.h
index f432d1b..5e08f30 100644
--- a/src/system.h
+++ b/src/system.h
@@ -2,7 +2,6 @@
 #include <cmath>
 #include <functional>
 #include <random>
-#include <vector>
 
 #include "alloc.h"
 #include "execution.h"
@@ -16,8 +15,8 @@ class System {
   index_t const max_tree_node_size;
   T const dt;
   T const constant;
-  std::vector<T, allocator<T>> m;
-  std::vector<vec<T, N>, allocator<vec<T, N>>> x, v, a, ao;
+  vector<T> m;
+  vector<vec<T, N>> x, v, a, ao;
 
   // random generation
   std::mt19937 gen{42};  // fix random generation