vrf_pbft/

engine.rs

use crate::error::{Error, Result};
use crate::node::Node;
use crate::pbft::{Message, Round};
use crate::types::{Block, Role, Stage};

pub struct Config {
    fault_tolerance: u64,
    // Per-role expected committee sizes for VRF sortition
    expected_proposers: u64,
    expected_validators: u64,
    expected_packers: u64,
}

impl Config {
    /// Create a config with the given fault tolerance `f`.
    /// The system tolerates up to `f` byzantine nodes.
    /// Requires at least `3f + 1` total nodes.
    pub fn new(fault_tolerance: u64) -> Self {
        Self {
            fault_tolerance,
            expected_proposers: 1,
            expected_validators: 3 * fault_tolerance + 1,
            expected_packers: fault_tolerance + 1,
        }
    }

    pub fn with_expected_committee(mut self, n: u64) -> Self {
        self.expected_validators = n;
        self
    }

    /// Minimum votes needed for quorum: 2f + 1
    pub fn threshold(&self) -> u64 {
        2 * self.fault_tolerance + 1
    }

    /// Minimum nodes needed: 3f + 1
    pub fn min_nodes(&self) -> u64 {
        3 * self.fault_tolerance + 1
    }
}

pub struct ConsensusEngine {
    nodes: Vec<Node>,
    config: Config,
    round: u64,
    committed: Vec<Block>,
    messages: Vec<Message>,
}

impl ConsensusEngine {
    pub fn new(config: Config) -> Self {
        Self {
            nodes: Vec::new(),
            config,
            round: 0,
            committed: Vec::new(),
            messages: Vec::new(),
        }
    }

    pub fn add_node(&mut self, node: Node) {
        self.nodes.push(node);
    }

    pub fn committed_blocks(&self) -> &[Block] {
        &self.committed
    }

    pub fn round(&self) -> u64 {
        self.round
    }

    pub fn nodes(&self) -> &[Node] {
        &self.nodes
    }

    pub fn messages(&self) -> &[Message] {
        &self.messages
    }

    fn total_weight(&self) -> u64 {
        self.nodes.iter().map(|n| n.weight()).sum()
    }

    /// Run a single consensus round through all PBFT phases.
    pub fn run_round(&mut self) -> Result<Block> {
        let min = self.config.min_nodes() as usize;
        if self.nodes.len() < min {
            return Err(Error::NotEnoughNodes {
                needed: min,
                have: self.nodes.len(),
            });
        }

        self.round += 1;
        self.messages.clear();
        let total_weight = self.total_weight();
        let role_expected = [
            (Role::Proposer, self.config.expected_proposers),
            (Role::Validator, self.config.expected_validators),
            (Role::Packer, self.config.expected_packers),
        ];

        // Phase 0: VRF role assignment
        for node in &mut self.nodes {
            node.assign_role(self.round, total_weight, &role_expected)?;
        }

        // Find proposer (first one wins if multiple)
        let proposer_idx = self
            .nodes
            .iter()
            .position(|n| n.role() == Role::Proposer)
            .ok_or(Error::NoProposer(self.round))?;

        // Collect validator indices
        let validator_indices: Vec<usize> = self
            .nodes
            .iter()
            .enumerate()
            .filter(|(_, n)| n.role() == Role::Validator)
            .map(|(i, _)| i)
            .collect();

        // Adaptive threshold: 2/3+1 of actual committee, capped at 2f+1
        let effective_threshold = if validator_indices.is_empty() {
            self.config.threshold()
        } else {
            let adaptive = (validator_indices.len() as u64 * 2 / 3) + 1;
            adaptive.min(self.config.threshold())
        };

        // Protocol-level round tracks votes and stage transitions
        let mut protocol = Round::new(self.round, Role::Proposer, effective_threshold);

        // Phase 1: PrePrepare -- proposer creates block
        let block = self.nodes[proposer_idx].propose_block(self.round)?;
        let block_hash = block.hash();
        protocol.set_block(block.clone());
        self.messages.push(Message::pre_prepare(
            self.nodes[proposer_idx].id(),
            self.round,
            block.clone(),
        ));

        protocol.advance(Stage::Prepare)?;

        // Phase 2: Prepare -- validators vote on the block
        for &idx in &validator_indices {
            let approve = self.nodes[idx].validate_block(&block);
            if approve {
                protocol.add_prepare_vote(self.nodes[idx].id());
            }
            self.messages.push(Message::prepare(
                self.nodes[idx].id(),
                self.round,
                block_hash,
                approve,
            ));
        }

        if !protocol.has_prepare_quorum() {
            return Err(Error::InsufficientVotes {
                needed: effective_threshold,
                got: protocol.prepare_count(),
            });
        }

        protocol.advance(Stage::Commit)?;

        // Phase 3: Commit -- validators confirm
        for &idx in &validator_indices {
            let approve = self.nodes[idx].validate_block(&block);
            if approve {
                protocol.add_commit_vote(self.nodes[idx].id());
            }
            self.messages.push(Message::commit(
                self.nodes[idx].id(),
                self.round,
                block_hash,
                approve,
            ));
        }

        if !protocol.has_commit_quorum() {
            return Err(Error::InsufficientVotes {
                needed: effective_threshold,
                got: protocol.commit_count(),
            });
        }

        protocol.advance(Stage::Reply)?;

        // Phase 4: Reply -- proposer broadcasts final block, all nodes commit
        self.messages.push(Message::reply(
            self.nodes[proposer_idx].id(),
            self.round,
            block.clone(),
        ));

        for node in &mut self.nodes {
            node.commit_block(block.clone());
        }

        self.committed.push(block.clone());
        Ok(block)
    }

    /// Run multiple rounds. Skips rounds where VRF doesn't elect a proposer.
    pub fn run(&mut self, rounds: u64) -> Result<Vec<Block>> {
        let mut blocks = Vec::new();
        for _ in 0..rounds {
            match self.run_round() {
                Ok(block) => blocks.push(block),
                Err(Error::NoProposer(_)) => continue,
                Err(Error::InsufficientVotes { .. }) => continue,
                Err(e) => return Err(e),
            }
        }
        Ok(blocks)
    }
}