pg_orrery/src/gist_tle.c

/*
 * gist_tle.c -- GiST operator class for 2-D orbital indexing on TLE
 *
 * Every TLE defines an orbit whose perigee/apogee altitudes and
 * inclination form a 2-D bounding box in (altitude, inclination) space.
 * Two orbits can only be in proximity if their altitude bands overlap
 * AND their inclination ranges overlap -- both necessary but not
 * sufficient conditions for conjunction.
 *
 * The GiST index stores [alt_low, alt_high] x [inc_low, inc_high] keys,
 * enabling fast coarse filtering.  For leaf entries inc_low == inc_high
 * (exact value); internal nodes widen to the bounding box of children.
 *
 * The && (overlap) operator is always rechecked: real conjunction
 * screening requires propagation.
 *
 * Semi-major axis from Kepler's third law using WGS-72 constants:
 *   a = (KE / n)^(2/3)       [earth radii]
 *   perigee_km = a*(1-e)*AE - AE
 *   apogee_km  = a*(1+e)*AE - AE
 */

#include "postgres.h"
#include "fmgr.h"
#include "access/gist.h"
#include "access/stratnum.h"
#include "utils/float.h"
#include "norad.h"
#include "types.h"
#include <math.h>
#include <float.h>

PG_FUNCTION_INFO_V1(tle_overlap);
PG_FUNCTION_INFO_V1(tle_alt_distance);
PG_FUNCTION_INFO_V1(gist_tle_compress);
PG_FUNCTION_INFO_V1(gist_tle_decompress);
PG_FUNCTION_INFO_V1(gist_tle_consistent);
PG_FUNCTION_INFO_V1(gist_tle_union);
PG_FUNCTION_INFO_V1(gist_tle_penalty);
PG_FUNCTION_INFO_V1(gist_tle_picksplit);
PG_FUNCTION_INFO_V1(gist_tle_same);
PG_FUNCTION_INFO_V1(gist_tle_distance);

/* Floating-point comparison tolerance (km and radians) */
#define KEY_EPSILON  1.0e-9

/* Domain widths for normalizing penalty across dimensions */
#define ALT_DOMAIN   50000.0   /* km — covers GEO + HEO margins */
#define INC_DOMAIN   M_PI      /* radians — full inclination range */

/* sizeof(pg_tle) == 112, matching INTERNALLENGTH in CREATE TYPE. */

/*
 * 2-D orbital key extracted from a TLE's mean elements.
 * Altitude band (perigee/apogee) plus inclination range.
 * This is the GiST internal key -- much cheaper to compare
 * than propagating two full state vectors.
 *
 * For leaf entries: inc_low == inc_high (exact inclination).
 * For internal nodes: bounding box over all children.
 */
typedef struct tle_orbital_key
{
	double	alt_low;	/* perigee altitude, km */
	double	alt_high;	/* apogee altitude, km */
	double	inc_low;	/* inclination lower bound, radians */
	double	inc_high;	/* inclination upper bound, radians */
} tle_orbital_key;


/* ----------------------------------------------------------------
 * tle_to_orbital_key -- compute [perigee, apogee] x [inc, inc]
 *
 * Uses WGS-72 KE and AE (the only constants valid for SGP4 elements).
 * Degenerate TLEs with n <= 0 map to zero-width ranges at 0.
 * Inclination is stored in radians (same as pg_tle).
 * ----------------------------------------------------------------
 */
