1use crate::challenge::{Challenge, ChallengeGenerator};
4use crate::mml_path::MMLPathValidator;
5use crate::neural_path::NeuralPathValidator;
6use ai3_lib::{AI3Engine, EngineStats, TensorEngine};
7use pot_o_core::TribeResult;
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10use std::time::Instant;
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct PotOProof {
15 pub challenge_id: String,
16 pub challenge_hash: String,
17 pub tensor_result_hash: String,
18 pub mml_score: f64,
19 pub path_signature: String,
20 pub path_distance: u32,
21 pub computation_nonce: u64,
22 pub computation_hash: String,
23 pub miner_pubkey: String,
24 pub timestamp: chrono::DateTime<chrono::Utc>,
25}
26
27#[derive(Debug, Clone, Serialize, Deserialize)]
29pub struct ProofPayload {
30 pub proof: PotOProof,
31 pub signature: Vec<u8>,
32}
33
34pub struct PotOConsensus {
37 pub engine: Box<dyn TensorEngine>,
38 pub challenge_gen: ChallengeGenerator,
39 pub mml_validator: MMLPathValidator,
40 pub neural_validator: NeuralPathValidator,
41}
42
43impl PotOConsensus {
44 pub fn new(difficulty: u64, max_tensor_dim: usize) -> Self {
45 Self {
46 engine: Box::new(AI3Engine::new()),
47 challenge_gen: ChallengeGenerator::new(difficulty, max_tensor_dim),
48 mml_validator: MMLPathValidator::default(),
49 neural_validator: NeuralPathValidator::default(),
50 }
51 }
52
53 pub fn generate_challenge(&self, slot: u64, slot_hash: &str) -> TribeResult<Challenge> {
55 self.challenge_gen.generate(slot, slot_hash)
56 }
57
58 pub fn mine(
61 &self,
62 challenge: &Challenge,
63 miner_pubkey: &str,
64 max_iterations: u64,
65 ) -> TribeResult<Option<PotOProof>> {
66 let start = Instant::now();
67 let task = challenge.to_mining_task(miner_pubkey);
68
69 let output_tensor = self.engine.execute_task(&task)?;
70 let mml_score = self
71 .mml_validator
72 .compute_mml_score(&challenge.input_tensor, &output_tensor)?;
73
74 for nonce in 0..max_iterations {
75 let actual_path = self
76 .neural_validator
77 .compute_actual_path(&output_tensor, nonce)?;
78 let expected = self.neural_validator.expected_path_signature(&challenge.id);
79 let min_len = actual_path.len().min(expected.len());
80 let distance = NeuralPathValidator::hamming_distance(
81 &actual_path[..min_len],
82 &expected[..min_len],
83 );
84
85 if distance <= challenge.path_distance_max
86 && self
87 .mml_validator
88 .validate(mml_score, challenge.mml_threshold)
89 {
90 let tensor_result_hash = output_tensor.calculate_hash();
91 let path_sig = NeuralPathValidator::path_to_hex(&actual_path);
92 let computation_hash = Self::compute_proof_hash(
93 &challenge.id,
94 &tensor_result_hash,
95 mml_score,
96 &path_sig,
97 nonce,
98 );
99
100 let elapsed = start.elapsed();
101 self.engine.record_result(true, elapsed);
102
103 return Ok(Some(PotOProof {
104 challenge_id: challenge.id.clone(),
105 challenge_hash: challenge.slot_hash.clone(),
106 tensor_result_hash,
107 mml_score,
108 path_signature: path_sig,
109 path_distance: distance,
110 computation_nonce: nonce,
111 computation_hash,
112 miner_pubkey: miner_pubkey.to_string(),
113 timestamp: chrono::Utc::now(),
114 }));
115 }
116 }
117
118 self.engine.record_result(false, start.elapsed());
119 Ok(None)
120 }
121
122 pub fn verify_proof(&self, proof: &PotOProof, challenge: &Challenge) -> TribeResult<bool> {
125 let expected_hash = Self::compute_proof_hash(
127 &proof.challenge_id,
128 &proof.tensor_result_hash,
129 proof.mml_score,
130 &proof.path_signature,
131 proof.computation_nonce,
132 );
133 if expected_hash != proof.computation_hash {
134 return Ok(false);
135 }
136
137 if proof.mml_score < 0.0 {
139 return Ok(false);
140 }
141
142 let task = challenge.to_mining_task(&proof.miner_pubkey);
144 let output_tensor = self.engine.execute_task(&task)?;
145 let recomputed_hash = output_tensor.calculate_hash();
146 if recomputed_hash != proof.tensor_result_hash {
147 return Ok(false);
148 }
149
150 let recomputed_mml = self
152 .mml_validator
153 .compute_mml_score(&challenge.input_tensor, &output_tensor)?;
154 if (recomputed_mml - proof.mml_score).abs() > f64::EPSILON {
155 return Ok(false);
156 }
157
158 if !self
160 .mml_validator
161 .validate(proof.mml_score, challenge.mml_threshold)
162 {
163 return Ok(false);
164 }
165
166 let actual_path = self
168 .neural_validator
169 .compute_actual_path(&output_tensor, proof.computation_nonce)?;
170 let expected_path = self.neural_validator.expected_path_signature(&challenge.id);
171 let min_len = actual_path.len().min(expected_path.len());
172 let recomputed_distance = NeuralPathValidator::hamming_distance(
173 &actual_path[..min_len],
174 &expected_path[..min_len],
175 );
176 if recomputed_distance != proof.path_distance {
177 return Ok(false);
178 }
179
180 if proof.path_distance > challenge.path_distance_max {
182 return Ok(false);
183 }
184
185 let recomputed_path_sig = NeuralPathValidator::path_to_hex(&actual_path);
187 if recomputed_path_sig != proof.path_signature {
188 return Ok(false);
189 }
190
191 Ok(true)
192 }
193
194 pub fn expected_paths_and_calcs(&self, challenge: &Challenge) -> (u64, u64) {
198 let expected_paths = self
199 .neural_validator
200 .expected_path_signature(&challenge.id)
201 .len() as u64;
202 let expected_calcs = 1 + challenge.difficulty;
203 (expected_paths, expected_calcs)
204 }
205
206 pub fn engine_stats(&self) -> EngineStats {
208 self.engine.get_stats()
209 }
210
211 pub fn compute_proof_hash(
213 challenge_id: &str,
214 tensor_result_hash: &str,
215 mml_score: f64,
216 path_signature: &str,
217 nonce: u64,
218 ) -> String {
219 let mut hasher = Sha256::new();
220 hasher.update(challenge_id.as_bytes());
221 hasher.update(tensor_result_hash.as_bytes());
222 hasher.update(mml_score.to_le_bytes());
223 hasher.update(path_signature.as_bytes());
224 hasher.update(nonce.to_le_bytes());
225 hex::encode(hasher.finalize())
226 }
227}
228
229#[cfg(test)]
230mod tests {
231 use super::*;
232
233 #[test]
234 fn test_consensus_lifecycle() {
235 let consensus = PotOConsensus::new(1, 8);
236 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
237 let challenge = consensus.generate_challenge(100, hash).unwrap();
238 assert!(!challenge.id.is_empty());
239
240 let result = consensus
242 .mine(&challenge, "test_miner_pubkey", 1000)
243 .unwrap();
244 assert!(result.is_some(), "Should find a proof with low difficulty");
245
246 let proof = result.unwrap();
247 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
248 assert!(valid, "Mined proof should verify");
249 }
250
251 #[test]
252 fn test_proof_hash_deterministic() {
253 let h1 = PotOConsensus::compute_proof_hash("chal", "tensor", 0.5, "path", 42);
254 let h2 = PotOConsensus::compute_proof_hash("chal", "tensor", 0.5, "path", 42);
255 assert_eq!(h1, h2);
256 }
257
258 #[test]
259 fn test_expected_paths_and_calcs() {
260 let consensus = PotOConsensus::new(2, 8);
261 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
262 let challenge = consensus.generate_challenge(100, hash).unwrap();
263 let (expected_paths, expected_calcs) = consensus.expected_paths_and_calcs(&challenge);
264 assert!(expected_paths > 0, "expected_paths should be positive");
265 assert!(expected_calcs > 0, "expected_calcs should be positive");
266 let path_len = consensus
267 .neural_validator
268 .expected_path_signature(&challenge.id)
269 .len() as u64;
270 assert_eq!(
271 expected_paths, path_len,
272 "expected_paths should match path signature length"
273 );
274 assert_eq!(expected_calcs, 1 + challenge.difficulty);
275 }
276
277 #[test]
278 fn test_tampered_tensor_result_hash_rejected() {
279 let consensus = PotOConsensus::new(1, 8);
280 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
281 let challenge = consensus.generate_challenge(100, hash).unwrap();
282 let mut proof = consensus
283 .mine(&challenge, "test_miner_pubkey", 1000)
284 .unwrap()
285 .expect("should mine");
286
287 proof.tensor_result_hash = "deadbeef".repeat(8);
288 proof.computation_hash = PotOConsensus::compute_proof_hash(
289 &proof.challenge_id,
290 &proof.tensor_result_hash,
291 proof.mml_score,
292 &proof.path_signature,
293 proof.computation_nonce,
294 );
295
296 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
297 assert!(!valid, "Tampered tensor_result_hash must be rejected");
298 }
299
300 #[test]
301 fn test_negative_mml_score_rejected() {
302 let consensus = PotOConsensus::new(1, 8);
303 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
304 let challenge = consensus.generate_challenge(100, hash).unwrap();
305 let mut proof = consensus
306 .mine(&challenge, "test_miner_pubkey", 1000)
307 .unwrap()
308 .expect("should mine");
309
310 proof.mml_score = -1.0;
311 proof.computation_hash = PotOConsensus::compute_proof_hash(
312 &proof.challenge_id,
313 &proof.tensor_result_hash,
314 proof.mml_score,
315 &proof.path_signature,
316 proof.computation_nonce,
317 );
318
319 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
320 assert!(!valid, "Negative MML score must be rejected");
321 }
322
323 #[test]
324 fn test_tampered_mml_score_rejected() {
325 let consensus = PotOConsensus::new(1, 8);
326 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
327 let challenge = consensus.generate_challenge(100, hash).unwrap();
328 let mut proof = consensus
329 .mine(&challenge, "test_miner_pubkey", 1000)
330 .unwrap()
331 .expect("should mine");
332
333 proof.mml_score = 999.0;
334 proof.computation_hash = PotOConsensus::compute_proof_hash(
335 &proof.challenge_id,
336 &proof.tensor_result_hash,
337 proof.mml_score,
338 &proof.path_signature,
339 proof.computation_nonce,
340 );
341
342 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
343 assert!(!valid, "Tampered MML score must be rejected");
344 }
345
346 #[test]
347 fn test_tampered_path_distance_rejected() {
348 let consensus = PotOConsensus::new(1, 8);
349 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
350 let challenge = consensus.generate_challenge(100, hash).unwrap();
351 let mut proof = consensus
352 .mine(&challenge, "test_miner_pubkey", 1000)
353 .unwrap()
354 .expect("should mine");
355
356 proof.path_distance = proof.path_distance.wrapping_add(100);
357
358 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
359 assert!(!valid, "Tampered path_distance must be rejected");
360 }
361
362 #[test]
363 fn test_tampered_path_signature_rejected() {
364 let consensus = PotOConsensus::new(1, 8);
365 let hash = "abcdef0123456789abcdef0123456789abcdef0123456789abcdef0123456789";
366 let challenge = consensus.generate_challenge(100, hash).unwrap();
367 let mut proof = consensus
368 .mine(&challenge, "test_miner_pubkey", 1000)
369 .unwrap()
370 .expect("should mine");
371
372 proof.path_signature = "ff".repeat(64);
373 proof.computation_hash = PotOConsensus::compute_proof_hash(
374 &proof.challenge_id,
375 &proof.tensor_result_hash,
376 proof.mml_score,
377 &proof.path_signature,
378 proof.computation_nonce,
379 );
380
381 let valid = consensus.verify_proof(&proof, &challenge).unwrap();
382 assert!(!valid, "Tampered path_signature must be rejected");
383 }
384}