grib-core 0.8.0

//! Shared binary encoding primitives used by GRIB section parsers and writers.

use crate::error::{Error, Result};

pub const U24_MAX: u32 = 0x00ff_ffff;

const WMO_I8_MAX: i64 = 0x7f;
const WMO_I16_MAX: i64 = 0x7fff;
const WMO_I24_MAX: i64 = 0x7f_ffff;
const WMO_I32_MAX: i64 = 0x7fff_ffff;
const IBM_FRACTION_SCALE: f64 = 16_777_216.0;
const IEEE_BINARY128_EXPONENT_BIAS: i32 = 16_383;
const IEEE_BINARY128_FRACTION_BITS: u32 = 112;
const IEEE_BINARY128_FRACTION_MASK: u128 = (1_u128 << IEEE_BINARY128_FRACTION_BITS) - 1;

pub fn write_u8_be(out: &mut Vec<u8>, value: u8) -> Result<()> {
    reserve_bytes(out, 1)?;
    out.push(value);
    Ok(())
}

pub fn write_u16_be(out: &mut Vec<u8>, value: u16) -> Result<()> {
    reserve_bytes(out, 2)?;
    out.extend_from_slice(&value.to_be_bytes());
    Ok(())
}

pub fn write_u24_be(out: &mut Vec<u8>, value: u32) -> Result<()> {
    if value > U24_MAX {
        return Err(Error::ValueOutOfRange(format!(
            "value {value} does not fit in an unsigned 24-bit integer"
        )));
    }

    reserve_bytes(out, 3)?;
    out.extend_from_slice(&[
        ((value >> 16) & 0xff) as u8,
        ((value >> 8) & 0xff) as u8,
        (value & 0xff) as u8,
    ]);
    Ok(())
}

pub fn write_u32_be(out: &mut Vec<u8>, value: u32) -> Result<()> {
    reserve_bytes(out, 4)?;
    out.extend_from_slice(&value.to_be_bytes());
    Ok(())
}

pub fn write_u64_be(out: &mut Vec<u8>, value: u64) -> Result<()> {
    reserve_bytes(out, 8)?;
    out.extend_from_slice(&value.to_be_bytes());
    Ok(())
}

pub fn read_u24_be(bytes: &[u8]) -> Option<u32> {
    let bytes = bytes.get(..3)?;
    Some((u32::from(bytes[0]) << 16) | (u32::from(bytes[1]) << 8) | u32::from(bytes[2]))
}

pub fn decode_wmo_i8(byte: u8) -> i16 {
    let magnitude = i16::from(byte & 0x7f);
    if byte & 0x80 == 0 {
        magnitude
    } else {
        -magnitude
    }
}

pub fn decode_wmo_i16(bytes: &[u8]) -> Option<i16> {
    let raw = u16::from_be_bytes(bytes.get(..2)?.try_into().ok()?);
    let magnitude = (raw & 0x7fff) as i16;
    Some(if raw & 0x8000 == 0 {
        magnitude
    } else {
        -magnitude
    })
}

pub fn decode_wmo_i24(bytes: &[u8]) -> Option<i32> {
    let raw = read_u24_be(bytes)?;
    let magnitude = (raw & 0x7f_ffff) as i32;
    Some(if raw & 0x80_0000 == 0 {
        magnitude
    } else {
        -magnitude
    })
}

pub fn decode_wmo_i32(bytes: &[u8]) -> Option<i32> {
    let raw = u32::from_be_bytes(bytes.get(..4)?.try_into().ok()?);
    let magnitude = (raw & 0x7fff_ffff) as i32;
    Some(if raw & 0x8000_0000 == 0 {
        magnitude
    } else {
        -magnitude
    })
}

pub fn encode_wmo_i8(value: i16) -> Option<u8> {
    let magnitude = checked_magnitude(value, WMO_I8_MAX)?;
    Some(if value < 0 {
        0x80 | magnitude as u8
    } else {
        magnitude as u8
    })
}

pub fn encode_wmo_i16(value: i16) -> Option<[u8; 2]> {
    let magnitude = checked_magnitude(value, WMO_I16_MAX)? as u16;
    Some(if value < 0 {
        (0x8000 | magnitude).to_be_bytes()
    } else {
        magnitude.to_be_bytes()
    })
}

