Node Graph Database

node_graph_db.c

#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include <lmdb.h>

#ifdef _MSC_VER
  #include <intrin.h>
  #pragma intrinsic(_BitScanReverse)
#endif

/* --- Configuration --- */
#define MAX_LVLS        16          
#define MIN_CELL_SHIFT  8           // lvl 0 Cell = 256 width
#define MIN_CELL_SIZE   (1 << MIN_CELL_SHIFT)
#define DB_MAP_SIZE     (2ULL * 1024 * 1024 * 1024) // 2GB Map
#define BATCH_LIMIT     1000        
#define DB_QRY_LIMIT    (8*1024)   // Max items to return in a query
#define COORD_OFFSET    (2147483648ULL) 
/* LOD THRESHOLD:
 * If a viewport spans more than this many cells of a specific lvl,
 * that lvl is considered "too detailed" (sub-pixel) and is skipped.
 * Lower = More aggressive optimization (objects disappear sooner).
 * Higher = More detail preserved when zoomed out.
 */
#define LOD_TILE_LIMIT  64

typedef __uint128_t objid_t; 
struct spatial_row {
  objid_t id;          // 16 bytes
  int32_t minx, maxx;  // 8 bytes
  int32_t miny, maxy;  // 8 bytes
}; // Total 32 bytes
_Static_assert(sizeof(struct spatial_row) == 32, "struct spatial_row must be exactly 32 bytes");

