text_block_permutation_optimizer/

algo.rs

use std::cmp;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Instant;

use rapidfuzz::distance::levenshtein;
use rayon::prelude::*;
use tracing::info;

// Levensthein Distance,
// how strong should shuffling be penalized.
const JUMP_STATIC_PENALTY: f64 = 2.0;
const JUMP_OFFSET_PENALTY: f64 = 0.2;

// Bigger blocks are slower, but enhance quality of output
const BLOCK_SIZES: &[usize] = &[1, 2, 3, 4, 6, 9, 12, 15];
// Base shift window; scaled up proportionally for larger blocks
const BASE_BLOCK_SHIFT: usize = 30;
// Scaling factor: max_shift = BASE_BLOCK_SHIFT + block_size * SHIFT_PER_BLOCK_SIZE
const SHIFT_PER_BLOCK_SIZE: usize = 8;

type IndexedBlock<'a> = (usize, &'a str);

#[derive(Debug, Clone)]
pub struct MatchingLoss {
    pub fuzzy: f64,
    pub jump_distance: f64,
    pub geo_distance: f64,
}

impl MatchingLoss {
    pub fn abs(&self) -> f64 {
        self.fuzzy + self.jump_distance + self.geo_distance
    }
}

/// Convert a non-negative f64 to u64 preserving ordering.
/// Positive IEEE 754 floats have the same ordering as their bit patterns.
#[inline]
fn f64_to_ord_u64(v: f64) -> u64 {
    debug_assert!(v >= 0.0);
    v.to_bits()
}

#[inline]
fn ord_u64_to_f64(v: u64) -> f64 {
    f64::from_bits(v)
}

