ahda 0.1.0

// ahda: Pseudoalignment compression and conversion between formats.
//
// Copyright 2025 Tommi Mäklin [tommi@maklin.fi].
//
// Copyrights in this project are retained by contributors. No copyright assignment
// is required to contribute to this project.
//
// Except as otherwise noted (below and/or in individual files), this
// project is licensed under the Apache License, Version 2.0
// <LICENSE-APACHE> or <http://www.apache.org/licenses/LICENSE-2.0> or
// the MIT license, <LICENSE-MIT> or <http://opensource.org/licenses/MIT>,
// at your option.
//

//! Block header and block flags used in the encoded format.

use crate::compression::gzwrapper::deflate_bytes;
use crate::compression::gzwrapper::inflate_bytes;

use std::io::Read;

use bincode::{Encode, Decode};
use bincode::encode_into_std_write;
use bincode::decode_from_slice;

use crate::compression::MetadataCompression;

type E = Box<dyn std::error::Error>;

/// Block header for encoded data
///
/// Always the first 32 bytes at the beginning of a .ahda block.
///
/// Must always conform to this format.
///
/// The BlockHeader contains the following data:
/// - Number of records stored in this block.
/// - The metadata compression method used for [BlockFlags].
/// - Type of bitmap stored in this block.
/// - A 1 byte unused placeholder value.
/// - The number of bytes in the block contents that follow the [BlockFlags] bytes.
/// - Number of bytes containing the [BlockFlags] that follow the header.
/// - Fields that are present in the [BlockFlags].
/// - A 2 byte unused placeholder value.
/// - A 8 byte unused placeholder value.
///
#[derive(Clone, Debug, Decode, Encode, PartialEq)]
pub struct BlockHeader {
    /// Number of records (queries) stored in this block
    pub num_records: u32,

    /// Compression method used for [BlockFlags], see [MetadataCompression].
    pub metadata_compression: u8,

    /// Bitmap type used to encode this block, see [BitmapType](crate::compression::BitmapType) for details.
    ///
    /// The address space of the block bitmap type must fit within the file
    /// bitmap type specified in [crate::headers::file::FileHeader]. Given that this is satisfied, the
    /// block bitmap type may differ from the file bitmap type.
    pub bitmap_type: u16,

    /// Placeholder:
    pub placeholder1: u8,

    /// Number of bytes in the block contents that follow the flags bytes.
    pub block_len: u32,

    /// Number of bytes in [BlockFlags] that follow the header bytes.
    pub flags_len: u64,

    /// [BlockFlags] fields that are present for records in this block.
    ///
    /// Every block must contain at least the fields given in the [crate::headers::file::FileHeader]
    /// `fields_present`. A block may also contain extra fields that are not
    /// specified in [crate::headers::file::FileHeader].
    pub fields_present: u16,

    /// Placeholder
    pub placeholder2: u16,

    /// Placeholder
    pub placeholder3: u64,
}

impl BlockHeader {
    /// Does the header require the encoded data to contain the query names?
    pub fn promises_query_names(
        &self,
    ) -> bool {
        (self.fields_present & crate::MASK_QUERIES) != 0
    }

    /// Does the header require the encoded data to contain the query indexes?
    pub fn promises_query_ids(
        &self,
    ) -> bool {
        (self.fields_present & crate::MASK_QUERY_IDS) != 0
    }
}

impl Default for BlockHeader {
    /// Initialize a header with default values.
    ///
    /// **Note**: This header is not a valid ahda header for most cases. Update
    /// the fields accordingly before using.
    ///
    /// If you need to encode individual ahda blocks, prefer using
    /// [pack_records](crate::compression::pack_records).
    fn default() -> BlockHeader {
        BlockHeader {
            num_records: 0_u32,
            metadata_compression: MetadataCompression::default().to_u8(),
            bitmap_type: crate::compression::BitmapType::default().to_u16(),
            block_len: 0_u32,
            flags_len: 0_u64,
            fields_present: crate::MASK_QUERY_IDS | crate::MASK_QUERIES,
            placeholder1: 0,
            placeholder2: 0,
            placeholder3: 0,
        }
    }
}

/// Data about the records in this block
///
/// Variable length, use [BlockHeader].flags_len to get size
///
/// May contain these fields:
/// - `queries`: Names of the query sequences in the original query file.
/// - `query_ids`: Indexes of the query sequences in the original query file.
///
/// If both `queries` and `query_ids` are given, they must be in the same order,
/// ie. the data in both arrays at the same index corresponds to the same
/// record.
///
/// The current implementation of ahda assumes that `query_ids` is always
/// present and filled.
///
#[derive(Clone, Decode, Debug, Encode, PartialEq)]
pub struct BlockFlags {
    /// Names of query records in the original query file.
    pub queries: Option<Vec<Vec<u8>>>,
    /// Indexes of the query records in the original query file.
    pub query_ids: Option<Vec<u32>>,
}