pub fn encode_wmo_i24(value: i32) -> Option<[u8; 3]> {
    let magnitude = checked_magnitude(value, WMO_I24_MAX)? as u32;
    let raw = if value < 0 {
        0x80_0000 | magnitude
    } else {
        magnitude
    };
    Some([
        ((raw >> 16) & 0xff) as u8,
        ((raw >> 8) & 0xff) as u8,
        (raw & 0xff) as u8,
    ])
}

pub fn encode_wmo_i32(value: i32) -> Option<[u8; 4]> {
    let magnitude = checked_magnitude(value, WMO_I32_MAX)? as u32;
    Some(if value < 0 {
        (0x8000_0000 | magnitude).to_be_bytes()
    } else {
        magnitude.to_be_bytes()
    })
}

pub fn decode_ibm_f32(bytes: [u8; 4]) -> f32 {
    if bytes == [0, 0, 0, 0] {
        return 0.0;
    }

    let sign = if bytes[0] & 0x80 == 0 { 1.0 } else { -1.0 };
    let exponent = i32::from(bytes[0] & 0x7f) - 64;
    let mantissa = (u32::from(bytes[1]) << 16) | (u32::from(bytes[2]) << 8) | u32::from(bytes[3]);
    let value = sign * f64::from(mantissa) / IBM_FRACTION_SCALE * 16f64.powi(exponent);
    value as f32
}

pub fn decode_ibm_f32_slice(bytes: &[u8]) -> Option<f32> {
    Some(decode_ibm_f32(bytes.get(..4)?.try_into().ok()?))
}

pub fn encode_ibm_f32(value: f32) -> Option<[u8; 4]> {
    if value == 0.0 {
        return Some([0, 0, 0, 0]);
    }
    if !value.is_finite() {
        return None;
    }

    let sign = if value.is_sign_negative() { 0x80 } else { 0x00 };
    let mut fraction = f64::from(value.abs());
    let mut exponent = 64i32;

    while fraction < 0.0625 {
        fraction *= 16.0;
        exponent -= 1;
        if exponent < 0 {
            return None;
        }
    }

    while fraction >= 1.0 {
        fraction /= 16.0;
        exponent += 1;
        if exponent > 0x7f {
            return None;
        }
    }

    let mut mantissa = (fraction * IBM_FRACTION_SCALE).round() as u32;
    if mantissa == 0 {
        return Some([0, 0, 0, 0]);
    }
    if mantissa >= 0x0100_0000 {
        mantissa >>= 4;
        exponent += 1;
        if exponent > 0x7f {
            return None;
        }
    }

    Some([
        sign | exponent as u8,
        ((mantissa >> 16) & 0xff) as u8,
        ((mantissa >> 8) & 0xff) as u8,
        (mantissa & 0xff) as u8,
    ])
}

