helpers_codec/common/

codec.rs

// MIT License
//
// Copyright (c) 2026 Raja Lehtihet & Wael El Oraiby
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all
// copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

//! Example message and packing helpers.
//!
//! The routines in this module are intentionally simple and geared toward examples.
//! They provide:
//! - byte <-> ring-element encoding for scaled binary messages,
//! - coefficient/ring-element compression and decompression helpers.

use nc_polynomial::{PolynomialError, RingContext, RingElem};

/// Errors returned by example codec helpers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CodecError {
    /// Message needs more bits than ring coefficient capacity.
    MessageTooLong {
        message_bits: usize,
        capacity_bits: usize,
    },
    /// Decoding requested too many bytes for the ring capacity.
    DecodeLengthTooLong {
        requested_bits: usize,
        capacity_bits: usize,
    },
    /// Compression bit width must be in the supported range.
    InvalidCompressionBits(u8),
    /// Packed coefficient vector length does not match ring size.
    CompressedLengthMismatch { expected: usize, actual: usize },
    /// Ring element construction or ring arithmetic failed.
    Polynomial(PolynomialError),
}

impl core::fmt::Display for CodecError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::MessageTooLong {
                message_bits,
                capacity_bits,
            } => write!(
                f,
                "message requires {message_bits} bits but ring capacity is {capacity_bits} bits"
            ),
            Self::DecodeLengthTooLong {
                requested_bits,
                capacity_bits,
            } => write!(
                f,
                "decode requested {requested_bits} bits but ring capacity is {capacity_bits} bits"
            ),
            Self::InvalidCompressionBits(bits) => write!(
                f,
                "invalid compression width {bits}, expected a value in 1..=16"
            ),
            Self::CompressedLengthMismatch { expected, actual } => write!(
                f,
                "compressed coefficient length mismatch: expected {expected}, got {actual}"
            ),
            Self::Polynomial(err) => write!(f, "polynomial operation failed: {err}"),
        }
    }
}

impl std::error::Error for CodecError {}

impl From<PolynomialError> for CodecError {
    fn from(value: PolynomialError) -> Self {
        Self::Polynomial(value)
    }
}

/// Encodes bytes into a ring element by mapping message bits to `{0, floor((q+1)/2)}`.
///
/// Bit order is little-endian within each byte.
pub fn encode_message_scaled_bits_example(
    ctx: &RingContext,
    message: &[u8],
) -> Result<RingElem, CodecError> {
    // In quotient rings like `Z_q[x]/(x^n+1)`, we treat `n` as the usable bit capacity.
    // Coefficient index `i` stores bit `i`.
    let capacity_bits = ctx.max_degree();
    let message_bits = message.len().saturating_mul(8);
    if message_bits > capacity_bits {
        return Err(CodecError::MessageTooLong {
            message_bits,
            capacity_bits,
        });
    }

    // Allocate full polynomial width (`n + 1` slots in this crate's dense representation).
    let mut coeffs = vec![0_u64; ctx.max_degree() + 1];
    // Encoding for bit=1 is midpoint-ish so threshold decoding is robust to moderate noise.
    let one_value = (ctx.modulus() + 1) / 2;

    // Little-endian bit order per byte:
    // - bit 0 goes to lower coefficient index,
    // - bit 7 goes to higher coefficient index.
    for (byte_idx, &byte) in message.iter().enumerate() {
        for bit_idx in 0..8 {
            if (byte >> bit_idx) & 1 == 1 {
                coeffs[byte_idx * 8 + bit_idx] = one_value;
            }
        }
    }

    // Canonicalize into ring representation.
    Ok(ctx.element(&coeffs)?)
}