static void
tle_to_orbital_key(const pg_tle *tle, tle_orbital_key *key)
{
	double n = tle->mean_motion;	/* rad/min */
	double e = tle->eccentricity;
	double a_er;					/* semi-major axis, earth radii */

	if (n <= 0.0)
	{
		key->alt_low  = 0.0;
		key->alt_high = 0.0;
		key->inc_low  = 0.0;
		key->inc_high = 0.0;
		return;
	}

	a_er = pow(WGS72_KE / n, 2.0 / 3.0);
	key->alt_low  = a_er * (1.0 - e) * WGS72_AE - WGS72_AE;
	key->alt_high = a_er * (1.0 + e) * WGS72_AE - WGS72_AE;

	/* Guard against numerical inversion from near-zero eccentricity */
	if (key->alt_low > key->alt_high)
	{
		double tmp = key->alt_low;
		key->alt_low = key->alt_high;
		key->alt_high = tmp;
	}

	/* Leaf entry: exact inclination (radians) */
	key->inc_low  = tle->inclination;
	key->inc_high = tle->inclination;
}


/* ----------------------------------------------------------------
 * key_overlaps -- do two orbital keys overlap in BOTH dimensions?
 *
 * Altitude bands AND inclination ranges must both overlap.
 * ----------------------------------------------------------------
 */
static inline bool
key_overlaps(const tle_orbital_key *a, const tle_orbital_key *b)
{
	return (a->alt_low <= b->alt_high) && (b->alt_low <= a->alt_high)
		&& (a->inc_low <= b->inc_high) && (b->inc_low <= a->inc_high);
}


/* ----------------------------------------------------------------
 * key_contains -- does outer fully contain inner in both dimensions?
 * ----------------------------------------------------------------
 */
static inline bool
key_contains(const tle_orbital_key *outer, const tle_orbital_key *inner)
{
	return (outer->alt_low <= inner->alt_low)
		&& (inner->alt_high <= outer->alt_high)
		&& (outer->inc_low <= inner->inc_low)
		&& (inner->inc_high <= outer->inc_high);
}


/* ----------------------------------------------------------------
 * key_merge -- expand dst bounding box to encompass src
 * ----------------------------------------------------------------
 */
static inline void
key_merge(tle_orbital_key *dst, const tle_orbital_key *src)
{
	if (src->alt_low < dst->alt_low)
		dst->alt_low = src->alt_low;
	if (src->alt_high > dst->alt_high)
		dst->alt_high = src->alt_high;
	if (src->inc_low < dst->inc_low)
		dst->inc_low = src->inc_low;
	if (src->inc_high > dst->inc_high)
		dst->inc_high = src->inc_high;
}


/* ----------------------------------------------------------------
 * alt_separation -- minimum altitude gap between two keys
 *
 * Returns 0 if the altitude bands overlap.
 * Used as one component of the 2-D orbital distance metric.
 * ----------------------------------------------------------------
 */
static inline double
alt_separation(const tle_orbital_key *a, const tle_orbital_key *b)
{
	if (a->alt_high < b->alt_low)
		return b->alt_low - a->alt_high;
	if (b->alt_high < a->alt_low)
		return a->alt_low - b->alt_high;
	return 0.0;
}


/* ----------------------------------------------------------------
 * inc_separation -- minimum inclination gap between two keys
 *
 * Analogous to alt_separation, but for the inclination dimension.
 * Returns 0 if the inclination ranges overlap.
 * ----------------------------------------------------------------
 */
static inline double
inc_separation(const tle_orbital_key *a, const tle_orbital_key *b)
{
	if (a->inc_high < b->inc_low)
		return b->inc_low - a->inc_high;
	if (b->inc_high < a->inc_low)
		return a->inc_low - b->inc_high;
	return 0.0;
}


/* ----------------------------------------------------------------
 * orbital_distance -- 2-D distance combining altitude and inclination
 *
 * Converts the inclination gap (radians) to km via Earth radius,
 * then returns L2 norm.  At Earth's surface, 1 radian of inclination
 * difference corresponds to ~6378 km of cross-track separation.
 *
 * For internal nodes, both alt_separation and inc_separation are
 * lower bounds on the gap to any child.  Since sqrt(a^2 + b^2) is
 * monotonically increasing, the L2 combination is also a valid
 * lower bound — satisfying GiST's distance contract.
 * ----------------------------------------------------------------
 */
