fastalp 0.1.7

use crate::{
  bitpack::bitpack_encoded,
  constants::AlpFloat,
  sampler::{BestParams, find_best_params},
};

pub const ALP_MAGIC: [u8; 2] = *b"AP";
pub const TYPE_F64: u8 = 1;
pub const TYPE_F32: u8 = 2;

/// 单个异常值记录
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Exception<R> {
  pub pos: u16,
  pub bits: R,
}

pub type ExceptionF64 = Exception<u64>;
pub type ExceptionF32 = Exception<u32>;

/// 通用压缩浮点数组并直接写入 `dst` 缓冲区
pub fn compress_into<F: AlpFloat>(data: &[F], dst: &mut Vec<u8>) {
  let count = data.len().min(u16::MAX as usize) as u16;
  if count == 0 {
    dst.extend_from_slice(&ALP_MAGIC);
    dst.push(F::TYPE_BYTE);
    dst.extend_from_slice(&0u16.to_le_bytes());
    dst.push(0); // exp
    dst.push(0); // fac
    dst.push(0); // bit_width
    F::write_base(F::ZERO_INT, dst); // base
    dst.extend_from_slice(&0u16.to_le_bytes()); // exceptions count
    return;
  }

  let slice = &data[..count as usize];
  let BestParams { exp, fac } = find_best_params(slice);

  let exp_factor = F::exp_factor(exp, fac);
  let fac_int = F::fac_int(fac);
  let frac_exp = F::frac_exp(exp);

  let mut encoded_ints = Vec::with_capacity(slice.len());
  let mut exceptions = Vec::new();
  let mut min_val = F::MAX_INT;
  let mut max_val = F::MIN_INT;

  for (i, &val) in slice.iter().enumerate() {
    match F::try_encode_fast(val, exp_factor, fac_int, frac_exp) {
      Some(enc) => {
        encoded_ints.push(enc);
        min_val = min_val.min(enc);
        max_val = max_val.max(enc);
      }
      None => {
        encoded_ints.push(F::ZERO_INT);
        exceptions.push(Exception {
          pos: i as u16,
          bits: val.to_raw_bits(),
        });
      }
    }
  }

  let base = if min_val <= max_val {
    min_val
  } else {
    F::ZERO_INT
  };
  let max_offset = if min_val <= max_val {
    F::calc_range(min_val, max_val)
  } else {
    0
  };

  if !exceptions.is_empty() {
    for exc in &exceptions {
      // SAFETY: exc.pos 是在上方遍历 slice (0..slice.len()) 时记录的索引，encoded_ints 的长度与 slice.len() 完全一致，因此 exc.pos as usize 严格小于 encoded_ints.len()，索引安全有效。
      unsafe {
        *encoded_ints.get_unchecked_mut(exc.pos as usize) = base;
      }
    }
  }

  let bit_width = F::bits_needed(max_offset);

  // 1. Header (8B)
  dst.extend_from_slice(&ALP_MAGIC);
  dst.push(F::TYPE_BYTE);
  dst.extend_from_slice(&count.to_le_bytes());
  dst.push(exp);
  dst.push(fac);
  dst.push(bit_width);

  // 2. Base
  F::write_base(base, dst);

  // 3. Bitpacked data
  bitpack_encoded::<F>(&encoded_ints, base, bit_width, dst);

  // 4. Exceptions
  let exc_count = exceptions.len() as u16;
  dst.extend_from_slice(&exc_count.to_le_bytes());
  for exc in exceptions {
    F::write_exception(exc.pos, exc.bits, dst);
  }
}

/// 压缩 f64 浮点数组并直接写入 `dst` 缓冲区
#[inline]
pub fn compress_f64_into(data: &[f64], dst: &mut Vec<u8>) {
  compress_into(data, dst);
}

/// 压缩 f32 浮点数组并直接写入 `dst` 缓冲区
#[inline]
pub fn compress_f32_into(data: &[f32], dst: &mut Vec<u8>) {
  compress_into(data, dst);
}

/// 通用压缩浮点数切片
#[inline]
pub fn compress<F: AlpFloat>(data: &[F]) -> Vec<u8> {
  let mut dst = Vec::with_capacity(data.len() * 2 + 16);
  compress_into(data, &mut dst);
  dst
}

/// 压缩 f64 浮点数切片
#[inline]
pub fn compress_f64(data: &[f64]) -> Vec<u8> {
  compress(data)
}

/// 压缩 f32 浮点数切片
#[inline]
pub fn compress_f32(data: &[f32]) -> Vec<u8> {
  compress(data)
}