ab_farmer_components/

proving.rs

//! Utilities for turning solution candidates (from auditing) into solutions (proving)
//!
//! Solutions generated by [`auditing`](crate::auditing) need to be converted into actual solutions
//! before they can be sent to the node and this is exactly what this module is about.

use crate::auditing::ChunkCandidate;
use crate::reading::{ReadingError, read_record_metadata, read_sector_record_chunks};
use crate::sector::{
    RecordMetadata, SectorContentsMap, SectorContentsMapFromBytesError, SectorMetadataChecksummed,
};
use crate::shard_commitment::{ShardCommitmentsRootsCache, derive_solution_shard_commitment};
use crate::{ReadAt, ReadAtSync};
use ab_core_primitives::hashes::Blake3Hash;
use ab_core_primitives::pieces::{ExtendedRecord, PieceOffset, Record, RecordChunk, RecordRoot};
use ab_core_primitives::pos::PosSeed;
use ab_core_primitives::sectors::{SBucket, SectorId};
use ab_core_primitives::shard::NumShards;
use ab_core_primitives::solutions::{
    ChunkProof, ShardMembershipEntropy, Solution, SolutionDistance,
};
use ab_erasure_coding::ErasureCoding;
use ab_merkle_tree::blake3::BalancedMerkleTree;
use ab_proof_of_space::PosProofs;
use futures::FutureExt;
use std::collections::VecDeque;
use std::io;
use thiserror::Error;

/// Solutions that can be proven if necessary.
///
/// Solutions are generated on demand during iteration.
pub trait ProvableSolutions: ExactSizeIterator {
    /// Best solution distance found, `None` in case there are no solutions
    fn best_solution_distance(&self) -> Option<SolutionDistance>;
}

/// Errors that happen during proving
#[derive(Debug, Error)]
pub enum ProvingError {
    /// Failed to create polynomial for record
    #[error("Failed to create polynomial for record at offset {piece_offset}: {error}")]
    FailedToCreatePolynomialForRecord {
        /// Piece offset
        piece_offset: PieceOffset,
        /// Lower-level error
        error: String,
    },
    /// Failed to decode sector contents map
    #[error("Failed to decode sector contents map: {0}")]
    FailedToDecodeSectorContentsMap(#[from] SectorContentsMapFromBytesError),
    /// Sector contents map and sector metadata disagree on the size of the audited s-bucket
    #[error(
        "S-bucket {s_bucket} has {actual} chunks in sector contents map, but {expected} in sector \
        metadata"
    )]
    SBucketSizeMismatch {
        /// S-bucket
        s_bucket: SBucket,
        /// Number of chunks according to sector metadata
        expected: u16,
        /// Number of chunks according to sector contents map
        actual: usize,
    },
    /// I/O error occurred
    #[error("Proving I/O error: {0}")]
    Io(#[from] io::Error),
    /// Record reading error
    #[error("Record reading error: {0}")]
    RecordReadingError(#[from] ReadingError),
}

impl ProvingError {
    /// Whether this error is fatal and makes farm unusable
    pub fn is_fatal(&self) -> bool {
        #[expect(
            clippy::rest_pattern_accessible_field,
            reason = "Do not care about fields"
        )]
        match self {
            ProvingError::FailedToCreatePolynomialForRecord { .. } => false,
            ProvingError::FailedToDecodeSectorContentsMap(_) => false,
            ProvingError::SBucketSizeMismatch { .. } => false,
            ProvingError::Io(_) => true,
            ProvingError::RecordReadingError(error) => error.is_fatal(),
        }
    }
}

#[derive(Debug, Clone)]
struct WinningChunk {
    /// Piece offset in a sector
    piece_offset: PieceOffset,
    /// Solution distance of this chunk
    solution_distance: SolutionDistance,
}

