use chrono::{DateTime, Utc};
use fnv::FnvHashMap;
use rust_decimal::Decimal;
use rustrade_instrument::instrument::name::InstrumentNameExchange;
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, fmt::Debug};
#[non_exhaustive]
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Shortability {
pub shortable: Option<bool>,
pub easy_to_borrow: Option<bool>,
pub available: Option<Decimal>,
pub fee_rate: Option<Decimal>,
pub delisting_at: Option<DateTime<Utc>>,
}
impl Shortability {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_shortable(mut self, shortable: bool) -> Self {
self.shortable = Some(shortable);
self
}
#[must_use]
pub fn with_easy_to_borrow(mut self, easy_to_borrow: bool) -> Self {
self.easy_to_borrow = Some(easy_to_borrow);
self
}
#[must_use]
pub fn with_available(mut self, available: Decimal) -> Self {
self.available = Some(available);
self
}
#[must_use]
pub fn with_fee_rate(mut self, fee_rate: Decimal) -> Self {
self.fee_rate = Some(fee_rate);
self
}
#[must_use]
pub fn with_delisting_at(mut self, delisting_at: DateTime<Utc>) -> Self {
self.delisting_at = Some(delisting_at);
self
}
}
pub trait ShortabilityProvider: Debug + Send + Sync {
fn shortability(
&self,
instrument: &InstrumentNameExchange,
time: DateTime<Utc>,
) -> Shortability;
fn short_sale_restricted(
&self,
_instrument: &InstrumentNameExchange,
_time: DateTime<Utc>,
) -> bool {
false
}
}
#[derive(Debug, Clone, Default)]
pub struct ShortabilityTable {
shortability: FnvHashMap<InstrumentNameExchange, BTreeMap<DateTime<Utc>, Shortability>>,
restrictions: FnvHashMap<InstrumentNameExchange, Vec<(DateTime<Utc>, DateTime<Utc>)>>,
}
impl ShortabilityTable {
pub fn new() -> Self {
Self::default()
}
#[must_use]
pub fn with_shortability(
mut self,
instrument: InstrumentNameExchange,
from: DateTime<Utc>,
shortability: Shortability,
) -> Self {
self.shortability
.entry(instrument)
.or_default()
.insert(from, shortability);
self
}
#[must_use]
pub fn with_short_sale_restriction(
mut self,
instrument: InstrumentNameExchange,
from: DateTime<Utc>,
until: DateTime<Utc>,
) -> Self {
assert!(
from < until,
"a short-sale restriction on {instrument} must end after it starts: {from} to {until}"
);
self.restrictions
.entry(instrument)
.or_default()
.push((from, until));
self
}
}
impl ShortabilityProvider for ShortabilityTable {
fn shortability(
&self,
instrument: &InstrumentNameExchange,
time: DateTime<Utc>,
) -> Shortability {
self.shortability
.get(instrument)
.and_then(|series| series.range(..=time).next_back())
.map(|(_, shortability)| shortability.clone())
.unwrap_or_default()
}
fn short_sale_restricted(
&self,
instrument: &InstrumentNameExchange,
time: DateTime<Utc>,
) -> bool {
self.restrictions.get(instrument).is_some_and(|windows| {
windows
.iter()
.any(|(from, until)| *from <= time && time < *until)
})
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)] mod tests {
use super::*;
use rust_decimal_macros::dec;
fn utc(text: &str) -> DateTime<Utc> {
DateTime::parse_from_rfc3339(text).unwrap().to_utc()
}
fn name() -> InstrumentNameExchange {
InstrumentNameExchange::new("AAPL")
}
#[test]
fn shortability_starts_unknown_and_each_setter_fills_one_field() {
assert_eq!(
Shortability::new(),
Shortability {
shortable: None,
easy_to_borrow: None,
available: None,
fee_rate: None,
delisting_at: None,
}
);
let full = Shortability::new()
.with_shortable(true)
.with_easy_to_borrow(false)
.with_available(dec!(1500))
.with_fee_rate(dec!(0.12))
.with_delisting_at(utc("2026-10-20T00:00:00Z"));
assert_eq!(full.shortable, Some(true));
assert_eq!(full.easy_to_borrow, Some(false));
assert_eq!(full.available, Some(dec!(1500)));
assert_eq!(full.fee_rate, Some(dec!(0.12)));
assert_eq!(full.delisting_at, Some(utc("2026-10-20T00:00:00Z")));
}
#[test]
fn shortability_round_trips_through_serde_and_reads_missing_fields_as_unknown() {
let full = Shortability::new()
.with_shortable(true)
.with_available(dec!(10))
.with_delisting_at(utc("2026-10-20T00:00:00Z"));
let json = serde_json::to_string(&full).unwrap();
assert_eq!(serde_json::from_str::<Shortability>(&json).unwrap(), full);
let sparse: Shortability = serde_json::from_str(r#"{"shortable":false}"#).unwrap();
assert_eq!(sparse, Shortability::new().with_shortable(false));
}
#[test]
fn a_table_steps_at_each_entry_and_knows_nothing_before_the_first() {
let hard = Shortability::new()
.with_shortable(true)
.with_easy_to_borrow(false);
let closed = Shortability::new().with_shortable(false);
let table = ShortabilityTable::new()
.with_shortability(name(), utc("2026-01-03T00:00:00Z"), closed.clone())
.with_shortability(name(), utc("2026-01-01T00:00:00Z"), hard.clone());
let at = |text| table.shortability(&name(), utc(text));
assert_eq!(at("2025-12-31T23:59:59Z"), Shortability::new());
assert_eq!(at("2026-01-01T00:00:00Z"), hard);
assert_eq!(at("2026-01-02T23:59:59Z"), hard);
assert_eq!(at("2026-01-03T00:00:00Z"), closed);
assert_eq!(
table.shortability(
&InstrumentNameExchange::new("MSFT"),
utc("2026-01-03T00:00:00Z")
),
Shortability::new()
);
}
#[test]
fn a_later_entry_at_the_same_instant_replaces_the_earlier() {
let instant = utc("2026-01-01T00:00:00Z");
let table = ShortabilityTable::new()
.with_shortability(name(), instant, Shortability::new().with_shortable(true))
.with_shortability(name(), instant, Shortability::new().with_shortable(false));
assert_eq!(table.shortability(&name(), instant).shortable, Some(false));
}
#[test]
fn a_restriction_is_in_effect_from_its_start_until_just_before_its_end() {
let table = ShortabilityTable::new()
.with_short_sale_restriction(
name(),
utc("2026-01-01T15:00:00Z"),
utc("2026-01-02T21:00:00Z"),
)
.with_short_sale_restriction(
name(),
utc("2026-01-05T15:00:00Z"),
utc("2026-01-06T21:00:00Z"),
);
let restricted = |text| table.short_sale_restricted(&name(), utc(text));
assert!(!restricted("2026-01-01T14:59:59Z"));
assert!(restricted("2026-01-01T15:00:00Z"));
assert!(restricted("2026-01-02T20:59:59Z"));
assert!(!restricted("2026-01-02T21:00:00Z"));
assert!(restricted("2026-01-06T00:00:00Z"));
assert!(!table.short_sale_restricted(
&InstrumentNameExchange::new("MSFT"),
utc("2026-01-01T16:00:00Z")
));
}
#[test]
#[should_panic(expected = "must end after it starts")]
fn a_restriction_that_ends_before_it_starts_panics() {
let _ = ShortabilityTable::new().with_short_sale_restriction(
name(),
utc("2026-01-02T00:00:00Z"),
utc("2026-01-01T00:00:00Z"),
);
}
#[test]
fn a_provider_is_unrestricted_unless_it_says_otherwise() {
#[derive(Debug)]
struct Constant;
impl ShortabilityProvider for Constant {
fn shortability(&self, _: &InstrumentNameExchange, _: DateTime<Utc>) -> Shortability {
Shortability::new()
}
}
assert!(!Constant.short_sale_restricted(&name(), utc("2026-01-01T00:00:00Z")));
}
}