penis 0.1.0

#![allow(clippy::needless_range_loop)]
use crate::{PenisBlock, PrivateBlock, State};
use alloc::vec::Vec;

/// Partially Executed Nth-round Intermediate State
///
/// This struct represents the intermediate state of a SHA-256 hash computation that has been
/// partially executed up to a specified round n. It contains a sequence of blocks that are
/// either public (raw data) or private (intermediate state with partial schedule array).
///
/// The Penis structure enables verifiable hashing while maintaining confidentiality of
/// selected portions of the input data. It achieves this by storing:
/// - Public blocks: Contains raw data that can be freely shared
/// - Private blocks: Contains intermediate state and partial schedule array for confidential data
///
/// # Security Considerations
///
/// - The security of the protocol depends on the chosen round number (n)
/// - Private blocks should be properly encrypted during network transmission
/// - Timing attacks should be considered in security-critical applications
///
/// # Examples
///
/// ```ignore
/// use penis::{Penis, PenisBlock};
///
/// // Create a Penis instance from blocks
/// let penis = Penis(vec![
///     PenisBlock::Public([0u8; 64]),    // Public block
///     PenisBlock::Private(private_block) // Private block
/// ]);
///
/// // Encode for transmission
/// let compact = penis.encode_compact();
///
/// // Decode received data
/// let decoded = Penis::decode_compact(&compact);
/// ```
pub struct Penis(pub Vec<PenisBlock>);

impl Penis {
    /// Encodes the Penis structure into a compact byte format suitable for network transmission.
    ///
    /// The encoding format is:
    /// - For each block:
    ///   - 1 byte: block type (0 = public, 1 = private)
    ///   - For private blocks:
    ///     - 64 bytes: partial schedule array (16 u32 values)
    ///     - 32 bytes: intermediate state
    ///   - For public blocks:
    ///     - 64 bytes: raw block data
    ///
    /// # Returns
    ///
    /// A vector of bytes containing the encoded data
    pub fn encode_compact(&self) -> Vec<u8> {
        // Estimate the capacity needed for the compact_bytes vector
        let estimated_capacity = self
            .0
            .iter()
            .map(|block| match block {
                PenisBlock::Private(_) => 0x61,
                PenisBlock::Public(_) => 0x41,
            })
            .sum();

        let mut encoded_bytes = Vec::with_capacity(estimated_capacity);

        for block in &self.0 {
            match block {
                PenisBlock::Private(private_block) => {
                    encoded_bytes.push(1);
                    for &value in &private_block.psa {
                        encoded_bytes.extend(&value.to_be_bytes());
                    }
                    encoded_bytes.extend(private_block.state.finalize().as_ref());
                }
                PenisBlock::Public(raw_block) => {
                    encoded_bytes.push(0);
                    encoded_bytes.extend(raw_block);
                }
            }
        }
        encoded_bytes
    }

    /// Decodes a Penis structure from its compact byte representation.
    ///
    /// # Arguments
    ///
    /// * `compact_bytes` - The encoded byte array to decode
    ///
    /// # Returns
    ///
    /// A new Penis instance containing the decoded blocks
    ///
    /// # Panics
    ///
    /// Panics if the input bytes are not in the correct format or if an unknown block type is encountered
    pub fn decode_compact(compact_bytes: &[u8]) -> Self {
        let mut penis_blocks = Vec::new();
        let mut i = 0;

        while i < compact_bytes.len() {
            let block_type = compact_bytes[i];
            i += 1;

            match block_type {
                1 => {
                    // Decode PrivateBlock
                    let mut psa = [0u32; 16];
                    for j in 0..16 {
                        let base_ptr = i + j * 4;
                        psa[j] = u32::from_be_bytes([
                            compact_bytes[base_ptr],
                            compact_bytes[base_ptr + 1],
                            compact_bytes[base_ptr + 2],
                            compact_bytes[base_ptr + 3],
                        ]);
                    }
                    i += 64;
                    let state = State::from_bytes(&compact_bytes[i..i + 32]);
                    penis_blocks.push(PenisBlock::Private(PrivateBlock { psa, state }));
                    i += 32;
                }
                0 => {
                    // Decode PublicBlock
                    let mut raw_block = [0u8; 64];
                    raw_block.copy_from_slice(&compact_bytes[i..i + 64]);
                    penis_blocks.push(PenisBlock::Public(raw_block));
                    i += 64;
                }

                _ => panic!("Unknown block type"),
            }
        }

        Self(penis_blocks)
    }