impl BlockFlags {
    /// Get the `fields_present` value to store in this block's [BlockHeader].
    ///
    /// ## Usage
    ///
    /// ```rust
    /// use ahda::headers::block::BlockFlags;
    ///
    /// let mut flags = BlockFlags {
    ///     queries: None,
    ///     query_ids: None,
    /// };
    ///
    /// assert_eq!(flags.fields_present(), 0_u16);
    ///
    /// flags.queries = Some(vec![]);
    /// assert_eq!(flags.fields_present(), 1_u16);
    ///
    /// flags.query_ids = Some(vec![]);
    /// assert_eq!(flags.fields_present(), 3_u16);
    ///
    /// flags.queries = None;
    /// assert_eq!(flags.fields_present(), 2_u16);
    /// ```
    pub fn fields_present(
        &self,
    ) -> u16 {
        let mut flags = 0_u16;
        if self.queries.is_some() {
            flags |= crate::MASK_QUERIES;
        }
        if self.query_ids.is_some() {
            flags |= crate::MASK_QUERY_IDS;
        }
        flags
    }
}

/// Encode a [BlockHeader].
///
/// ## Usage
///
/// ```rust
/// use ahda::headers::block::encode_block_header;
/// use ahda::headers::block::BlockHeader;
///
/// let data = BlockHeader {
///     num_records: 31,
///     placeholder1: 0,
///     block_len: 65511,
///     flags_len: 921,
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0, metadata_compression: 0
/// };
///
/// let expected: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
///
/// let got = encode_block_header(&data);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected);
/// ```
pub fn encode_block_header(
    header: &BlockHeader,
) -> Result<Vec<u8>, E> {
    let mut bytes: Vec<u8> = Vec::new();
    let nbytes = encode_into_std_write(
        header,
        &mut bytes,
        bincode::config::standard().with_fixed_int_encoding(),
    )?;
    assert_eq!(nbytes, 32);
    Ok(bytes)
}

/// Decode a [BlockHeader] from an u8 array.
///
/// ## Usage
///
/// ```rust
/// use ahda::headers::block::decode_block_header;
/// use ahda::headers::block::BlockHeader;
///
/// let expected_header = BlockHeader {
///     num_records: 31,
///     placeholder1: 0,
///     block_len: 65511,
///     flags_len: 921,
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0, metadata_compression: 0
/// };
///
/// let data: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
///
/// let got = decode_block_header(&data);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected_header);
/// ```
pub fn decode_block_header(
    header_bytes: &[u8],
) -> Result<BlockHeader, E> {
    Ok(decode_from_slice(header_bytes, bincode::config::standard().with_fixed_int_encoding())?.0)
}

/// Read a [BlockHeader] from a [Read].
///
/// ## Usage
///
/// ```rust
/// use ahda::headers::block::read_block_header;
/// use ahda::headers::block::BlockHeader;
/// use std::io::Cursor;
///
/// let expected_header = BlockHeader {
///     num_records: 31,
///     placeholder1: 0,
///     block_len: 65511,
///     flags_len: 921,
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0, metadata_compression: 0
/// };
///
/// let data_bytes: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
/// let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);
///
/// let got = read_block_header(&mut data);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected_header);
/// ```
pub fn read_block_header<R: Read>(
    conn: &mut R,
) -> Result<BlockHeader, E> {
    let mut header_bytes: [u8; 32] = [0_u8; 32];
    conn.read_exact(&mut header_bytes)?;
    let res = decode_block_header(&header_bytes)?;
    Ok(res)
}

/// Read [BlockFlags] from a [Read].
///
/// ## Usage
///
/// ```rust
/// use ahda::compression::MetadataCompression;
/// use ahda::headers::block::read_block_flags;
/// use ahda::headers::block::BlockFlags;
/// use ahda::headers::block::BlockHeader;
/// use std::io::Cursor;
///
/// let header = BlockHeader {
///     num_records: 0,
///     placeholder1: 0,
///     block_len: 0,
///     flags_len: 33,
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0,
///     metadata_compression: MetadataCompression::Flate2.to_u8(),
/// };
///
/// let expected_flags = BlockFlags {
///     queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]),
///     query_ids: Some(vec![1, 0, 2]),
/// };
///
/// let data_bytes: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];
/// let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);
///
/// let got = read_block_flags(&header, &mut data);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected_flags);
/// ```
pub fn read_block_flags<R: Read>(
    header: &BlockHeader,
    conn: &mut R,
) -> Result<BlockFlags, E> {
    let mut flags_bytes: Vec<u8> = vec![0; header.flags_len as usize];
    conn.read_exact(&mut flags_bytes)?;
    let metadata_compression = MetadataCompression::from_u8(header.metadata_compression)?;
    let res = decode_block_flags(&flags_bytes, &metadata_compression)?;
    Ok(res)
}