/// Container for solution candidates.
///
/// [`SolutionCandidates::into_solutions`] is used to get an iterator over proven solutions that are
/// generated on demand during iteration.
#[derive(Debug)]
pub struct SolutionCandidates<'a, Sector>
where
    Sector: 'a,
{
    public_key_hash: &'a Blake3Hash,
    sector_id: SectorId,
    shard_commitments_roots_cache: &'a ShardCommitmentsRootsCache,
    shard_membership_entropy: ShardMembershipEntropy,
    num_shards: NumShards,
    s_bucket: SBucket,
    sector: Sector,
    sector_metadata: &'a SectorMetadataChecksummed,
    chunk_candidates: VecDeque<ChunkCandidate>,
}

impl<'a, Sector> Clone for SolutionCandidates<'a, Sector>
where
    Sector: Clone + 'a,
{
    fn clone(&self) -> Self {
        Self {
            public_key_hash: self.public_key_hash,
            sector_id: self.sector_id,
            shard_commitments_roots_cache: self.shard_commitments_roots_cache,
            shard_membership_entropy: self.shard_membership_entropy,
            num_shards: self.num_shards,
            s_bucket: self.s_bucket,
            sector: self.sector.clone(),
            sector_metadata: self.sector_metadata,
            chunk_candidates: self.chunk_candidates.clone(),
        }
    }
}

impl<'a, Sector> SolutionCandidates<'a, Sector>
where
    Sector: ReadAtSync + 'a,
{
    #[expect(clippy::too_many_arguments, reason = "Private API")]
    pub(crate) fn new(
        public_key_hash: &'a Blake3Hash,
        sector_id: SectorId,
        shard_commitments_roots_cache: &'a ShardCommitmentsRootsCache,
        shard_membership_entropy: ShardMembershipEntropy,
        num_shards: NumShards,
        s_bucket: SBucket,
        sector: Sector,
        sector_metadata: &'a SectorMetadataChecksummed,
        chunk_candidates: VecDeque<ChunkCandidate>,
    ) -> Self {
        Self {
            public_key_hash,
            sector_id,
            shard_commitments_roots_cache,
            shard_membership_entropy,
            num_shards,
            s_bucket,
            sector,
            sector_metadata,
            chunk_candidates,
        }
    }

    /// Total number of candidates
    pub fn len(&self) -> usize {
        self.chunk_candidates.len()
    }

    /// Returns true if no candidates inside
    pub fn is_empty(&self) -> bool {
        self.chunk_candidates.is_empty()
    }

    /// Turn solution candidates into actual solutions
    pub fn into_solutions<PosProofGenerator>(
        self,
        erasure_coding: &'a ErasureCoding,
        table_generator: PosProofGenerator,
    ) -> Result<impl ProvableSolutions<Item = MaybeSolution> + 'a, ProvingError>
    where
        PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
    {
        SolutionsIterator::<'a, _, _>::new(
            self.public_key_hash,
            self.sector_id,
            self.shard_commitments_roots_cache,
            self.shard_membership_entropy,
            self.num_shards,
            self.s_bucket,
            self.sector,
            self.sector_metadata,
            erasure_coding,
            self.chunk_candidates,
            table_generator,
        )
    }
}

type MaybeSolution = Result<Solution, ProvingError>;

struct SolutionsIterator<'a, PosProofGenerator, Sector>
where
    Sector: ReadAtSync + 'a,
    PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
{
    public_key_hash: &'a Blake3Hash,
    sector_id: SectorId,
    shard_commitments_roots_cache: &'a ShardCommitmentsRootsCache,
    shard_membership_entropy: ShardMembershipEntropy,
    num_shards: NumShards,
    s_bucket: SBucket,
    sector_metadata: &'a SectorMetadataChecksummed,
    s_bucket_offsets: Box<[u32; Record::NUM_S_BUCKETS]>,
    erasure_coding: &'a ErasureCoding,
    sector_contents_map: SectorContentsMap,
    sector: ReadAt<Sector, !>,
    winning_chunks: VecDeque<WinningChunk>,
    count: usize,
    best_solution_distance: Option<SolutionDistance>,
    table_generator: PosProofGenerator,
}