static inline double
orbital_distance(const tle_orbital_key *a, const tle_orbital_key *b)
{
	double alt_gap = alt_separation(a, b);
	double inc_gap = inc_separation(a, b);
	double inc_km  = inc_gap * WGS72_AE;	/* radians → km via Earth radius */

	return sqrt(alt_gap * alt_gap + inc_km * inc_km);
}


/* ================================================================
 * SQL-callable operators
 * ================================================================
 */


/*
 * tle_overlap(tle, tle) -> bool     [the && operator]
 *
 * True if two TLEs overlap in both altitude AND inclination.
 * This is a fast pre-filter: overlapping keys are necessary
 * but not sufficient for actual conjunction.
 */
Datum
tle_overlap(PG_FUNCTION_ARGS)
{
	pg_tle		   *a = (pg_tle *) PG_GETARG_POINTER(0);
	pg_tle		   *b = (pg_tle *) PG_GETARG_POINTER(1);
	tle_orbital_key	ka, kb;

	tle_to_orbital_key(a, &ka);
	tle_to_orbital_key(b, &kb);

	PG_RETURN_BOOL(key_overlaps(&ka, &kb));
}


/*
 * tle_alt_distance(tle, tle) -> float8     [the <-> operator]
 *
 * 2-D orbital distance in km: combines altitude-band separation
 * with inclination gap (converted to km via Earth radius).
 * Returns 0 only if both altitude bands AND inclination ranges overlap.
 *
 * The C symbol name is kept as tle_alt_distance to avoid SQL migration
 * churn — the SQL FUNCTION and OPERATOR definitions are unchanged.
 */
Datum
tle_alt_distance(PG_FUNCTION_ARGS)
{
	pg_tle		   *a = (pg_tle *) PG_GETARG_POINTER(0);
	pg_tle		   *b = (pg_tle *) PG_GETARG_POINTER(1);
	tle_orbital_key	ka, kb;

	tle_to_orbital_key(a, &ka);
	tle_to_orbital_key(b, &kb);

	PG_RETURN_FLOAT8(orbital_distance(&ka, &kb));
}


/* ================================================================
 * GiST support functions
 * ================================================================
 */


/*
 * gist_tle_compress -- extract orbital key from a leaf TLE
 *
 * Leaf entries carry the full pg_tle; we compress to tle_orbital_key.
 * Internal entries are already tle_orbital_key from union operations.
 *
 * The allocation must be sizeof(pg_tle) bytes — which matches
 * INTERNALLENGTH — not sizeof(tle_orbital_key).  GiST's
 * index_form_tuple() copies typlen bytes from the datum pointer.
 */
Datum
gist_tle_compress(PG_FUNCTION_ARGS)
{
	GISTENTRY  *entry = (GISTENTRY *) PG_GETARG_POINTER(0);
	GISTENTRY  *retval;

	if (entry->leafkey)
	{
		pg_tle			*tle = (pg_tle *) DatumGetPointer(entry->key);
		tle_orbital_key	*key = (tle_orbital_key *) palloc0(sizeof(pg_tle));

		tle_to_orbital_key(tle, key);

		retval = (GISTENTRY *) palloc(sizeof(GISTENTRY));
		gistentryinit(*retval, PointerGetDatum(key),
					  entry->rel, entry->page, entry->offset, false);
	}
	else
	{
		/* Internal node: already a tle_orbital_key */
		retval = entry;
	}

	PG_RETURN_POINTER(retval);
}


/*
 * gist_tle_decompress -- identity (we operate on compressed keys directly)
 */
Datum
gist_tle_decompress(PG_FUNCTION_ARGS)
{
	PG_RETURN_POINTER(PG_GETARG_POINTER(0));
}


/*
 * gist_tle_consistent -- can this subtree contain matches for the query?
 *
 * Checks overlap in both altitude AND inclination dimensions.
 *
 * For leaf entries, recheck = false: the orbital key is computed
 * identically to the SQL operator, so the GiST check is exact.
 * For internal nodes, recheck is irrelevant (GiST ignores it).
 */
