diff --git a/crates/fidc-core/src/numeric_factors.rs b/crates/fidc-core/src/numeric_factors.rs index 81b985d..320339e 100644 --- a/crates/fidc-core/src/numeric_factors.rs +++ b/crates/fidc-core/src/numeric_factors.rs @@ -19,6 +19,7 @@ pub struct NumericFactorMap { #[derive(Clone)] enum Storage { + Empty, Owned(Vec<(CompactString, f64)>), Shared(SharedRow), } @@ -50,12 +51,13 @@ fn compact_key(key: Cow<'static, str>) -> CompactString { impl NumericFactorMap { pub const fn new() -> Self { Self { - storage: Storage::Owned(Vec::new()), + storage: Storage::Empty, } } pub fn len(&self) -> usize { match &self.storage { + Storage::Empty => 0, Storage::Owned(entries) => entries.len(), Storage::Shared(row) => row.len(), } @@ -65,6 +67,7 @@ impl NumericFactorMap { } pub fn clear(&mut self) { match &mut self.storage { + Storage::Empty => {}, Storage::Owned(entries) => entries.clear(), Storage::Shared(_) => *self = Self::new(), } @@ -77,6 +80,7 @@ impl NumericFactorMap { pub(crate) fn has_normalized_finite_entries(&self) -> bool { match &self.storage { + Storage::Empty => true, Storage::Owned(entries) => entries.iter().all(|(name, value)| { normalized_name(name) && value.is_finite() }), @@ -85,18 +89,23 @@ impl NumericFactorMap { } fn owned_entries(&mut self) -> &mut Vec<(CompactString, f64)> { + if matches!(self.storage, Storage::Empty) { + self.storage = Storage::Owned(Vec::new()); + } if matches!(self.storage, Storage::Shared(_)) { let entries = self.iter().map(|(key, value)| (key.clone(), *value)).collect(); self.storage = Storage::Owned(entries); } match &mut self.storage { Storage::Owned(entries) => entries, + Storage::Empty => unreachable!("empty storage was initialized"), Storage::Shared(_) => unreachable!("shared row was materialized"), } } pub fn get(&self, key: &str) -> Option<&f64> { match &self.storage { + Storage::Empty => None, Storage::Owned(entries) => entries .binary_search_by(|(name, _)| name.as_str().cmp(key)) .ok() @@ -106,7 +115,9 @@ impl NumericFactorMap { } pub fn get_mut(&mut self, key: &str) -> Option<&mut f64> { - if !self.contains_key(key) { + if matches!(self.storage, Storage::Empty) + || matches!(self.storage, Storage::Shared(_)) && !self.contains_key(key) + { return None; } let entries = self.owned_entries(); @@ -145,7 +156,9 @@ impl NumericFactorMap { } pub fn remove(&mut self, key: &str) -> Option { - if !self.contains_key(key) { + if matches!(self.storage, Storage::Empty) + || matches!(self.storage, Storage::Shared(_)) && !self.contains_key(key) + { return None; } let entries = self.owned_entries(); @@ -161,6 +174,7 @@ impl NumericFactorMap { pub fn iter(&self) -> Iter<'_> { match &self.storage { + Storage::Empty => Iter(IterStorage::Owned([].iter())), Storage::Owned(entries) => Iter(IterStorage::Owned(entries.iter())), Storage::Shared(row) => Iter(IterStorage::Shared(row.iter())), } @@ -228,6 +242,7 @@ impl IntoIterator for NumericFactorMap { type IntoIter = std::vec::IntoIter; fn into_iter(self) -> Self::IntoIter { match self.storage { + Storage::Empty => Vec::new().into_iter(), Storage::Owned(entries) => entries.into_iter(), Storage::Shared(row) => row.iter().map(|(key, value)| (key.clone(), *value)) .collect::>().into_iter(), @@ -255,7 +270,11 @@ impl FromIterator<(CompactString, f64)> for NumericFactorMap { false } }); - Self { storage: Storage::Owned(entries) } + if entries.is_empty() { + Self::new() + } else { + Self { storage: Storage::Owned(entries) } + } } } impl Extend<(Cow<'static, str>, f64)> for NumericFactorMap { @@ -367,7 +386,7 @@ mod tests { let normalized = rows.iter().map(NumericFactorMap::has_normalized_finite_entries).collect::>(); assert_eq!(NumericFactorMap::share_rows(rows.iter_mut()), rows.len()); for (index, row) in rows.iter().enumerate() { - assert!(matches!(row.storage, Storage::Shared(_))); + assert!(row.is_empty() || matches!(row.storage, Storage::Shared(_))); assert_eq!(row_bits(row), before[index]); assert_eq!(serde_json::to_string(row).unwrap(), serialized[index]); assert_eq!(row.has_normalized_finite_entries(), normalized[index]); diff --git a/crates/fidc-core/src/numeric_factors/shared_rows.rs b/crates/fidc-core/src/numeric_factors/shared_rows.rs index 263bdd3..c8c16f4 100644 --- a/crates/fidc-core/src/numeric_factors/shared_rows.rs +++ b/crates/fidc-core/src/numeric_factors/shared_rows.rs @@ -100,7 +100,9 @@ pub(super) fn share<'a>(rows: impl IntoIterator return 0; } let Some(owned_bytes) = rows.iter().try_fold(0usize, |sum, row| { - let Storage::Owned(entries) = &row.storage else { return None }; + let Storage::Owned(entries) = &row.storage else { + return matches!(row.storage, Storage::Empty).then_some(sum); + }; sum.checked_add(entries.capacity().checked_mul(std::mem::size_of::<(CompactString, f64)>())?) }) else { return 0 }; let mut names = BTreeMap::<&str, &CompactString>::new(); @@ -166,7 +168,11 @@ pub(super) fn share<'a>(rows: impl IntoIterator }); let count = rows.len(); for (index, row) in rows.into_iter().enumerate() { - row.storage = Storage::Shared(SharedRow { data: Arc::clone(&data), index }); + row.storage = if data.lengths[index] == 0 { + Storage::Empty + } else { + Storage::Shared(SharedRow { data: Arc::clone(&data), index }) + }; } count }