1use rust_decimal::Decimal;
37use serde::{Deserialize, Serialize};
38use std::collections::HashMap;
39
40#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)]
42#[serde(rename_all = "snake_case")]
43pub enum StrategyType {
44 #[default]
46 Balancer,
47 Hold,
49 Scalper,
51 Dca,
53 Custom(String),
55}
56
57impl StrategyType {
58 #[must_use]
60 pub fn from_name(name: &str) -> Self {
61 match name.to_lowercase().as_str() {
62 "balancer" | "balance" | "rebalance" => Self::Balancer,
63 "hold" | "holding" | "hodl" => Self::Hold,
64 "scalp" | "scalper" | "scalping" => Self::Scalper,
65 "dca" | "dollar_cost_average" => Self::Dca,
66 _ => Self::Custom(name.to_string()),
67 }
68 }
69
70 #[must_use]
72 pub fn as_str(&self) -> &str {
73 match self {
74 Self::Balancer => "balancer",
75 Self::Hold => "hold",
76 Self::Scalper => "scalper",
77 Self::Dca => "dca",
78 Self::Custom(name) => name,
79 }
80 }
81}
82
83#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
85pub struct AllocationPortion {
86 pub percentage: Decimal,
88 pub allocation: Allocation,
90}
91
92impl AllocationPortion {
93 #[must_use]
95 pub fn new(percentage: Decimal, allocation: Allocation) -> Self {
96 Self {
97 percentage,
98 allocation,
99 }
100 }
101
102 #[must_use]
104 pub fn hold(ticker: impl Into<String>, percentage: Decimal) -> Self {
105 Self {
106 percentage,
107 allocation: Allocation::hold(ticker),
108 }
109 }
110}
111
112#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)]
119pub struct Allocation {
120 #[serde(default)]
122 pub strategy: StrategyType,
123
124 #[serde(skip_serializing_if = "Option::is_none")]
126 pub ticker: Option<String>,
127
128 #[serde(default, skip_serializing_if = "Vec::is_empty")]
130 pub portions: Vec<AllocationPortion>,
131
132 #[serde(default, skip_serializing_if = "HashMap::is_empty")]
134 pub params: HashMap<String, String>,
135}
136
137impl Allocation {
138 #[must_use]
140 pub fn new() -> Self {
141 Self::default()
142 }
143
144 #[must_use]
146 pub fn balancer(portions: Vec<AllocationPortion>) -> Self {
147 Self {
148 strategy: StrategyType::Balancer,
149 ticker: None,
150 portions,
151 params: HashMap::new(),
152 }
153 }
154
155 #[must_use]
157 pub fn hold(ticker: impl Into<String>) -> Self {
158 Self {
159 strategy: StrategyType::Hold,
160 ticker: Some(ticker.into()),
161 portions: Vec::new(),
162 params: HashMap::new(),
163 }
164 }
165
166 #[must_use]
168 pub fn scalper(ticker: impl Into<String>) -> Self {
169 Self {
170 strategy: StrategyType::Scalper,
171 ticker: Some(ticker.into()),
172 portions: Vec::new(),
173 params: HashMap::new(),
174 }
175 }
176
177 #[must_use]
189 pub fn from_simple(allocations: &[(String, Decimal)]) -> Self {
190 let portions = allocations
191 .iter()
192 .map(|(ticker, pct)| AllocationPortion::hold(ticker, *pct))
193 .collect();
194
195 Self::balancer(portions)
196 }
197
198 #[must_use]
200 pub fn from_map(allocations: HashMap<String, Decimal>) -> Self {
201 let portions = allocations
202 .into_iter()
203 .map(|(ticker, pct)| AllocationPortion::hold(ticker, pct))
204 .collect();
205
206 Self::balancer(portions)
207 }
208
209 pub fn set_param(&mut self, key: impl Into<String>, value: impl Into<String>) {
211 self.params.insert(key.into(), value.into());
212 }
213
214 #[must_use]
216 pub fn get_param(&self, key: &str) -> Option<&str> {
217 self.params.get(key).map(|s| s.as_str())
218 }
219
220 pub fn add_portion(&mut self, portion: AllocationPortion) {
222 self.portions.push(portion);
223 }
224
225 #[must_use]
227 pub fn is_hold(&self) -> bool {
228 matches!(self.strategy, StrategyType::Hold) && self.ticker.is_some()
229 }
230
231 #[must_use]
233 pub fn is_balancer(&self) -> bool {
234 matches!(self.strategy, StrategyType::Balancer) && !self.portions.is_empty()
235 }
236
237 #[must_use]
239 pub fn all_tickers(&self) -> Vec<String> {
240 let mut tickers = Vec::new();
241 self.collect_tickers(&mut tickers);
242 tickers
243 }
244
245 fn collect_tickers(&self, tickers: &mut Vec<String>) {
246 if let Some(ref ticker) = self.ticker {
247 if !tickers.contains(ticker) {
248 tickers.push(ticker.clone());
249 }
250 }
251 for portion in &self.portions {
252 portion.allocation.collect_tickers(tickers);
253 }
254 }
255
256 #[must_use]
258 pub fn total_percentage(&self) -> Decimal {
259 self.portions.iter().map(|p| p.percentage).sum()
260 }
261
262 #[must_use]
264 pub fn is_valid(&self) -> bool {
265 if self.portions.is_empty() {
266 return self.ticker.is_some();
268 }
269
270 let total = self.total_percentage();
271 let diff = (total - Decimal::from(100)).abs();
272 diff < Decimal::ONE }
274
275 pub fn normalize(&mut self) {
277 let total = self.total_percentage();
278 if total.is_zero() {
279 return;
280 }
281
282 let factor = Decimal::from(100) / total;
283 for portion in &mut self.portions {
284 portion.percentage *= factor;
285 }
286 }
287
288 #[must_use]
293 pub fn flatten(&self) -> HashMap<String, Decimal> {
294 let mut result = HashMap::new();
295 self.flatten_recursive(Decimal::from(100), &mut result);
296 result
297 }
298
299 fn flatten_recursive(&self, parent_pct: Decimal, result: &mut HashMap<String, Decimal>) {
300 if let Some(ref ticker) = self.ticker {
302 let entry = result.entry(ticker.clone()).or_insert(Decimal::ZERO);
303 *entry += parent_pct;
304 return;
305 }
306
307 for portion in &self.portions {
309 let effective_pct = parent_pct * portion.percentage / Decimal::from(100);
310 portion.allocation.flatten_recursive(effective_pct, result);
311 }
312 }
313}
314
315impl From<super::DesiredAllocation> for Allocation {
317 fn from(desired: super::DesiredAllocation) -> Self {
318 let map: HashMap<String, Decimal> = desired.iter().map(|(k, v)| (k.clone(), *v)).collect();
319 Self::from_map(map)
320 }
321}
322
323#[cfg(test)]
324mod tests {
325 use super::*;
326 use rust_decimal_macros::dec;
327
328 #[test]
329 fn test_simple_allocation() {
330 let alloc = Allocation::from_simple(&[
331 ("SBER".to_string(), dec!(30)),
332 ("LKOH".to_string(), dec!(30)),
333 ("GAZP".to_string(), dec!(40)),
334 ]);
335
336 assert!(alloc.is_balancer());
337 assert!(alloc.is_valid());
338 assert_eq!(alloc.portions.len(), 3);
339
340 let flat = alloc.flatten();
341 assert_eq!(flat.get("SBER"), Some(&dec!(30)));
342 assert_eq!(flat.get("LKOH"), Some(&dec!(30)));
343 assert_eq!(flat.get("GAZP"), Some(&dec!(40)));
344 }
345
346 #[test]
347 fn test_hold_allocation() {
348 let alloc = Allocation::hold("BTC");
349 assert!(alloc.is_hold());
350 assert_eq!(alloc.ticker, Some("BTC".to_string()));
351 assert_eq!(alloc.all_tickers(), vec!["BTC".to_string()]);
352 }
353
354 #[test]
355 fn test_nested_allocation() {
356 let stocks = Allocation::balancer(vec![
360 AllocationPortion::hold("SBER", dec!(50)),
361 AllocationPortion::hold("LKOH", dec!(50)),
362 ]);
363
364 let crypto = Allocation::balancer(vec![
365 AllocationPortion::hold("BTC", dec!(70)),
366 AllocationPortion::hold("ETH", dec!(30)),
367 ]);
368
369 let root = Allocation::balancer(vec![
370 AllocationPortion::new(dec!(60), stocks),
371 AllocationPortion::new(dec!(40), crypto),
372 ]);
373
374 assert!(root.is_valid());
375
376 let flat = root.flatten();
377 assert_eq!(flat.get("SBER"), Some(&dec!(30)));
382 assert_eq!(flat.get("LKOH"), Some(&dec!(30)));
383 assert_eq!(flat.get("BTC"), Some(&dec!(28)));
384 assert_eq!(flat.get("ETH"), Some(&dec!(12)));
385
386 let all_tickers = root.all_tickers();
387 assert_eq!(all_tickers.len(), 4);
388 }
389
390 #[test]
391 fn test_strategy_types() {
392 assert_eq!(StrategyType::from_name("balancer"), StrategyType::Balancer);
393 assert_eq!(StrategyType::from_name("HOLD"), StrategyType::Hold);
394 assert_eq!(StrategyType::from_name("scalper"), StrategyType::Scalper);
395 assert_eq!(
396 StrategyType::from_name("custom_strat"),
397 StrategyType::Custom("custom_strat".to_string())
398 );
399 }
400
401 #[test]
402 fn test_normalize_allocation() {
403 let mut alloc = Allocation::balancer(vec![
404 AllocationPortion::hold("A", dec!(30)),
405 AllocationPortion::hold("B", dec!(20)),
406 ]);
407
408 alloc.normalize();
410
411 assert_eq!(alloc.portions[0].percentage, dec!(60)); assert_eq!(alloc.portions[1].percentage, dec!(40)); }
414
415 #[test]
416 fn test_deeply_nested_allocation() {
417 let level3 = Allocation::balancer(vec![
419 AllocationPortion::hold("X", dec!(50)),
420 AllocationPortion::hold("Y", dec!(50)),
421 ]);
422
423 let level2 = Allocation::balancer(vec![AllocationPortion::new(dec!(100), level3)]);
424
425 let root = Allocation::balancer(vec![
426 AllocationPortion::new(dec!(50), level2),
427 AllocationPortion::hold("Z", dec!(50)),
428 ]);
429
430 let flat = root.flatten();
431 assert_eq!(flat.get("X"), Some(&dec!(25)));
435 assert_eq!(flat.get("Y"), Some(&dec!(25)));
436 assert_eq!(flat.get("Z"), Some(&dec!(50)));
437 }
438}