warriors/trisha/rs/warrior.rs

use rayon::prelude::*;
use triton_vm::prelude::*;
use triton_vm::proof::Claim;

use trident::runtime::{
    Deployer, ExecutionResult, GuessResult, Guesser, ProgramBundle, ProgramInput, ProofData,
    Prover, Runner, Verifier,
};

use crate::convert;

pub struct ProveResult {
    pub proof_data: ProofData,
    pub cycle_count: u64,
    pub padded_height: u64,
}

pub struct Warrior;

impl Warrior {
    pub fn new() -> Self {
        eprintln!("Backend: cpu");
        Warrior
    }

    pub fn prove_full(
        &self,
        bundle: &ProgramBundle,
        input: &ProgramInput,
    ) -> Result<ProveResult, String> {
        let program = Program::from_code(&bundle.assembly)
            .map_err(|e| format!("TASM parse error: {}", e))?;

        let (pub_in, non_det) = convert::to_triton_inputs(input);

        let op_count = bundle.assembly.lines().count();
        eprintln!("Proving {} ({} ops)...", bundle.name, op_count);

        let (aet, output) =
            VM::trace_execution(program.clone(), pub_in.clone(), non_det)
                .map_err(|e| format!("execution error: {}", e))?;

        let cycle_count = aet.processor_trace.nrows() as u64;
        let padded_height = aet.padded_height() as u64;

        let claim = Claim::about_program(&program)
            .with_input(pub_in.individual_tokens)
            .with_output(output);

        let proof = Stark::default()
            .prove(&claim, &aet)
            .map_err(|e| format!("proving error: {}", e))?;

        eprintln!(
            "Proof generated ({} cycles, padded height {})",
            cycle_count, padded_height
        );

        Ok(ProveResult {
            proof_data: ProofData {
                claim: convert::to_trident_claim(&claim),
                proof_bytes: convert::proof_to_bytes(&proof),
                format: "stark-triton-v2".to_string(),
            },
            cycle_count,
            padded_height,
        })
    }
}

impl Runner for Warrior {
    fn run(&self, bundle: &ProgramBundle, input: &ProgramInput) -> Result<ExecutionResult, String> {
        let program =
            Program::from_code(&bundle.assembly).map_err(|e| format!("TASM parse error: {}", e))?;

        let (pub_in, non_det) = convert::to_triton_inputs(input);

        let op_count = bundle.assembly.lines().count();
        eprintln!("Executing {} ({} ops)...", bundle.name, op_count);

        let (aet, output) = VM::trace_execution(program, pub_in, non_det)
            .map_err(|e| format!("execution error: {}", e))?;
        let cycle_count = aet.processor_trace.nrows() as u64;
        Ok(convert::to_execution_result(&output, cycle_count))
    }
}

impl Prover for Warrior {
    fn prove(&self, bundle: &ProgramBundle, input: &ProgramInput) -> Result<ProofData, String> {
        self.prove_full(bundle, input).map(|r| r.proof_data)
    }
}

impl Verifier for Warrior {
    fn verify(&self, proof_data: &ProofData) -> Result<bool, String> {
        let claim = convert::to_triton_claim_native(
            &proof_data.claim.program_hash,
            &proof_data.claim.public_input,
            &proof_data.claim.public_output,
        );
        let proof = convert::bytes_to_proof(&proof_data.proof_bytes)?;
        let stark = Stark::default();
        Ok(triton_vm::verify(stark, &claim, &proof))
    }
}

impl Deployer for Warrior {
    fn deploy(&self, bundle: &ProgramBundle, proof: Option<&ProofData>) -> Result<String, String> {
        let program =
            Program::from_code(&bundle.assembly).map_err(|e| format!("TASM parse error: {}", e))?;
        let digest = program.hash();
        let digest_u64s = convert::digest_to_u64s(&digest);
        let digest_str = digest_u64s
            .iter()
            .map(|v| v.to_string())
            .collect::<Vec<_>>()
            .join(":");

        eprintln!("Program:   {}", bundle.name);
        eprintln!("Digest:    {}", digest_str);
        eprintln!(
            "Proof:     {}",
            if proof.is_some() { "attached" } else { "none" }
        );
        eprintln!();
        eprintln!("On-chain deployment requires a running Neptune node.");
        eprintln!("Neptune RPC is not yet available in this release.");
        eprintln!("The program is ready for deployment — use the digest");
        eprintln!("to construct a LockScript when Neptune SDK is available.");

        Ok(digest_str)
    }
}

impl Guesser for Warrior {
    fn guess(
        &self,
        bundle: &ProgramBundle,
        _input: &ProgramInput,
        difficulty: u64,
        max_attempts: u64,
    ) -> Result<GuessResult, String> {
        let program =
            Program::from_code(&bundle.assembly).map_err(|e| format!("TASM parse error: {}", e))?;

        let base: [BFieldElement; 5] = program.hash().0;

        let threads = rayon::current_num_threads();
        eprintln!(
            "Mining {} (difficulty {}, max {} attempts, {} threads)...",
            bundle.name, difficulty, max_attempts, threads
        );

        let start = std::time::Instant::now();

        let found = (0u64..max_attempts).into_par_iter().find_any(|&nonce| {
            let input = [base[0], base[1], base[2], base[3], base[4], BFieldElement::new(nonce)];
            Tip5::hash_varlen(&input).0[0].value() < difficulty
        });

        match found {
            Some(nonce) => {
                let input = [base[0], base[1], base[2], base[3], base[4], BFieldElement::new(nonce)];
                let hash = Tip5::hash_varlen(&input);
                let elapsed = start.elapsed();
                let rate = nonce as f64 / elapsed.as_secs_f64();
                eprintln!(
                    "Found nonce {} in {} attempts ({:.0} H/s, {:.1}s)",
                    nonce,
                    nonce + 1,
                    rate,
                    elapsed.as_secs_f64()
                );
                Ok(GuessResult {
                    nonce,
                    digest: hash.0.iter().map(|b| b.value()).collect(),
                    attempts: nonce + 1,
                })
            }
            None => Err(format!(
                "no solution found within {} attempts",
                max_attempts
            )),
        }
    }
}