/// Core function: Optimize the matching order of a list of strings against ground truth.
///
/// # Arguments
///
/// * `ground` - The ground truth string to match against
/// * `blocks` - Slice of strings to be matched/ordered
///
/// # Returns
///
/// Tuple of (MatchingLoss, Vec of (index, text) tuples representing optimal order)
pub fn optimize<'a>(ground: &str, blocks: &[&'a str]) -> (MatchingLoss, Vec<IndexedBlock<'a>>) {
    let ground_normalized = normalize_whitespace(ground);
    let ground_words: Vec<&str> = ground_normalized.split_whitespace().collect();
    let ground_concat = ground_words.join(" ");

    // Precompute the BatchComparator for the ground truth (builds pattern match vector once).
    let scorer = levenshtein::BatchComparator::new(ground_concat.bytes());
    let ground_len = ground_concat.len();

    // Pre-compute total byte length for buffer allocation
    let total_block_bytes: usize =
        blocks.iter().map(|b| b.len()).sum::<usize>() + blocks.len().saturating_sub(1);
    let mut main_buf = Vec::<u8>::with_capacity(total_block_bytes);

    // Convert blocks slice to indexed tuples
    let mut matching: Vec<IndexedBlock<'a>> =
        blocks.iter().enumerate().map(|(a, b)| (a, *b)).collect();

    let mut optimum_loss = score_matching(&scorer, ground_len, &matching, &mut main_buf, None);
    info!("Start fuzzy matching with loss of {:?}", optimum_loss);

    let shift_count = AtomicU64::new(0);
    let start_time = Instant::now();
    let n = matching.len();

    for &block_size in BLOCK_SIZES.iter().rev() {
        let level_start = Instant::now();
        let mut level_improvements = 0u64;
        let max_shift = BASE_BLOCK_SHIFT + block_size * SHIFT_PER_BLOCK_SIZE;
        info!(
            "Optimizing on block level {} (max_shift={})",
            block_size, max_shift
        );

        loop {
            // Shared atomic cutoff: when any thread finds a better solution,
            // all other threads immediately use the tighter cutoff for pruning.
            let shared_cutoff = AtomicU64::new(f64_to_ord_u64(optimum_loss.abs()));
            let snapshot = matching.clone();

            let best = (0..n)
                .into_par_iter()
                .map_init(
                    || {
                        (
                            snapshot.clone(),
                            Vec::<u8>::with_capacity(total_block_bytes),
                        )
                    },
                    |(local_matching, buf), start_pos| {
                        let end_pos = cmp::min(start_pos + block_size, n);
                        let actual_bs = end_pos - start_pos;

                        if !is_consecutive_block(local_matching, start_pos, actual_bs) {
                            return None;
                        }

                        let lo = start_pos.saturating_sub(max_shift);
                        let hi = cmp::min(n - actual_bs + 1, start_pos + max_shift);

                        let mut best_local: Option<(f64, usize)> = None;
                        let mut local_count = 0u64;

                        for target in lo..hi {
                            if target == start_pos {
                                continue;
                            }
                            local_count += 1;

                            // Read shared cutoff — tightens as other threads find improvements
                            let cutoff = ord_u64_to_f64(shared_cutoff.load(Ordering::Relaxed));

                            apply_shift(local_matching, start_pos, actual_bs, target);

                            let loss = score_matching(
                                &scorer,
                                ground_len,
                                local_matching,
                                buf,
                                Some(cutoff),
                            );

                            let loss_val = loss.abs();
                            if loss_val < cutoff {
                                best_local = Some((loss_val, target));
                                // Publish tighter cutoff to all threads
                                shared_cutoff
                                    .fetch_min(f64_to_ord_u64(loss_val), Ordering::Relaxed);
                            }

                            undo_shift(local_matching, start_pos, actual_bs, target);
                        }

                        shift_count.fetch_add(local_count, Ordering::Relaxed);
                        best_local.map(|(loss, target)| (loss, start_pos, target))
                    },
                )
                .flatten()
                .min_by(|a, b| a.0.partial_cmp(&b.0).unwrap());

            if let Some((_best_loss, best_start, best_target)) = best {
                let end_pos = cmp::min(best_start + block_size, n);
                let actual_bs = end_pos - best_start;
                apply_shift(&mut matching, best_start, actual_bs, best_target);
                let old = optimum_loss.abs();
                optimum_loss = score_matching(&scorer, ground_len, &matching, &mut main_buf, None);
                level_improvements += 1;
                info!(
                    "  bs={:02} loss={:.1} -> {:.1} | {} trials, {:.2?}",
                    block_size,
                    old,
                    optimum_loss.abs(),
                    shift_count.load(Ordering::Relaxed),
                    start_time.elapsed()
                );
            } else {
                break;
            }

            if shift_count.load(Ordering::Relaxed) > 4_000_000 {
                break;
            }
        }
        info!(
            "  level {} done: {} improvements, {:.2?} elapsed (total {:.2?})",
            block_size,
            level_improvements,
            level_start.elapsed(),
            start_time.elapsed()
        );
    }

    info!(
        "Returning optimum with {:?} after {} trials ({:.2?})",
        optimum_loss,
        shift_count.load(Ordering::Relaxed),
        start_time.elapsed()
    );
    (optimum_loss, matching)
}

/// Check that the original indexes in a block form a consecutive ascending sequence.
#[inline]
fn is_consecutive_block(matching: &[IndexedBlock], start: usize, len: usize) -> bool {
    for i in 1..len {
        if matching[start + i].0 != matching[start + i - 1].0 + 1 {
            return false;
        }
    }
    true
}

/// Shift block at `[src..src+len]` to position `dst` in-place.
#[inline]
fn apply_shift<T>(slice: &mut [T], src: usize, len: usize, dst: usize) {
    if dst < src {
        slice[dst..src + len].rotate_right(len);
    } else {
        slice[src..dst + len].rotate_left(len);
    }
}

/// Undo a shift (reverse of apply_shift)
#[inline]
fn undo_shift<T>(slice: &mut [T], src: usize, len: usize, dst: usize) {
    if dst < src {
        slice[dst..src + len].rotate_left(len);
    } else {
        slice[src..dst + len].rotate_right(len);
    }
}