/// Read a [BlockHeader] and [BlockFlags] from a [Read].
///
/// ## Usage
///
/// ```rust
/// use ahda::compression::MetadataCompression;
/// use ahda::headers::block::read_block_header_and_flags;
/// use ahda::headers::block::BlockHeader;
/// use ahda::headers::block::BlockFlags;
/// use std::io::Cursor;
///
/// let expected_header = BlockHeader {
///     num_records: 31,
///     placeholder1: 0,
///     block_len: 65511,
///     flags_len: 13,
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0,
///     metadata_compression: 0,
/// };
/// let expected_flags = BlockFlags {
///     queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]),
///     query_ids: Some(vec![1, 0, 2])
/// };
///
/// let data_bytes: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3, 1, 97, 1, 98, 1, 99, 1, 3, 1, 0, 2];
/// let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);
///
/// let got = read_block_header_and_flags(&mut data);
/// assert!(got.is_ok());
/// let (header, flags) = got.unwrap();
/// assert_eq!(header, expected_header);
/// assert_eq!(flags, expected_flags);
/// ```
pub fn read_block_header_and_flags<R: Read>(
    conn: &mut R,
) -> Result<(BlockHeader, BlockFlags), E> {
    let header = read_block_header(conn)?;
    let flags = read_block_flags(&header, conn)?;
    Ok((header, flags))
}

/// Encode [BlockFlags].
///
/// ## Usage
///
/// ```rust
/// use ahda::compression::MetadataCompression;
/// use ahda::headers::block::encode_block_flags;
/// use ahda::headers::block::BlockFlags;
///
/// let flags = BlockFlags {
///     queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]),
///     query_ids: Some(vec![1, 0, 2]),
/// };
///
/// let expected: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];
///
/// let got = encode_block_flags(&flags, &MetadataCompression::Flate2);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected);
/// ```
pub fn encode_block_flags(
    flags: &BlockFlags,
    compression: &MetadataCompression,
) -> Result<Vec<u8>, E> {
    let mut bytes: Vec<u8> = Vec::new();

    match compression {
        MetadataCompression::BincodeStandard => {
            let _ = encode_into_std_write(
                flags,
                &mut bytes,
                bincode::config::standard(),
            )?;
        },
        MetadataCompression::Flate2 => {
            let _ = encode_into_std_write(
                flags,
                &mut bytes,
                bincode::config::standard(),
            )?;
            bytes = deflate_bytes(&bytes)?;
        },
    }

    Ok(bytes)
}

/// Decode [BlockFlags] from an u8 array.
///
/// ## Usage
///
/// ```rust
/// use ahda::compression::MetadataCompression;
/// use ahda::headers::block::decode_block_flags;
/// use ahda::headers::block::BlockFlags;
///
/// let expected_flags = BlockFlags {
///     queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]),
///     query_ids: Some(vec![1, 0, 2]),
/// };
///
/// let data: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];
///
/// let got = decode_block_flags(&data, &MetadataCompression::Flate2);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected_flags);
/// ```
pub fn decode_block_flags(
    bytes: &[u8],
    compression: &MetadataCompression,
) -> Result<BlockFlags, E> {
    let flags: BlockFlags = match compression {
        MetadataCompression::BincodeStandard => {
            decode_from_slice(
                bytes,
                bincode::config::standard(),
            )?.0
        },
        MetadataCompression::Flate2 => {
            let inflated = inflate_bytes(bytes)?;
            decode_from_slice(
                &inflated,
                bincode::config::standard(),
            )?.0
        },
    };


    Ok(flags)
}

/// Encode a [BlockHeader] and [BlockFlags].
///
/// ## Usage
///
/// ```rust
/// use ahda::compression::MetadataCompression;
/// use ahda::headers::block::encode_block_header_and_flags;
/// use ahda::headers::block::BlockHeader;
/// use ahda::headers::block::BlockFlags;
/// use std::io::Cursor;
///
/// let mut header = BlockHeader {
///     num_records: 31,
///     placeholder1: 0,
///     block_len: 65511,
///     flags_len: 0, // Note we don't need to set this
///     fields_present: 0,
///     placeholder2: 0,
///     placeholder3: 0,
///     bitmap_type: 0,
///     metadata_compression: 0,
/// };
/// let flags = BlockFlags {
///     queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]),
///     query_ids: Some(vec![1, 0, 2])
/// };
///
/// let expected: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3, 1, 97, 1, 98, 1, 99, 1, 3, 1, 0, 2];
///
/// let got = encode_block_header_and_flags(&mut header, &flags);
/// assert!(got.is_ok());
/// assert_eq!(got.unwrap(), expected);
/// ```
pub fn encode_block_header_and_flags(
    header: &mut BlockHeader,
    flags: &BlockFlags,
) -> Result<Vec<u8>, E> {
    let mut flags_bytes = encode_block_flags(flags, &MetadataCompression::from_u8(header.metadata_compression)?)?;
    header.flags_len = flags_bytes.len() as u64;
    let mut bytes = encode_block_header(header)?;
    bytes.append(&mut flags_bytes);
    Ok(bytes)
}