/// Decodes bytes from a scaled-bit ring element using quarter-range thresholding.
///
/// Coefficients in `[q/4, 3q/4)` decode to bit `1`, and the rest decode to `0`.
pub fn decode_message_scaled_bits_example(
    element: &RingElem,
    output_len_bytes: usize,
) -> Result<Vec<u8>, CodecError> {
    // Decode limit follows the same bit-capacity convention used at encode time.
    let capacity_bits = element.params().max_degree();
    let requested_bits = output_len_bytes.saturating_mul(8);
    if requested_bits > capacity_bits {
        return Err(CodecError::DecodeLengthTooLong {
            requested_bits,
            capacity_bits,
        });
    }

    let q = element.params().modulus();
    // Threshold window near middle of modulus range.
    // Values in `[q/4, 3q/4)` decode to 1; others decode to 0.
    let lower = q / 4;
    let upper = (3 * q) / 4;

    let mut out = vec![0_u8; output_len_bytes];
    for bit_index in 0..requested_bits {
        let coeff = element.coefficients()[bit_index];
        let bit = coeff >= lower && coeff < upper;
        if bit {
            // Restore little-endian bit position.
            out[bit_index / 8] |= 1_u8 << (bit_index % 8);
        }
    }

    Ok(out)
}

/// Compresses one coefficient into `bits` bits.
pub fn compress_coefficient_example(value: u64, modulus: u64, bits: u8) -> Result<u16, CodecError> {
    validate_compression_bits(bits)?;

    // Quantization levels = 2^bits.
    let levels = 1_u128 << bits;
    let q = modulus as u128;
    let v = (value % modulus) as u128;

    // Rounded scaling from `[0, q)` into `[0, 2^bits)`.
    let compressed = ((v * levels + (q / 2)) / q) % levels;
    Ok(compressed as u16)
}

/// Decompresses one `bits`-bit coefficient back into `Z_q`.
pub fn decompress_coefficient_example(
    compressed: u16,
    modulus: u64,
    bits: u8,
) -> Result<u64, CodecError> {
    validate_compression_bits(bits)?;

    let levels = 1_u128 << bits;
    let q = modulus as u128;
    let c = (compressed as u128) % levels;

    // Rounded inverse scaling from `[0, 2^bits)` back into `[0, q)`.
    let decompressed = (c * q + (levels / 2)) / levels;
    Ok((decompressed % q) as u64)
}

/// Compresses all coefficients of a ring element into `bits` bits each.
pub fn compress_ring_elem_example(element: &RingElem, bits: u8) -> Result<Vec<u16>, CodecError> {
    validate_compression_bits(bits)?;

    let q = element.params().modulus();
    // Compress each coefficient independently.
    let mut packed = Vec::with_capacity(element.coefficients().len());
    for &coeff in element.coefficients() {
        packed.push(compress_coefficient_example(coeff, q, bits)?);
    }
    Ok(packed)
}

/// Decompresses packed coefficients into a canonical ring element.
pub fn decompress_ring_elem_example(
    ctx: &RingContext,
    packed: &[u16],
    bits: u8,
) -> Result<RingElem, CodecError> {
    validate_compression_bits(bits)?;

    // Packed vector must match dense polynomial width.
    let expected = ctx.max_degree() + 1;
    if packed.len() != expected {
        return Err(CodecError::CompressedLengthMismatch {
            expected,
            actual: packed.len(),
        });
    }

    // Decompress each coefficient independently and rebuild ring element.
    let mut coeffs = vec![0_u64; expected];
    for (index, &coeff) in packed.iter().enumerate() {
        coeffs[index] = decompress_coefficient_example(coeff, ctx.modulus(), bits)?;
    }

    Ok(ctx.element(&coeffs)?)
}

