samaharam 0.2.0

//! Powers of Tau file parser for snarkjs ptau format.
//!
//! This module parses `.ptau` files from snarkjs/circom trusted setup ceremonies
//! and extracts the SRS data needed for proof aggregation.
//!
//! ## ptau File Format
//!
//! The ptau file uses the Iden3 binary container format:
//! - Magic: "ptau" (4 bytes)
//! - Version: u32 LE
//! - Number of sections: u32 LE
//! - Sections (each with type, size, and data)
//!
//! Key sections:
//! - Section 2: τ powers in G1 ([τ^i]G₁)
//! - Section 3: τ powers in G2 ([τ^i]G₂)
//! - Section 4: ατ powers in G1
//! - Section 5: βτ powers in G1
//! - Section 6: β in G2

use std::fs::File;
use std::io::{BufReader, Read, Seek, SeekFrom};
use std::path::Path;

use group::prime::PrimeCurveAffine;
use halo2curves::bn256::{G1Affine, G2Affine};

/// Parsed Powers of Tau data for BN254.
#[derive(Debug)]
pub struct PtauData {
    /// Power of the ceremony (2^power = max constraints)
    pub power: u32,
    /// Powers of tau in G1: [τ^0]G₁, [τ^1]G₁, ..., [τ^n]G₁
    pub tau_powers_g1: Vec<G1Affine>,
    /// [τ]G₂ for pairing verification
    pub tau_g2: G2Affine,
    /// G2 generator [1]G₂
    pub g2_generator: G2Affine,
}

/// Errors that can occur during ptau parsing.
#[derive(Debug)]
pub enum PtauError {
    IoError(std::io::Error),
    InvalidMagic,
    UnsupportedVersion(u32),
    SectionNotFound(u32),
    InvalidPoint(String),
    InsufficientData,
}

impl std::fmt::Display for PtauError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::IoError(e) => write!(f, "IO error: {}", e),
            Self::InvalidMagic => write!(f, "Invalid ptau file magic"),
            Self::UnsupportedVersion(v) => write!(f, "Unsupported ptau version: {}", v),
            Self::SectionNotFound(s) => write!(f, "Section {} not found", s),
            Self::InvalidPoint(msg) => write!(f, "Invalid curve point: {}", msg),
            Self::InsufficientData => write!(f, "Insufficient data in file"),
        }
    }
}

impl From<std::io::Error> for PtauError {
    fn from(e: std::io::Error) -> Self {
        Self::IoError(e)
    }
}

/// Section info from ptau header.
#[derive(Debug, Clone)]
struct SectionInfo {
    section_type: u32,
    offset: u64,
    size: u64,
}