#[cfg(test)]
mod tests {

    #[test]
    fn encode_block_header() {
        use super::encode_block_header;
        use super::BlockHeader;

        let data = BlockHeader{ num_records: 31, placeholder1: 0, block_len: 65511, flags_len: 921, fields_present: 0, placeholder2: 0, placeholder3: 0, bitmap_type: 0, metadata_compression: 0 };
        let expected: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];

        let got = encode_block_header(&data).unwrap();
        assert_eq!(got, expected);
    }

    #[test]
    fn decode_block_header() {
        use super::decode_block_header;
        use super::BlockHeader;

        let expected = BlockHeader{ num_records: 31, placeholder1: 0, block_len: 65511, flags_len: 921, fields_present: 0, placeholder2: 0, placeholder3: 0, bitmap_type: 0, metadata_compression: 0 };
        let data: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];

        let got = decode_block_header(&data).unwrap();
        assert_eq!(got, expected);
    }

    #[test]
    fn read_block_header() {
        use super::read_block_header;
        use super::BlockHeader;

        use std::io::Cursor;

        let expected = BlockHeader{ num_records: 31, placeholder1: 0, block_len: 65511, flags_len: 921, fields_present: 0, placeholder2: 0, placeholder3: 0, bitmap_type: 0, metadata_compression: 0 };
        let data_bytes: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 153, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0];
        let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);

        let got = read_block_header(&mut data).unwrap();
        assert_eq!(got, expected);
    }

    #[test]
    fn encode_block_flags() {
        use super::encode_block_flags;
        use super::BlockFlags;
        use crate::MetadataCompression;

        let data = BlockFlags{ queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]), query_ids: Some(vec![1, 0, 2]) };
        let expected: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];

        let got = encode_block_flags(&data, &MetadataCompression::Flate2).unwrap();
        assert_eq!(got, expected);
    }

    #[test]
    fn decode_block_flags() {
        use super::decode_block_flags;
        use super::BlockFlags;
        use crate::compression::MetadataCompression;

        let expected = BlockFlags{ queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]), query_ids: Some(vec![1, 0, 2]) };
        let data: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];

        let got = decode_block_flags(&data, &MetadataCompression::Flate2).unwrap();
        assert_eq!(got, expected);
    }

    #[test]
    fn read_block_flags() {
        use super::read_block_flags;
        use super::BlockFlags;
        use super::BlockHeader;

        use std::io::Cursor;

        let expected = BlockFlags{ queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]), query_ids: Some(vec![1, 0, 2]) };
        let data_bytes: Vec<u8> = vec![31, 139, 8, 0, 0, 0, 0, 0, 0, 255, 99, 100, 102, 76, 100, 76, 98, 76, 102, 100, 102, 100, 96, 2, 0, 249, 181, 108, 55, 13, 0, 0, 0];
        let header = BlockHeader {
            num_records: 31,
            placeholder1: 0,
            block_len: 65511,
            flags_len: data_bytes.len() as u64,
            fields_present: 0,
            placeholder2: 0,
            placeholder3: 0,
            bitmap_type: 0,
            metadata_compression: 1,
        };
        let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);

        let got = read_block_flags(&header, &mut data).unwrap();
        assert_eq!(got, expected);
    }


    #[test]
    fn read_block_header_and_flags() {
        use super::read_block_header_and_flags;
        use super::BlockFlags;
        use super::BlockHeader;

        use std::io::Cursor;

        let data_bytes: Vec<u8> = vec![31, 0, 0, 0, 0, 0, 0, 0, 231, 255, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 3, 1, 97, 1, 98, 1, 99, 1, 3, 1, 0, 2];
        let mut data: Cursor<Vec<u8>> = Cursor::new(data_bytes);

        let expected_header = BlockHeader{ num_records: 31, placeholder1: 0, block_len: 65511, flags_len: 13, fields_present: 0, placeholder2: 0, placeholder3: 0, bitmap_type: 0, metadata_compression: 0 };
        let expected_flags = BlockFlags{ queries: Some(vec!["a".as_bytes().to_vec(), "b".as_bytes().to_vec(), "c".as_bytes().to_vec()]), query_ids: Some(vec![1, 0, 2]) };

        let (got_header, got_flags) = read_block_header_and_flags(&mut data).unwrap();
        assert_eq!(got_header, expected_header);
        assert_eq!(got_flags, expected_flags);
    }
}