/// Converts a big-endian IEEE 754 binary128 value to binary64.
///
/// Values outside the binary64 range become signed infinity or signed zero.
/// Finite values are rounded to nearest, with ties to even.
pub fn decode_ieee_binary128(bytes: [u8; 16]) -> f64 {
    let bits = u128::from_be_bytes(bytes);
    let sign = ((bits >> 127) as u64) << 63;
    let exponent = ((bits >> IEEE_BINARY128_FRACTION_BITS) & 0x7fff) as u16;
    let fraction = bits & IEEE_BINARY128_FRACTION_MASK;

    if exponent == 0x7fff {
        if fraction == 0 {
            return f64::from_bits(sign | (0x7ff_u64 << 52));
        }
        let mut payload = (fraction >> (IEEE_BINARY128_FRACTION_BITS - 52)) as u64;
        if payload == 0 {
            payload = 1;
        }
        return f64::from_bits(sign | (0x7ff_u64 << 52) | payload);
    }

    let (significand, binary_exponent) = if exponent == 0 {
        (fraction, 1 - IEEE_BINARY128_EXPONENT_BIAS - 112)
    } else {
        (
            (1_u128 << IEEE_BINARY128_FRACTION_BITS) | fraction,
            i32::from(exponent) - IEEE_BINARY128_EXPONENT_BIAS - 112,
        )
    };
    if significand == 0 {
        return f64::from_bits(sign);
    }

    let leading_bit = 127 - significand.leading_zeros() as i32;
    let mut result_exponent = leading_bit + binary_exponent;
    if result_exponent > 1023 {
        return f64::from_bits(sign | (0x7ff_u64 << 52));
    }

    if result_exponent >= -1022 {
        let shift = u32::try_from(leading_bit - 52)
            .expect("binary128 normal significand always exceeds binary64 precision");
        let mut rounded = round_shift_right_to_even(significand, shift);
        if rounded == 1_u128 << 53 {
            rounded >>= 1;
            result_exponent += 1;
            if result_exponent > 1023 {
                return f64::from_bits(sign | (0x7ff_u64 << 52));
            }
        }

        let exponent_bits = u64::try_from(result_exponent + 1023)
            .expect("normal binary64 exponent is non-negative");
        let fraction_bits = (rounded as u64) & ((1_u64 << 52) - 1);
        return f64::from_bits(sign | (exponent_bits << 52) | fraction_bits);
    }

    let subnormal_shift = u32::try_from(-(binary_exponent + 1074))
        .expect("binary64 subnormal conversion requires a right shift");
    let rounded = round_shift_right_to_even(significand, subnormal_shift);
    if rounded >= 1_u128 << 52 {
        return f64::from_bits(sign | (1_u64 << 52));
    }
    f64::from_bits(sign | rounded as u64)
}

/// Encodes a binary64 value exactly as big-endian IEEE 754 binary128.
pub fn encode_ieee_binary128(value: f64) -> [u8; 16] {
    let bits = value.to_bits();
    let sign = u128::from(bits >> 63) << 127;
    let exponent = ((bits >> 52) & 0x7ff) as u16;
    let fraction = bits & ((1_u64 << 52) - 1);

    let encoded = match exponent {
        0 if fraction == 0 => sign,
        0 => {
            let leading_bit = 63 - fraction.leading_zeros();
            let unbiased_exponent = leading_bit as i32 - 1074;
            let binary128_exponent =
                u128::try_from(unbiased_exponent + IEEE_BINARY128_EXPONENT_BIAS)
                    .expect("every binary64 subnormal is normal in binary128");
            let trailing = u128::from(fraction ^ (1_u64 << leading_bit));
            let binary128_fraction = trailing << (IEEE_BINARY128_FRACTION_BITS - leading_bit);
            sign | (binary128_exponent << IEEE_BINARY128_FRACTION_BITS) | binary128_fraction
        }
        0x7ff => {
            sign | (0x7fff_u128 << IEEE_BINARY128_FRACTION_BITS) | (u128::from(fraction) << 60)
        }
        _ => {
            let unbiased_exponent = i32::from(exponent) - 1023;
            let binary128_exponent =
                u128::try_from(unbiased_exponent + IEEE_BINARY128_EXPONENT_BIAS)
                    .expect("every finite binary64 value is representable in binary128");
            sign | (binary128_exponent << IEEE_BINARY128_FRACTION_BITS)
                | (u128::from(fraction) << 60)
        }
    };

    encoded.to_be_bytes()
}

fn round_shift_right_to_even(value: u128, shift: u32) -> u128 {
    match shift {
        0 => value,
        1..=127 => {
            let quotient = value >> shift;
            let remainder_mask = (1_u128 << shift) - 1;
            let remainder = value & remainder_mask;
            let halfway = 1_u128 << (shift - 1);
            quotient
                + u128::from(remainder > halfway || (remainder == halfway && quotient & 1 != 0))
        }
        128 => u128::from(value > (1_u128 << 127)),
        _ => 0,
    }
}

fn reserve_bytes(out: &mut Vec<u8>, additional: usize) -> Result<()> {
    out.try_reserve(additional)
        .map_err(|error| Error::allocation("output bytes", additional, error))
}

fn checked_magnitude<T>(value: T, max: i64) -> Option<i64>
where
    T: Into<i64>,
{
    let value = value.into();
    let magnitude = value.checked_abs()?;
    (magnitude <= max).then_some(magnitude)
}

