neo_frizbee 0.8.1

Fast fuzzy matching via SIMD smith waterman, similar algorithm to FZF/FZY
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112

use std::{
    arch::x86_64::*,
    simd::{Simd, num::SimdUint},
};

#[target_feature(enable = "avx2", enable = "avx512f", enable = "avx512bw")]
pub fn interleave<const W: usize>(strs: [&str; 32]) -> [Simd<u16, 32>; W] {
    let str_bytes: [&[u8]; 32] = std::array::from_fn(|i| strs[i].as_bytes());
    let str_lens: [usize; 32] = std::array::from_fn(|i| str_bytes[i].len());

    let mut interleaved = [Simd::splat(0); W];
    for offset in (0..W).step_by(32) {
        let simds = to_simd(str_bytes, str_lens, offset);
        let interleaved_chunk = interleave_chunk(simds).map(|s| s.cast::<u16>());

        let copy_len = (W - offset).min(32);
        interleaved[offset..offset + copy_len].copy_from_slice(&interleaved_chunk[0..copy_len]);
    }

    interleaved
}

#[inline(always)]
fn to_simd(str_bytes: [&[u8]; 32], str_lens: [usize; 32], offset: usize) -> [__m256i; 32] {
    unsafe {
        std::array::from_fn(|i| {
            let len = str_lens[i];
            if offset >= len {
                // Beyond string length - return zeros
                return _mm256_setzero_si256();
            }

            let remaining = len - offset;
            let load_len = remaining.min(32);

            if load_len == 32 {
                _mm256_loadu_si256(str_bytes[i][offset..].as_ptr() as *const __m256i)
            } else {
                let mut temp = _mm256_setzero_si256();
                std::ptr::copy_nonoverlapping(
                    str_bytes[i][offset..].as_ptr(),
                    &mut temp as *mut __m256i as *mut u8,
                    load_len,
                );
                temp
            }
        })
    }
}

#[inline(always)]
pub fn interleave_chunk(mut simds: [__m256i; 32]) -> [Simd<u8, 32>; 32] {
    unsafe {
        // distance = 16
        for i in 0..16 {
            let (lo, hi) = interleave_u8x32(simds[i], simds[i + 16]);
            simds[i] = lo;
            simds[i + 16] = hi;
        }

        // distance = 8
        for base in (0..32).step_by(16) {
            for i in 0..8 {
                let (lo, hi) = interleave_u8x32(simds[base + i], simds[base + i + 8]);
                simds[base + i] = lo;
                simds[base + i + 8] = hi;
            }
        }

        // distance = 4
        for base in (0..32).step_by(8) {
            for i in 0..4 {
                let (lo, hi) = interleave_u8x32(simds[base + i], simds[base + i + 4]);
                simds[base + i] = lo;
                simds[base + i + 4] = hi;
            }
        }

        // distance = 2
        for base in (0..32).step_by(4) {
            for i in 0..2 {
                let (lo, hi) = interleave_u8x32(simds[base + i], simds[base + i + 2]);
                simds[base + i] = lo;
                simds[base + i + 2] = hi;
            }
        }

        // distance = 1
        for base in (0..32).step_by(2) {
            let (lo, hi) = interleave_u8x32(simds[base], simds[base + 1]);
            simds[base] = lo;
            simds[base + 1] = hi;
        }

        std::mem::transmute::<[__m256i; 32], [Simd<u8, 32>; 32]>(simds)
    }
}

#[inline(always)]
unsafe fn interleave_u8x32(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
    unsafe {
        // Use vpunpcklwd and vpunpckhwd for 16-bit interleaving
        let lo = _mm256_unpacklo_epi8(a, b);
        let hi = _mm256_unpackhi_epi8(a, b);

        // Fix the lane crossing issue in AVX2
        let lo_fixed = _mm256_permute4x64_epi64(lo, 0b11011000); // 0xD8
        let hi_fixed = _mm256_permute4x64_epi64(hi, 0b11011000);

        (lo_fixed, hi_fixed)
    }
}