tetcore_runtime/
model_registry.rs

1use crate::RuntimeError;
2use serde::{Deserialize, Serialize};
3use std::collections::HashMap;
4use tetcore_primitives::{Address, Hash32};
5
6pub const REVENUE_SPLIT_BASIS_POINTS: u32 = 10000;
7
8#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
9pub enum ModelState {
10    Registered,
11    Active,
12    Paused,
13    Deprecated,
14    Disabled,
15}
16
17#[derive(Clone, Debug, Serialize, Deserialize)]
18pub struct PricingPolicy {
19    pub mode: PricingMode,
20    pub base_fee: u64,
21    pub token_per_char: u64,
22}
23
24#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)]
25pub enum PricingMode {
26    Owner,
27    Market,
28    Hybrid,
29}
30
31#[derive(Clone, Debug, Serialize, Deserialize)]
32pub struct RevenueSplit {
33    pub model_owner_bps: u16,
34    pub operator_bps: u16,
35    pub shard_provider_bps: u16,
36    pub validator_bps: u16,
37    pub vault_bps: u16,
38}
39
40impl Default for RevenueSplit {
41    fn default() -> Self {
42        Self {
43            model_owner_bps: 3000,
44            operator_bps: 3000,
45            shard_provider_bps: 2000,
46            validator_bps: 1000,
47            vault_bps: 1000,
48        }
49    }
50}
51
52impl RevenueSplit {
53    pub fn is_valid(&self) -> bool {
54        let total = self.model_owner_bps as u32
55            + self.operator_bps as u32
56            + self.shard_provider_bps as u32
57            + self.validator_bps as u32
58            + self.vault_bps as u32;
59        total == REVENUE_SPLIT_BASIS_POINTS
60    }
61}
62
63#[derive(Clone, Debug, Serialize, Deserialize)]
64pub struct Model {
65    pub model_id: Hash32,
66    pub owner: Address,
67    pub version: u32,
68    pub shard_root: Hash32,
69    pub shard_count: u32,
70    pub state: ModelState,
71    pub pricing_policy: PricingPolicy,
72    pub revenue_split: RevenueSplit,
73    pub vault_enabled: bool,
74}
75
76pub struct ModelRegistryModule {
77    models: HashMap<Hash32, Model>,
78    model_versions: HashMap<Hash32, Vec<Hash32>>,
79}
80
81impl ModelRegistryModule {
82    pub fn new() -> Self {
83        Self {
84            models: HashMap::new(),
85            model_versions: HashMap::new(),
86        }
87    }
88
89    pub fn register_model(&mut self, model: Model) -> Result<(), RuntimeError> {
90        if !model.revenue_split.is_valid() {
91            return Err(RuntimeError::InvalidState);
92        }
93
94        self.models.insert(model.model_id.clone(), model);
95        Ok(())
96    }
97
98    pub fn get_model(&self, model_id: &Hash32) -> Option<&Model> {
99        self.models.get(model_id)
100    }
101
102    pub fn get_model_mut(&mut self, model_id: &Hash32) -> Option<&mut Model> {
103        self.models.get_mut(model_id)
104    }
105
106    pub fn update_version(
107        &mut self,
108        model_id: &Hash32,
109        new_version: u32,
110        new_shard_root: Hash32,
111    ) -> Result<(), RuntimeError> {
112        let model = self
113            .models
114            .get_mut(model_id)
115            .ok_or(RuntimeError::InvalidState)?;
116        model.version = new_version;
117        model.shard_root = new_shard_root;
118        Ok(())
119    }
120
121    pub fn set_model_state(
122        &mut self,
123        model_id: &Hash32,
124        state: ModelState,
125    ) -> Result<(), RuntimeError> {
126        let model = self
127            .models
128            .get_mut(model_id)
129            .ok_or(RuntimeError::InvalidState)?;
130        model.state = state;
131        Ok(())
132    }
133
134    pub fn update_pricing(
135        &mut self,
136        model_id: &Hash32,
137        policy: PricingPolicy,
138    ) -> Result<(), RuntimeError> {
139        let model = self
140            .models
141            .get_mut(model_id)
142            .ok_or(RuntimeError::InvalidState)?;
143        model.pricing_policy = policy;
144        Ok(())
145    }
146
147    pub fn update_revenue_split(
148        &mut self,
149        model_id: &Hash32,
150        split: RevenueSplit,
151    ) -> Result<(), RuntimeError> {
152        if !split.is_valid() {
153            return Err(RuntimeError::InvalidState);
154        }
155        let model = self
156            .models
157            .get_mut(model_id)
158            .ok_or(RuntimeError::InvalidState)?;
159        model.revenue_split = split;
160        Ok(())
161    }
162
163    pub fn transfer_ownership(
164        &mut self,
165        model_id: &Hash32,
166        new_owner: Address,
167    ) -> Result<(), RuntimeError> {
168        let model = self
169            .models
170            .get_mut(model_id)
171            .ok_or(RuntimeError::InvalidState)?;
172        model.owner = new_owner;
173        Ok(())
174    }
175
176    pub fn all_models(&self) -> &HashMap<Hash32, Model> {
177        &self.models
178    }
179
180    pub fn is_active(&self, model_id: &Hash32) -> bool {
181        self.models
182            .get(model_id)
183            .map(|m| m.state == ModelState::Active)
184            .unwrap_or(false)
185    }
186}
187
188impl Default for ModelRegistryModule {
189    fn default() -> Self {
190        Self::new()
191    }
192}
tetcore_runtime/model_registry.rs

tetcore_runtime/
model_registry.rs
