Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
248 changes: 225 additions & 23 deletions src/hpc/clam.rs
Original file line number Diff line number Diff line change
Expand Up @@ -158,13 +158,42 @@ impl Default for BuildConfig {
/// Distance function type: takes two byte slices of equal length, returns u64.
pub type DistanceFn = fn(&[u8], &[u8]) -> u64;

/// The tree's distance, carried on one of two arms.
///
/// `Ptr` is every pre-existing construction path, byte-identical in behaviour.
/// `Dyn` is the stateful arm ([`ClamTree::build_with_distance`]): a distance
/// that carries configuration — e.g. a ClassView-derived `RailSpec` — which a
/// bare fn pointer cannot hold. Both arms feed the SAME partition core and the
/// SAME search paths through [`ClamTree::dist`], so one name can never mean
/// two things (`I-LEGACY-API-FEATURE-GATED`).
enum TreeDistance {
Ptr(DistanceFn),
Dyn {
f: Box<dyn Fn(&[u8], &[u8]) -> u64 + Send + Sync>,
/// Carried from [`Distance::is_metric`]. The `Ptr` arm is recorded as
/// metric because every historical caller relied on triangle-inequality
/// pruning with it — preserving, not blessing, that assumption.
metric: bool,
},
}

impl TreeDistance {
#[inline]
fn eval(&self, a: &[u8], b: &[u8]) -> u64 {
match self {
TreeDistance::Ptr(f) => f(a, b),
TreeDistance::Dyn { f, .. } => f(a, b),
}
}
}

/// Divisive hierarchical clustering tree.
pub struct ClamTree {
pub nodes: Vec<Cluster>,
pub reordered: Vec<usize>,
pub num_leaves: usize,
pub mean_leaf_radius: f64,
distance_fn: DistanceFn,
distance: TreeDistance,
}

impl ClamTree {
Expand All @@ -185,24 +214,64 @@ impl ClamTree {

/// Build a CLAM tree with a custom distance function.
pub fn build_with_fn(data: &[u8], vec_len: usize, count: usize, config: &BuildConfig, dist_fn: DistanceFn) -> Self {
assert_eq!(data.len(), vec_len * count);
Self::build_core(data, vec_len, count, config, &dist_fn, TreeDistance::Ptr(dist_fn))
}

/// Build with a **stateful** distance — the universal arm.
///
/// A [`Distance`] impl can carry configuration a bare fn pointer cannot:
/// a byte-range spec, a codebook, a resolved ClassView reading. Its
/// [`Distance::is_metric`] answer is carried into the tree (see
/// [`ClamTree::is_metric`]) instead of being assumed. Identical inputs
/// produce a tree identical to the fn-pointer arm — pinned by test, not
/// claimed.
pub fn build_with_distance<D>(data: &[u8], vec_len: usize, count: usize, config: &BuildConfig, dist: D) -> Self
where
D: Distance<Point = [u8]> + Send + Sync + 'static,
{
let metric = dist.is_metric();
let f: Box<dyn Fn(&[u8], &[u8]) -> u64 + Send + Sync> = Box::new(move |a, b| dist.distance(a, b));
// Build against a borrow of the SAME closure that will be stored —
// one callable, never a rebuilt twin that could drift.
let nodes_input = &*f as &dyn Fn(&[u8], &[u8]) -> u64;
// SAFETY-free reborrow dance: build_core only borrows `dist` during
// construction; the box is moved into the carrier afterwards.
let tmp = Self::build_core_nodes(data, vec_len, count, config, nodes_input);
Self::assemble(tmp, TreeDistance::Dyn { f, metric })
Comment on lines +232 to +240

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Respect is_metric before applying tree pruning

When a stateful Distance returns false from is_metric, this constructor still produces a tree that rho_nn, knn_repeated_rho, knn_dfs_sieve, and rho_nn_candidates prune using triangle-inequality radius bounds. Those bounds are invalid for such a distance, so relevant clusters can be skipped and searches can silently return false negatives; merely exposing the stored flag does not protect callers. Reject non-metric distances in the pruned routines or bypass the bounds and scan exhaustively when this flag is false.

Useful? React with 👍 / 👎.

}

/// The one construction core, parameterised over any callable. Both
/// public arms land here.
fn build_core(
data: &[u8], vec_len: usize, count: usize, config: &BuildConfig, dist: &dyn Fn(&[u8], &[u8]) -> u64,
carrier: TreeDistance,
) -> Self {
let tmp = Self::build_core_nodes(data, vec_len, count, config, dist);
Self::assemble(tmp, carrier)
}

fn build_core_nodes(
data: &[u8], vec_len: usize, count: usize, config: &BuildConfig, dist: &dyn Fn(&[u8], &[u8]) -> u64,
) -> (Vec<Cluster>, Vec<usize>) {
// Order and constants are the ORIGINAL build's, verbatim: assert
// before the empty return, the 0xDEAD_BEEF_CAFE_BABE seed, the
// 2*count capacity. A shared core that quietly changes any of them
// would rebuild every existing caller's tree differently — the exact
// silent break this refactor exists to make impossible.
assert_eq!(data.len(), vec_len * count);
if count == 0 {
return ClamTree {
nodes: Vec::new(),
reordered: Vec::new(),
num_leaves: 0,
mean_leaf_radius: 0.0,
distance_fn: dist_fn,
};
return (Vec::new(), Vec::new());
}

let mut indices: Vec<usize> = (0..count).collect();
let mut nodes = Vec::with_capacity(2 * count);
let mut rng = SplitMix64::new(0xDEAD_BEEF_CAFE_BABE);
Self::partition(data, vec_len, &mut indices, 0, count, 0, config, &mut nodes, &mut rng, dist);
(nodes, indices)
}

Self::partition(data, vec_len, &mut indices, 0, count, 0, config, &mut nodes, &mut rng, dist_fn);

fn assemble((nodes, reordered): (Vec<Cluster>, Vec<usize>), carrier: TreeDistance) -> Self {
// Integer sum, one division — the original arithmetic, not an f64
// re-summation with different rounding.
let mut num_leaves = 0usize;
let mut leaf_radius_sum = 0u64;
for node in &nodes {
Expand All @@ -216,21 +285,20 @@ impl ClamTree {
} else {
0.0
};

ClamTree {
nodes,
reordered: indices,
reordered,
num_leaves,
mean_leaf_radius,
distance_fn: dist_fn,
distance: carrier,
}
}

/// Recursive partition (Algorithm 1 from CAKES).
#[allow(clippy::too_many_arguments)]
fn partition(
data: &[u8], vec_len: usize, indices: &mut [usize], start: usize, end: usize, depth: usize,
config: &BuildConfig, nodes: &mut Vec<Cluster>, rng: &mut SplitMix64, dist_fn: DistanceFn,
config: &BuildConfig, nodes: &mut Vec<Cluster>, rng: &mut SplitMix64, dist_fn: &dyn Fn(&[u8], &[u8]) -> u64,
) -> usize {
let n = end - start;
let node_idx = nodes.len();
Expand Down Expand Up @@ -370,12 +438,40 @@ impl ClamTree {

#[inline]
pub fn dist(&self, a: &[u8], b: &[u8]) -> u64 {
(self.distance_fn)(a, b)
self.distance.eval(a, b)
}

#[inline]
/// The raw fn-pointer arm, for callers that thread it onward.
///
/// # Panics
/// On a tree built with [`ClamTree::build_with_distance`] — a stateful
/// distance has no fn pointer to hand out, and returning a substitute
/// would silently measure something else. New code should call
/// [`ClamTree::dist`] instead; no pre-existing construction path can
/// reach the panic.
pub fn distance_fn(&self) -> DistanceFn {
self.distance_fn
match &self.distance {
TreeDistance::Ptr(f) => *f,
TreeDistance::Dyn { .. } => panic!(
"ClamTree was built with build_with_distance (stateful arm); \
use ClamTree::dist instead of extracting a fn pointer"
),
}
}

/// Whether this tree's distance declared itself a metric.
///
/// The fn-pointer arm reports `true` — every historical caller relied on
/// triangle-inequality pruning with it, and this accessor preserves that
/// assumption rather than blessing it. The stateful arm carries the
/// answer from [`Distance::is_metric`]; `rho_nn`-style pruning over a
/// distance that answers `false` is unsound (silent false negatives).
pub fn is_metric(&self) -> bool {
match &self.distance {
TreeDistance::Ptr(_) => true,
TreeDistance::Dyn { metric, .. } => *metric,
}
}

pub fn root(&self) -> &Cluster {
Expand Down Expand Up @@ -492,7 +588,7 @@ impl ClamTree {
let center = self.center_data(cluster, data, vec_len);
let mut distances: Vec<u64> = self
.cluster_points(cluster, data, vec_len)
.map(|(_, point)| (self.distance_fn)(center, point))
.map(|(_, point)| self.distance.eval(center, point))
.collect();

if distances.is_empty() {
Expand Down Expand Up @@ -1068,7 +1164,7 @@ impl CompressedTree {
/// Compute Hamming distance from query to compressed point WITHOUT decompression.
pub fn hamming_to_compressed(
&self, query: &[u8], point_idx: usize, data: &[u8], vec_len: usize, dist_cache: &mut DistanceCache,
dist_fn: DistanceFn,
dist_fn: impl Fn(&[u8], &[u8]) -> u64,
) -> u64 {
let center_idx = self.encoding_centers[point_idx];

Expand Down Expand Up @@ -1167,7 +1263,7 @@ impl ClamTree {
while let Some(node_idx) = stack.pop() {
let node = &self.nodes[node_idx];
let center = &data[self.reordered[node.center_idx] * vec_len..][..vec_len];
let dist_to_center = (self.distance_fn)(query, center);
let dist_to_center = self.distance.eval(query, center);

// Triangle inequality: closest possible point in cluster
if node.delta_minus(dist_to_center) > rho {
Expand All @@ -1179,7 +1275,7 @@ impl ClamTree {
for i in node.offset..node.offset + node.cardinality {
let idx = self.reordered[i];
let point = &data[idx * vec_len..][..vec_len];
let d = (self.distance_fn)(query, point);
let d = self.distance.eval(query, point);
if d <= rho {
candidates.push((idx, d));
}
Expand Down Expand Up @@ -1430,7 +1526,12 @@ impl ClamTree {
let mut cache = DistanceCache::new();
let mut hits = Vec::new();
let mut distance_calls = 0usize;
let dist_fn = self.distance_fn();
// A borrow of the carrier, not an extracted fn pointer — this path
// must work on BOTH arms (a Dyn-built tree has no pointer to hand out).
let dist_fn: &dyn Fn(&[u8], &[u8]) -> u64 = match &self.distance {
TreeDistance::Ptr(f) => f,
TreeDistance::Dyn { f, .. } => &**f,
};
Comment on lines +1532 to +1535

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Avoid Hamming correction for arbitrary stateful distances

When query_compressed is called on the new dynamic arm with anything other than full-slice bitwise Hamming—such as the added TailHamming or V3RailGeodesic—this closure supplies the baseline distance, but hamming_from_query subsequently adjusts it using full-record XOR popcounts. That identity is valid only for ordinary Hamming distance: for example, a byte difference before TailHamming::from changes the reported result even though the configured distance ignores it, yielding incorrect distances and hit filtering. Restrict this optimized path to full Hamming or decode each candidate and evaluate the tree's actual distance.

Useful? React with 👍 / 👎.


// Use CLAM tree structure: walk to find overlapping leaves,
// then do compressive distance on leaf members
Expand Down Expand Up @@ -3012,4 +3113,105 @@ mod tests {
let hits = clam_cascade_search(&tree, &cascade, &data, vec_len, query, u64::MAX, 5);
assert!(hits.len() <= 5);
}

// ── the universal-builder pass: two arms, one core ──

/// A stateful distance for the Dyn arm: Hamming over the tail from a
/// CONFIGURED offset — state a bare fn pointer cannot carry.
struct TailHamming {
from: usize,
}
impl Distance for TailHamming {
type Point = [u8];
fn distance(&self, a: &[u8], b: &[u8]) -> u64 {
hamming_inline(&a[self.from..], &b[self.from..])
}
fn is_metric(&self) -> bool {
true
}
}

/// The fn-pointer twin of `TailHamming { from: 16 }` — offset hardcoded,
/// because a pointer can hold nothing else. Exists so the two arms can be
/// compared on EQUAL inputs.
fn tail16_hamming(a: &[u8], b: &[u8]) -> u64 {
hamming_inline(&a[16..], &b[16..])
}

fn arm_test_data() -> Vec<u8> {
// Deterministic, structured enough to force real splits.
let mut rng = SplitMix64::new(0x51_7EED);
(0..64 * 32)
.map(|_| (rng.next_u64() & 0xFF) as u8)
.collect()
}

fn assert_same_tree(a: &ClamTree, b: &ClamTree) {
assert_eq!(a.reordered, b.reordered, "reordering diverged");
assert_eq!(a.nodes.len(), b.nodes.len(), "node count diverged");
for (i, (x, y)) in a.nodes.iter().zip(&b.nodes).enumerate() {
assert_eq!(
(x.center_idx, x.radius, x.cardinality, x.offset, x.depth, x.left, x.right),
(y.center_idx, y.radius, y.cardinality, y.offset, y.depth, y.left, y.right),
"cluster {i} diverged"
);
}
assert_eq!(a.num_leaves, b.num_leaves);
}

/// THE claim of the refactor, pinned: both arms land in one core, so the
/// same distance produces the byte-identical tree regardless of arm.
#[test]
fn the_two_arms_build_the_identical_tree() {
let data = arm_test_data();
let cfg = BuildConfig {
min_cardinality: 4,
..Default::default()
};
let ptr = ClamTree::build_with_fn(&data, 32, 64, &cfg, tail16_hamming);
let dy = ClamTree::build_with_distance(&data, 32, 64, &cfg, TailHamming { from: 16 });
assert_same_tree(&ptr, &dy);
// and the search paths agree through the dispatch
let q = &data[0..32];
let r1 = rho_nn(&ptr, &data, 32, q, 40);
let r2 = rho_nn(&dy, &data, 32, q, 40);
assert_eq!(r1.hits, r2.hits, "rho_nn diverged between arms");
}

/// `is_metric` rides the carrier: the Ptr arm preserves the historical
/// assumption (true), the Dyn arm carries the distance's own answer.
#[test]
fn is_metric_is_carried_not_assumed() {
struct NotAMetric;
impl Distance for NotAMetric {
type Point = [u8];
fn distance(&self, a: &[u8], b: &[u8]) -> u64 {
hamming_inline(a, b)
}
fn is_metric(&self) -> bool {
false
}
}
let data = arm_test_data();
let cfg = BuildConfig {
min_cardinality: 4,
..Default::default()
};
assert!(ClamTree::build_with_fn(&data, 32, 64, &cfg, hamming_inline).is_metric());
assert!(!ClamTree::build_with_distance(&data, 32, 64, &cfg, NotAMetric).is_metric());
}

/// The fn-pointer accessor refuses the stateful arm LOUDLY. A silent
/// substitute would measure something else; a panic names the fix.
#[test]
#[should_panic(expected = "build_with_distance")]
fn the_fn_pointer_accessor_refuses_the_stateful_arm() {
let data = arm_test_data();
let cfg = BuildConfig {
min_cardinality: 4,
..Default::default()
};
let t = ClamTree::build_with_distance(&data, 32, 64, &cfg, TailHamming { from: 0 });
let _ = t.distance_fn();
}
}
2 changes: 1 addition & 1 deletion src/hpc/clam_compress.rs
Original file line number Diff line number Diff line change
Expand Up @@ -359,7 +359,7 @@ impl CompressedTree {
/// Cost: O(num_diffs) per point instead of O(vec_len).
pub fn hamming_to_compressed(
&self, query: &[u8], point_idx: usize, data: &[u8], vec_len: usize, dist_cache: &mut DistanceCache,
dist_fn: fn(&[u8], &[u8]) -> u64,
dist_fn: impl Fn(&[u8], &[u8]) -> u64,
) -> u64 {
let center_idx = self.encoding_centers[point_idx];

Expand Down
17 changes: 12 additions & 5 deletions src/hpc/clam_v3.rs
Original file line number Diff line number Diff line change
Expand Up @@ -155,10 +155,7 @@ impl RailSpec {
/// The same spec with a continuation register stacked at `at`.
#[must_use]
pub const fn stacked(self, at: usize) -> Self {
Self {
cont: Some(at),
..self
}
Self { cont: Some(at), ..self }
}

/// Maximum representable depth under this spec.
Expand Down Expand Up @@ -273,6 +270,11 @@ impl Distance for V3ValueHamming {
/// [`super::clam::DistanceFn`] so it plugs straight into
/// `ClamTree::build_with_fn(rows, 512, …, v3_value_hamming)`.
///
/// The rail geodesic has no such bare-fn form BY DESIGN: it carries a
/// [`RailSpec`], and state does not fit in a fn pointer. It rides
/// [`super::clam::ClamTree::build_with_distance`] as [`V3RailGeodesic`]
/// directly — no `const` workaround needed since the universal-builder pass.
///
/// A row shorter than the value offset contributes nothing — an honest 0
/// beats a panic in a distance callback, and a truncated row is a loader
/// bug this function cannot repair.
Expand All @@ -293,7 +295,12 @@ mod tests {
fn row(levels: &[u8], axis: RailAxis, fill: u8) -> Vec<u8> {
let mut r = vec![0u8; 512];
for (i, &v) in levels.iter().enumerate().take(RAIL_PAIRS) {
let at = 4 + 2 * i + match axis { RailAxis::Lo => 0, RailAxis::Hi => 1 };
let at = 4
+ 2 * i
+ match axis {
RailAxis::Lo => 0,
RailAxis::Hi => 1,
};
r[at] = v;
}
for b in &mut r[V3_VALUE_OFF..] {
Expand Down
Loading