1use std::collections::{HashMap, HashSet};
2use std::time::{Duration, Instant, SystemTime};
3
4use crate::options;
5use crate::stocks::{self, Adjustment, BarPoint, Currency, DataFeed, TimeFrame};
6use crate::symbols::option_contract_symbol;
7
8pub(crate) type BarsMap = HashMap<String, Vec<BarPoint>>;
9
10#[derive(Debug, Clone)]
11pub struct CachedEntry<T> {
12 pub value: T,
13 pub stored_at: Instant,
14}
15
16impl<T> CachedEntry<T> {
17 pub fn new(value: T) -> Self {
18 Self {
19 value,
20 stored_at: Instant::now(),
21 }
22 }
23
24 pub fn is_fresh(&self, ttl: Duration, now: Instant) -> bool {
25 is_timestamp_fresh(self.stored_at, ttl, now)
26 }
27}
28
29pub fn is_timestamp_fresh(stored_at: Instant, ttl: Duration, now: Instant) -> bool {
30 now.saturating_duration_since(stored_at) < ttl
31}
32
33#[derive(Debug, Clone)]
34pub struct StockBarsRequest {
35 pub key: String,
36 pub symbols: Vec<String>,
37 pub timeframe: TimeFrame,
38 pub start: Option<String>,
39 pub end: Option<String>,
40 pub limit: u32,
41 pub adjustment: Option<Adjustment>,
42 pub feed: Option<DataFeed>,
43 pub currency: Option<Currency>,
44 pub chunk_size: usize,
45}
46
47impl StockBarsRequest {
48 pub(crate) fn normalized(mut self) -> Self {
49 self.symbols = normalize_stock_list(&self.symbols);
50 self
51 }
52
53 pub(crate) fn merge_from(&mut self, next: &Self) {
54 self.symbols = merge_values(&self.symbols, &next.symbols);
55 self.timeframe = next.timeframe.clone();
56 self.start = next.start.clone();
57 self.end = next.end.clone();
58 self.limit = next.limit;
59 self.adjustment = next.adjustment.clone();
60 self.feed = next.feed;
61 self.currency = next.currency.clone();
62 self.chunk_size = next.chunk_size;
63 }
64}
65
66#[derive(Debug, Default)]
67pub(crate) struct SnapshotCache<T> {
68 pub subscribed: HashSet<String>,
69 pub values: HashMap<String, CachedEntry<T>>,
70 pub empty: HashMap<String, Instant>,
71 pub updated_at: Option<SystemTime>,
72}
73
74impl<T: Clone> SnapshotCache<T> {
75 pub(crate) fn reconcile(
76 &mut self,
77 requested: &[String],
78 fetched: &HashMap<String, T>,
79 updated_at: SystemTime,
80 ) -> usize {
81 let mut count = 0;
82 let mut seen = HashSet::new();
83 let stored_at = Instant::now();
84
85 for key in requested {
86 if !seen.insert(key) {
87 continue;
88 }
89
90 self.subscribed.insert(key.clone());
91 if let Some(value) = fetched.get(key) {
92 self.values.insert(
93 key.clone(),
94 CachedEntry {
95 value: value.clone(),
96 stored_at,
97 },
98 );
99 self.empty.remove(key);
100 count += 1;
101 } else {
102 self.values.remove(key);
103 self.empty.insert(key.clone(), stored_at);
104 }
105 }
106
107 if !seen.is_empty() {
108 self.updated_at = Some(updated_at);
109 }
110 count
111 }
112}
113
114#[derive(Debug, Default)]
115pub(crate) struct StockBarsCache {
116 pub requests: HashMap<String, StockBarsRequest>,
117 pub values: HashMap<String, HashMap<String, CachedEntry<Vec<BarPoint>>>>,
118 pub empty: HashMap<String, HashMap<String, Instant>>,
119 pub updated_at: HashMap<String, SystemTime>,
120}
121
122#[derive(Debug, Default)]
123pub(crate) struct CacheState {
124 pub stocks: SnapshotCache<stocks::Snapshot>,
125 pub options: SnapshotCache<options::Snapshot>,
126 pub bars: StockBarsCache,
127}
128
129pub(crate) fn normalize_values<S: AsRef<str>>(values: &[S]) -> Vec<String> {
130 let mut normalized = Vec::new();
131 let mut seen = HashSet::new();
132 for value in values {
133 let value = value.as_ref().trim();
134 if !value.is_empty() && seen.insert(value.to_string()) {
135 normalized.push(value.to_string());
136 }
137 }
138 normalized
139}
140
141pub(crate) fn normalize_stock_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<(String, String)> {
142 normalize_values(symbols)
143 .into_iter()
144 .map(|symbol| {
145 let resolved = stocks::display_stock_symbol(&symbol);
146 (symbol, resolved)
147 })
148 .collect()
149}
150
151pub(crate) fn normalize_option_symbols<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
152 let normalized = normalize_values(symbols);
153 let mut values = Vec::new();
154 let mut seen = HashSet::new();
155 for symbol in normalized {
156 let symbol = option_contract_symbol(&symbol);
157 if !symbol.is_empty() && seen.insert(symbol.clone()) {
158 values.push(symbol);
159 }
160 }
161 values
162}
163
164pub fn collect_cached_hits<T: Clone>(
165 requested: &[String],
166 cached: &HashMap<String, CachedEntry<T>>,
167 empty: &HashMap<String, Instant>,
168 ttl: Duration,
169 now: Instant,
170) -> (HashMap<String, T>, Vec<String>) {
171 let mut hits = HashMap::new();
172 let mut missing = Vec::new();
173 for key in requested {
174 if let Some(entry) = cached.get(key) {
175 if entry.is_fresh(ttl, now) {
176 hits.insert(key.clone(), entry.value.clone());
177 } else {
178 missing.push(key.clone());
179 }
180 } else if empty
181 .get(key)
182 .is_some_and(|stored_at| is_timestamp_fresh(*stored_at, ttl, now))
183 {
184 continue;
185 } else {
186 missing.push(key.clone());
187 }
188 }
189 (hits, missing)
190}
191
192pub(crate) fn unwrap_bars_map(
193 cached: &HashMap<String, CachedEntry<Vec<BarPoint>>>,
194) -> BarsMap {
195 cached
196 .iter()
197 .map(|(symbol, entry)| (symbol.clone(), entry.value.clone()))
198 .collect()
199}
200
201pub(crate) fn missing_bar_symbols(
202 requested: &[String],
203 cached: Option<&HashMap<String, CachedEntry<Vec<BarPoint>>>>,
204 empty: Option<&HashMap<String, Instant>>,
205 ttl: Duration,
206 now: Instant,
207) -> Vec<String> {
208 requested
209 .iter()
210 .filter(|symbol| {
211 if cached
212 .and_then(|values| values.get(*symbol))
213 .is_some_and(|entry| entry.is_fresh(ttl, now))
214 {
215 return false;
216 }
217 if empty
218 .and_then(|values| values.get(*symbol))
219 .is_some_and(|stored_at| is_timestamp_fresh(*stored_at, ttl, now))
220 {
221 return false;
222 }
223 true
224 })
225 .cloned()
226 .collect()
227}
228
229fn normalize_stock_list<S: AsRef<str>>(symbols: &[S]) -> Vec<String> {
230 normalize_stock_symbols(symbols)
231 .into_iter()
232 .map(|(_, resolved)| resolved)
233 .collect()
234}
235
236fn merge_values(current: &[String], next: &[String]) -> Vec<String> {
237 let mut merged = Vec::new();
238 let mut seen = HashSet::new();
239 for value in current.iter().chain(next.iter()) {
240 if seen.insert(value.clone()) {
241 merged.push(value.clone());
242 }
243 }
244 merged
245}
246
247#[cfg(test)]
248mod tests {
249 use super::*;
250
251 #[test]
252 fn expired_price_entries_are_treated_as_missing() {
253 let now = Instant::now();
254 let stale = now
255 .checked_sub(Duration::from_secs(20))
256 .expect("ttl test clock");
257 let mut cached = HashMap::new();
258 cached.insert(
259 "AAPL".to_string(),
260 CachedEntry {
261 value: 1,
262 stored_at: stale,
263 },
264 );
265 cached.insert(
266 "MSFT".to_string(),
267 CachedEntry {
268 value: 2,
269 stored_at: now,
270 },
271 );
272
273 let (hits, missing) = collect_cached_hits(
274 &[
275 "AAPL".to_string(),
276 "MSFT".to_string(),
277 "GOOG".to_string(),
278 ],
279 &cached,
280 &HashMap::new(),
281 Duration::from_secs(15),
282 now,
283 );
284
285 assert_eq!(hits.get("MSFT"), Some(&2));
286 assert!(!hits.contains_key("AAPL"));
287 assert_eq!(missing, vec!["AAPL".to_string(), "GOOG".to_string()]);
288 }
289
290 #[test]
291 fn expired_empty_entries_are_retried() {
292 let now = Instant::now();
293 let stale = now
294 .checked_sub(Duration::from_secs(20))
295 .expect("ttl test clock");
296 let mut empty = HashMap::new();
297 empty.insert("AAPL".to_string(), stale);
298 empty.insert("MSFT".to_string(), now);
299
300 let (hits, missing) = collect_cached_hits(
301 &["AAPL".to_string(), "MSFT".to_string()],
302 &HashMap::<String, CachedEntry<i32>>::new(),
303 &empty,
304 Duration::from_secs(15),
305 now,
306 );
307
308 assert!(hits.is_empty());
309 assert_eq!(missing, vec!["AAPL".to_string()]);
310 }
311}