#[cfg(test)]
mod tests {
    use super::{
        decode_ibm_f32, decode_ibm_f32_slice, decode_ieee_binary128, decode_wmo_i16,
        decode_wmo_i24, decode_wmo_i32, decode_wmo_i8, encode_ibm_f32, encode_ieee_binary128,
        encode_wmo_i16, encode_wmo_i24, encode_wmo_i32, encode_wmo_i8, read_u24_be, write_u16_be,
        write_u24_be, write_u32_be, write_u64_be, write_u8_be,
    };

    #[test]
    fn writes_big_endian_unsigned_integers() {
        let mut out = Vec::new();
        write_u8_be(&mut out, 0x12).unwrap();
        write_u16_be(&mut out, 0x3456).unwrap();
        write_u24_be(&mut out, 0x789abc).unwrap();
        write_u32_be(&mut out, 0xdef0_1234).unwrap();
        write_u64_be(&mut out, 0x5678_9abc_def0_1234).unwrap();

        assert_eq!(
            out,
            [
                0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc, 0xde, 0xf0, 0x12, 0x34, 0x56, 0x78, 0x9a, 0xbc,
                0xde, 0xf0, 0x12, 0x34
            ]
        );
    }

    #[test]
    fn rejects_unsigned_24_bit_overflow() {
        let mut out = Vec::new();
        let error = write_u24_be(&mut out, 0x0100_0000).unwrap_err();
        assert!(matches!(error, crate::Error::ValueOutOfRange(_)));
        assert!(out.is_empty());
    }

    #[test]
    fn reads_unsigned_24_bit_integers() {
        assert_eq!(read_u24_be(&[0x12, 0x34, 0x56]), Some(0x12_3456));
        assert_eq!(read_u24_be(&[0x12, 0x34]), None);
    }

    #[test]
    fn roundtrips_wmo_signed_integers() {
        assert_eq!(decode_wmo_i8(encode_wmo_i8(-5).unwrap()), -5);
        assert_eq!(decode_wmo_i16(&encode_wmo_i16(-500).unwrap()), Some(-500));
        assert_eq!(
            decode_wmo_i24(&encode_wmo_i24(-500_000).unwrap()),
            Some(-500_000)
        );
        assert_eq!(
            decode_wmo_i32(&encode_wmo_i32(-500_000_000).unwrap()),
            Some(-500_000_000)
        );
    }

    #[test]
    fn matches_known_wmo_signed_byte_patterns() {
        assert_eq!(decode_wmo_i8(0x85), -5);
        assert_eq!(decode_wmo_i16(&[0x80, 0x05]), Some(-5));
        assert_eq!(decode_wmo_i24(&[0x80, 0x00, 0x05]), Some(-5));
        assert_eq!(decode_wmo_i32(&[0x80, 0x00, 0x00, 0x05]), Some(-5));
        assert_eq!(encode_wmo_i8(5), Some(0x05));
        assert_eq!(encode_wmo_i16(5), Some([0x00, 0x05]));
        assert_eq!(encode_wmo_i24(5), Some([0x00, 0x00, 0x05]));
        assert_eq!(encode_wmo_i32(5), Some([0x00, 0x00, 0x00, 0x05]));
    }

    #[test]
    fn rejects_wmo_signed_integer_overflow() {
        assert_eq!(encode_wmo_i8(128), None);
        assert_eq!(encode_wmo_i16(i16::MIN), None);
        assert_eq!(encode_wmo_i24(0x80_0000), None);
        assert_eq!(encode_wmo_i32(i32::MIN), None);
        assert_eq!(decode_wmo_i16(&[0x00]), None);
        assert_eq!(decode_wmo_i24(&[0x00, 0x00]), None);
        assert_eq!(decode_wmo_i32(&[0x00, 0x00, 0x00]), None);
    }