Datum
gist_tle_consistent(PG_FUNCTION_ARGS)
{
	GISTENTRY		*entry    = (GISTENTRY *) PG_GETARG_POINTER(0);
	pg_tle			*query    = (pg_tle *) PG_GETARG_POINTER(1);
	StrategyNumber	 strategy = (StrategyNumber) PG_GETARG_UINT16(2);
	/* arg 3 is the query type OID, unused */
	bool			*recheck  = (bool *) PG_GETARG_POINTER(4);
	tle_orbital_key	*key      = (tle_orbital_key *) DatumGetPointer(entry->key);
	tle_orbital_key	 query_key;
	bool			 result;

	tle_to_orbital_key(query, &query_key);

	/*
	 * Leaf keys are exact (same tle_to_orbital_key as the operator),
	 * so no recheck needed.  For internal nodes PostgreSQL ignores
	 * the flag, but we set true by convention.
	 */
	*recheck = !GIST_LEAF(entry);

	switch (strategy)
	{
		case RTOverlapStrategyNumber:		/* && */
			result = key_overlaps(key, &query_key);
			break;

		default:
			elog(ERROR, "gist_tle_consistent: unrecognized strategy %d",
				 strategy);
			result = false;
			break;
	}

	PG_RETURN_BOOL(result);
}


/*
 * gist_tle_union -- compute 2-D bounding box for a set of entries
 *
 * The union is [min(alt_low), max(alt_high)] x [min(inc_low), max(inc_high)].
 *
 * The entry vector is 0-based: valid indices are 0 .. entryvec->n - 1.
 * This differs from picksplit's 1-based convention.
 */
Datum
gist_tle_union(PG_FUNCTION_ARGS)
{
	GistEntryVector	*entryvec = (GistEntryVector *) PG_GETARG_POINTER(0);
	int				*sizep    = (int *) PG_GETARG_POINTER(1);
	int				 i;
	tle_orbital_key	*result;
	tle_orbital_key	*cur;

	result = (tle_orbital_key *) palloc0(sizeof(pg_tle));
	cur = (tle_orbital_key *) DatumGetPointer(entryvec->vector[0].key);
	*result = *cur;

	for (i = 1; i < entryvec->n; i++)
	{
		cur = (tle_orbital_key *) DatumGetPointer(entryvec->vector[i].key);
		key_merge(result, cur);
	}

	*sizep = sizeof(pg_tle);

	PG_RETURN_POINTER(result);
}


/*
 * gist_tle_penalty -- cost of inserting a new entry into an existing subtree
 *
 * Uses margin (half-perimeter) instead of area.  Leaf entries have
 * inc_low == inc_high, giving zero area -- area-based penalty would
 * degenerate to 0 for every insertion, making the tree random.
 * Margin remains non-zero for degenerate (zero-width) bounding boxes.
 */
Datum
gist_tle_penalty(PG_FUNCTION_ARGS)
{
	GISTENTRY		*origentry = (GISTENTRY *) PG_GETARG_POINTER(0);
	GISTENTRY		*newentry  = (GISTENTRY *) PG_GETARG_POINTER(1);
	float			*penalty   = (float *) PG_GETARG_POINTER(2);
	tle_orbital_key	*orig = (tle_orbital_key *) DatumGetPointer(origentry->key);
	tle_orbital_key	*add  = (tle_orbital_key *) DatumGetPointer(newentry->key);
	double			 orig_margin;
	double			 merged_margin;

	orig_margin = (orig->alt_high - orig->alt_low) / ALT_DOMAIN
				+ (orig->inc_high - orig->inc_low) / INC_DOMAIN;
	merged_margin = (fmax(orig->alt_high, add->alt_high) - fmin(orig->alt_low, add->alt_low)) / ALT_DOMAIN
				  + (fmax(orig->inc_high, add->inc_high) - fmin(orig->inc_low, add->inc_low)) / INC_DOMAIN;

	*penalty = (float)(merged_margin - orig_margin);

	PG_RETURN_POINTER(penalty);
}