/// Score a candidate matching against the ground truth.
fn score_matching(
    scorer: &levenshtein::BatchComparator<u8>,
    ground_len: usize,
    matching: &[IndexedBlock],
    buf: &mut Vec<u8>,
    best_known: Option<f64>,
) -> MatchingLoss {
    // Compute jump distance first (very cheap pre-filter).
    let jump_distance = calculate_jump_distance(matching);
    if let Some(best) = best_known {
        if jump_distance >= best {
            return MatchingLoss {
                fuzzy: f64::MAX,
                jump_distance,
                geo_distance: 0.0,
            };
        }
    }

    // Build concatenated bytes into reusable buffer
    buf.clear();
    for (i, (_, text)) in matching.iter().enumerate() {
        if i > 0 {
            buf.push(b' ');
        }
        buf.extend_from_slice(text.as_bytes());
    }

    // Use score_cutoff for early termination in levenshtein computation
    let lev_dist = if let Some(best) = best_known {
        let max_fuzzy = (best - jump_distance).max(0.0) as usize;
        let args = levenshtein::Args::default()
            .score_hint(max_fuzzy)
            .score_cutoff(max_fuzzy);
        match scorer.distance_with_args(buf.iter().copied(), &args) {
            Some(d) => d as f64,
            None => {
                return MatchingLoss {
                    fuzzy: (ground_len + buf.len()) as f64,
                    jump_distance,
                    geo_distance: 0.0,
                };
            }
        }
    } else {
        scorer.distance(buf.iter().copied()) as f64
    };

    MatchingLoss {
        fuzzy: lev_dist,
        jump_distance,
        geo_distance: 0.0,
    }
}

fn calculate_jump_distance(matching: &[IndexedBlock]) -> f64 {
    let mut j = 0.0;
    for i in 1..matching.len() {
        let left = matching[i - 1];
        let right = matching[i];
        let expected_right = left.0 + 1;
        let distance = (expected_right as i32 - right.0 as i32).abs() as usize;
        if distance == 0 {
            continue;
        }
        j += JUMP_STATIC_PENALTY + (distance as f64) * JUMP_OFFSET_PENALTY;
    }
    j
}

fn normalize_whitespace(text: &str) -> String {
    text.split_whitespace().collect::<Vec<_>>().join(" ")
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_levenshtein_distance() {
        assert_eq!(levenshtein::distance("abc".chars(), "abc".chars()), 0);
        assert_eq!(levenshtein::distance("abc".chars(), "ab".chars()), 1);
        assert_eq!(levenshtein::distance("ab".chars(), "abc".chars()), 1);
        assert_eq!(levenshtein::distance("abc".chars(), "def".chars()), 3);
        assert_eq!(levenshtein::distance("".chars(), "abc".chars()), 3);
        assert_eq!(levenshtein::distance("abc".chars(), "".chars()), 3);
    }

    #[test]
    fn test_normalize_whitespace() {
        assert_eq!(normalize_whitespace("  hello   world  "), "hello world");
        assert_eq!(normalize_whitespace("a\nb\tc"), "a b c");
    }

    #[test]
    fn test_optimize_simple() {
        let ground = "hello world";
        let blocks = &["world", "hello"];
        let (loss, result) = optimize(ground, blocks);
        assert!(loss.abs() >= 0.0);
        assert_eq!(result.len(), 2);
    }

    #[test]
    fn test_apply_undo_shift() {
        let mut v = vec![0, 1, 2, 3, 4, 5, 6, 7];
        apply_shift(&mut v, 2, 2, 5);
        assert_eq!(v, vec![0, 1, 4, 5, 6, 2, 3, 7]);
        undo_shift(&mut v, 2, 2, 5);
        assert_eq!(v, vec![0, 1, 2, 3, 4, 5, 6, 7]);

        apply_shift(&mut v, 5, 2, 1);
        assert_eq!(v, vec![0, 5, 6, 1, 2, 3, 4, 7]);
        undo_shift(&mut v, 5, 2, 1);
        assert_eq!(v, vec![0, 1, 2, 3, 4, 5, 6, 7]);
    }

    #[test]
    fn test_f64_ord_roundtrip() {
        let vals = [0.0, 1.0, 100.5, 315.6, 1000.0];
        for &v in &vals {
            assert_eq!(v, ord_u64_to_f64(f64_to_ord_u64(v)));
        }
        assert!(f64_to_ord_u64(1.0) < f64_to_ord_u64(2.0));
        assert!(f64_to_ord_u64(100.0) < f64_to_ord_u64(315.6));
    }
}