support dynamic staking system and ability to automatically increase stake

Author: JannikSt

shared/artifacts/abi/compute_registry.json

@@ -510,6 +510,25 @@
     ],
     "stateMutability": "view"
   },
+  {
+    "type": "function",
+    "name": "getProviderTotalCompute",
+    "inputs": [
+      {
+        "name": "provider",
+        "type": "address",
+        "internalType": "address"
+      }
+    ],
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "stateMutability": "view"
+  },
   {
     "type": "function",
     "name": "getProviderTotalNodes",
@@ -714,6 +733,25 @@
     ],
     "stateMutability": "view"
   },
+  {
+    "type": "function",
+    "name": "providerTotalCompute",
+    "inputs": [
+      {
+        "name": "",
+        "type": "address",
+        "internalType": "address"
+      }
+    ],
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "stateMutability": "view"
+  },
   {
     "type": "function",
     "name": "providers",

shared/artifacts/abi/prime_network.json

@@ -113,6 +113,30 @@
     "outputs": [],
     "stateMutability": "nonpayable"
   },
+  {
+    "type": "function",
+    "name": "calculateMinimumStake",
+    "inputs": [
+      {
+        "name": "provider",
+        "type": "address",
+        "internalType": "address"
+      },
+      {
+        "name": "computeUnits",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "outputs": [
+      {
+        "name": "",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "stateMutability": "view"
+  },
   {
     "type": "function",
     "name": "computePool",
@@ -317,6 +341,19 @@
     ],
     "stateMutability": "view"
   },
+  {
+    "type": "function",
+    "name": "increaseStake",
+    "inputs": [
+      {
+        "name": "additionalStake",
+        "type": "uint256",
+        "internalType": "uint256"
+      }
+    ],
+    "outputs": [],
+    "stateMutability": "nonpayable"
+  },
   {
     "type": "function",
     "name": "invalidateNode",

shared/src/web3/contracts/core/builder.rs

@@ -4,10 +4,7 @@ use crate::web3::{
     contracts::{
         core::error::ContractError, // Using custom error ContractError
         implementations::{
-            ai_token_contract::AIToken, compute_pool_contract::ComputePool,
-            compute_registry_contract::ComputeRegistryContract,
-            prime_network_contract::PrimeNetworkContract,
-            work_validators::synthetic_data_validator::SyntheticDataWorkValidator,
+            ai_token_contract::AIToken, compute_pool_contract::ComputePool, compute_registry_contract::ComputeRegistryContract, prime_network_contract::PrimeNetworkContract, stake_manager::StakeManagerContract, work_validators::synthetic_data_validator::SyntheticDataWorkValidator
         },
     },
     wallet::Wallet,
@@ -20,6 +17,7 @@ pub struct Contracts {
     pub ai_token: AIToken,
     pub prime_network: PrimeNetworkContract,
     pub compute_pool: ComputePool,
+    pub stake_manager: Option<StakeManagerContract>,
     pub synthetic_data_validator: Option<SyntheticDataWorkValidator>,
 }
 
@@ -29,6 +27,7 @@ pub struct ContractBuilder<'a> {
     ai_token: Option<AIToken>,
     prime_network: Option<PrimeNetworkContract>,
     compute_pool: Option<ComputePool>,
+    stake_manager: Option<StakeManagerContract>,
     synthetic_data_validator: Option<SyntheticDataWorkValidator>,
 }
 
@@ -40,6 +39,7 @@ impl<'a> ContractBuilder<'a> {
             ai_token: None,
             prime_network: None,
             compute_pool: None,
+            stake_manager: None,
             synthetic_data_validator: None,
         }
     }
@@ -76,6 +76,11 @@ impl<'a> ContractBuilder<'a> {
         self
     }
 
+    pub fn with_stake_manager(mut self) -> Self {
+        self.stake_manager = Some(StakeManagerContract::new(self.wallet, "stake_manager.json"));
+        self
+    }
+
     // TODO: This is not ideal yet - now you have to init all contracts all the time
     pub fn build(self) -> Result<Contracts, ContractError> {
         // Using custom error ContractError
@@ -101,6 +106,7 @@ impl<'a> ContractBuilder<'a> {
                 None => return Err(ContractError::Other("PrimeNetwork not initialized".into())), // Custom error handling
             },
             synthetic_data_validator: self.synthetic_data_validator,
+            stake_manager: self.stake_manager,
         })
     }
 }

