1  // Copyright (c) 2025 ADnet Contributors
  2  // SPDX-License-Identifier: Apache-2.0
  3  
  4  //! DX Staking System (F-G16)
  5  //!
  6  //! Implements DX staking for dividend eligibility and governance participation.
  7  //! Staking is **default ON** (opt-out model) - all wallets are staked by default.
  8  //!
  9  //! # Key Properties
 10  //!
 11  //! - Default state: `staking_enabled = true` (opt-out model)
 12  //! - Staked DX cannot be transferred (must unstake first)
 13  //! - Received DX inherits recipient wallet's staking state
 14  //! - Unstaking takes effect at epoch transition (daily)
 15  //! - Only staked DX earns dividends
 16  //!
 17  //! # Epoch Transition Priority
 18  //!
 19  //! 1. SLASHING (always first - no escape from governance penalties)
 20  //! 2. DIVIDENDS (calculated on current staked balances before changes)
 21  //! 3. STAKING (pending stake requests take effect)
 22  //! 4. UNSTAKING (pending unstake requests take effect last)
 23  
 24  use crate::TokenCommitment;
 25  use anyhow::{ensure, Result};
 26  use serde::{Deserialize, Serialize};
 27  use sha2::{Digest, Sha256};
 28  use std::collections::HashMap;
 29  
 30  // ============================================================================
 31  // Constants
 32  // ============================================================================
 33  
 34  /// Default staking state for new wallets (true = opt-out model)
 35  pub const DEFAULT_STAKING_STATE: bool = true;
 36  
 37  /// Minimum DX required for governance eligibility (can only be decreased via governance)
 38  pub const GOVERNANCE_THRESHOLD_DX: u64 = 10_000;
 39  
 40  // ============================================================================
 41  // Staking State
 42  // ============================================================================
 43  
 44  /// Staking state for a single wallet
 45  #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
 46  pub struct StakingState {
 47      /// Whether staking is enabled (default: true)
 48      pub is_staking: bool,
 49      /// Epoch when unstake was requested (None if not unstaking)
 50      pub pending_unstake_epoch: Option<u64>,
 51      /// Epoch when stake was requested (None if not re-staking)
 52      pub pending_stake_epoch: Option<u64>,
 53      /// Commitment to staked balance (for privacy)
 54      pub staked_balance_commitment: Option<TokenCommitment>,
 55      /// Last known staked balance (for dividend calculation)
 56      /// Note: Actual balance is ZK-shielded, this is for internal tracking
 57      staked_balance: u64,
 58  }
 59  
 60  impl Default for StakingState {
 61      fn default() -> Self {
 62          Self {
 63              is_staking: DEFAULT_STAKING_STATE,
 64              pending_unstake_epoch: None,
 65              pending_stake_epoch: None,
 66              staked_balance_commitment: None,
 67              staked_balance: 0,
 68          }
 69      }
 70  }
 71  
 72  impl StakingState {
 73      /// Create a new staking state with default settings (staking ON)
 74      pub fn new() -> Self {
 75          Self::default()
 76      }
 77  
 78      /// Create a staking state with specific initial balance
 79      pub fn with_balance(balance: u64) -> Self {
 80          Self {
 81              is_staking: DEFAULT_STAKING_STATE,
 82              pending_unstake_epoch: None,
 83              pending_stake_epoch: None,
 84              staked_balance_commitment: None,
 85              staked_balance: balance,
 86          }
 87      }
 88  
 89      /// Check if wallet is currently staking
 90      pub fn is_staking(&self) -> bool {
 91          self.is_staking
 92      }
 93  
 94      /// Check if there's a pending unstake request
 95      pub fn has_pending_unstake(&self) -> bool {
 96          self.pending_unstake_epoch.is_some()
 97      }
 98  
 99      /// Check if there's a pending stake request
100      pub fn has_pending_stake(&self) -> bool {
101          self.pending_stake_epoch.is_some()
102      }
103  
104      /// Get staked balance (returns 0 if not staking)
105      pub fn staked_balance(&self) -> u64 {
106          if self.is_staking {
107              self.staked_balance
108          } else {
109              0
110          }
111      }
112  
113      /// Get raw balance regardless of staking state
114      ///
115      /// Use this for checking actual token holdings after unstaking.
116      pub fn balance(&self) -> u64 {
117          self.staked_balance
118      }
119  
120      /// Update staked balance
121      pub fn set_balance(&mut self, balance: u64) {
122          self.staked_balance = balance;
123      }
124  
125      /// Add to staked balance (when receiving DX while staked)
126      pub fn add_balance(&mut self, amount: u64) {
127          self.staked_balance = self.staked_balance.saturating_add(amount);
128      }
129  
130      /// Subtract from staked balance (for slashing)
131      pub fn subtract_balance(&mut self, amount: u64) {
132          self.staked_balance = self.staked_balance.saturating_sub(amount);
133      }
134  }
135  
136  // ============================================================================
137  // Pending Operations
138  // ============================================================================
139  
140  /// A pending staking operation
141  #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
142  pub struct PendingOperation {
143      /// DX public key (wallet address hash)
144      pub dx_pubkey: [u8; 32],
145      /// Epoch when the operation was requested
146      pub requested_epoch: u64,
147  }
148  
149  // ============================================================================
150  // Staking Registry
151  // ============================================================================
152  
153  /// Registry tracking all wallet staking states
154  #[derive(Clone, Debug, Default)]
155  pub struct StakingRegistry {
156      /// Wallet states by DX pubkey hash
157      wallet_states: HashMap<[u8; 32], StakingState>,
158      /// Pending unstake requests (to be processed at epoch transition)
159      pending_unstakes: Vec<PendingOperation>,
160      /// Pending stake requests (to be processed at epoch transition)
161      pending_stakes: Vec<PendingOperation>,
162      /// Total staked DX across all wallets (aggregate, for explorer)
163      total_staked: u64,
164      /// Total number of staked wallets (for explorer)
165      staked_wallet_count: u64,
166      /// Current epoch
167      current_epoch: u64,
168  }
169  
170  impl StakingRegistry {
171      /// Create a new empty registry
172      pub fn new() -> Self {
173          Self::default()
174      }
175  
176      /// Create registry at a specific epoch
177      pub fn at_epoch(epoch: u64) -> Self {
178          Self {
179              current_epoch: epoch,
180              ..Default::default()
181          }
182      }
183  
184      /// Get or create wallet state (new wallets default to staking ON)
185      pub fn get_or_create_wallet(&mut self, dx_pubkey: &[u8; 32]) -> &mut StakingState {
186          if !self.wallet_states.contains_key(dx_pubkey) {
187              self.wallet_states.insert(*dx_pubkey, StakingState::new());
188              self.staked_wallet_count += 1;
189          }
190          self.wallet_states.get_mut(dx_pubkey).unwrap()
191      }
192  
193      /// Get wallet state (read-only)
194      pub fn get_wallet(&self, dx_pubkey: &[u8; 32]) -> Option<&StakingState> {
195          self.wallet_states.get(dx_pubkey)
196      }
197  
198      /// Check if a wallet is currently staking
199      ///
200      /// Returns true if:
201      /// - Wallet exists and has is_staking = true
202      /// - Wallet doesn't exist (new wallets default to staking)
203      pub fn is_staked(&self, dx_pubkey: &[u8; 32]) -> bool {
204          self.wallet_states
205              .get(dx_pubkey)
206              .map(|s| s.is_staking)
207              .unwrap_or(DEFAULT_STAKING_STATE)
208      }
209  
210      /// Check if wallet can transfer (must be unstaked)
211      pub fn can_transfer(&self, dx_pubkey: &[u8; 32]) -> bool {
212          !self.is_staked(dx_pubkey)
213      }
214  
215      /// Request to unstake (takes effect at next epoch transition)
216      ///
217      /// # Arguments
218      /// * `dx_pubkey` - The wallet's DX public key
219      ///
220      /// # Returns
221      /// Ok if unstake request was queued, Err if already unstaked
222      pub fn request_unstake(&mut self, dx_pubkey: &[u8; 32]) -> Result<()> {
223          let current_epoch = self.current_epoch;
224          let state = self.get_or_create_wallet(dx_pubkey);
225  
226          ensure!(state.is_staking, "Wallet is already unstaked");
227          ensure!(
228              state.pending_unstake_epoch.is_none(),
229              "Unstake already pending"
230          );
231  
232          state.pending_unstake_epoch = Some(current_epoch);
233  
234          self.pending_unstakes.push(PendingOperation {
235              dx_pubkey: *dx_pubkey,
236              requested_epoch: current_epoch,
237          });
238  
239          Ok(())
240      }
241  
242      /// Request to stake (re-enable staking after unstaking)
243      ///
244      /// Takes effect at next epoch transition.
245      pub fn request_stake(&mut self, dx_pubkey: &[u8; 32]) -> Result<()> {
246          // Copy epoch before mutable borrow
247          let current_epoch = self.current_epoch;
248          let state = self.get_or_create_wallet(dx_pubkey);
249  
250          ensure!(!state.is_staking, "Wallet is already staking");
251          ensure!(state.pending_stake_epoch.is_none(), "Stake already pending");
252  
253          state.pending_stake_epoch = Some(current_epoch);
254  
255          self.pending_stakes.push(PendingOperation {
256              dx_pubkey: *dx_pubkey,
257              requested_epoch: current_epoch,
258          });
259  
260          Ok(())
261      }
262  
263      /// Process pending stake requests (Priority 3 in epoch transition)
264      ///
265      /// Called during epoch transition AFTER dividends are calculated.
266      pub fn process_pending_stakes(&mut self, new_epoch: u64) -> Result<Vec<[u8; 32]>> {
267          let mut processed = Vec::new();
268  
269          for op in self.pending_stakes.drain(..) {
270              if let Some(state) = self.wallet_states.get_mut(&op.dx_pubkey) {
271                  if state.pending_stake_epoch.is_some() {
272                      state.is_staking = true;
273                      state.pending_stake_epoch = None;
274                      self.staked_wallet_count += 1;
275                      self.total_staked = self.total_staked.saturating_add(state.staked_balance);
276                      processed.push(op.dx_pubkey);
277                  }
278              }
279          }
280  
281          self.current_epoch = new_epoch;
282          Ok(processed)
283      }
284  
285      /// Process pending unstake requests (Priority 4 in epoch transition - LAST)
286      ///
287      /// Called during epoch transition AFTER stake requests are processed.
288      pub fn process_pending_unstakes(&mut self, new_epoch: u64) -> Result<Vec<[u8; 32]>> {
289          let mut processed = Vec::new();
290  
291          for op in self.pending_unstakes.drain(..) {
292              if let Some(state) = self.wallet_states.get_mut(&op.dx_pubkey) {
293                  if state.pending_unstake_epoch.is_some() {
294                      state.is_staking = false;
295                      state.pending_unstake_epoch = None;
296                      self.staked_wallet_count = self.staked_wallet_count.saturating_sub(1);
297                      self.total_staked = self.total_staked.saturating_sub(state.staked_balance);
298                      processed.push(op.dx_pubkey);
299                  }
300              }
301          }
302  
303          self.current_epoch = new_epoch;
304          Ok(processed)
305      }
306  
307      /// Register a new wallet with initial balance
308      ///
309      /// New wallets are staked by default.
310      pub fn register_wallet(&mut self, dx_pubkey: &[u8; 32], initial_balance: u64) {
311          let state = self.get_or_create_wallet(dx_pubkey);
312          state.set_balance(initial_balance);
313          self.total_staked = self.total_staked.saturating_add(initial_balance);
314      }
315  
316      /// Update wallet balance (for tracking)
317      pub fn update_balance(&mut self, dx_pubkey: &[u8; 32], new_balance: u64) {
318          if let Some(state) = self.wallet_states.get_mut(dx_pubkey) {
319              let old_balance = state.staked_balance;
320              state.set_balance(new_balance);
321  
322              if state.is_staking {
323                  self.total_staked = self.total_staked.saturating_sub(old_balance);
324                  self.total_staked = self.total_staked.saturating_add(new_balance);
325              }
326          }
327      }
328  
329      /// Apply slashing to a wallet (reduces staked balance)
330      ///
331      /// Returns the actual amount slashed.
332      pub fn apply_slashing(&mut self, dx_pubkey: &[u8; 32], slash_amount: u64) -> u64 {
333          if let Some(state) = self.wallet_states.get_mut(dx_pubkey) {
334              let actual_slash = slash_amount.min(state.staked_balance);
335              state.subtract_balance(actual_slash);
336              self.total_staked = self.total_staked.saturating_sub(actual_slash);
337              actual_slash
338          } else {
339              0
340          }
341      }
342  
343      /// Get staked balance for a wallet
344      pub fn get_staked_balance(&self, dx_pubkey: &[u8; 32]) -> u64 {
345          self.wallet_states
346              .get(dx_pubkey)
347              .map(|s| s.staked_balance())
348              .unwrap_or(0)
349      }
350  
351      /// Get total staked DX across all wallets (public aggregate)
352      pub fn total_staked(&self) -> u64 {
353          self.total_staked
354      }
355  
356      /// Get count of staked wallets (public aggregate)
357      pub fn staked_wallet_count(&self) -> u64 {
358          self.staked_wallet_count
359      }
360  
361      /// Get current epoch
362      pub fn current_epoch(&self) -> u64 {
363          self.current_epoch
364      }
365  
366      /// Advance to new epoch
367      pub fn set_epoch(&mut self, epoch: u64) {
368          self.current_epoch = epoch;
369      }
370  
371      /// Get all staked wallets with their balances (for dividend calculation)
372      pub fn get_staked_wallets(&self) -> Vec<([u8; 32], u64)> {
373          self.wallet_states
374              .iter()
375              .filter(|(_, state)| state.is_staking)
376              .map(|(pubkey, state)| (*pubkey, state.staked_balance))
377              .collect()
378      }
379  
380      /// Get iterator over all wallet states
381      pub fn iter_wallets(&self) -> impl Iterator<Item = (&[u8; 32], &StakingState)> {
382          self.wallet_states.iter()
383      }
384  }
385  
386  // ============================================================================
387  // Staking Proof (for privacy-preserving staking operations)
388  // ============================================================================
389  
390  /// Zero-knowledge proof of valid staking operation
391  #[derive(Clone, Debug, Serialize, Deserialize)]
392  pub struct StakingProof {
393      /// DX public key
394      pub dx_pubkey: [u8; 32],
395      /// Operation type
396      pub operation: StakingOperation,
397      /// Commitment to balance (hides actual amount)
398      pub balance_commitment: [u8; 32],
399      /// Proof data
400      pub proof_data: Vec<u8>,
401      /// Epoch of operation
402      pub epoch: u64,
403  }
404  
405  /// Type of staking operation
406  #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
407  pub enum StakingOperation {
408      /// Request to unstake
409      Unstake,
410      /// Request to re-stake
411      Stake,
412  }
413  
414  impl StakingProof {
415      /// Create a new staking proof
416      pub fn new(
417          dx_pubkey: [u8; 32],
418          operation: StakingOperation,
419          balance: u64,
420          randomness: u64,
421          epoch: u64,
422      ) -> Self {
423          // Create balance commitment
424          let mut hasher = Sha256::new();
425          hasher.update(b"STAKING_BALANCE_COMMIT");
426          hasher.update(balance.to_le_bytes());
427          hasher.update(randomness.to_le_bytes());
428          let mut balance_commitment = [0u8; 32];
429          balance_commitment.copy_from_slice(&hasher.finalize());
430  
431          // Create proof data
432          let mut hasher = Sha256::new();
433          hasher.update(b"STAKING_PROOF");
434          hasher.update(dx_pubkey);
435          hasher.update([operation.clone() as u8]);
436          hasher.update(balance_commitment);
437          hasher.update(epoch.to_le_bytes());
438          let proof_data = hasher.finalize().to_vec();
439  
440          Self {
441              dx_pubkey,
442              operation,
443              balance_commitment,
444              proof_data,
445              epoch,
446          }
447      }
448  
449      /// Verify the proof
450      pub fn verify(&self) -> bool {
451          // Verify proof data is not empty
452          if self.proof_data.is_empty() {
453              return false;
454          }
455  
456          // Verify balance commitment is not zero
457          if self.balance_commitment == [0u8; 32] {
458              return false;
459          }
460  
461          // In production: verify ZK proof
462          true
463      }
464  }
465  
466  // ============================================================================
467  // Tests
468  // ============================================================================
469  
470  #[cfg(test)]
471  mod tests {
472      use super::*;
473  
474      #[test]
475      fn test_default_staking_state() {
476          let state = StakingState::new();
477          assert!(state.is_staking());
478          assert!(!state.has_pending_unstake());
479          assert!(!state.has_pending_stake());
480      }
481  
482      #[test]
483      fn test_staking_registry_new_wallet() {
484          let mut registry = StakingRegistry::new();
485          let pubkey = [1u8; 32];
486  
487          // New wallet should be staking by default
488          assert!(registry.is_staked(&pubkey));
489  
490          // Getting or creating should maintain default state
491          let state = registry.get_or_create_wallet(&pubkey);
492          assert!(state.is_staking());
493      }
494  
495      #[test]
496      fn test_request_unstake() {
497          let mut registry = StakingRegistry::at_epoch(100);
498          let pubkey = [2u8; 32];
499  
500          registry.register_wallet(&pubkey, 10_000);
501  
502          // Should be staking initially
503          assert!(registry.is_staked(&pubkey));
504          assert!(!registry.can_transfer(&pubkey));
505  
506          // Request unstake
507          registry.request_unstake(&pubkey).unwrap();
508  
509          // Still staking until epoch transition
510          assert!(registry.is_staked(&pubkey));
511  
512          // Process epoch transition (unstaking is LAST)
513          registry.process_pending_unstakes(101).unwrap();
514  
515          // Now should be unstaked
516          assert!(!registry.is_staked(&pubkey));
517          assert!(registry.can_transfer(&pubkey));
518      }
519  
520      #[test]
521      fn test_request_stake_after_unstake() {
522          let mut registry = StakingRegistry::at_epoch(100);
523          let pubkey = [3u8; 32];
524  
525          registry.register_wallet(&pubkey, 10_000);
526  
527          // Unstake
528          registry.request_unstake(&pubkey).unwrap();
529          registry.process_pending_unstakes(101).unwrap();
530          assert!(!registry.is_staked(&pubkey));
531  
532          // Re-stake
533          registry.request_stake(&pubkey).unwrap();
534          registry.process_pending_stakes(102).unwrap();
535          assert!(registry.is_staked(&pubkey));
536      }
537  
538      #[test]
539      fn test_already_unstaked_error() {
540          let mut registry = StakingRegistry::at_epoch(100);
541          let pubkey = [4u8; 32];
542  
543          registry.register_wallet(&pubkey, 10_000);
544          registry.request_unstake(&pubkey).unwrap();
545          registry.process_pending_unstakes(101).unwrap();
546  
547          // Should error when trying to unstake again
548          assert!(registry.request_unstake(&pubkey).is_err());
549      }
550  
551      #[test]
552      fn test_already_staking_error() {
553          let mut registry = StakingRegistry::new();
554          let pubkey = [5u8; 32];
555  
556          registry.register_wallet(&pubkey, 10_000);
557  
558          // Should error when trying to stake while already staking
559          assert!(registry.request_stake(&pubkey).is_err());
560      }
561  
562      #[test]
563      fn test_total_staked_tracking() {
564          let mut registry = StakingRegistry::at_epoch(100);
565  
566          let pubkey1 = [6u8; 32];
567          let pubkey2 = [7u8; 32];
568  
569          registry.register_wallet(&pubkey1, 5_000);
570          registry.register_wallet(&pubkey2, 3_000);
571  
572          assert_eq!(registry.total_staked(), 8_000);
573          assert_eq!(registry.staked_wallet_count(), 2);
574  
575          // Unstake one
576          registry.request_unstake(&pubkey1).unwrap();
577          registry.process_pending_unstakes(101).unwrap();
578  
579          assert_eq!(registry.total_staked(), 3_000);
580          assert_eq!(registry.staked_wallet_count(), 1);
581      }
582  
583      #[test]
584      fn test_apply_slashing() {
585          let mut registry = StakingRegistry::new();
586          let pubkey = [8u8; 32];
587  
588          registry.register_wallet(&pubkey, 10_000);
589          assert_eq!(registry.total_staked(), 10_000);
590  
591          // Slash 1% (100 DX)
592          let slashed = registry.apply_slashing(&pubkey, 100);
593          assert_eq!(slashed, 100);
594          assert_eq!(registry.get_staked_balance(&pubkey), 9_900);
595          assert_eq!(registry.total_staked(), 9_900);
596      }
597  
598      #[test]
599      fn test_slashing_capped_at_balance() {
600          let mut registry = StakingRegistry::new();
601          let pubkey = [9u8; 32];
602  
603          registry.register_wallet(&pubkey, 100);
604  
605          // Try to slash more than balance
606          let slashed = registry.apply_slashing(&pubkey, 1_000);
607          assert_eq!(slashed, 100);
608          assert_eq!(registry.get_staked_balance(&pubkey), 0);
609      }
610  
611      #[test]
612      fn test_get_staked_wallets() {
613          let mut registry = StakingRegistry::at_epoch(100);
614  
615          let pubkey1 = [10u8; 32];
616          let pubkey2 = [11u8; 32];
617  
618          registry.register_wallet(&pubkey1, 5_000);
619          registry.register_wallet(&pubkey2, 3_000);
620  
621          // Unstake one
622          registry.request_unstake(&pubkey2).unwrap();
623          registry.process_pending_unstakes(101).unwrap();
624  
625          let staked = registry.get_staked_wallets();
626          assert_eq!(staked.len(), 1);
627          assert_eq!(staked[0].0, pubkey1);
628          assert_eq!(staked[0].1, 5_000);
629      }
630  
631      #[test]
632      fn test_staking_proof() {
633          let pubkey = [12u8; 32];
634          let proof = StakingProof::new(pubkey, StakingOperation::Unstake, 10_000, 12345, 100);
635  
636          assert!(proof.verify());
637          assert_eq!(proof.dx_pubkey, pubkey);
638          assert_eq!(proof.operation, StakingOperation::Unstake);
639      }
640  
641      #[test]
642      fn test_epoch_transition_priority() {
643          // This test verifies the correct order of operations
644          let mut registry = StakingRegistry::at_epoch(100);
645  
646          let pubkey1 = [13u8; 32]; // Will unstake
647          let pubkey2 = [14u8; 32]; // Will re-stake
648  
649          registry.register_wallet(&pubkey1, 5_000);
650          registry.register_wallet(&pubkey2, 3_000);
651  
652          // First, unstake pubkey2 so we can re-stake it
653          registry.request_unstake(&pubkey2).unwrap();
654          registry.process_pending_unstakes(101).unwrap();
655  
656          // Now queue operations for next epoch
657          registry.set_epoch(101);
658          registry.request_unstake(&pubkey1).unwrap();
659          registry.request_stake(&pubkey2).unwrap();
660  
661          // Before processing: pubkey1 staked, pubkey2 unstaked
662          assert!(registry.is_staked(&pubkey1));
663          assert!(!registry.is_staked(&pubkey2));
664  
665          // Process in priority order:
666          // 1. Slashing would go here
667          // 2. Dividends would be calculated here (on current state)
668          // 3. Process stakes
669          registry.process_pending_stakes(102).unwrap();
670          // pubkey2 is now staked again
671          assert!(registry.is_staked(&pubkey2));
672  
673          // 4. Process unstakes (LAST)
674          registry.process_pending_unstakes(102).unwrap();
675          // pubkey1 is now unstaked
676          assert!(!registry.is_staked(&pubkey1));
677      }
678  }