struct index_row {
  uint64_t spatial_key; 
  struct spatial_row row;    
};
struct heap_item {
  objid_t id;
  double dist_sq;
};
struct db_ctx {
  MDB_env *env;
  MDB_dbi spatial_dbi;
  MDB_dbi index_dbi;
  MDB_txn *batch_txn;
  int batch_count;
  struct heap_item qry_heap[DB_QRY_LIMIT]; 
};
static inline uint64_t
morton_encode(uint32_t x, uint32_t y) {
#if defined(ARCH_X86) && defined(__BMI2__)
  return _pdep_u64(x, 0x5555555555555555ULL) | 
         (_pdep_u64(y, 0x5555555555555555ULL) << 1);
#else
  uint64_t sx = x, sy = y;
  sx = (sx | (sx << 16)) & 0x0000FFFF0000FFFFULL;
  sx = (sx | (sx << 8))  & 0x00FF00FF00FF00FFULL;
  sx = (sx | (sx << 4))  & 0x0F0F0F0F0F0F0F0FULL;
  sx = (sx | (sx << 2))  & 0x3333333333333333ULL;
  sx = (sx | (sx << 1))  & 0x5555555555555555ULL;
  
  sy = (sy | (sy << 16)) & 0x0000FFFF0000FFFFULL;
  sy = (sy | (sy << 8))  & 0x00FF00FF00FF00FFULL;
  sy = (sy | (sy << 4))  & 0x0F0F0F0F0F0F0F0FULL;
  sy = (sy | (sy << 2))  & 0x3333333333333333ULL;
  sy = (sy | (sy << 1))  & 0x5555555555555555ULL;
  return sx | (sy << 1);
#endif
}
static inline int
get_size_lvl(int32_t w, int32_t h) {
  uint32_t max_dim = (uint32_t)((w > h) ? w : h);
  if (max_dim <= MIN_CELL_SIZE) {
    return 0;
  }
  uint32_t v = max_dim - 1u;
  unsigned int msb;
#if defined(_MSC_VER)
  unsigned long index;
  _BitScanReverse(&index, v);
  msb = (unsigned int)index + 1u;
#elif defined(__GNUC__) || defined(__clang__)
  msb = 32u - __builtin_clz(v);
#else
  msb = 0;
  while (v >>= 1) msb++;
  msb++;
#endif
  int lvl = (int)(msb - MIN_CELL_SHIFT);
  return (lvl < 0) ? 0 : (lvl >= MAX_LVLS) ? MAX_LVLS - 1 : lvl;
}
static int
cmp_uint128(const MDB_val *a, const MDB_val *b) {
  uint64_t va[2], vb[2];
  memcpy(va, a->mv_data, 16);
  memcpy(vb, b->mv_data, 16);
#if __BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__
  uint64_t ha = va[1], la = va[0];
  uint64_t hb = vb[1], lb = vb[0];
#else
  uint64_t ha = va[0], la = va[1];
  uint64_t hb = vb[0], lb = vb[1];
#endif
  if (ha != hb) {
    return (ha > hb) - (ha < hb);
  }
  return (la > lb) - (la < lb);
}
static inline uint32_t
world_to_grid(int32_t val, int lvl) {
  uint64_t shifted = (int64_t)val + COORD_OFFSET; 
  return (uint32_t)(shifted >> (MIN_CELL_SHIFT + lvl));
}
static uint64_t
get_center_key(int32_t x, int32_t y, int32_t w, int32_t h) {
  int lvl = get_size_lvl(w, h);
  int32_t cx = x + (w >> 1);
  int32_t cy = y + (h >> 1);
  uint32_t gx = world_to_grid(cx, lvl);
  uint32_t gy = world_to_grid(cy, lvl);
  uint64_t m_code = morton_encode(gx, gy);
  return ((uint64_t)lvl << 56) | m_code;
}
static struct db_ctx*
db_init(const char *path) {
  struct db_ctx *ctx = calloc(1, sizeof(struct db_ctx));
  mdb_env_create(&ctx->env);
  mdb_env_set_maxdbs(ctx->env, 2);
  mdb_env_set_mapsize(ctx->env, DB_MAP_SIZE);
  mdb_env_open(ctx->env, path, MDB_NOTLS | MDB_NOSUBDIR, 0664);

  MDB_txn *txn;
  mdb_txn_begin(ctx->env, NULL, 0, &txn);
  mdb_dbi_open(txn, "spatial", MDB_CREATE | MDB_INTEGERKEY | MDB_DUPSORT | MDB_DUPFIXED, &ctx->spatial_dbi);
  mdb_dbi_open(txn, "index", MDB_CREATE, &ctx->index_dbi);
  mdb_set_compare(txn, ctx->index_dbi, cmp_uint128);
  mdb_txn_commit(txn);
  return ctx;
}
static void
db_close(struct db_ctx *ctx) {
  if (ctx->batch_txn) {
    mdb_txn_commit(ctx->batch_txn);
  }
  mdb_dbi_close(ctx->env, ctx->spatial_dbi);
  mdb_dbi_close(ctx->env, ctx->index_dbi);
  mdb_env_close(ctx->env);
  free(ctx);
}
static MDB_txn*
get_write_txn(struct db_ctx *ctx, int *is_local) {
  if (ctx->batch_txn) {
    *is_local = 0;
    return ctx->batch_txn;
  }
  MDB_txn *txn;
  mdb_txn_begin(ctx->env, NULL, 0, &txn);
  *is_local = 1;
  return txn;
}
static void
check_batch(struct db_ctx *ctx) {
  if (ctx->batch_txn) {
    ctx->batch_count++;
    if (ctx->batch_count >= BATCH_LIMIT) {
      mdb_txn_commit(ctx->batch_txn);
      mdb_txn_begin(ctx->env, NULL, 0, &ctx->batch_txn);
      ctx->batch_count = 0;
    }
  }
}
static void
db_batch_begin(struct db_ctx *ctx) {
  if (ctx->batch_txn) {
    return;
  }
  mdb_txn_begin(ctx->env, NULL, 0, &ctx->batch_txn);
  ctx->batch_count = 0;
}
static void
db_batch_end(struct db_ctx *ctx) {
  if (ctx->batch_txn) {
    mdb_txn_commit(ctx->batch_txn);
    ctx->batch_txn = NULL;
  }
}
static void
db_upsert(struct db_ctx *ctx, objid_t id, int32_t x, int32_t y, int32_t w, int32_t h) {
  int local = 0;
  MDB_txn *txn = get_write_txn(ctx, &local);
  MDB_val k_id = { sizeof(objid_t), &id };
  MDB_val v_idx;

  // 1. Lookup & Cleanup Old Position
  if (mdb_get(txn, ctx->index_dbi, &k_id, &v_idx) == 0) {
    struct index_row *old_rec = (struct index_row*)v_idx.mv_data;
    MDB_val k_sp = { sizeof(uint64_t), &old_rec->spatial_key };
    MDB_val v_sp = { sizeof(struct spatial_row), &old_rec->row };
    mdb_del(txn, ctx->spatial_dbi, &k_sp, &v_sp);
  }
  // 2. Prepare New Record
  uint64_t new_sp_key = get_center_key(x, y, w, h);

  struct spatial_row sp_row;
  memset(&sp_row, 0, sizeof(struct spatial_row)); 
  sp_row.id = id;
  sp_row.minx = x; sp_row.miny = y;
  sp_row.maxx = x + w; sp_row.maxy = y + h;

  struct index_row idx_row;
  idx_row.spatial_key = new_sp_key;
  idx_row.row = sp_row;

  // 3. Zero-Copy Insert (Spatial)
  MDB_val k_sp_new = { sizeof(uint64_t), &new_sp_key };
  MDB_val v_sp_new = { sizeof(struct spatial_row), NULL };
  
  // MDB_RESERVE writes directly to DB map -> Fastest possible path
  if (mdb_put(txn, ctx->spatial_dbi, &k_sp_new, &v_sp_new, MDB_RESERVE) == 0) {
    memcpy(v_sp_new.mv_data, &sp_row, sizeof(struct spatial_row));
  }
  // 4. Update Index
  MDB_val v_idx_new = { sizeof(struct index_row), &idx_row };
  mdb_put(txn, ctx->index_dbi, &k_id, &v_idx_new, 0);
  if (local) {
    mdb_txn_commit(txn);
  } else {
    check_batch(ctx);
  }
}
static void
db_obj_del(struct db_ctx *ctx, objid_t id) {
  int local = 0;
  MDB_txn *txn = get_write_txn(ctx, &local);
  MDB_val k_id = { sizeof(objid_t), &id };
  MDB_val v_idx;

  if (mdb_get(txn, ctx->index_dbi, &k_id, &v_idx) == 0) {
    struct index_row *rec = (struct index_row*)v_idx.mv_data;
    MDB_val k_sp = { sizeof(uint64_t), &rec->spatial_key };
    MDB_val v_sp = { sizeof(struct spatial_row), &rec->row };
    
    mdb_del(txn, ctx->spatial_dbi, &k_sp, &v_sp);
    mdb_del(txn, ctx->index_dbi, &k_id, 0);
  }
  if (local) {
    mdb_txn_commit(txn);
  } else {
    check_batch(ctx);
  }
}
static void
heap_push(struct heap_item *h, int *n, struct heap_item val) {
  int i = (*n)++;
  h[i] = val;
  while (i > 0) {
    int p = (i - 1) >> 1;
    if (h[p].dist_sq >= h[i].dist_sq) break;
    struct heap_item t = h[i]; h[i] = h[p]; h[p] = t;
    i = p;
  }
}
static inline void
heap_replace(struct heap_item *h, int n, struct heap_item val) {
  /* 
   * OPTIMIZATION: Branchless Child Selection 
   * Eliminates branch mispredictions during the "Bubble Down" phase.
   */
  // "Half" optimization: Don't write 'val' to memory in every loop. 
  // Just find the hole, write once at the end.
  int i = 0;
  while (1) {
    int l = (i << 1) + 1;
    int r = (i << 1) + 2;
    int swap_idx = i;
    // Find largest child
    if (l < n && h[l].dist_sq > h[swap_idx].dist_sq) swap_idx = l;
    if (r < n && h[r].dist_sq > h[swap_idx].dist_sq) swap_idx = r;
    if (swap_idx == i) break; 
    h[i] = h[swap_idx]; // Move child up
    i = swap_idx;       // Move index down
  }
  h[i] = val; // Final write
}
static void
heap_sort_results(struct heap_item *h, int n) {
  /*
   * OPTIMIZATION: Final Sort
   * The query collects a Max-Heap (Index 0 is furthest).
   * Users usually want [Closest .... Furthest].
   * We use a standard Heap Sort to reverse it in-place.
   */
  // Pop elements one by one from the heap
  // This naturally moves the largest items to the end of the array.
  // Result: [Smallest, ..., Largest] (Exactly what we want)
  int original_count = n;
  while (n > 1) {
    struct heap_item max_val = h[0];
    struct heap_item last_val = h[n - 1];
    n--;
    // Move Max to end
    h[n] = max_val;
    // Restore heap property for the remaining
    heap_replace(h, n, last_val);
  }
}
static int
db_qry(struct db_ctx *ctx, int32_t vx, int32_t vy, int32_t vw, int32_t vh,
       objid_t *out_ids, int max_count) {
  
  int heap_cnt = 0;
  double cx = vx + vw * 0.5;
  double cy = vy + vh * 0.5;

  MDB_txn *txn;
  mdb_txn_begin(ctx->env, 0, MDB_RDONLY, &txn);
  MDB_cursor *cur;
  mdb_cursor_open(txn, ctx->spatial_dbi, &cur);

  // Pre-calculate squared cutoff to avoid SQRT or checks if heap is full
  double max_dist_sq = 1e300; // Infinity
  for (int lvl = MAX_LVLS - 1; lvl >= 0; lvl--) {
    int32_t cell_size = MIN_CELL_SIZE << lvl;
    int32_t half_cell = cell_size >> 1;

    int32_t search_min_x = vx - half_cell;
    int32_t search_min_y = vy - half_cell;
    int32_t search_max_x = vx + vw + half_cell; 
    int32_t search_max_y = vy + vh + half_cell;

    uint32_t t_min_x = world_to_grid(search_min_x, lvl);
    uint32_t t_min_y = world_to_grid(search_min_y, lvl);
    uint32_t t_max_x = world_to_grid(search_max_x, lvl);
    uint32_t t_max_y = world_to_grid(search_max_y, lvl);

    /* Granular LOD Skip */
    uint32_t span_x = t_max_x - t_min_x;
    uint32_t span_y = t_max_y - t_min_y;
    if (span_x > LOD_TILE_LIMIT || span_y > LOD_TILE_LIMIT) {
      break;
    }
    for (uint32_t ty = t_min_y; ty <= t_max_y; ty++) {
      for (uint32_t tx = t_min_x; tx <= t_max_x; tx++) {
        uint64_t m_code = morton_encode(tx, ty);
        uint64_t key = ((uint64_t)lvl << 56) | m_code;
        MDB_val v, k = { sizeof(key), &key };
        if (mdb_cursor_get(cur, &k, &v, MDB_SET) == 0) {
          do {
            struct spatial_row obj;
            memcpy(&obj, v.mv_data, sizeof(struct spatial_row));
            // 1. AABB Intersection (Integer math only)
            if (obj.maxx >= vx && obj.minx <= vx + vw &&
                obj.maxy >= vy && obj.miny <= vy + vh) {
              
              // 2. Distance Calculation (Floating point)
              double ocx = (obj.minx + obj.maxx) * 0.5;
              double ocy = (obj.miny + obj.maxy) * 0.5;
              double dist_sq = (ocx - cx)*(ocx - cx) + (ocy - cy)*(ocy - cy);

              // 3. Heap Logic
              if (heap_cnt < max_count) {
                heap_push(ctx->qry_heap, &heap_cnt, (struct heap_item){obj.id, dist_sq});
                // Track max distance if heap just got full
                if (heap_cnt == max_count) {
                   max_dist_sq = ctx->qry_heap[0].dist_sq; 
                }
              }  else if (dist_sq < max_dist_sq) {
                // Optimization: Check dist against cached max before calling function
                heap_replace(ctx->qry_heap, heap_cnt, (struct heap_item){obj.id, dist_sq});
                // Update cached max distance (Root is always largest)
                max_dist_sq = ctx->qry_heap[0].dist_sq;
              }
            }
          } while (mdb_cursor_get(cur, &k, &v, MDB_NEXT_DUP) == 0);
        }
      }
    }
  }
  mdb_cursor_close(cur);
  mdb_txn_abort(txn);
  // Sort the result
  // Turns [Furthest ... Random] into [Closest ... Furthest]
  heap_sort_results(ctx->qry_heap, heap_cnt);
  for (int i = 0; i < heap_cnt; i++) {
    out_ids[i] = ctx->qry_heap[i].id;
  }
  return heap_cnt;
}