shared/src/web3/contracts/core/contract.rs

@@ -45,6 +45,7 @@ impl Contract {
             "synthetic_data_work_validator.json" => {
                 include_abi!("../../../../artifacts/abi/synthetic_data_work_validator.json")
             }
+            "stake_manager.json" => include_abi!("../../../../artifacts/abi/stake_manager.json"),
             _ => panic!("Unknown ABI file: {}", path),
         };
 

shared/src/web3/contracts/implementations/compute_registry_contract.rs

@@ -5,7 +5,7 @@ use super::{
 use crate::web3::contracts::helpers::utils::get_selector;
 use crate::web3::wallet::Wallet;
 use alloy::dyn_abi::DynSolValue;
-use alloy::primitives::Address;
+use alloy::primitives::{Address, U256};
 
 pub struct ComputeRegistryContract {
     instance: Contract,
@@ -41,6 +41,20 @@ impl ComputeRegistryContract {
         Ok(provider)
     }
 
+    pub async fn get_provider_total_compute(
+        &self,
+        address: Address,
+    ) -> Result<U256, Box<dyn std::error::Error>> {
+        let provider_response = self
+            .instance
+            .instance()
+            .function("getProviderTotalCompute", &[address.into()])?
+            .call()
+            .await?;
+
+        Ok(U256::from(provider_response.first().unwrap().as_uint().unwrap().0))
+    }
+
     pub async fn get_node(
         &self,
         #[allow(unused_variables)] provider_address: Address,

shared/src/web3/contracts/implementations/mod.rs

@@ -3,3 +3,4 @@ pub mod compute_pool_contract;
 pub mod compute_registry_contract;
 pub mod prime_network_contract;
 pub mod work_validators;
+pub mod stake_manager;

shared/src/web3/contracts/implementations/prime_network_contract.rs

@@ -32,6 +32,22 @@ impl PrimeNetworkContract {
         Ok(register_tx)
     }
 
+    pub async fn stake(
+        &self,
+        additional_stake: U256,
+    ) -> Result<FixedBytes<32>, Box<dyn std::error::Error>> {
+        let stake_tx = self
+            .instance
+            .instance()
+            .function("increaseStake", &[additional_stake.into()])?
+            .send()
+            .await?
+            .watch()
+            .await?;
+
+        Ok(stake_tx)
+    }
+
     pub async fn add_compute_node(
         &self,
         node_address: Address,

shared/src/web3/contracts/implementations/stake_manager.rs

@@ -0,0 +1,54 @@
+use crate::web3::contracts::constants::addresses::STAKE_MANAGER_ADDRESS;
+use crate::web3::contracts::core::contract::Contract;
+use crate::web3::wallet::Wallet;
+use alloy::primitives::{Address, U256};
+
+pub struct StakeManagerContract {
+    instance: Contract,
+}
+
+impl StakeManagerContract {
+    pub fn new(wallet: &Wallet, abi_file_path: &str) -> Self {
+        let instance = Contract::new(STAKE_MANAGER_ADDRESS, wallet, abi_file_path);
+        Self { instance }
+    }
+
+    pub async fn get_stake_minimum(&self) -> Result<U256, Box<dyn std::error::Error>> {
+        let result = self
+            .instance
+            .instance()
+            .function("getStakeMinimum", &[])?
+            .call()
+            .await?;
+        
+        let minimum: U256 = result
+            .into_iter()
+            .next()
+            .map(|value| value.as_uint().unwrap_or_default())
+            .unwrap_or_default()
+            .0;
+        Ok(minimum)
+    }
+    pub async fn get_stake(&self, staker: Address) -> Result<U256, Box<dyn std::error::Error>> {
+        let result = self
+            .instance
+            .instance()
+            .function("getStake", &[staker.into()])?
+            .call()
+            .await?;
+        println!("Result: {:?}", result);
+
+        Ok(result[0].as_uint().unwrap_or_default().0)
+    }
+
+    pub async fn calculate_stake(
+        &self,
+        compute_units: U256,
+        provider_total_compute: U256,
+    ) -> Result<U256, Box<dyn std::error::Error>> {
+        let min_stake_per_unit = self.get_stake_minimum().await?;
+        let total_compute = provider_total_compute + compute_units;
+        let required_stake = total_compute * min_stake_per_unit;
+        Ok(required_stake)
+    }
+}