1use chrono::{DateTime, Utc};
11use fnv::FnvHashMap;
12use rust_decimal::Decimal;
13use rustrade_instrument::instrument::name::InstrumentNameExchange;
14use serde::{Deserialize, Serialize};
15use std::{collections::BTreeMap, fmt::Debug};
16
17#[non_exhaustive]
29#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
30pub struct Shortability {
31 pub shortable: Option<bool>,
33 pub easy_to_borrow: Option<bool>,
35 pub available: Option<Decimal>,
38 pub fee_rate: Option<Decimal>,
45}
46
47impl Shortability {
48 pub fn new() -> Self {
50 Self::default()
51 }
52
53 #[must_use]
55 pub fn with_shortable(mut self, shortable: bool) -> Self {
56 self.shortable = Some(shortable);
57 self
58 }
59
60 #[must_use]
62 pub fn with_easy_to_borrow(mut self, easy_to_borrow: bool) -> Self {
63 self.easy_to_borrow = Some(easy_to_borrow);
64 self
65 }
66
67 #[must_use]
69 pub fn with_available(mut self, available: Decimal) -> Self {
70 self.available = Some(available);
71 self
72 }
73
74 #[must_use]
76 pub fn with_fee_rate(mut self, fee_rate: Decimal) -> Self {
77 self.fee_rate = Some(fee_rate);
78 self
79 }
80}
81
82pub trait ShortabilityProvider: Debug + Send + Sync {
90 fn shortability(
92 &self,
93 instrument: &InstrumentNameExchange,
94 time: DateTime<Utc>,
95 ) -> Shortability;
96
97 fn short_sale_restricted(
104 &self,
105 _instrument: &InstrumentNameExchange,
106 _time: DateTime<Utc>,
107 ) -> bool {
108 false
109 }
110}
111
112#[derive(Debug, Clone, Default)]
119pub struct ShortabilityTable {
120 shortability: FnvHashMap<InstrumentNameExchange, BTreeMap<DateTime<Utc>, Shortability>>,
121 restrictions: FnvHashMap<InstrumentNameExchange, Vec<(DateTime<Utc>, DateTime<Utc>)>>,
122}
123
124impl ShortabilityTable {
125 pub fn new() -> Self {
127 Self::default()
128 }
129
130 #[must_use]
133 pub fn with_shortability(
134 mut self,
135 instrument: InstrumentNameExchange,
136 from: DateTime<Utc>,
137 shortability: Shortability,
138 ) -> Self {
139 self.shortability
140 .entry(instrument)
141 .or_default()
142 .insert(from, shortability);
143 self
144 }
145
146 #[must_use]
153 pub fn with_short_sale_restriction(
154 mut self,
155 instrument: InstrumentNameExchange,
156 from: DateTime<Utc>,
157 until: DateTime<Utc>,
158 ) -> Self {
159 assert!(
160 from < until,
161 "a short-sale restriction on {instrument} must end after it starts: {from} to {until}"
162 );
163 self.restrictions
164 .entry(instrument)
165 .or_default()
166 .push((from, until));
167 self
168 }
169}
170
171impl ShortabilityProvider for ShortabilityTable {
172 fn shortability(
173 &self,
174 instrument: &InstrumentNameExchange,
175 time: DateTime<Utc>,
176 ) -> Shortability {
177 self.shortability
178 .get(instrument)
179 .and_then(|series| series.range(..=time).next_back())
180 .map(|(_, shortability)| shortability.clone())
181 .unwrap_or_default()
182 }
183
184 fn short_sale_restricted(
185 &self,
186 instrument: &InstrumentNameExchange,
187 time: DateTime<Utc>,
188 ) -> bool {
189 self.restrictions.get(instrument).is_some_and(|windows| {
190 windows
191 .iter()
192 .any(|(from, until)| *from <= time && time < *until)
193 })
194 }
195}
196
197#[cfg(test)]
198#[allow(clippy::unwrap_used)] mod tests {
200 use super::*;
201 use rust_decimal_macros::dec;
202
203 fn utc(text: &str) -> DateTime<Utc> {
204 DateTime::parse_from_rfc3339(text).unwrap().to_utc()
205 }
206
207 fn name() -> InstrumentNameExchange {
208 InstrumentNameExchange::new("AAPL")
209 }
210
211 #[test]
212 fn shortability_starts_unknown_and_each_setter_fills_one_field() {
213 assert_eq!(
214 Shortability::new(),
215 Shortability {
216 shortable: None,
217 easy_to_borrow: None,
218 available: None,
219 fee_rate: None,
220 }
221 );
222
223 let full = Shortability::new()
224 .with_shortable(true)
225 .with_easy_to_borrow(false)
226 .with_available(dec!(1500))
227 .with_fee_rate(dec!(0.12));
228 assert_eq!(full.shortable, Some(true));
229 assert_eq!(full.easy_to_borrow, Some(false));
230 assert_eq!(full.available, Some(dec!(1500)));
231 assert_eq!(full.fee_rate, Some(dec!(0.12)));
232 }
233
234 #[test]
235 fn shortability_round_trips_through_serde_and_reads_missing_fields_as_unknown() {
236 let full = Shortability::new()
237 .with_shortable(true)
238 .with_available(dec!(10));
239 let json = serde_json::to_string(&full).unwrap();
240 assert_eq!(serde_json::from_str::<Shortability>(&json).unwrap(), full);
241
242 let sparse: Shortability = serde_json::from_str(r#"{"shortable":false}"#).unwrap();
243 assert_eq!(sparse, Shortability::new().with_shortable(false));
244 }
245
246 #[test]
247 fn a_table_steps_at_each_entry_and_knows_nothing_before_the_first() {
248 let hard = Shortability::new()
249 .with_shortable(true)
250 .with_easy_to_borrow(false);
251 let closed = Shortability::new().with_shortable(false);
252 let table = ShortabilityTable::new()
254 .with_shortability(name(), utc("2026-01-03T00:00:00Z"), closed.clone())
255 .with_shortability(name(), utc("2026-01-01T00:00:00Z"), hard.clone());
256
257 let at = |text| table.shortability(&name(), utc(text));
258 assert_eq!(at("2025-12-31T23:59:59Z"), Shortability::new());
259 assert_eq!(at("2026-01-01T00:00:00Z"), hard);
260 assert_eq!(at("2026-01-02T23:59:59Z"), hard);
261 assert_eq!(at("2026-01-03T00:00:00Z"), closed);
262 assert_eq!(
263 table.shortability(
264 &InstrumentNameExchange::new("MSFT"),
265 utc("2026-01-03T00:00:00Z")
266 ),
267 Shortability::new()
268 );
269 }
270
271 #[test]
272 fn a_later_entry_at_the_same_instant_replaces_the_earlier() {
273 let instant = utc("2026-01-01T00:00:00Z");
274 let table = ShortabilityTable::new()
275 .with_shortability(name(), instant, Shortability::new().with_shortable(true))
276 .with_shortability(name(), instant, Shortability::new().with_shortable(false));
277
278 assert_eq!(table.shortability(&name(), instant).shortable, Some(false));
279 }
280
281 #[test]
282 fn a_restriction_is_in_effect_from_its_start_until_just_before_its_end() {
283 let table = ShortabilityTable::new()
284 .with_short_sale_restriction(
285 name(),
286 utc("2026-01-01T15:00:00Z"),
287 utc("2026-01-02T21:00:00Z"),
288 )
289 .with_short_sale_restriction(
290 name(),
291 utc("2026-01-05T15:00:00Z"),
292 utc("2026-01-06T21:00:00Z"),
293 );
294
295 let restricted = |text| table.short_sale_restricted(&name(), utc(text));
296 assert!(!restricted("2026-01-01T14:59:59Z"));
297 assert!(restricted("2026-01-01T15:00:00Z"));
298 assert!(restricted("2026-01-02T20:59:59Z"));
299 assert!(!restricted("2026-01-02T21:00:00Z"));
300 assert!(restricted("2026-01-06T00:00:00Z"));
301 assert!(!table.short_sale_restricted(
302 &InstrumentNameExchange::new("MSFT"),
303 utc("2026-01-01T16:00:00Z")
304 ));
305 }
306
307 #[test]
308 #[should_panic(expected = "must end after it starts")]
309 fn a_restriction_that_ends_before_it_starts_panics() {
310 let _ = ShortabilityTable::new().with_short_sale_restriction(
311 name(),
312 utc("2026-01-02T00:00:00Z"),
313 utc("2026-01-01T00:00:00Z"),
314 );
315 }
316
317 #[test]
318 fn a_provider_is_unrestricted_unless_it_says_otherwise() {
319 #[derive(Debug)]
320 struct Constant;
321 impl ShortabilityProvider for Constant {
322 fn shortability(&self, _: &InstrumentNameExchange, _: DateTime<Utc>) -> Shortability {
323 Shortability::new()
324 }
325 }
326
327 assert!(!Constant.short_sale_restricted(&name(), utc("2026-01-01T00:00:00Z")));
328 }
329}