impl<'a, PosProofGenerator, Sector> ExactSizeIterator
    for SolutionsIterator<'a, PosProofGenerator, Sector>
where
    Sector: ReadAtSync + 'a,
    PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
{
}

impl<'a, PosProofGenerator, Sector> Iterator for SolutionsIterator<'a, PosProofGenerator, Sector>
where
    Sector: ReadAtSync + 'a,
    PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
{
    type Item = MaybeSolution;

    fn next(&mut self) -> Option<Self::Item> {
        let WinningChunk {
            piece_offset,
            solution_distance: _,
        } = self.winning_chunks.pop_front()?;

        self.count -= 1;

        // Derive PoSpace proofs
        let pos_proofs =
            (self.table_generator)(&self.sector_id.derive_evaluation_seed(piece_offset));

        let maybe_solution = try {
            let sector_record_chunks_fut = read_sector_record_chunks(
                piece_offset,
                self.sector_metadata.pieces_in_sector,
                &self.s_bucket_offsets,
                &self.sector_contents_map,
                &pos_proofs,
                &self.sector,
            );
            let mut sector_record_chunks = sector_record_chunks_fut
                .now_or_never()
                .expect("Sync reader; qed")
                .map_err(ProvingError::RecordReadingError)?;

            let chunk = sector_record_chunks.extended_record[self.s_bucket];

            let ExtendedRecord { source, parity } = &mut *sector_record_chunks.extended_record;
            self.erasure_coding
                .recover_all(source, parity, &sector_record_chunks.present)
                .map_err(|error| ReadingError::FailedToErasureDecodeRecord {
                    piece_offset,
                    error,
                })
                .map_err(ProvingError::RecordReadingError)?;

            let record_merkle_tree = BalancedMerkleTree::<{ Record::NUM_S_BUCKETS }>::new_boxed(
                RecordChunk::array_to_repr(sector_record_chunks.extended_record.as_s_buckets()),
            );

            // NOTE: We do not check plot consistency using checksum because it is more
            // expensive and consensus will verify validity of the proof anyway
            let record_metadata_fut = read_record_metadata(
                piece_offset,
                self.sector_metadata.pieces_in_sector,
                &self.sector,
            );
            let RecordMetadata {
                piece_header,
                piece_checksum: _,
            } = record_metadata_fut
                .now_or_never()
                .expect("Sync reader; qed")
                .map_err(ProvingError::RecordReadingError)?;

            let proof_of_space = pos_proofs.for_s_bucket(self.s_bucket).expect(
                "Reading checked that the map matches the proofs, and the map has a chunk of this \
                record in this s-bucket; qed",
            );

            let chunk_proof = record_merkle_tree
                .all_proofs()
                .nth(usize::from(self.s_bucket))
                .expect("Chunk offset is valid, hence corresponding proof exists; qed");

            let history_size = self.sector_metadata.history_size;
            let shard_commitment = derive_solution_shard_commitment(
                self.public_key_hash,
                &self.shard_commitments_roots_cache.shard_commitments_seed(),
                &self.shard_commitments_roots_cache.get(history_size),
                history_size,
                &self.shard_membership_entropy,
                self.num_shards,
            );

            Solution {
                public_key_hash: *self.public_key_hash,
                shard_commitment,
                piece_local_segment_index: piece_header.local_segment_index.as_inner(),
                piece_super_segment_index: piece_header.super_segment_index.as_inner(),
                segment_root: piece_header.segment_root,
                segment_proof: piece_header.segment_proof,
                record_root: RecordRoot::from(record_merkle_tree.root()),
                record_proof: piece_header.record_proof,
                chunk,
                chunk_proof: ChunkProof::from(chunk_proof),
                proof_of_space,
                history_size,
                sector_index: self.sector_metadata.sector_index,
                piece_offset,
                segment_position: piece_header.segment_position.as_inner(),
                piece_shard_index: piece_header.shard_index.as_inner(),
                padding: [0; _],
            }
        };

        match maybe_solution {
            Ok(solution) => Some(Ok(solution)),
            Err(error) => Some(Err(error)),
        }
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.count, Some(self.count))
    }
}

