NAC_Blockchain/nac-cbpp/src/consensus.rs

//! CBPP共识引擎

use crate::block::Block;
use crate::validator::ValidatorSet;
use crate::vote::{Vote, VoteSet, VoteType};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// 共识状态
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConsensusState {
    NewHeight,    // 新高度
    Propose,      // 提议阶段
    Prevote,      // 预投票阶段
    Precommit,    // 预提交阶段
    Commit,       // 提交阶段
}

/// CBPP共识引擎
#[derive(Debug)]
pub struct ConsensusEngine {
    state: ConsensusState,
    height: u64,
    round: u32,
    validator_set: ValidatorSet,
    prevotes: HashMap<String, VoteSet>,
    precommits: HashMap<String, VoteSet>,
    locked_block: Option<Block>,
    locked_round: Option<u32>,
}

impl ConsensusEngine {
    pub fn new() -> Self {
        ConsensusEngine {
            state: ConsensusState::NewHeight,
            height: 0,
            round: 0,
            validator_set: ValidatorSet::new(),
            prevotes: HashMap::new(),
            precommits: HashMap::new(),
            locked_block: None,
            locked_round: None,
        }
    }
    
    pub fn is_initialized(&self) -> bool {
        true
    }
    
    /// 设置验证者集合
    pub fn set_validator_set(&mut self, validator_set: ValidatorSet) {
        self.validator_set = validator_set;
    }
    
    /// 开始新高度
    pub fn start_new_height(&mut self, height: u64) {
        self.height = height;
        self.round = 0;
        self.state = ConsensusState::NewHeight;
        self.prevotes.clear();
        self.precommits.clear();
        self.locked_block = None;
        self.locked_round = None;
    }
    
    /// 进入提议阶段
    pub fn enter_propose(&mut self) {
        self.state = ConsensusState::Propose;
    }
    
    /// 处理提议
    pub fn handle_proposal(&mut self, block: Block) -> bool {
        if self.state != ConsensusState::Propose {
            return false;
        }
        
        // 验证区块
        if !self.validate_block(&block) {
            return false;
        }
        
        // 进入预投票阶段
        self.state = ConsensusState::Prevote;
        true
    }
    
    /// 处理预投票
    pub fn handle_prevote(&mut self, vote: Vote) -> bool {
        if vote.vote_type != VoteType::Prevote {
            return false;
        }
        
        let vote_set = self.prevotes
            .entry(vote.block_hash.clone())
            .or_insert_with(|| VoteSet::new(self.validator_set.total_voting_power()));
        
        vote_set.add_vote(vote);
        
        // 检查是否达到2/3+多数
        if self.check_prevote_majority() {
            self.state = ConsensusState::Precommit;
            return true;
        }
        
        false
    }
    
    /// 处理预提交
    pub fn handle_precommit(&mut self, vote: Vote) -> bool {
        if vote.vote_type != VoteType::Precommit {
            return false;
        }
        
        let vote_set = self.precommits
            .entry(vote.block_hash.clone())
            .or_insert_with(|| VoteSet::new(self.validator_set.total_voting_power()));
        
        vote_set.add_vote(vote);
        
        // 检查是否达到2/3+多数
        if self.check_precommit_majority() {
            self.state = ConsensusState::Commit;
            return true;
        }
        
        false
    }
    
    /// 提交区块
    pub fn commit_block(&mut self) -> Option<Block> {
        if self.state != ConsensusState::Commit {
            return None;
        }
        
        let block = self.locked_block.take();
        
        // 进入新高度
        self.start_new_height(self.height + 1);
        
        block
    }
    
    /// 验证区块
    fn validate_block(&self, _block: &Block) -> bool {
        // 简化实现，实际应该验证：
        // 1. 区块签名
        // 2. 交易有效性
        // 3. 状态转换
        // 4. Merkle根
        true
    }
    
    /// 检查预投票是否达到多数
    fn check_prevote_majority(&self) -> bool {
        for vote_set in self.prevotes.values() {
            let voting_power = vote_set.len() as u64 * 1000; // 简化计算
            if vote_set.has_two_thirds_majority(voting_power) {
                return true;
            }
        }
        false
    }
    
    /// 检查预提交是否达到多数
    fn check_precommit_majority(&self) -> bool {
        for vote_set in self.precommits.values() {
            let voting_power = vote_set.len() as u64 * 1000; // 简化计算
            if vote_set.has_two_thirds_majority(voting_power) {
                return true;
            }
        }
        false
    }
    
    /// 获取当前状态
    pub fn state(&self) -> ConsensusState {
        self.state
    }
    
    /// 获取当前高度
    pub fn height(&self) -> u64 {
        self.height
    }
    
    /// 获取当前轮次
    pub fn round(&self) -> u32 {
        self.round
    }
}

impl Default for ConsensusEngine {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::validator::Validator;

    #[test]
    fn test_consensus_engine_creation() {
        let engine = ConsensusEngine::new();
        assert_eq!(engine.state(), ConsensusState::NewHeight);
        assert_eq!(engine.height(), 0);
        assert_eq!(engine.round(), 0);
    }

    #[test]
    fn test_start_new_height() {
        let mut engine = ConsensusEngine::new();
        engine.start_new_height(1);
        
        assert_eq!(engine.height(), 1);
        assert_eq!(engine.round(), 0);
        assert_eq!(engine.state(), ConsensusState::NewHeight);
    }

    #[test]
    fn test_consensus_flow() {
        let mut engine = ConsensusEngine::new();
        
        // 设置验证者
        let mut validator_set = ValidatorSet::new();
        validator_set.add_validator(Validator::new("v1".to_string(), 1000));
        validator_set.add_validator(Validator::new("v2".to_string(), 1000));
        validator_set.add_validator(Validator::new("v3".to_string(), 1000));
        engine.set_validator_set(validator_set);
        
        // 开始新高度
        engine.start_new_height(1);
        assert_eq!(engine.state(), ConsensusState::NewHeight);
        
        // 进入提议阶段
        engine.enter_propose();
        assert_eq!(engine.state(), ConsensusState::Propose);
        
        // 处理提议
        let block = Block::new(1, "genesis".to_string(), "v1".to_string());
        assert!(engine.handle_proposal(block));
        assert_eq!(engine.state(), ConsensusState::Prevote);
    }
}