fn validate_compression_bits(bits: u8) -> Result<(), CodecError> {
    if !(1..=16).contains(&bits) {
        return Err(CodecError::InvalidCompressionBits(bits));
    }
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::{
        CodecError, compress_coefficient_example, compress_ring_elem_example,
        decode_message_scaled_bits_example, decompress_coefficient_example,
        decompress_ring_elem_example, encode_message_scaled_bits_example,
    };
    use nc_polynomial::RingContext;

    fn ctx() -> RingContext {
        let max_degree = 32;
        let mut modulus_poly = vec![0_u64; max_degree + 1];
        modulus_poly[0] = 1;
        modulus_poly[max_degree] = 1;

        RingContext::from_parts(max_degree, 998_244_353, &modulus_poly, 3)
            .expect("context should build")
    }

    #[test]
    fn message_scaled_bits_round_trip() {
        let ctx = ctx();
        let message = [0xA5_u8, 0x5A, 0xC3, 0x3C];

        // End-to-end encode/decode should preserve payload exactly.
        let encoded =
            encode_message_scaled_bits_example(&ctx, &message).expect("encoding should work");
        let decoded = decode_message_scaled_bits_example(&encoded, message.len())
            .expect("decoding should work");

        assert_eq!(decoded, message);
    }

    #[test]
    fn message_encoding_rejects_oversized_payload() {
        let ctx = ctx();
        let oversized = [0_u8; 5]; // 40 bits, ring has capacity 32 bits.

        let err = encode_message_scaled_bits_example(&ctx, &oversized).expect_err("expected error");
        assert_eq!(
            err,
            CodecError::MessageTooLong {
                message_bits: 40,
                capacity_bits: 32,
            }
        );
    }

    #[test]
    fn message_decoding_rejects_oversized_request() {
        let ctx = ctx();
        let encoded = encode_message_scaled_bits_example(&ctx, &[0xAA, 0x55, 0x11, 0xEE])
            .expect("encoding should work");

        let err = decode_message_scaled_bits_example(&encoded, 5).expect_err("expected error");
        assert_eq!(
            err,
            CodecError::DecodeLengthTooLong {
                requested_bits: 40,
                capacity_bits: 32,
            }
        );
    }

    #[test]
    fn coefficient_compression_round_trips_with_bounded_error() {
        let modulus = 998_244_353_u64;
        let bits = 10_u8;
        // Quantization step size.
        let tolerance = modulus / (1_u64 << bits);

        for value in [
            0_u64,
            1,
            2,
            17,
            31_337,
            modulus / 3,
            modulus - 2,
            modulus - 1,
        ] {
            let compressed =
                compress_coefficient_example(value, modulus, bits).expect("compress should work");
            let decompressed = decompress_coefficient_example(compressed, modulus, bits)
                .expect("decompress should work");

            // Distance in cyclic modular space.
            let direct_distance = value.abs_diff(decompressed);
            let wraparound_distance = modulus - direct_distance;
            let distance = direct_distance.min(wraparound_distance);
            assert!(distance <= tolerance + 1);
        }
    }

    #[test]
    fn ring_elem_compression_round_trip_shape_and_domain() {
        let ctx = ctx();
        let elem = ctx
            .element(&[
                1, 17, 999, 1_337, 7, 5, 3, 2, 11, 13, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61,
                67, 71, 73, 79, 83, 89, 97, 101, 103, 107, 109, 113,
            ])
            .expect("element should build");

        // Compression keeps vector width fixed.
        let packed = compress_ring_elem_example(&elem, 11).expect("compression should work");
        assert_eq!(packed.len(), ctx.max_degree() + 1);

        // Decompression should still produce valid residues in `Z_q`.
        let decoded =
            decompress_ring_elem_example(&ctx, &packed, 11).expect("decompression should work");
        assert_eq!(decoded.coefficients().len(), ctx.max_degree() + 1);
        assert!(decoded.coefficients().iter().all(|&c| c < ctx.modulus()));
    }

    #[test]
    fn compression_rejects_invalid_bit_width() {
        let err = compress_coefficient_example(5, 17, 0).expect_err("expected error");
        assert_eq!(err, CodecError::InvalidCompressionBits(0));

        let err = decompress_coefficient_example(5, 17, 17).expect_err("expected error");
        assert_eq!(err, CodecError::InvalidCompressionBits(17));
    }
}