/*
 * Comparison callback for qsort in picksplit.
 * Sorts entries by a sort key chosen at call time.
 */
typedef struct
{
	int		index;		/* position in the original entry vector */
	double	sortval;	/* midpoint in the chosen dimension */
} picksplit_item;

static int
picksplit_cmp(const void *a, const void *b)
{
	double ma = ((const picksplit_item *) a)->sortval;
	double mb = ((const picksplit_item *) b)->sortval;

	if (ma < mb)
		return -1;
	if (ma > mb)
		return 1;
	return 0;
}


/*
 * gist_tle_picksplit -- split an overfull page into two groups
 *
 * Standard R-tree approach: compute spread in both dimensions, split
 * along whichever dimension has the greater spread.  This prevents
 * the tree from becoming biased toward one dimension.
 *
 * GiST convention for picksplit: entryvec->vector[] is 1-based
 * (FirstOffsetNumber), vector[0] is unused/uninitialized.
 * entryvec->n includes the unused slot, so the actual entry count
 * is (entryvec->n - 1) and valid indices are
 * FirstOffsetNumber .. entryvec->n - 1.  The OffsetNumbers placed
 * into spl_left[] and spl_right[] must also be 1-based.
 */
Datum
gist_tle_picksplit(PG_FUNCTION_ARGS)
{
	GistEntryVector	*entryvec = (GistEntryVector *) PG_GETARG_POINTER(0);
	GIST_SPLITVEC	*splitvec = (GIST_SPLITVEC *) PG_GETARG_POINTER(1);
	OffsetNumber	 maxoff = entryvec->n - 1;
	int				 nentries = maxoff - FirstOffsetNumber + 1;
	picksplit_item	*items;
	tle_orbital_key	*left_union;
	tle_orbital_key	*right_union;
	tle_orbital_key	*cur;
	int				 split_at;
	int				 i;
	OffsetNumber	 off;
	double			 alt_min, alt_max, inc_min, inc_max;
	double			 alt_spread, inc_spread;
	bool			 split_on_alt;

	/* First pass: compute spread in both dimensions */
	cur = (tle_orbital_key *) DatumGetPointer(entryvec->vector[FirstOffsetNumber].key);
	alt_min = (cur->alt_low + cur->alt_high) / 2.0;
	alt_max = alt_min;
	inc_min = (cur->inc_low + cur->inc_high) / 2.0;
	inc_max = inc_min;

	for (off = FirstOffsetNumber + 1; off <= maxoff; off++)
	{
		double alt_mid, inc_mid;

		cur = (tle_orbital_key *) DatumGetPointer(entryvec->vector[off].key);
		alt_mid = (cur->alt_low + cur->alt_high) / 2.0;
		inc_mid = (cur->inc_low + cur->inc_high) / 2.0;

		if (alt_mid < alt_min) alt_min = alt_mid;
		if (alt_mid > alt_max) alt_max = alt_mid;
		if (inc_mid < inc_min) inc_min = inc_mid;
		if (inc_mid > inc_max) inc_max = inc_mid;
	}

	/*
	 * Normalize spreads so they're comparable.  Altitude is km above
	 * surface (range ~0-35786), inclination in radians (range 0-pi).
	 * Divide each by its domain width to get a 0-1 fraction.
	 */
	alt_spread = (alt_max - alt_min) / 35786.0;	/* GEO altitude above surface */
	inc_spread = (inc_max - inc_min) / M_PI;
	split_on_alt = (alt_spread >= inc_spread);

	/* Second pass: compute sort values in the chosen dimension */
	items = (picksplit_item *) palloc(sizeof(picksplit_item) * nentries);
	for (i = 0, off = FirstOffsetNumber; off <= maxoff; i++, off++)
	{
		cur = (tle_orbital_key *) DatumGetPointer(entryvec->vector[off].key);
		items[i].index = off;	/* store 1-based OffsetNumber directly */
		if (split_on_alt)
			items[i].sortval = (cur->alt_low + cur->alt_high) / 2.0;
		else
			items[i].sortval = (cur->inc_low + cur->inc_high) / 2.0;
	}

	qsort(items, nentries, sizeof(picksplit_item), picksplit_cmp);

	split_at = nentries / 2;

	/* Allocate offset arrays */
	splitvec->spl_left  = (OffsetNumber *) palloc(sizeof(OffsetNumber) * nentries);
	splitvec->spl_right = (OffsetNumber *) palloc(sizeof(OffsetNumber) * nentries);
	splitvec->spl_nleft  = 0;
	splitvec->spl_nright = 0;

	/* Compute union keys and assign entries */
	left_union  = (tle_orbital_key *) palloc0(sizeof(pg_tle));
	right_union = (tle_orbital_key *) palloc0(sizeof(pg_tle));

	/* Seed the unions from the first entry in each half */
	cur = (tle_orbital_key *) DatumGetPointer(
			entryvec->vector[items[0].index].key);
	*left_union = *cur;

	cur = (tle_orbital_key *) DatumGetPointer(
			entryvec->vector[items[split_at].index].key);
	*right_union = *cur;

	for (i = 0; i < nentries; i++)
	{
		OffsetNumber idx = items[i].index;	/* already 1-based */

		cur = (tle_orbital_key *) DatumGetPointer(
				entryvec->vector[idx].key);

		if (i < split_at)
		{
			splitvec->spl_left[splitvec->spl_nleft++] = idx;
			key_merge(left_union, cur);
		}
		else
		{
			splitvec->spl_right[splitvec->spl_nright++] = idx;
			key_merge(right_union, cur);
		}
	}

	splitvec->spl_ldatum = PointerGetDatum(left_union);
	splitvec->spl_rdatum = PointerGetDatum(right_union);

	pfree(items);

	PG_RETURN_POINTER(splitvec);
}


