ethrex-common 17.0.0

use bytes::Bytes;
use serde::{Deserialize, Serialize};
use std::fmt::Display;

/// Enum used to identify the different proving systems.
#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProverType {
    Exec,
    RISC0,
    SP1,
    TDX,
}

impl From<ProverType> for u32 {
    fn from(value: ProverType) -> u32 {
        match value {
            ProverType::Exec => 0,
            ProverType::RISC0 => 1,
            ProverType::SP1 => 2,
            ProverType::TDX => 3,
        }
    }
}

impl ProverType {
    /// Used to iterate through all the possible proving systems
    pub fn all() -> impl Iterator<Item = ProverType> {
        [
            ProverType::Exec,
            ProverType::RISC0,
            ProverType::SP1,
            ProverType::TDX,
        ]
        .into_iter()
    }
}

impl Display for ProverType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Exec => write!(f, "Exec"),
            Self::RISC0 => write!(f, "RISC0"),
            Self::SP1 => write!(f, "SP1"),
            Self::TDX => write!(f, "TDX"),
        }
    }
}

/// Raw proof bytes with a prover identifier.
#[derive(PartialEq, Serialize, Deserialize, Clone, Debug)]
pub struct ProofBytes {
    pub prover_type: ProverType,
    pub proof: Vec<u8>,
}

/// Output produced by a prover backend.
///
/// - [`Proof`](ProverOutput::Proof): just the proof bytes — for L1 and L2
///   on-chain verification where public values are unnecessary or recoverable.
/// - [`ProofWithPublicValues`](ProverOutput::ProofWithPublicValues): proof
///   bytes together with public values — needed by the Aligned Layer
///   verification path.
#[derive(PartialEq, Serialize, Deserialize, Clone, Debug)]
pub enum ProverOutput {
    Proof(ProofBytes),
    ProofWithPublicValues {
        proof_bytes: ProofBytes,
        public_values: Vec<u8>,
    },
}

impl ProverOutput {
    pub fn prover_type(&self) -> ProverType {
        match self {
            Self::Proof(p) => p.prover_type,
            Self::ProofWithPublicValues { proof_bytes, .. } => proof_bytes.prover_type,
        }
    }

    pub fn proof_bytes(&self) -> &ProofBytes {
        match self {
            Self::Proof(p) => p,
            Self::ProofWithPublicValues { proof_bytes, .. } => proof_bytes,
        }
    }

    pub fn into_proof_bytes(self) -> ProofBytes {
        match self {
            Self::Proof(p) => p,
            Self::ProofWithPublicValues { proof_bytes, .. } => proof_bytes,
        }
    }

    pub fn public_values(&self) -> Option<&[u8]> {
        match self {
            Self::Proof(_) => None,
            Self::ProofWithPublicValues { public_values, .. } => Some(public_values),
        }
    }
}

/// Indicates the prover which proof *format* to generate
#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default)]
pub enum ProofFormat {
    #[default]
    /// A compressed proof wrapped over groth16. EVM friendly.
    Groth16,
    /// Fixed size STARK execution proof.
    Compressed,
}

/// Generic enum for the ProverServer <--> ProverClient Communication Protocol.
///
/// The type parameter `I` represents the input type sent from server to prover:
/// - L2 uses `ProverInputData`
/// - L1 uses `ProgramInput`
#[allow(clippy::large_enum_variant)]
#[derive(Serialize, Deserialize)]
pub enum ProofData<I> {
    /// The client performs any needed setup steps (e.g. key registration).
    ProverSetup {
        prover_type: ProverType,
        payload: Bytes,
    },

    /// The Server acknowledges the receipt of the setup and its completion.
    ProverSetupACK,

    /// The Client requests work from the coordinator.
    /// The commit hash ensures client and server are compatible.
    /// The prover_type tells the coordinator which backend the client runs,
    /// so it can skip work items that already have a proof for that type.
    InputRequest {
        commit_hash: String,
        prover_type: ProverType,
    },

    /// The Server responds with VersionMismatch when the prover's code version
    /// does not match the version needed to prove the next work item.
    VersionMismatch,

    /// The Server responds with ProverTypeNotNeeded when the connecting prover's
    /// backend type is not in the set of required proof types for this deployment.
    ProverTypeNotNeeded { prover_type: ProverType },

    /// The Server responds with an InputResponse containing the input data.
    /// If all fields are None, the Client knows no work is available.
    InputResponse {
        id: Option<u64>,
        input: Option<I>,
        format: Option<ProofFormat>,
    },

    /// The Client submits the proof generated by the prover for the specified id.
    ProofSubmit { id: u64, proof: ProverOutput },

    /// The Server acknowledges the receipt of the proof and updates its state.
    ProofSubmitACK { id: u64 },
}

impl<I> ProofData<I> {
    pub fn prover_setup(prover_type: ProverType, payload: Bytes) -> Self {
        ProofData::ProverSetup {
            prover_type,
            payload,
        }
    }

    pub fn prover_setup_ack() -> Self {
        ProofData::ProverSetupACK
    }

    pub fn input_request(commit_hash: String, prover_type: ProverType) -> Self {
        ProofData::InputRequest {
            commit_hash,
            prover_type,
        }
    }

    pub fn version_mismatch() -> Self {
        ProofData::VersionMismatch
    }

    pub fn input_response(id: u64, input: I, format: ProofFormat) -> Self {
        ProofData::InputResponse {
            id: Some(id),
            input: Some(input),
            format: Some(format),
        }
    }

    pub fn empty_input_response() -> Self {
        ProofData::InputResponse {
            id: None,
            input: None,
            format: None,
        }
    }

    pub fn proof_submit(id: u64, proof: ProverOutput) -> Self {
        ProofData::ProofSubmit { id, proof }
    }

    pub fn proof_submit_ack(id: u64) -> Self {
        ProofData::ProofSubmitACK { id }
    }
}