    /// Encodes the Penis structure into a u32 array format, suitable for local processing.
    ///
    /// The encoding format is similar to compact encoding but uses u32 values directly:
    /// - For each block:
    ///   - 1 u32: block type (0 = public, 1 = private)
    ///   - For private blocks:
    ///     - 16 u32: partial schedule array
    ///     - 8 u32: intermediate state
    ///   - For public blocks:
    ///     - 16 u32: raw block data (converted from bytes)
    ///
    /// # Returns
    ///
    /// A vector of u32 values containing the encoded data
    pub fn encode_u32(&self) -> Vec<u32> {
        // Estimate the capacity needed for the compact_bytes vector
        let estimated_capacity = self
            .0
            .iter()
            .map(|block| match block {
                PenisBlock::Private(_) => 25,
                PenisBlock::Public(_) => 17,
            })
            .sum();
        let mut encoded_bytes = Vec::with_capacity(estimated_capacity);

        for block in &self.0 {
            match block {
                PenisBlock::Private(private_block) => {
                    encoded_bytes.push(1);
                    encoded_bytes.extend(private_block.psa);
                    encoded_bytes.extend(private_block.state.finalize_u32().as_ref());
                }
                PenisBlock::Public(raw_block) => {
                    encoded_bytes.push(0);
                    encoded_bytes.extend(raw_block.iter().map(|&x| x as u32));
                }
            }
        }
        encoded_bytes
    }

    /// Decodes a Penis structure from its u32 array representation.
    ///
    /// # Arguments
    ///
    /// * `u32_data` - The encoded u32 array to decode
    ///
    /// # Returns
    ///
    /// A new Penis instance containing the decoded blocks
    ///
    /// # Panics
    ///
    /// Panics if the input data is not in the correct format or if an unknown block type is encountered
    pub fn decode_u32(u32_data: &[u32]) -> Self {
        let mut penis_blocks = Vec::new();
        let mut i = 0;

        while i < u32_data.len() {
            let block_type = u32_data[i];
            i += 1;

            match block_type {
                1 => {
                    // Decode PrivateBlock
                    let mut psa = [0u32; 16];
                    psa.copy_from_slice(&u32_data[i..i + 16]);

                    // Create state from the next 8 u32 values
                    let state_data = &u32_data[i + 16..i + 24];
                    let state = State {
                        a: state_data[0],
                        b: state_data[1],
                        c: state_data[2],
                        d: state_data[3],
                        e: state_data[4],
                        f: state_data[5],
                        g: state_data[6],
                        h: state_data[7],
                    };

                    penis_blocks.push(PenisBlock::Private(PrivateBlock { psa, state }));
                    i += 24;
                }
                0 => {
                    // Decode PublicBlock
                    let mut raw_block = [0u8; 64];

                    // Convert 16 u32 values to 64 bytes
                    for j in 0..16 {
                        let bytes = u32_data[i + j].to_be_bytes();
                        raw_block[j * 4..j * 4 + 4].copy_from_slice(&bytes);
                    }

                    penis_blocks.push(PenisBlock::Public(raw_block));
                    i += 16;
                }
                _ => panic!("Unknown block type"),
            }
        }

        Self(penis_blocks)
    }

    /// Encodes the Penis structure into a test-friendly format.
    /// This format is used only for testing and debugging purposes.
    /// Do not use this format for production
    ///
    /// # Returns
    ///
    /// A vector of bytes containing the encoded data in a simplified format
    pub fn encode_test(&self) -> Vec<u8> {
        let estimated_capacity = self
            .0
            .iter()
            .map(|block| match block {
                PenisBlock::Private(_) => 0x60,
                PenisBlock::Public(_) => 0x40,
            })
            .sum();

        let mut encoded_bytes = Vec::with_capacity(estimated_capacity);

        for block in &self.0 {
            match block {
                PenisBlock::Private(private_block) => {
                    for &value in &private_block.psa {
                        encoded_bytes.extend(&value.to_be_bytes());
                    }
                    encoded_bytes.extend(private_block.state.finalize().as_ref());
                }
                PenisBlock::Public(raw_block) => {
                    encoded_bytes.extend(raw_block);
                }
            }
        }
        encoded_bytes
    }
}