    #[test]
    fn decodes_known_ibm_float_patterns() {
        assert_eq!(decode_ibm_f32([0x00, 0x00, 0x00, 0x00]), 0.0);
        assert_eq!(decode_ibm_f32([0x41, 0x10, 0x00, 0x00]), 1.0);
        assert_eq!(decode_ibm_f32([0xc1, 0x20, 0x00, 0x00]), -2.0);
        assert_eq!(decode_ibm_f32([0x41, 0xa0, 0x00, 0x00]), 10.0);
        assert_eq!(decode_ibm_f32_slice(&[0x41, 0x10, 0x00, 0x00]), Some(1.0));
        assert_eq!(decode_ibm_f32_slice(&[0x41, 0x10, 0x00]), None);
    }

    #[test]
    fn roundtrips_finite_ibm_float_values() {
        for value in [0.0_f32, 0.25, 0.5, 1.0, -2.0, 10.0, 16.0, 1234.5] {
            let encoded = encode_ibm_f32(value).unwrap();
            let decoded = decode_ibm_f32(encoded);
            let tolerance = value.abs().max(1.0) * 1.0e-6;
            assert!(
                (decoded - value).abs() <= tolerance,
                "value {value} encoded as {encoded:02x?} decoded as {decoded}"
            );
        }
    }

    #[test]
    fn encodes_known_ibm_float_patterns() {
        assert_eq!(encode_ibm_f32(0.0), Some([0x00, 0x00, 0x00, 0x00]));
        assert_eq!(encode_ibm_f32(1.0), Some([0x41, 0x10, 0x00, 0x00]));
        assert_eq!(encode_ibm_f32(-2.0), Some([0xc1, 0x20, 0x00, 0x00]));
        assert_eq!(encode_ibm_f32(10.0), Some([0x41, 0xa0, 0x00, 0x00]));
    }

    #[test]
    fn rejects_non_finite_ibm_float_values() {
        assert_eq!(encode_ibm_f32(f32::NAN), None);
        assert_eq!(encode_ibm_f32(f32::INFINITY), None);
        assert_eq!(encode_ibm_f32(f32::NEG_INFINITY), None);
    }

    #[test]
    fn roundtrips_binary64_values_through_binary128() {
        for value in [
            0.0,
            -0.0,
            1.0,
            -2.5,
            f64::MIN_POSITIVE,
            f64::from_bits(1),
            f64::MAX,
            f64::INFINITY,
            f64::NEG_INFINITY,
        ] {
            let decoded = decode_ieee_binary128(encode_ieee_binary128(value));
            assert_eq!(decoded.to_bits(), value.to_bits(), "failed for {value:?}");
        }

        assert!(decode_ieee_binary128(encode_ieee_binary128(f64::NAN)).is_nan());
    }

    #[test]
    fn decodes_known_binary128_values() {
        assert_eq!(
            decode_ieee_binary128(0x3fff_0000_0000_0000_0000_0000_0000_0000_u128.to_be_bytes()),
            1.0
        );
        assert_eq!(
            decode_ieee_binary128(0xc000_0000_0000_0000_0000_0000_0000_0000_u128.to_be_bytes()),
            -2.0
        );
    }

    #[test]
    fn rounds_binary128_ties_to_even_binary64_values() {
        let exponent = 0x3fff_u128 << 112;
        let halfway_above_one = exponent | (1_u128 << 59);
        let above_halfway = halfway_above_one | 1;
        let halfway_above_odd = exponent | (3_u128 << 59);

        assert_eq!(
            decode_ieee_binary128(halfway_above_one.to_be_bytes()).to_bits(),
            1.0_f64.to_bits()
        );
        assert_eq!(
            decode_ieee_binary128(above_halfway.to_be_bytes()).to_bits(),
            1.0_f64.to_bits() + 1
        );
        assert_eq!(
            decode_ieee_binary128(halfway_above_odd.to_be_bytes()).to_bits(),
            1.0_f64.to_bits() + 2
        );
    }

    #[test]
    fn rounds_binary128_at_binary64_subnormal_boundary() {
        let half_minimum_subnormal = (15_308_u128 << 112).to_be_bytes();
        let above_half_minimum_subnormal = ((15_308_u128 << 112) | 1).to_be_bytes();

        assert_eq!(decode_ieee_binary128(half_minimum_subnormal).to_bits(), 0);
        assert_eq!(
            decode_ieee_binary128(above_half_minimum_subnormal).to_bits(),
            1
        );
    }
}