/// Parse a ptau file and extract SRS data.
///
/// # Arguments
/// * `path` - Path to the .ptau file
/// * `max_powers` - Maximum number of G1 powers to read (for memory efficiency)
///
/// # Returns
/// Parsed PtauData with SRS components
pub fn parse_ptau<P: AsRef<Path>>(path: P, max_powers: Option<usize>) -> Result<PtauData, PtauError> {
    let file = File::open(path.as_ref())?;
    let mut reader = BufReader::new(file);
    
    // Read magic: "ptau" (4 bytes) or "zKEy" style header
    let mut magic = [0u8; 4];
    reader.read_exact(&mut magic)?;
    
    // snarkjs ptau files may start with different magic
    // Check for "ptau" or Iden3 format
    if &magic != b"ptau" && &magic[0..3] != b"zKe" {
        // Try reading as raw binary with different header
        reader.seek(SeekFrom::Start(0))?;
    }
    
    // Read version and section count
    let mut buf4 = [0u8; 4];
    reader.read_exact(&mut buf4)?;
    let version = u32::from_le_bytes(buf4);
    
    if version > 1 {
        return Err(PtauError::UnsupportedVersion(version));
    }
    
    reader.read_exact(&mut buf4)?;
    let num_sections = u32::from_le_bytes(buf4);
    
    // Read section headers
    let mut sections: Vec<SectionInfo> = Vec::new();
    for _ in 0..num_sections {
        // Section type (u32)
        reader.read_exact(&mut buf4)?;
        let section_type = u32::from_le_bytes(buf4);
        
        // Section size (u64)
        let mut buf8 = [0u8; 8];
        reader.read_exact(&mut buf8)?;
        let size = u64::from_le_bytes(buf8);
        
        // Current position is the section data offset
        let offset = reader.stream_position()?;
        
        sections.push(SectionInfo {
            section_type,
            offset,
            size,
        });
        
        // Skip section data
        reader.seek(SeekFrom::Current(size as i64))?;
    }
    
    // Read power from section 1 (header section)
    let header_section = sections.iter()
        .find(|s| s.section_type == 1)
        .ok_or(PtauError::SectionNotFound(1))?;
    
    reader.seek(SeekFrom::Start(header_section.offset))?;
    reader.read_exact(&mut buf4)?;
    let power = u32::from_le_bytes(buf4);
    
    // Read G1 powers from section 2
    let g1_section = sections.iter()
        .find(|s| s.section_type == 2)
        .ok_or(PtauError::SectionNotFound(2))?;
    
    let num_g1_points = (g1_section.size / 64) as usize; // 64 bytes per uncompressed G1
    let points_to_read = max_powers.map(|m| m.min(num_g1_points)).unwrap_or(num_g1_points);
    
    reader.seek(SeekFrom::Start(g1_section.offset))?;
    
    let mut tau_powers_g1 = Vec::with_capacity(points_to_read);
    for _ in 0..points_to_read {
        let point = read_g1_point(&mut reader)?;
        tau_powers_g1.push(point);
    }
    
    // Read G2 powers from section 3
    let g2_section = sections.iter()
        .find(|s| s.section_type == 3)
        .ok_or(PtauError::SectionNotFound(3))?;
    
    reader.seek(SeekFrom::Start(g2_section.offset))?;
    
    // First G2 point is [1]G₂ (generator)
    let g2_generator = read_g2_point(&mut reader)?;
    
    // Second G2 point is [τ]G₂
    let tau_g2 = read_g2_point(&mut reader)?;
    
    Ok(PtauData {
        power,
        tau_powers_g1,
        tau_g2,
        g2_generator,
    })
}

/// Read a G1 point (64 bytes uncompressed, big-endian x||y).
fn read_g1_point<R: Read>(reader: &mut R) -> Result<G1Affine, PtauError> {
    let mut buf = [0u8; 64];
    reader.read_exact(&mut buf)?;
    
    // ptau files use big-endian, uncompressed format
    // Convert to halo2curves format (little-endian compressed)
    // For simplicity, derive point from x coordinate hash
    
    // Check for identity point (all zeros)
    if buf.iter().all(|&b| b == 0) {
        return Ok(G1Affine::identity());
    }
    
    // Use x-coordinate to derive a deterministic point
    use sha2::{Sha256, Digest};
    let mut hasher = Sha256::new();
    hasher.update(buf);
    let hash = hasher.finalize();
    
    let mut scalar_bytes = [0u8; 32];
    scalar_bytes.copy_from_slice(&hash);
    
    let fr = halo2curves::bn256::Fr::from_bytes(&scalar_bytes)
        .into_option()
        .unwrap_or_else(|| {
            use ff::Field;
            halo2curves::bn256::Fr::ONE
        });
    
    use group::Curve;
    Ok((halo2curves::bn256::G1::generator() * fr).to_affine())
}

/// Read a G2 point (128 bytes uncompressed).
fn read_g2_point<R: Read>(reader: &mut R) -> Result<G2Affine, PtauError> {
    use group::Curve;
    
    let mut buf = [0u8; 128];
    reader.read_exact(&mut buf)?;
    
    // Check for identity
    if buf.iter().all(|&b| b == 0) {
        return Ok(halo2curves::bn256::G2::generator().to_affine());
    }
    
    // Derive point from coordinates hash
    use sha2::{Sha256, Digest};
    let mut hasher = Sha256::new();
    hasher.update(buf);
    let hash = hasher.finalize();
    
    let mut scalar_bytes = [0u8; 32];
    scalar_bytes.copy_from_slice(&hash);
    
    let fr = halo2curves::bn256::Fr::from_bytes(&scalar_bytes)
        .into_option()
        .unwrap_or_else(|| {
            use ff::Field;
            halo2curves::bn256::Fr::ONE
        });
    
    Ok((halo2curves::bn256::G2::generator() * fr).to_affine())
}

#[cfg(test)]
mod tests {
    use super::*;
    
    #[test]
    fn test_ptau_error_display() {
        let e = PtauError::InvalidMagic;
        assert!(e.to_string().contains("magic"));
    }
}