impl<'a, PosProofGenerator, Sector> ProvableSolutions
    for SolutionsIterator<'a, PosProofGenerator, Sector>
where
    Sector: ReadAtSync + 'a,
    PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
{
    fn best_solution_distance(&self) -> Option<SolutionDistance> {
        self.best_solution_distance
    }
}

impl<'a, PosProofGenerator, Sector> SolutionsIterator<'a, PosProofGenerator, Sector>
where
    Sector: ReadAtSync + 'a,
    PosProofGenerator: (FnMut(&PosSeed) -> Box<PosProofs>) + 'a,
{
    // TODO: Struct for arguments
    #[expect(clippy::too_many_arguments, reason = "TODO")]
    fn new(
        public_key_hash: &'a Blake3Hash,
        sector_id: SectorId,
        shard_commitments_roots_cache: &'a ShardCommitmentsRootsCache,
        shard_membership_entropy: ShardMembershipEntropy,
        num_shards: NumShards,
        s_bucket: SBucket,
        sector: Sector,
        sector_metadata: &'a SectorMetadataChecksummed,
        erasure_coding: &'a ErasureCoding,
        chunk_candidates: VecDeque<ChunkCandidate>,
        table_generator: PosProofGenerator,
    ) -> Result<Self, ProvingError> {
        let sector_contents_map = {
            let mut sector_contents_map_bytes =
                vec![0; SectorContentsMap::encoded_size(sector_metadata.pieces_in_sector)];

            sector.read_at(&mut sector_contents_map_bytes, 0)?;

            SectorContentsMap::from_bytes(
                &sector_contents_map_bytes,
                sector_metadata.pieces_in_sector,
            )?
        };

        let s_bucket_piece_offsets = sector_contents_map
            .iter_s_bucket_piece_offsets(s_bucket)
            .expect("S-bucket audit index is guaranteed to be in range; qed")
            .collect::<Vec<_>>();
        let s_bucket_size = sector_metadata
            .s_bucket_sizes
            .get(usize::from(s_bucket))
            .copied()
            .expect("S-bucket audit index is guaranteed to be in range; qed");
        // Auditing found candidates among as many chunks as sector metadata says the s-bucket has
        if s_bucket_piece_offsets.len() != usize::from(s_bucket_size) {
            return Err(ProvingError::SBucketSizeMismatch {
                s_bucket,
                expected: s_bucket_size,
                actual: s_bucket_piece_offsets.len(),
            });
        }
        let winning_chunks = chunk_candidates
            .into_iter()
            .map(move |chunk_candidate| {
                let piece_offset = s_bucket_piece_offsets
                    .get(chunk_candidate.chunk_offset as usize)
                    .expect(
                        "Auditing only creates candidates within the s-bucket size from metadata, \
                        which matches the map as checked above; qed",
                    );

                WinningChunk {
                    piece_offset: *piece_offset,
                    solution_distance: chunk_candidate.solution_distance,
                }
            })
            .collect::<VecDeque<_>>();

        let best_solution_distance = winning_chunks
            .front()
            .map(|winning_chunk| winning_chunk.solution_distance);

        let s_bucket_offsets = sector_metadata.s_bucket_offsets();

        let count = winning_chunks.len();

        Ok(Self {
            public_key_hash,
            sector_id,
            shard_commitments_roots_cache,
            shard_membership_entropy,
            num_shards,
            s_bucket,
            sector_metadata,
            s_bucket_offsets,
            erasure_coding,
            sector_contents_map,
            sector: ReadAt::from_sync(sector),
            winning_chunks,
            count,
            best_solution_distance,
            table_generator,
        })
    }
}