Skip to content

Commit 1373068

Browse files
authored
Add: L2 distances in Rust (#237)
Closes #202
1 parent 876fe64 commit 1373068

File tree

2 files changed

+132
-1
lines changed

2 files changed

+132
-1
lines changed

javascript/simsimd.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ export const sqeuclidean = (
3333
* @brief Computes the Euclidean distance between two vectors.
3434
* @param {Float64Array|Float32Array|Int8Array|Uint8Array} a - The first vector.
3535
* @param {Float64Array|Float32Array|Int8Array|Uint8Array} b - The second vector.
36-
* @returns {number} The squared Euclidean distance between vectors a and b.
36+
* @returns {number} The Euclidean distance between vectors a and b.
3737
*/
3838
export const euclidean = (
3939
a: Float64Array | Float32Array | Int8Array | Uint8Array,

rust/lib.rs

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,12 @@ extern "C" {
8080
fn simsimd_l2sq_f32(a: *const f32, b: *const f32, c: usize, d: *mut Distance);
8181
fn simsimd_l2sq_f64(a: *const f64, b: *const f64, c: usize, d: *mut Distance);
8282

83+
fn simsimd_l2_i8(a: *const i8, b: *const i8, c: usize, d: *mut Distance);
84+
fn simsimd_l2_f16(a: *const u16, b: *const u16, c: usize, d: *mut Distance);
85+
fn simsimd_l2_bf16(a: *const u16, b: *const u16, c: usize, d: *mut Distance);
86+
fn simsimd_l2_f32(a: *const f32, b: *const f32, c: usize, d: *mut Distance);
87+
fn simsimd_l2_f64(a: *const f64, b: *const f64, c: usize, d: *mut Distance);
88+
8389
fn simsimd_hamming_b8(a: *const u8, b: *const u8, c: usize, d: *mut Distance);
8490
fn simsimd_jaccard_b8(a: *const u8, b: *const u8, c: usize, d: *mut Distance);
8591

@@ -215,13 +221,27 @@ where
215221
/// between corresponding elements of the two slices.
216222
fn l2sq(a: &[Self], b: &[Self]) -> Option<Distance>;
217223

224+
/// Computes the Euclidean distance between two slices.
225+
/// The Euclidean distance is the square root of
226+
// sum of the squared differences between corresponding
227+
/// elements of the two slices.
228+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance>;
229+
218230
/// Computes the squared Euclidean distance between two slices.
219231
/// The squared Euclidean distance is the sum of the squared differences
220232
/// between corresponding elements of the two slices.
221233
fn sqeuclidean(a: &[Self], b: &[Self]) -> Option<Distance> {
222234
SpatialSimilarity::l2sq(a, b)
223235
}
224236

237+
/// Computes the Euclidean distance between two slices.
238+
/// The Euclidean distance is the square root of the
239+
/// sum of the squared differences between corresponding
240+
/// elements of the two slices.
241+
fn euclidean(a: &[Self], b: &[Self]) -> Option<Distance> {
242+
SpatialSimilarity::l2(a, b)
243+
}
244+
225245
/// Computes the squared Euclidean distance between two slices.
226246
/// The squared Euclidean distance is the sum of the squared differences
227247
/// between corresponding elements of the two slices.
@@ -347,6 +367,16 @@ impl SpatialSimilarity for i8 {
347367
unsafe { simsimd_l2sq_i8(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
348368
Some(distance_value)
349369
}
370+
371+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance> {
372+
if a.len() != b.len() {
373+
return None;
374+
}
375+
let mut distance_value: Distance = 0.0;
376+
let distance_ptr: *mut Distance = &mut distance_value as *mut Distance;
377+
unsafe { simsimd_l2_i8(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
378+
Some(distance_value)
379+
}
350380
}
351381

352382
impl SpatialSimilarity for f16 {
@@ -391,6 +421,20 @@ impl SpatialSimilarity for f16 {
391421
unsafe { simsimd_l2sq_f16(a_ptr, b_ptr, a.len(), distance_ptr) };
392422
Some(distance_value)
393423
}
424+
425+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance> {
426+
427+
if a.len() != b.len() {
428+
return None;
429+
}
430+
// Explicitly cast `*const f16` to `*const u16`
431+
let a_ptr = a.as_ptr() as *const u16;
432+
let b_ptr = b.as_ptr() as *const u16;
433+
let mut distance_value: Distance = 0.0;
434+
let distance_ptr: *mut Distance = &mut distance_value as *mut Distance;
435+
unsafe { simsimd_l2_f16(a_ptr, b_ptr, a.len(), distance_ptr) };
436+
Some(distance_value)
437+
}
394438
}
395439

396440
impl SpatialSimilarity for bf16 {
@@ -435,6 +479,19 @@ impl SpatialSimilarity for bf16 {
435479
unsafe { simsimd_l2sq_bf16(a_ptr, b_ptr, a.len(), distance_ptr) };
436480
Some(distance_value)
437481
}
482+
483+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance> {
484+
if a.len() != b.len() {
485+
return None;
486+
}
487+
// Explicitly cast `*const bf16` to `*const u16`
488+
let a_ptr = a.as_ptr() as *const u16;
489+
let b_ptr = b.as_ptr() as *const u16;
490+
let mut distance_value: Distance = 0.0;
491+
let distance_ptr: *mut Distance = &mut distance_value as *mut Distance;
492+
unsafe { simsimd_l2sq_bf16(a_ptr, b_ptr, a.len(), distance_ptr) };
493+
Some(distance_value)
494+
}
438495
}
439496

440497
impl SpatialSimilarity for f32 {
@@ -467,6 +524,16 @@ impl SpatialSimilarity for f32 {
467524
unsafe { simsimd_l2sq_f32(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
468525
Some(distance_value)
469526
}
527+
528+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance> {
529+
if a.len() != b.len() {
530+
return None;
531+
}
532+
let mut distance_value: Distance = 0.0;
533+
let distance_ptr: *mut Distance = &mut distance_value as *mut Distance;
534+
unsafe { simsimd_l2_f32(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
535+
Some(distance_value)
536+
}
470537
}
471538

472539
impl SpatialSimilarity for f64 {
@@ -499,6 +566,16 @@ impl SpatialSimilarity for f64 {
499566
unsafe { simsimd_l2sq_f64(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
500567
Some(distance_value)
501568
}
569+
570+
fn l2(a: &[Self], b: &[Self]) -> Option<Distance> {
571+
if a.len() != b.len() {
572+
return None;
573+
}
574+
let mut distance_value: Distance = 0.0;
575+
let distance_ptr: *mut Distance = &mut distance_value as *mut Distance;
576+
unsafe { simsimd_l2_f64(a.as_ptr(), b.as_ptr(), a.len(), distance_ptr) };
577+
Some(distance_value)
578+
}
502579
}
503580

504581
impl ProbabilitySimilarity for f16 {
@@ -850,6 +927,60 @@ mod tests {
850927
}
851928
}
852929

930+
#[test]
931+
fn test_l2_f32() {
932+
let a: &[f32; 3] = &[1.0, 2.0, 3.0];
933+
let b: &[f32; 3] = &[4.0, 5.0, 6.0];
934+
if let Some(result) = SpatialSimilarity::euclidean(a, b) {
935+
println!("The result of l2_f32 is {:.8}", result);
936+
assert_almost_equal(5.2, result, 0.01);
937+
}
938+
}
939+
940+
#[test]
941+
fn test_l2_f64() {
942+
let a: &[f64; 3] = &[1.0, 2.0, 3.0];
943+
let b: &[f64; 3] = &[4.0, 5.0, 6.0];
944+
if let Some(result) = SpatialSimilarity::euclidean(a, b) {
945+
println!("The result of l2_f64 is {:.8}", result);
946+
assert_almost_equal(5.2, result, 0.01);
947+
}
948+
}
949+
950+
#[test]
951+
fn test_l2_f16() {
952+
let a_half: Vec<HalfF16> = vec![1.0, 2.0, 3.0]
953+
.iter()
954+
.map(|&x| HalfF16::from_f32(x))
955+
.collect();
956+
let b_half: Vec<HalfF16> = vec![4.0, 5.0, 6.0]
957+
.iter()
958+
.map(|&x| HalfF16::from_f32(x))
959+
.collect();
960+
961+
962+
let a_simsimd: &[f16] =
963+
unsafe { std::slice::from_raw_parts(a_half.as_ptr() as *const f16, a_half.len()) };
964+
let b_simsimd: &[f16] =
965+
unsafe { std::slice::from_raw_parts(b_half.as_ptr() as *const f16, b_half.len()) };
966+
967+
if let Some(result) = SpatialSimilarity::euclidean(&a_simsimd, &b_simsimd) {
968+
println!("The result of l2_f16 is {:.8}", result);
969+
assert_almost_equal(5.2, result, 0.01);
970+
}
971+
972+
}
973+
974+
#[test]
975+
fn test_l2_i8() {
976+
let a = &[1, 2, 3];
977+
let b = &[4, 5, 6];
978+
979+
if let Some(result) = SpatialSimilarity::euclidean(a, b) {
980+
println!("The result of l2_i8 is {:.8}", result);
981+
assert_almost_equal(5.2, result, 0.01);
982+
}
983+
}
853984
// Adding new tests for bit-level distances
854985
#[test]
855986
fn test_hamming_u8() {

0 commit comments

Comments
 (0)