/*
 * gist_tle_same -- equality test on compressed keys
 *
 * Two orbital keys are "same" if all four bounds match within
 * tolerance.  This lets the GiST machinery detect duplicate keys.
 */
Datum
gist_tle_same(PG_FUNCTION_ARGS)
{
	tle_orbital_key	*a = (tle_orbital_key *) PG_GETARG_POINTER(0);
	tle_orbital_key	*b = (tle_orbital_key *) PG_GETARG_POINTER(1);
	bool			*result = (bool *) PG_GETARG_POINTER(2);

	*result = (fabs(a->alt_low  - b->alt_low)  < KEY_EPSILON
			&& fabs(a->alt_high - b->alt_high) < KEY_EPSILON
			&& fabs(a->inc_low  - b->inc_low)  < KEY_EPSILON
			&& fabs(a->inc_high - b->inc_high) < KEY_EPSILON);

	PG_RETURN_POINTER(result);
}


/*
 * gist_tle_distance -- GiST distance function for KNN ordering
 *
 * Returns the 2-D orbital distance: L2 norm of altitude separation
 * and inclination gap (in km).  For entries where both dimensions
 * overlap this is 0, making the entry a candidate.
 *
 * Both alt_separation and inc_separation are lower bounds for
 * internal nodes, so the L2 combination is also a valid lower
 * bound — satisfying GiST's distance contract for correct KNN.
 */
Datum
gist_tle_distance(PG_FUNCTION_ARGS)
{
	GISTENTRY		*entry = (GISTENTRY *) PG_GETARG_POINTER(0);
	pg_tle			*query = (pg_tle *) PG_GETARG_POINTER(1);
	/* strategy number at arg 2, unused for single-distance class */
	tle_orbital_key	*key   = (tle_orbital_key *) DatumGetPointer(entry->key);
	tle_orbital_key	 query_key;

	tle_to_orbital_key(query, &query_key);

	PG_RETURN_FLOAT8(orbital_distance(key, &query_key));
}