diff --git a/src/frame/base.rs b/src/frame/base.rs index df4ff40..49b873d 100644 --- a/src/frame/base.rs +++ b/src/frame/base.rs @@ -1,15 +1,13 @@ use crate::matrix::Matrix; - use chrono::NaiveDate; use std::collections::HashMap; - use std::fmt; use std::ops::{Index, IndexMut, Not, Range}; -// --- Helper Enums and Structs for Indexing --- +// Helper enums and structs for indexing support /// Represents the different types of row indices a Frame can have. -#[derive(Debug, Clone, PartialEq, Eq)] // Added Eq for HashMaps etc. +#[derive(Debug, Clone, PartialEq, Eq)] // Derive Eq to enable usage as HashMap keys pub enum RowIndex { /// Integer-based index (e.g., 0, 1, 2, ...). Values must be unique. Int(Vec), @@ -21,24 +19,6 @@ pub enum RowIndex { impl RowIndex { /// Returns the number of elements in the index. - /// - /// # Examples - /// ``` - /// # use rustframe::frame::RowIndex; - /// # use chrono::NaiveDate; - /// # fn d(y: i32, m: u32, d: u32) -> NaiveDate { NaiveDate::from_ymd_opt(y,m,d).unwrap() } - /// let idx_int = RowIndex::Int(vec![10, 20, 5]); - /// assert_eq!(idx_int.len(), 3); - /// - /// let idx_date = RowIndex::Date(vec![d(2024,1,1), d(2024,1,2)]); - /// assert_eq!(idx_date.len(), 2); - /// - /// let idx_range = RowIndex::Range(0..5); - /// assert_eq!(idx_range.len(), 5); - /// - /// let idx_empty_int = RowIndex::Int(vec![]); - /// assert_eq!(idx_empty_int.len(), 0); - /// ``` pub fn len(&self) -> usize { match self { RowIndex::Int(v) => v.len(), @@ -48,36 +28,22 @@ impl RowIndex { } /// Checks if the index is empty. - /// - /// # Examples - /// ``` - /// # use rustframe::frame::RowIndex; - /// # use chrono::NaiveDate; - /// let idx_int = RowIndex::Int(vec![10, 20, 5]); - /// assert!(!idx_int.is_empty()); - /// - /// let idx_range = RowIndex::Range(0..0); - /// assert!(idx_range.is_empty()); - /// - /// let idx_empty_date = RowIndex::Date(vec![]); - /// assert!(idx_empty_date.is_empty()); - /// ``` pub fn is_empty(&self) -> bool { self.len() == 0 } } /// Internal helper for fast lookups from index value to physical row position. -#[derive(Debug, Clone, PartialEq, Eq)] // Added Eq +#[derive(Debug, Clone, PartialEq, Eq)] enum RowIndexLookup { Int(HashMap), Date(HashMap), - None, // Used for Range index + None, // Variant for default range-based indexing } -// --- Frame Struct Definition --- +// Frame struct definition and associated implementations -/// A data frame – a Matrix with string-identified columns and a typed row index. +/// A data frame - a Matrix with string-identified columns and a typed row index. /// /// `Frame` extends the concept of a `Matrix` by adding named columns /// and an index for rows, which can be integers, dates, or a default range. @@ -142,11 +108,10 @@ enum RowIndexLookup { /// /// // --- Example 4: Element-wise Multiplication --- /// let f_prod = &f1 * &f2; -/// // Use approx comparison for floats if necessary -/// assert!((f_prod["C1"][0] - 0.1_f64).abs() < 1e-9); // 1.0 * 0.1 -/// assert!((f_prod["C1"][1] - 0.4_f64).abs() < 1e-9); // 2.0 * 0.2 -/// assert!((f_prod["C2"][0] - 0.9_f64).abs() < 1e-9); // 3.0 * 0.3 -/// assert!((f_prod["C2"][1] - 1.6_f64).abs() < 1e-9); // 4.0 * 0.4 +/// assert!((f_prod["C1"][0] - 0.1_f64).abs() < 1e-9); +/// assert!((f_prod["C1"][1] - 0.4_f64).abs() < 1e-9); +/// assert!((f_prod["C2"][0] - 0.9_f64).abs() < 1e-9); +/// assert!((f_prod["C2"][1] - 1.6_f64).abs() < 1e-9); /// /// // --- Example 5: Column Manipulation and Sorting --- /// let mut frame_manip = Frame::new( @@ -163,13 +128,14 @@ enum RowIndexLookup { /// assert_eq!(frame_manip["DataX"], vec![3, 4]); // X has A's original data /// let deleted_c = frame_manip.delete_column("DataC"); /// assert_eq!(deleted_c, vec![1, 2]); -/// assert_eq!(frame_manip.columns(), &["DataX", "DataB"]); // Order after delete -/// frame_manip.sort_columns(); // Sorts ["DataX", "DataB"] -> ["DataB", "DataX"] +/// assert_eq!(frame_manip.columns(), &["DataX", "DataB"]); +/// frame_manip.sort_columns(); /// assert_eq!(frame_manip.columns(), &["DataB", "DataX"]); -/// assert_eq!(frame_manip["DataB"], vec![5, 6]); // B keeps its data -/// assert_eq!(frame_manip["DataX"], vec![3, 4]); // X keeps its data (originally A's) +/// assert_eq!(frame_manip["DataB"], vec![5, 6]); +/// assert_eq!(frame_manip["DataX"], vec![3, 4]); /// ``` -// Implement Debug manually for Frame + +// Custom Debug implementation for Frame impl fmt::Debug for Frame { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Frame") @@ -178,13 +144,13 @@ impl fmt::Debug for Frame { .field("matrix_dims", &(self.matrix.rows(), self.matrix.cols())) .field("col_lookup", &self.col_lookup) .field("index_lookup", &self.index_lookup) - // Optionally hide matrix data in Debug unless specifically requested + // Matrix data is omitted from Debug output by default // .field("matrix", &self.matrix) .finish() } } -#[derive(Clone, PartialEq)] // Removed Eq as T doesn't require Eq +#[derive(Clone, PartialEq)] pub struct Frame { /// Vector holding the column names in their current order. column_names: Vec, @@ -199,13 +165,13 @@ pub struct Frame { } impl Frame { - /* ---------- Constructors ---------- */ + // Constructors /// Creates a new [`Frame`] from a matrix, column names, and an optional row index. /// Panics if the underlying `Matrix` requirements are not met (e.g., rows>0, cols>0 for standard constructors), /// or if column name/index constraints are violated. pub fn new>(matrix: Matrix, names: Vec, index: Option) -> Self { - // --- Column Validation --- + // Validate that the number of column names matches the matrix's column count. if matrix.cols() != names.len() { panic!( "Frame::new: column name count mismatch (names: {}, matrix: {})", @@ -213,10 +179,9 @@ impl Frame { matrix.cols() ); } - // Note: Matrix constructors enforce rows > 0 and cols > 0 based on provided code. - // Therefore, we don't need to explicitly handle 0-row/0-col cases here, - // as they would have already panicked during Matrix creation. + // Matrix creation already enforces non-zero rows and columns. + // Build column name list and lookup map, ensuring no duplicates. let mut col_lookup = HashMap::with_capacity(names.len()); let column_names: Vec = names .into_iter() @@ -230,7 +195,7 @@ impl Frame { }) .collect(); - // --- Index Validation and Processing --- + // Validate and construct the row index and its lookup structure. let num_rows = matrix.rows(); let (index_values, index_lookup) = match index { Some(RowIndex::Int(vals)) => { @@ -241,6 +206,7 @@ impl Frame { num_rows ); } + // Build integer-to-physical-row mapping. let mut lookup = HashMap::with_capacity(num_rows); for (physical_row, index_val) in vals.iter().enumerate() { if lookup.insert(*index_val, physical_row).is_some() { @@ -257,6 +223,7 @@ impl Frame { num_rows ); } + // Build date-to-physical-row mapping. let mut lookup = HashMap::with_capacity(num_rows); for (physical_row, index_val) in vals.iter().enumerate() { if lookup.insert(*index_val, physical_row).is_some() { @@ -270,7 +237,10 @@ impl Frame { "Frame::new: Cannot explicitly provide a Range index. Use None for default range." ); } - None => (RowIndex::Range(0..num_rows), RowIndexLookup::None), + None => { + // Default to a sequential range index. + (RowIndex::Range(0..num_rows), RowIndexLookup::None) + } }; Self { @@ -282,48 +252,55 @@ impl Frame { } } - /* ---------- Accessors ---------- */ - /// Returns an immutable reference to the underlying Matrix. + // Accessors + + /// Returns an immutable reference to the underlying `Matrix`. #[inline] pub fn matrix(&self) -> &Matrix { &self.matrix } - /// Returns a mutable reference to the underlying Matrix. - /// Use with caution, as direct matrix manipulation bypasses Frame's name/index tracking. + + /// Returns a mutable reference to the underlying `Matrix`. + /// Use with caution: direct matrix edits bypass frame-level name/index consistency checks. #[inline] pub fn matrix_mut(&mut self) -> &mut Matrix { &mut self.matrix } - /// Returns a slice containing the current column names in order. + + /// Returns the list of column names in their current order. #[inline] pub fn columns(&self) -> &[String] { &self.column_names } - /// Returns a reference to the `RowIndex`. + + /// Returns a reference to the frame's row index. #[inline] pub fn index(&self) -> &RowIndex { &self.index } + /// Returns the number of rows in the frame. #[inline] pub fn rows(&self) -> usize { self.matrix.rows() } + /// Returns the number of columns in the frame. #[inline] pub fn cols(&self) -> usize { self.matrix.cols() } - /// Returns the physical column index for a given column name, if it exists. + + /// Returns the physical column index for the given column name, if present. #[inline] pub fn column_index(&self, name: &str) -> Option { self.col_lookup.get(name).copied() } - /// Internal helper to find the physical row index based on the index key and type. + /// Internal helper to compute the physical row index from a logical key. fn get_physical_row_index(&self, index_key: Idx) -> usize where - Self: RowIndexLookupHelper, // Uses the helper trait below + Self: RowIndexLookupHelper, // Requires `RowIndexLookupHelper` for `Idx` { >::lookup_row_index( index_key, @@ -332,16 +309,17 @@ impl Frame { ) } - /// Returns an immutable slice representing the data in the specified column. - /// Panics if the column name does not exist. + /// Returns an immutable slice of the specified column's data. + /// Panics if the column name is not found. pub fn column(&self, name: &str) -> &[T] { let idx = self .column_index(name) .unwrap_or_else(|| panic!("Frame::column: unknown column label: '{}'", name)); self.matrix.column(idx) } - /// Returns a mutable slice representing the data in the specified column. - /// Panics if the column name does not exist. + + /// Returns a mutable slice of the specified column's data. + /// Panics if the column name is not found. pub fn column_mut(&mut self, name: &str) -> &mut [T] { let idx = self .column_index(name) @@ -349,10 +327,10 @@ impl Frame { self.matrix.column_mut(idx) } - /* ---------- Row Access Methods ---------- */ - /// Returns an immutable view of the row corresponding to the given integer index key. - /// Panics if the key is not found in the `RowIndex::Int` or `RowIndex::Range`. - /// Panics if the index is `RowIndex::Date`. + // Row access methods + + /// Returns an immutable view of the row for the given integer key. + /// Panics if the key is invalid or if the index type is `Date`. pub fn get_row(&self, index_key: usize) -> FrameRowView<'_, T> { let idx = self.get_physical_row_index(index_key); FrameRowView { @@ -360,9 +338,9 @@ impl Frame { physical_row_idx: idx, } } - /// Returns a mutable view of the row corresponding to the given integer index key. - /// Panics if the key is not found in the `RowIndex::Int` or `RowIndex::Range`. - /// Panics if the index is `RowIndex::Date`. + + /// Returns a mutable view of the row for the given integer key. + /// Panics if the key is invalid or if the index type is `Date`. pub fn get_row_mut(&mut self, index_key: usize) -> FrameRowViewMut<'_, T> { let idx = self.get_physical_row_index(index_key); FrameRowViewMut { @@ -370,9 +348,9 @@ impl Frame { physical_row_idx: idx, } } - /// Returns an immutable view of the row corresponding to the given date index key. - /// Panics if the key is not found in the `RowIndex::Date`. - /// Panics if the index is `RowIndex::Int` or `RowIndex::Range`. + + /// Returns an immutable view of the row for the given date key. + /// Panics if the key is invalid or if the index type is not `Date`. pub fn get_row_date(&self, index_key: NaiveDate) -> FrameRowView<'_, T> { let idx = self.get_physical_row_index(index_key); FrameRowView { @@ -380,9 +358,9 @@ impl Frame { physical_row_idx: idx, } } - /// Returns a mutable view of the row corresponding to the given date index key. - /// Panics if the key is not found in the `RowIndex::Date`. - /// Panics if the index is `RowIndex::Int` or `RowIndex::Range`. + + /// Returns a mutable view of the row for the given date key. + /// Panics if the key is invalid or if the index type is not `Date`. pub fn get_row_date_mut(&mut self, index_key: NaiveDate) -> FrameRowViewMut<'_, T> { let idx = self.get_physical_row_index(index_key); FrameRowViewMut { @@ -391,41 +369,34 @@ impl Frame { } } - /* ---------- Column manipulation ---------- */ + // Column manipulation - /// Internal helper to swap two columns. Updates matrix, column names, and lookup map. - /// This is not intended for direct user consumption. Use `sort_columns`. + /// Internal helper to swap two columns, updating the matrix, name list, and lookup map. + /// Not intended for public use; prefer `sort_columns`. fn _swap_columns_internal(&mut self, a: &str, b: &str) { - // Avoid cloning strings if possible by getting indices first - let maybe_ia = self.column_index(a); - let maybe_ib = self.column_index(b); - - let ia = maybe_ia.unwrap_or_else(|| { + // Retrieve physical indices for the given column names. + let ia = self.column_index(a).unwrap_or_else(|| { panic!("Frame::_swap_columns_internal: unknown column label: {}", a) }); - let ib = maybe_ib.unwrap_or_else(|| { + let ib = self.column_index(b).unwrap_or_else(|| { panic!("Frame::_swap_columns_internal: unknown column label: {}", b) }); if ia == ib { - return; // No-op + return; // No action needed if columns are identical } - // 1. Swap data in the underlying matrix + // Swap the columns in the underlying matrix. self.matrix.swap_columns(ia, ib); - - // 2. Swap names in the ordered list + // Swap the names in the ordered list. self.column_names.swap(ia, ib); - - // 3. Update the lookup map to reflect the new physical indices - // The column originally named 'a' is now at physical index 'ib' + // Update lookup entries to reflect the new positions. self.col_lookup.insert(a.to_string(), ib); - // The column originally named 'b' is now at physical index 'ia' self.col_lookup.insert(b.to_string(), ia); } /// Renames an existing column. - /// Panics if the old name doesn't exist, or the new name already exists or is the same as the old name. + /// Panics if the source name is missing or if the target name is already in use. pub fn rename>(&mut self, old: &str, new: L) { let new_name = new.into(); if old == new_name { @@ -444,82 +415,63 @@ impl Frame { ); } - // Update lookup map + // Update lookup and ordered name list. self.col_lookup.remove(old); self.col_lookup.insert(new_name.clone(), idx); - - // Update ordered name list self.column_names[idx] = new_name; } - /// Adds a new column to the end of the frame. - /// Panics if the column name already exists or if the data length doesn't match the number of rows. + /// Adds a new column at the end of the frame. + /// Panics if the column name exists or if the data length mismatches the row count. pub fn add_column>(&mut self, name: L, column_data: Vec) { let name_str = name.into(); if self.col_lookup.contains_key(&name_str) { panic!("Frame::add_column: duplicate column label: {}", name_str); } - // Matrix::add_column checks length against self.rows(). - // This assumes self.rows() > 0 because Matrix constructors enforce this. + // Matrix enforces that the new column length matches the frame's row count. let new_col_idx = self.matrix.cols(); - self.matrix.add_column(new_col_idx, column_data); // Add to end + self.matrix.add_column(new_col_idx, column_data); - // Update frame metadata + // Update metadata to include the new column. self.column_names.push(name_str.clone()); self.col_lookup.insert(name_str, new_col_idx); } - /// Deletes a column from the frame by name. - /// Returns the data of the deleted column. - /// Panics if the column name does not exist. + /// Deletes a column by name and returns its data. + /// Panics if the column name is not found. pub fn delete_column(&mut self, name: &str) -> Vec { let idx = self .column_index(name) .unwrap_or_else(|| panic!("Frame::delete_column: unknown column label: '{}'", name)); - // Retrieve data before deleting. Requires Clone. - // Note: Matrix::delete_column might be more efficient if it returned the data. + // Clone out the data before removal. let deleted_data = self.matrix.column(idx).to_vec(); - - // Delete from matrix self.matrix.delete_column(idx); - // Remove from metadata + // Remove from metadata and rebuild lookup for shifted columns. self.column_names.remove(idx); - self.col_lookup.remove(name); // Remove the specific entry + self.col_lookup.remove(name); + self.rebuild_col_lookup(); - // Update indices in the lookup map for columns that shifted - // This is necessary because deleting column `idx` shifts all columns > `idx` one position to the left. - self.rebuild_col_lookup(); // Rebuild is simpler than manually adjusting indices - - // Handle index if last column removed - // Since Matrix must have cols >= 1, delete_column cannot result in 0 cols - // unless the matrix started with 1 col. Matrix must also have rows >= 1. - // So, after delete_column, the matrix will still have rows >= 1, but potentially 0 cols. - // This state (rows > 0, cols = 0) might be invalid depending on Matrix guarantees. - // Assuming Matrix *allows* this state after delete_column: + // If all columns are removed, reset default range index. if self.cols() == 0 { - // Ensure index matches row count. if let RowIndex::Range(_) = self.index { - // The range should reflect the number of rows remaining. self.index = RowIndex::Range(0..self.rows()); self.index_lookup = RowIndexLookup::None; } - // If the index was Int or Date, it remains unchanged as it tracks rows. } deleted_data } - /// Sorts the columns of the Frame alphabetically by name in place. - /// The data remains associated with its original column name. + /// Sorts columns alphabetically by name, preserving data associations. pub fn sort_columns(&mut self) { let n = self.column_names.len(); if n <= 1 { - return; // Already sorted (or empty) + return; // Nothing to sort } - // Simple selection sort based on column names. + // Selection sort on column names. for i in 0..n { let mut min_idx = i; for j in (i + 1)..n { @@ -528,18 +480,13 @@ impl Frame { } } if min_idx != i { - // Get names at current physical positions i and min_idx before swapping let col_i_name = self.column_names[i].clone(); let col_min_name = self.column_names[min_idx].clone(); - - // Use the internal swap function which handles matrix, names, and lookup self._swap_columns_internal(&col_i_name, &col_min_name); - // After swap, col_i_name is at min_idx, col_min_name is at i. - // The name list `self.column_names` is updated by the swap. } } - // Final check for internal consistency (optional, for debugging) + // Debug-only consistency check. #[cfg(debug_assertions)] { let mut temp_lookup = HashMap::with_capacity(self.cols()); @@ -548,14 +495,14 @@ impl Frame { } assert_eq!( self.col_lookup, temp_lookup, - "Internal col_lookup inconsistent after sort_columns" + "Inconsistent col_lookup after sort_columns" ); } } - /* ---------- Helpers ---------- */ - /// Rebuilds the `col_lookup` map based on the current `column_names` order. - /// Used internally after operations that change multiple column indices (like delete). + // Internal helpers + + /// Rebuilds the column lookup map to match the current `column_names` ordering. fn rebuild_col_lookup(&mut self) { self.col_lookup.clear(); for (i, name) in self.column_names.iter().enumerate() { @@ -564,13 +511,12 @@ impl Frame { } } -// --- Helper Trait for Row Index Lookup --- +// Trait for resolving logical to physical row indices /// Internal trait to abstract the logic for looking up physical row indices. trait RowIndexLookupHelper { fn lookup_row_index(key: Idx, index_values: &RowIndex, index_lookup: &RowIndexLookup) -> usize; } -/// Implementation for `usize` keys (used for `RowIndex::Int` and `RowIndex::Range`). impl RowIndexLookupHelper for Frame { fn lookup_row_index( key: usize, @@ -579,18 +525,16 @@ impl RowIndexLookupHelper for Frame { ) -> usize { match (index_values, index_lookup) { (RowIndex::Int(_), RowIndexLookup::Int(lookup)) => { - // Use the HashMap for O(1) average lookup + // Constant-time lookup using hash map *lookup.get(&key).unwrap_or_else(|| { panic!("Frame index: integer key {} not found in Int index", key) }) } (RowIndex::Range(range), RowIndexLookup::None) => { - // Direct mapping for Range, but check bounds + // Direct-range mapping with boundary check if range.contains(&key) { - // For a range S..E, the key `k` corresponds to physical row `k - S`. - // Frame::new ensures Range is always 0..N, so S=0. - debug_assert_eq!(range.start, 0, "Range index expected to start at 0"); - key // physical row = key - 0 + // Since Range always starts at 0, the physical index equals the key + key } else { panic!( "Frame index: integer key {} out of bounds for Range index {:?}", @@ -601,23 +545,17 @@ impl RowIndexLookupHelper for Frame { (RowIndex::Date(_), _) => { panic!("Frame index: incompatible key type usize for Date index") } - // Ensure state consistency + // Fallback panic on inconsistent internal index state #[allow(unreachable_patterns)] - (RowIndex::Int(_), RowIndexLookup::None) - | (RowIndex::Int(_), RowIndexLookup::Date(_)) - | (RowIndex::Date(_), RowIndexLookup::Int(_)) - | (RowIndex::Date(_), RowIndexLookup::None) - | (RowIndex::Range(_), RowIndexLookup::Int(_)) - | (RowIndex::Range(_), RowIndexLookup::Date(_)) => { + _ => { panic!( - "Frame index: inconsistent internal index state (lookup type mismatch for index type with usize key)" + "Frame index: inconsistent internal index state (lookup type mismatch for usize key)" ) } } } } -/// Implementation for `NaiveDate` keys (used for `RowIndex::Date`). impl RowIndexLookupHelper for Frame { fn lookup_row_index( key: NaiveDate, @@ -626,7 +564,7 @@ impl RowIndexLookupHelper for Frame { ) -> usize { match (index_values, index_lookup) { (RowIndex::Date(_), RowIndexLookup::Date(lookup)) => { - // Use the HashMap for O(1) average lookup + // Constant-time lookup via hash map *lookup.get(&key).unwrap_or_else(|| { panic!("Frame index: date key {} not found in Date index", key) }) @@ -634,21 +572,18 @@ impl RowIndexLookupHelper for Frame { (RowIndex::Int(_), _) | (RowIndex::Range(_), _) => { panic!("Frame index: incompatible key type NaiveDate for Int or Range index") } - // Ensure state consistency + // Fallback panic on inconsistent internal index state #[allow(unreachable_patterns)] - (RowIndex::Date(_), RowIndexLookup::None) - | (RowIndex::Date(_), RowIndexLookup::Int(_)) - | (RowIndex::Int(_), RowIndexLookup::Date(_)) - | (RowIndex::Range(_), RowIndexLookup::Date(_)) => { + _ => { panic!( - "Frame index: inconsistent internal index state (lookup type mismatch for index type with NaiveDate key)" + "Frame index: inconsistent internal index state (lookup type mismatch for NaiveDate key)" ) } } } } -// --- Row View Structs --- +// Row view types for frame rows /// An immutable view of a single row in a `Frame`. Allows access via `[]`. pub struct FrameRowView<'a, T: Clone + PartialEq> { @@ -658,7 +593,7 @@ pub struct FrameRowView<'a, T: Clone + PartialEq> { impl<'a, T: Clone + PartialEq + fmt::Debug> fmt::Debug for FrameRowView<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Collect references to the data for display + // Gather row data for debug output let row_data: Vec<&T> = (0..self.frame.cols()) .map(|c| &self.frame.matrix[(self.physical_row_idx, c)]) .collect(); @@ -671,28 +606,31 @@ impl<'a, T: Clone + PartialEq + fmt::Debug> fmt::Debug for FrameRowView<'a, T> { } impl<'a, T: Clone + PartialEq> FrameRowView<'a, T> { - /// Get an element by its physical column index. Panics if out of bounds (delegated to Matrix). + /// Returns a reference to the element at the given physical column index. + /// Panics with a descriptive message if `col_idx` is out of bounds. pub fn get_by_index(&self, col_idx: usize) -> &T { - // Add check here for better error message if desired if col_idx >= self.frame.cols() { panic!( - "FrameRowView::get_by_index: column index {} out of bounds for frame with {} columns", + "FrameRowView::get_by_index: column index {} out of bounds (frame has {} columns)", col_idx, self.frame.cols() ); } &self.frame.matrix[(self.physical_row_idx, col_idx)] } - /// Get an element by its column name. Panics if name not found. + + /// Returns a reference to the element in the column named `col_name`. + /// Panics if the column name is not found. pub fn get(&self, col_name: &str) -> &T { let idx = self .frame .column_index(col_name) - .unwrap_or_else(|| panic!("FrameRowView::get: column name '{}' not found", col_name)); + .unwrap_or_else(|| panic!("FrameRowView::get: unknown column '{}'", col_name)); self.get_by_index(idx) } } -// Indexing by column name (&str) + +// Immutable indexing by column name impl<'a, T: Clone + PartialEq> Index<&str> for FrameRowView<'a, T> { type Output = T; #[inline] @@ -700,7 +638,8 @@ impl<'a, T: Clone + PartialEq> Index<&str> for FrameRowView<'a, T> { self.get(col_name) } } -// Indexing by physical column index (usize) + +// Immutable indexing by physical column index impl<'a, T: Clone + PartialEq> Index for FrameRowView<'a, T> { type Output = T; #[inline] @@ -709,7 +648,8 @@ impl<'a, T: Clone + PartialEq> Index for FrameRowView<'a, T> { } } -/// A mutable view of a single row in a `Frame`. Allows access/mutation via `[]` or methods. +/// A mutable view of a single row in a `Frame`. +/// Supports indexed access and mutation via methods or `[]` operators. pub struct FrameRowViewMut<'a, T: Clone + PartialEq> { frame: &'a mut Frame, physical_row_idx: usize, @@ -717,102 +657,101 @@ pub struct FrameRowViewMut<'a, T: Clone + PartialEq> { impl<'a, T: Clone + PartialEq + fmt::Debug> fmt::Debug for FrameRowViewMut<'a, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - // Avoid borrowing frame immutably while already borrowed mutably + // Only display the row index and column names to avoid conflicting borrows f.debug_struct("FrameRowViewMut") .field("physical_row_idx", &self.physical_row_idx) - .field("columns", &self.frame.column_names) // Reading column names is fine - // Cannot easily display data without another borrow + .field("columns", &self.frame.column_names) .finish() } } impl<'a, T: Clone + PartialEq> FrameRowViewMut<'a, T> { - /// Get a mutable reference to an element by its physical column index. Panics if out of bounds (delegated to Matrix). + /// Returns a mutable reference to the element at the given physical column index. + /// Panics if `col_idx` is out of bounds. pub fn get_by_index_mut(&mut self, col_idx: usize) -> &mut T { - // Add check here for better error message if desired - let num_cols = self.frame.cols(); // Borrow checker friendly + let num_cols = self.frame.cols(); if col_idx >= num_cols { panic!( - "FrameRowViewMut::get_by_index_mut: column index {} out of bounds for frame with {} columns", + "FrameRowViewMut::get_by_index_mut: column index {} out of bounds (frame has {} columns)", col_idx, num_cols ); } &mut self.frame.matrix[(self.physical_row_idx, col_idx)] } - /// Get a mutable reference to an element by its column name. Panics if name not found. + + /// Returns a mutable reference to the element in the column named `col_name`. + /// Panics if the column name is not found. pub fn get_mut(&mut self, col_name: &str) -> &mut T { - // Need to get index first, as borrow checker prevents simultaneous borrow for index lookup and mutable access - // Clone name to avoid borrow conflict if name comes from frame itself (unlikely but possible) - let col_name_owned = col_name.to_string(); - let idx = self.frame.column_index(&col_name_owned).unwrap_or_else(|| { - panic!( - "FrameRowViewMut::get_mut: column name '{}' not found", - col_name_owned - ) - }); + let idx = self + .frame + .column_index(col_name) + .unwrap_or_else(|| panic!("FrameRowViewMut::get_mut: unknown column '{}'", col_name)); self.get_by_index_mut(idx) } - /// Set the value of an element by its physical column index. Panics if out of bounds. + + /// Sets the element at the given physical column index to `value`. + /// Panics if `col_idx` is out of bounds. pub fn set_by_index(&mut self, col_idx: usize, value: T) { - // get_by_index_mut already checks bounds *self.get_by_index_mut(col_idx) = value; } - /// Set the value of an element by its column name. Panics if name not found. + + /// Sets the element in the column named `col_name` to `value`. + /// Panics if the column name is not found. pub fn set(&mut self, col_name: &str, value: T) { - // get_mut already finds index and calls get_by_index_mut (which checks bounds) *self.get_mut(col_name) = value; } - // --- Read-only access needed for Index trait --- - /// Internal helper for immutable access by index (needed for Index impl). + // Read-only helpers for Index implementations + + /// Immutable reference to the element at the given physical column index. fn get_by_index_ref(&self, col_idx: usize) -> &T { - // Add check here for better error message if desired if col_idx >= self.frame.cols() { panic!( - "FrameRowViewMut::get_by_index_ref: column index {} out of bounds for frame with {} columns", + "FrameRowViewMut::get_by_index_ref: column index {} out of bounds (frame has {} columns)", col_idx, self.frame.cols() ); } &self.frame.matrix[(self.physical_row_idx, col_idx)] } - /// Internal helper for immutable access by name (needed for Index impl). + + /// Immutable reference to the element in the column named `col_name`. fn get_ref(&self, col_name: &str) -> &T { - let idx = self.frame.column_index(col_name).unwrap_or_else(|| { - panic!( - "FrameRowViewMut::get_ref: column name '{}' not found", - col_name - ) - }); + let idx = self + .frame + .column_index(col_name) + .unwrap_or_else(|| panic!("FrameRowViewMut::get_ref: unknown column '{}'", col_name)); self.get_by_index_ref(idx) } } -// Read-only indexing by column name (&str) + +// Immutable indexing by column name impl<'a, T: Clone + PartialEq> Index<&str> for FrameRowViewMut<'a, T> { type Output = T; #[inline] fn index(&self, col_name: &str) -> &Self::Output { - // Must use the read-only helper due to borrow rules self.get_ref(col_name) } } -// Read-only indexing by physical column index (usize) + +// Immutable indexing by physical column index impl<'a, T: Clone + PartialEq> Index for FrameRowViewMut<'a, T> { type Output = T; #[inline] fn index(&self, col_idx: usize) -> &Self::Output { - // Must use the read-only helper due to borrow rules self.get_by_index_ref(col_idx) } } -// Mutable indexing by column name (&str) + +// Mutable indexing by column name impl<'a, T: Clone + PartialEq> IndexMut<&str> for FrameRowViewMut<'a, T> { #[inline] fn index_mut(&mut self, col_name: &str) -> &mut Self::Output { self.get_mut(col_name) } } -// Mutable indexing by physical column index (usize) + +// Mutable indexing by physical column index impl<'a, T: Clone + PartialEq> IndexMut for FrameRowViewMut<'a, T> { #[inline] fn index_mut(&mut self, col_idx: usize) -> &mut Self::Output { @@ -820,8 +759,7 @@ impl<'a, T: Clone + PartialEq> IndexMut for FrameRowViewMut<'a, T> { } } -/* ---------- Frame Indexing Implementation ---------- */ -/// Allows accessing a column's data as a slice using `frame["col_name"]`. +/// Enables immutable access to a column's data via `frame["col_name"]`. impl Index<&str> for Frame { type Output = [T]; #[inline] @@ -829,7 +767,8 @@ impl Index<&str> for Frame { self.column(name) } } -/// Allows mutating a column's data as a slice using `frame["col_name"]`. + +/// Enables mutable access to a column's data via `frame["col_name"]`. impl IndexMut<&str> for Frame { #[inline] fn index_mut(&mut self, name: &str) -> &mut Self::Output { @@ -837,8 +776,9 @@ impl IndexMut<&str> for Frame { } } -/* ---------- Element-wise numerical ops ---------- */ -/// Macro to implement element-wise binary operations (+, -, *, /) for Frames. +/* ---------- Element-wise Arithmetic Operations ---------- */ +/// Generates implementations for element-wise arithmetic (`+`, `-`, `*`, `/`) on `Frame`. +/// Panics if column labels or row indices differ between operands. macro_rules! impl_elementwise_frame_op { ($OpTrait:ident, $method:ident) => { impl<'a, 'b, T> std::ops::$OpTrait<&'b Frame> for &'a Frame @@ -846,11 +786,12 @@ macro_rules! impl_elementwise_frame_op { T: Clone + PartialEq + std::ops::$OpTrait, { type Output = Frame; + fn $method(self, rhs: &'b Frame) -> Frame { - // 1. Check for compatibility + // Verify matching schema if self.column_names != rhs.column_names { panic!( - "Element-wise op ({}): column names mismatch. Left: {:?}, Right: {:?}", + "Element-wise {}: column names do not match. Left: {:?}, Right: {:?}", stringify!($method), self.column_names, rhs.column_names @@ -858,22 +799,22 @@ macro_rules! impl_elementwise_frame_op { } if self.index != rhs.index { panic!( - "Element-wise op ({}): row indices mismatch. Left: {:?}, Right: {:?}", + "Element-wise {}: row indices do not match. Left: {:?}, Right: {:?}", stringify!($method), self.index, rhs.index ); } - // 2. Perform operation on underlying matrices + // Apply the matrix operation let result_matrix = (&self.matrix).$method(&rhs.matrix); - // 3. Construct the new Frame - // Clone index unless it's Range, then pass None to use default construction + // Determine index for the result let new_index = match self.index { - RowIndex::Range(_) => None, // Frame::new handles None correctly - _ => Some(self.index.clone()), // Clone Int or Date index + RowIndex::Range(_) => None, + _ => Some(self.index.clone()), }; + Frame::new(result_matrix, self.column_names.clone(), new_index) } } @@ -884,17 +825,19 @@ impl_elementwise_frame_op!(Sub, sub); impl_elementwise_frame_op!(Mul, mul); impl_elementwise_frame_op!(Div, div); -/* ---------- Boolean-specific bitwise ops ---------- */ -/// Macro to implement element-wise binary bitwise operations (&, |, ^) for Frames of bool. +/* ---------- Boolean Bitwise Operations ---------- */ +/// Generates implementations for element-wise bitwise operations (`&`, `|`, `^`) on `Frame`. +/// Panics if column labels or row indices differ between operands. macro_rules! impl_bitwise_frame_op { ($OpTrait:ident, $method:ident) => { impl<'a, 'b> std::ops::$OpTrait<&'b Frame> for &'a Frame { type Output = Frame; + fn $method(self, rhs: &'b Frame) -> Frame { - // 1. Check for compatibility + // Verify matching schema if self.column_names != rhs.column_names { panic!( - "Bitwise op ({}): column names mismatch. Left: {:?}, Right: {:?}", + "Bitwise {}: column names do not match. Left: {:?}, Right: {:?}", stringify!($method), self.column_names, rhs.column_names @@ -902,21 +845,22 @@ macro_rules! impl_bitwise_frame_op { } if self.index != rhs.index { panic!( - "Bitwise op ({}): row indices mismatch. Left: {:?}, Right: {:?}", + "Bitwise {}: row indices do not match. Left: {:?}, Right: {:?}", stringify!($method), self.index, rhs.index ); } - // 2. Perform operation on underlying matrices + // Apply the matrix operation let result_matrix = (&self.matrix).$method(&rhs.matrix); - // 3. Construct the new Frame + // Determine index for the result let new_index = match self.index { RowIndex::Range(_) => None, _ => Some(self.index.clone()), }; + Frame::new(result_matrix, self.column_names.clone(), new_index) } } @@ -926,19 +870,21 @@ impl_bitwise_frame_op!(BitAnd, bitand); impl_bitwise_frame_op!(BitOr, bitor); impl_bitwise_frame_op!(BitXor, bitxor); -/// Implements element-wise logical NOT (!) for Frames of bool. Consumes the frame. +/// Implements logical NOT (`!`) for `Frame`, consuming the frame. impl Not for Frame { type Output = Frame; + fn not(self) -> Frame { - // Perform operation on underlying matrix (Matrix::not consumes the matrix) + // Apply NOT to the underlying matrix let result_matrix = !self.matrix; - // Construct the new Frame (index can be moved as self is consumed) + // Determine index for the result let new_index = match self.index { RowIndex::Range(_) => None, - _ => Some(self.index), // Move index + _ => Some(self.index), }; - Frame::new(result_matrix, self.column_names, new_index) // Move column names + + Frame::new(result_matrix, self.column_names, new_index) } } @@ -1399,7 +1345,7 @@ mod tests { let _ = row_view[2]; // Access column index 2 (out of bounds) } #[test] - #[should_panic(expected = "column name 'C' not found")] + #[should_panic(expected = "unknown column 'C'")] fn test_row_view_name_panic() { let frame = create_test_frame_f64(); let row_view = frame.get_row(0); @@ -1420,7 +1366,7 @@ mod tests { row_view_mut[2] = 0.0; // Access column index 2 (out of bounds) } #[test] - #[should_panic(expected = "column name 'C' not found")] + #[should_panic(expected = "unknown column 'C'")] fn test_row_view_mut_name_panic() { let mut frame = create_test_frame_f64(); let mut row_view_mut = frame.get_row_mut(0); @@ -1441,7 +1387,7 @@ mod tests { row_view_mut.set_by_index(3, 0.0); } #[test] - #[should_panic(expected = "column name 'C' not found")] // Panic from view set -> get_mut + #[should_panic(expected = "unknown column 'C'")] // Panic from view set -> get_mut fn test_row_view_mut_set_panic() { let mut frame = create_test_frame_f64(); let mut row_view_mut = frame.get_row_mut(0); @@ -1782,7 +1728,7 @@ mod tests { } #[test] - #[should_panic(expected = "row indices mismatch")] + #[should_panic(expected = "row indices do not match")] fn frame_elementwise_ops_panic_index() { let frame1 = create_test_frame_f64(); // Range index 0..3 let matrix2 = create_test_matrix_f64_alt(); // 3 rows @@ -1791,7 +1737,7 @@ mod tests { let _ = &frame1 + &frame2; // Should panic due to index mismatch } #[test] - #[should_panic(expected = "column names mismatch")] + #[should_panic(expected = "column names do not match")] fn frame_elementwise_ops_panic_cols() { let frame1 = create_test_frame_f64(); // Columns ["A", "B"] let matrix2 = create_test_matrix_f64_alt(); @@ -1799,7 +1745,7 @@ mod tests { let _ = &frame1 + &frame2; // Should panic due to column name mismatch } #[test] - #[should_panic(expected = "row indices mismatch")] + #[should_panic(expected = "row indices do not match")] fn frame_bitwise_ops_panic_index() { let frame1 = create_test_frame_bool(); // Range index 0..2 let matrix2 = create_test_matrix_bool_alt(); // 2 rows @@ -1808,7 +1754,7 @@ mod tests { let _ = &frame1 & &frame2; } #[test] - #[should_panic(expected = "column names mismatch")] + #[should_panic(expected = "column names do not match")] fn frame_bitwise_ops_panic_cols() { let frame1 = create_test_frame_bool(); // Cols P, Q let matrix2 = create_test_matrix_bool_alt(); diff --git a/src/matrix/boolops.rs b/src/matrix/boolops.rs index cb3c5e5..ddfa451 100644 --- a/src/matrix/boolops.rs +++ b/src/matrix/boolops.rs @@ -73,7 +73,6 @@ impl BoolOps for BoolMatrix { self.data().iter().filter(|&&v| v).count() } } -// use macros to generate the implementations for BitAnd, BitOr, BitXor, and Not #[cfg(test)] mod tests { diff --git a/src/matrix/mat.rs b/src/matrix/mat.rs index 8186195..5670267 100644 --- a/src/matrix/mat.rs +++ b/src/matrix/mat.rs @@ -85,20 +85,22 @@ impl Matrix { "column index out of bounds" ); if c1 == c2 { - return; // No-op if indices are the same + // Indices are equal; no operation required + return; } - // Loop through each row + // Iterate over each row to swap corresponding elements for r in 0..self.rows { - // Calculate the flat index for the element in row r, column c1 + // Compute the one-dimensional index for (row r, column c1) let idx1 = c1 * self.rows + r; - // Calculate the flat index for the element in row r, column c2 + // Compute the one-dimensional index for (row r, column c2) let idx2 = c2 * self.rows + r; - // Swap the elements directly in the data vector + // Exchange the two elements in the internal data buffer self.data.swap(idx1, idx2); } } + /// Deletes a column from the matrix. pub fn delete_column(&mut self, col: usize) { assert!(col < self.cols, "column index out of bounds"); @@ -144,35 +146,45 @@ impl Matrix { impl Index<(usize, usize)> for Matrix { type Output = T; + #[inline] fn index(&self, (r, c): (usize, usize)) -> &T { + // Validate that the requested indices are within bounds assert!(r < self.rows && c < self.cols, "index out of bounds"); + // Compute column-major offset and return reference &self.data[c * self.rows + r] } } + impl IndexMut<(usize, usize)> for Matrix { #[inline] fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut T { + // Validate that the requested indices are within bounds assert!(r < self.rows && c < self.cols, "index out of bounds"); + // Compute column-major offset and return mutable reference &mut self.data[c * self.rows + r] } } -/// A view of one row +/// Represents an immutable view of a single row in the matrix. pub struct MatrixRow<'a, T> { matrix: &'a Matrix, row: usize, } + impl<'a, T> MatrixRow<'a, T> { + /// Returns a reference to the element at the given column in this row. pub fn get(&self, c: usize) -> &T { &self.matrix[(self.row, c)] } + + /// Returns an iterator over all elements in this row. pub fn iter(&self) -> impl Iterator { (0..self.matrix.cols).map(move |c| &self.matrix[(self.row, c)]) } } -/// Macro to generate element‐wise impls for +, -, *, / +/// Generates element-wise arithmetic implementations for matrices. macro_rules! impl_elementwise_op { ($OpTrait:ident, $method:ident, $op:tt) => { impl<'a, 'b, T> std::ops::$OpTrait<&'b Matrix> for &'a Matrix @@ -182,8 +194,10 @@ macro_rules! impl_elementwise_op { type Output = Matrix; fn $method(self, rhs: &'b Matrix) -> Matrix { + // Ensure both matrices have identical dimensions assert_eq!(self.rows, rhs.rows, "row count mismatch"); assert_eq!(self.cols, rhs.cols, "col count mismatch"); + // Apply the operation element-wise and collect into a new matrix let data = self .data .iter() @@ -197,7 +211,7 @@ macro_rules! impl_elementwise_op { }; } -// invoke it 4 times: +// Instantiate element-wise addition, subtraction, multiplication, and division impl_elementwise_op!(Add, add, +); impl_elementwise_op!(Sub, sub, -); impl_elementwise_op!(Mul, mul, *); @@ -208,16 +222,17 @@ pub type BoolMatrix = Matrix; pub type IntMatrix = Matrix; pub type StringMatrix = Matrix; -// implement bit ops - and, or, xor, not -- using Macros - +/// Generates element-wise bitwise operations for boolean matrices. macro_rules! impl_bitwise_op { ($OpTrait:ident, $method:ident, $op:tt) => { impl<'a, 'b> std::ops::$OpTrait<&'b Matrix> for &'a Matrix { type Output = Matrix; fn $method(self, rhs: &'b Matrix) -> Matrix { + // Ensure both matrices have identical dimensions assert_eq!(self.rows, rhs.rows, "row count mismatch"); assert_eq!(self.cols, rhs.cols, "col count mismatch"); + // Apply the bitwise operation element-wise let data = self .data .iter() @@ -230,6 +245,8 @@ macro_rules! impl_bitwise_op { } }; } + +// Instantiate bitwise AND, OR, and XOR for boolean matrices impl_bitwise_op!(BitAnd, bitand, &); impl_bitwise_op!(BitOr, bitor, |); impl_bitwise_op!(BitXor, bitxor, ^); @@ -238,6 +255,7 @@ impl Not for Matrix { type Output = Matrix; fn not(self) -> Matrix { + // Invert each boolean element in the matrix let data = self.data.iter().map(|&v| !v).collect(); Matrix { rows: self.rows, @@ -247,12 +265,12 @@ impl Not for Matrix { } } -/// Axis along which to apply a reduction. +/// Specifies the axis along which to perform a reduction operation. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum Axis { - /// Operate column‑wise (vertical). + /// Apply reduction along columns (vertical axis). Col, - /// Operate row‑wise (horizontal). + /// Apply reduction along rows (horizontal axis). Row, } diff --git a/src/matrix/seriesops.rs b/src/matrix/seriesops.rs index 763614c..bc34a1f 100644 --- a/src/matrix/seriesops.rs +++ b/src/matrix/seriesops.rs @@ -1,6 +1,6 @@ use crate::matrix::{Axis, BoolMatrix, FloatMatrix}; -/// “Series–like” helpers that work along a single axis. +/// "Series-like" helpers that work along a single axis. /// /// *All* the old methods (`sum_*`, `prod_*`, `is_nan`, …) are exposed /// through this trait, so nothing needs to stay on an `impl Matrix`; @@ -100,7 +100,7 @@ impl SeriesOps for FloatMatrix { } fn cumsum_horizontal(&self) -> FloatMatrix { - // 1. Store row-wise cumulative sums temporarily + // Compute cumulative sums for each row and store in a temporary buffer let mut row_results: Vec> = Vec::with_capacity(self.rows()); for r in 0..self.rows() { let mut row_data = Vec::with_capacity(self.cols()); @@ -115,16 +115,15 @@ impl SeriesOps for FloatMatrix { row_results.push(row_data); } - // 2. Build the final data vector in column-major order + // Assemble the final data vector in column-major format let mut final_data = Vec::with_capacity(self.rows() * self.cols()); for c in 0..self.cols() { for r in 0..self.rows() { - // Get the element from row 'r', column 'c' of the row_results + // Extract the element at (r, c) from the temporary row-wise results final_data.push(row_results[r][c]); } } - // 3. Construct the matrix using the correctly ordered data FloatMatrix::from_vec(final_data, self.rows(), self.cols()) } @@ -146,7 +145,7 @@ impl SeriesOps for FloatMatrix { #[cfg(test)] mod tests { use super::*; - + // Helper function to create a FloatMatrix for SeriesOps testing fn create_float_test_matrix() -> FloatMatrix { // 3x3 matrix (column-major) with some NaNs @@ -361,4 +360,4 @@ mod tests { assert_eq!(matrix.count_nan_horizontal(), vec![2, 2]); assert_eq!(matrix.is_nan().data(), &[true, true, true, true]); } -} \ No newline at end of file +} diff --git a/src/utils/bdates.rs b/src/utils/bdates.rs index 4664d0f..c3818ba 100644 --- a/src/utils/bdates.rs +++ b/src/utils/bdates.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use std::error::Error; use std::hash::Hash; use std::result::Result; -use std::str::FromStr; // Import FromStr trait +use std::str::FromStr; /// Represents the frequency at which business dates should be generated. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -23,8 +23,8 @@ pub enum BDateFreq { /// is selected for the frequency. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum AggregationType { - Start, // Indicates picking the first valid business date in a group's period. - End, // Indicates picking the last valid business day in a group's period. + Start, // Select the first valid business day in the period + End, // Select the last valid business day in the period } impl BDateFreq { @@ -40,7 +40,7 @@ impl BDateFreq { /// /// Returns an error if the string does not match any known frequency. pub fn from_string(freq: String) -> Result> { - // Use the FromStr implementation directly + // Delegate parsing to the FromStr implementation freq.parse() } @@ -53,7 +53,7 @@ impl BDateFreq { BDateFreq::WeeklyMonday => "W", BDateFreq::MonthStart => "M", BDateFreq::QuarterStart => "Q", - BDateFreq::YearStart => "Y", // Changed to "Y" + BDateFreq::YearStart => "Y", BDateFreq::MonthEnd => "ME", BDateFreq::QuarterEnd => "QE", BDateFreq::WeeklyFriday => "WF", @@ -112,11 +112,11 @@ impl FromStr for BDateFreq { "W" | "WS" => BDateFreq::WeeklyMonday, "M" | "MS" => BDateFreq::MonthStart, "Q" | "QS" => BDateFreq::QuarterStart, - "Y" | "A" | "AS" | "YS" => BDateFreq::YearStart, // Added Y, YS, A, AS aliases + "Y" | "A" | "AS" | "YS" => BDateFreq::YearStart, // Support standard aliases for year start "ME" => BDateFreq::MonthEnd, "QE" => BDateFreq::QuarterEnd, "WF" => BDateFreq::WeeklyFriday, - "YE" | "AE" => BDateFreq::YearEnd, // Added AE alias + "YE" | "AE" => BDateFreq::YearEnd, // Include 'AE' alias for year end _ => return Err(format!("Invalid frequency specified: {}", freq).into()), }; Ok(r) @@ -131,21 +131,19 @@ pub struct BDatesList { start_date_str: String, end_date_str: String, freq: BDateFreq, - // Optional: Cache the generated list to avoid re-computation? - // For now, we recompute each time list(), count(), or groups() is called. + // TODO: cache the generated date list to reduce repeated computation. + // Currently, list(), count(), and groups() regenerate the list on every invocation. // cached_list: Option>, } -// Helper enum to represent the key for grouping dates into periods. -// Deriving traits for comparison and hashing allows using it as a HashMap key -// and for sorting groups chronologically. +// Enumeration of period keys used for grouping dates. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] enum GroupKey { - Daily(NaiveDate), // Group by the specific date (for Daily frequency) - Weekly(i32, u32), // Group by year and ISO week number - Monthly(i32, u32), // Group by year and month (1-12) - Quarterly(i32, u32), // Group by year and quarter (1-4) - Yearly(i32), // Group by year + Daily(NaiveDate), // Daily grouping: use the exact date + Weekly(i32, u32), // Weekly grouping: use year and ISO week number + Monthly(i32, u32), // Monthly grouping: use year and month (1-12) + Quarterly(i32, u32), // Quarterly grouping: use year and quarter (1-4) + Yearly(i32), // Yearly grouping: use year } /// Represents a collection of business dates generated according to specific rules. @@ -168,23 +166,23 @@ enum GroupKey { /// use rustframe::utils::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name /// /// fn main() -> Result<(), Box> { -/// let start_date = "2023-11-01".to_string(); // Wednesday -/// let end_date = "2023-11-07".to_string(); // Tuesday -/// let freq = BDateFreq::Daily; +/// let start_date = "2023-11-01".to_string(); // Wednesday +/// let end_date = "2023-11-07".to_string(); // Tuesday +/// let freq = BDateFreq::Daily; /// -/// let bdates = BDatesList::new(start_date, end_date, freq); +/// let bdates = BDatesList::new(start_date, end_date, freq); /// -/// let expected_dates = vec![ -/// NaiveDate::from_ymd_opt(2023, 11, 1).unwrap(), // Wed -/// NaiveDate::from_ymd_opt(2023, 11, 2).unwrap(), // Thu -/// NaiveDate::from_ymd_opt(2023, 11, 3).unwrap(), // Fri -/// NaiveDate::from_ymd_opt(2023, 11, 6).unwrap(), // Mon -/// NaiveDate::from_ymd_opt(2023, 11, 7).unwrap(), // Tue -/// ]; +/// let expected_dates = vec![ +/// NaiveDate::from_ymd_opt(2023, 11, 1).unwrap(), // Wed +/// NaiveDate::from_ymd_opt(2023, 11, 2).unwrap(), // Thu +/// NaiveDate::from_ymd_opt(2023, 11, 3).unwrap(), // Fri +/// NaiveDate::from_ymd_opt(2023, 11, 6).unwrap(), // Mon +/// NaiveDate::from_ymd_opt(2023, 11, 7).unwrap(), // Tue +/// ]; /// -/// assert_eq!(bdates.list()?, expected_dates); -/// assert_eq!(bdates.count()?, 5); -/// Ok(()) +/// assert_eq!(bdates.list()?, expected_dates); +/// assert_eq!(bdates.count()?, 5); +/// Ok(()) /// } /// ``` /// @@ -196,25 +194,25 @@ enum GroupKey { /// use rustframe::utils::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name /// /// fn main() -> Result<(), Box> { -/// let start_date = "2024-02-28".to_string(); // Wednesday -/// let freq = BDateFreq::WeeklyFriday; -/// let n_periods = 3; +/// let start_date = "2024-02-28".to_string(); // Wednesday +/// let freq = BDateFreq::WeeklyFriday; +/// let n_periods = 3; /// -/// let bdates = BDatesList::from_n_periods(start_date, freq, n_periods)?; +/// let bdates = BDatesList::from_n_periods(start_date, freq, n_periods)?; /// -/// // The first Friday on or after 2024-02-28 is Mar 1. -/// // The next two Fridays are Mar 8 and Mar 15. -/// let expected_dates = vec![ -/// NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(), -/// NaiveDate::from_ymd_opt(2024, 3, 8).unwrap(), -/// NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(), -/// ]; +/// // The first Friday on or after 2024-02-28 is Mar 1. +/// // The next two Fridays are Mar 8 and Mar 15. +/// let expected_dates = vec![ +/// NaiveDate::from_ymd_opt(2024, 3, 1).unwrap(), +/// NaiveDate::from_ymd_opt(2024, 3, 8).unwrap(), +/// NaiveDate::from_ymd_opt(2024, 3, 15).unwrap(), +/// ]; /// -/// assert_eq!(bdates.list()?, expected_dates); -/// assert_eq!(bdates.count()?, 3); -/// assert_eq!(bdates.start_date_str(), "2024-02-28"); // Keeps original start string -/// assert_eq!(bdates.end_date_str(), "2024-03-15"); // End date is the last generated date -/// Ok(()) +/// assert_eq!(bdates.list()?, expected_dates); +/// assert_eq!(bdates.count()?, 3); +/// assert_eq!(bdates.start_date_str(), "2024-02-28"); // Keeps original start string +/// assert_eq!(bdates.end_date_str(), "2024-03-15"); // End date is the last generated date +/// Ok(()) /// } /// ``` /// @@ -226,20 +224,20 @@ enum GroupKey { /// use rustframe::utils::{BDatesList, BDateFreq}; // Replace bdates with your actual crate/module name /// /// fn main() -> Result<(), Box> { -/// let start_date = "2023-11-20".to_string(); // Mon, Week 47 -/// let end_date = "2023-12-08".to_string(); // Fri, Week 49 -/// let freq = BDateFreq::WeeklyMonday; +/// let start_date = "2023-11-20".to_string(); // Mon, Week 47 +/// let end_date = "2023-12-08".to_string(); // Fri, Week 49 +/// let freq = BDateFreq::WeeklyMonday; /// -/// let bdates = BDatesList::new(start_date, end_date, freq); +/// let bdates = BDatesList::new(start_date, end_date, freq); /// -/// // Mondays in range: Nov 20, Nov 27, Dec 4 -/// let groups = bdates.groups()?; +/// // Mondays in range: Nov 20, Nov 27, Dec 4 +/// let groups = bdates.groups()?; /// -/// assert_eq!(groups.len(), 3); // One group per week containing a Monday -/// assert_eq!(groups[0], vec![NaiveDate::from_ymd_opt(2023, 11, 20).unwrap()]); // Week 47 -/// assert_eq!(groups[1], vec![NaiveDate::from_ymd_opt(2023, 11, 27).unwrap()]); // Week 48 -/// assert_eq!(groups[2], vec![NaiveDate::from_ymd_opt(2023, 12, 4).unwrap()]); // Week 49 -/// Ok(()) +/// assert_eq!(groups.len(), 3); // One group per week containing a Monday +/// assert_eq!(groups[0], vec![NaiveDate::from_ymd_opt(2023, 11, 20).unwrap()]); // Week 47 +/// assert_eq!(groups[1], vec![NaiveDate::from_ymd_opt(2023, 11, 27).unwrap()]); // Week 48 +/// assert_eq!(groups[2], vec![NaiveDate::from_ymd_opt(2023, 12, 4).unwrap()]); // Week 49 +/// Ok(()) /// } /// ``` impl BDatesList { @@ -286,19 +284,19 @@ impl BDatesList { let start_date = NaiveDate::parse_from_str(&start_date_str, "%Y-%m-%d")?; - // Use the generator to find all the dates + // Instantiate the date generator to compute the sequence of business dates. let generator = BDatesGenerator::new(start_date, freq, n_periods)?; let dates: Vec = generator.collect(); - // Should always have at least one date if n_periods > 0 and generator construction succeeded + // Confirm that the generator returned at least one date when n_periods > 0. let last_date = dates .last() - .ok_or("Generator failed to produce dates even though n_periods > 0")?; + .ok_or("Generator failed to produce dates for the specified periods")?; let end_date_str = last_date.format("%Y-%m-%d").to_string(); Ok(BDatesList { - start_date_str, // Keep the original start date string + start_date_str, end_date_str, freq, }) @@ -312,7 +310,7 @@ impl BDatesList { /// /// Returns an error if the start or end date strings cannot be parsed. pub fn list(&self) -> Result, Box> { - // Delegate the core logic to the internal helper function + // Retrieve the list of business dates via the shared helper function. get_bdates_list_with_freq(&self.start_date_str, &self.end_date_str, self.freq) } @@ -320,41 +318,33 @@ impl BDatesList { /// /// # Errors /// - /// Returns an error if the start or end date strings cannot be parsed (as it - /// calls `list` internally). + /// Returns an error if the start or end date strings cannot be parsed. pub fn count(&self) -> Result> { - // Get the list and return its length. Uses map to handle the Result elegantly. + // Compute the total number of business dates by invoking `list()` and returning its length. self.list().map(|list| list.len()) } /// Returns a list of date lists, where each inner list contains dates /// belonging to the same period (determined by frequency). /// - /// The outer list (groups) is sorted by period chronologically, and the - /// inner lists (dates within groups) are also sorted chronologically. - /// - /// For `Daily` frequency, each date forms its own group. For `Weekly` - /// frequencies, grouping is by ISO week number. For `Monthly`, `Quarterly`, - /// and `Yearly` frequencies, grouping is by the respective period. + /// The outer list (groups) is sorted chronologically by period, and the + /// inner lists (dates within each period) are also sorted. /// /// # Errors /// /// Returns an error if the start or end date strings cannot be parsed. pub fn groups(&self) -> Result>, Box> { - // Get the sorted list of all dates first. This sorted order is crucial - // for ensuring the inner vectors (dates within groups) are also sorted - // as we insert into the HashMap. + // Retrieve all business dates in chronological order. let dates = self.list()?; - // Use a HashMap to collect dates into their respective groups. + // Aggregate dates into buckets keyed by period. let mut groups: HashMap> = HashMap::new(); for date in dates { - // Determine the grouping key based on frequency. + // Derive the appropriate GroupKey for the current date based on the configured frequency. let key = match self.freq { BDateFreq::Daily => GroupKey::Daily(date), BDateFreq::WeeklyMonday | BDateFreq::WeeklyFriday => { - // Use ISO week number for consistent weekly grouping across year boundaries let iso_week = date.iso_week(); GroupKey::Weekly(iso_week.year(), iso_week.week()) } @@ -367,30 +357,19 @@ impl BDatesList { BDateFreq::YearStart | BDateFreq::YearEnd => GroupKey::Yearly(date.year()), }; - // Add the current date to the vector corresponding to the determined key. - // entry().or_insert_with() gets a mutable reference to the vector for the key, - // inserting a new empty vector if the key doesn't exist yet. - groups.entry(key).or_insert_with(Vec::new).push(date); // Using or_insert_with is slightly more idiomatic + // Append the date to its period group. + groups.entry(key).or_insert_with(Vec::new).push(date); } - // Convert the HashMap into a vector of (key, vector_of_dates) tuples. + // Transform the group map into a vector of (GroupKey, Vec) tuples. let mut sorted_groups: Vec<(GroupKey, Vec)> = groups.into_iter().collect(); - // Sort the vector of groups by the `GroupKey`. Since `GroupKey` derives `Ord`, - // this sorts the groups chronologically (Yearly < Quarterly < Monthly < Weekly < Daily, - // then by year, quarter, month, week, or date within each category). + // Sort groups chronologically using the derived `Ord` implementation on `GroupKey`. sorted_groups.sort_by(|(k1, _), (k2, _)| k1.cmp(k2)); - // The dates *within* each group (`Vec`) are already sorted - // because they were pushed in the order they appeared in the initially - // sorted `dates` vector obtained from `self.list()`. - // If the source `dates` wasn't guaranteed sorted, or for clarity, - // an inner sort could be added here: - // for (_, dates_in_group) in sorted_groups.iter_mut() { - // dates_in_group.sort(); - // } + // Note: Dates within each group remain sorted due to initial ordered input. - // Extract just the vectors of dates from the sorted tuples, discarding the keys. + // Discard group keys to return only the list of date vectors. let result_groups = sorted_groups.into_iter().map(|(_, dates)| dates).collect(); Ok(result_groups) @@ -435,7 +414,7 @@ impl BDatesList { } } -// --- Business Date Generator (Iterator) --- +// Business date iterator: generates a sequence of business dates for a given frequency and period count. /// An iterator that generates a sequence of business dates based on a start date, /// frequency, and a specified number of periods. @@ -451,22 +430,22 @@ impl BDatesList { /// ```rust /// use chrono::NaiveDate; /// use std::error::Error; -/// use rustframe::utils::{BDatesGenerator, BDateFreq}; +/// use rustframe::utils::{BDatesGenerator, BDateFreq}; /// /// fn main() -> Result<(), Box> { -/// let start = NaiveDate::from_ymd_opt(2023, 12, 28).unwrap(); // Thursday -/// let freq = BDateFreq::MonthEnd; -/// let n_periods = 4; // Dec '23, Jan '24, Feb '24, Mar '24 +/// let start = NaiveDate::from_ymd_opt(2023, 12, 28).unwrap(); // Thursday +/// let freq = BDateFreq::MonthEnd; +/// let n_periods = 4; // Dec '23, Jan '24, Feb '24, Mar '24 /// -/// let mut generator = BDatesGenerator::new(start, freq, n_periods)?; +/// let mut generator = BDatesGenerator::new(start, freq, n_periods)?; /// -/// // First month-end on or after 2023-12-28 is 2023-12-29 -/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2023, 12, 29).unwrap())); -/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 31).unwrap())); -/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 2, 29).unwrap())); // Leap year -/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 3, 29).unwrap())); // Mar 31 is Sun -/// assert_eq!(generator.next(), None); // Exhausted -/// Ok(()) +/// // First month-end on or after 2023-12-28 is 2023-12-29 +/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2023, 12, 29).unwrap())); +/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 1, 31).unwrap())); +/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 2, 29).unwrap())); // Leap year +/// assert_eq!(generator.next(), Some(NaiveDate::from_ymd_opt(2024, 3, 29).unwrap())); // Mar 31 is Sun +/// assert_eq!(generator.next(), None); // Exhausted +/// Ok(()) /// } /// ``` /// @@ -478,31 +457,30 @@ impl BDatesList { /// use rustframe::utils::{BDatesGenerator, BDateFreq}; // Replace bdates with your actual crate/module name /// /// fn main() -> Result<(), Box> { -/// let start = NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(); // Monday -/// let freq = BDateFreq::Daily; -/// let n_periods = 5; +/// let start = NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(); // Monday +/// let freq = BDateFreq::Daily; +/// let n_periods = 5; /// -/// let generator = BDatesGenerator::new(start, freq, n_periods)?; -/// let dates: Vec = generator.collect(); +/// let generator = BDatesGenerator::new(start, freq, n_periods)?; +/// let dates: Vec = generator.collect(); /// -/// let expected_dates = vec![ -/// NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(), // Mon -/// NaiveDate::from_ymd_opt(2024, 4, 30).unwrap(), // Tue -/// NaiveDate::from_ymd_opt(2024, 5, 1).unwrap(), // Wed -/// NaiveDate::from_ymd_opt(2024, 5, 2).unwrap(), // Thu -/// NaiveDate::from_ymd_opt(2024, 5, 3).unwrap(), // Fri -/// ]; +/// let expected_dates = vec![ +/// NaiveDate::from_ymd_opt(2024, 4, 29).unwrap(), // Mon +/// NaiveDate::from_ymd_opt(2024, 4, 30).unwrap(), // Tue +/// NaiveDate::from_ymd_opt(2024, 5, 1).unwrap(), // Wed +/// NaiveDate::from_ymd_opt(2024, 5, 2).unwrap(), // Thu +/// NaiveDate::from_ymd_opt(2024, 5, 3).unwrap(), // Fri +/// ]; /// -/// assert_eq!(dates, expected_dates); -/// Ok(()) +/// assert_eq!(dates, expected_dates); +/// Ok(()) /// } /// ``` #[derive(Debug, Clone)] pub struct BDatesGenerator { freq: BDateFreq, periods_remaining: usize, - // Stores the *next* date to be yielded by the iterator. - // This is None initially or when the iterator is exhausted. + // Next business date candidate to yield; None when iteration is complete. next_date_candidate: Option, } @@ -532,7 +510,8 @@ impl BDatesGenerator { let first_date = if n_periods > 0 { Some(find_first_bdate_on_or_after(start_date, freq)) } else { - None // No dates to generate if n_periods is 0 + // No dates when period count is zero. + None }; Ok(BDatesGenerator { @@ -546,29 +525,29 @@ impl BDatesGenerator { impl Iterator for BDatesGenerator { type Item = NaiveDate; - /// Returns the next business date in the sequence, or `None` if `n_periods` - /// dates have already been generated. + /// Returns the next business date in the sequence, or `None` if the specified + /// number of periods has been generated. fn next(&mut self) -> Option { - // Check if exhausted or if there was no initial date + // Terminate if no periods remain or no initial date is set. if self.periods_remaining == 0 || self.next_date_candidate.is_none() { return None; } - // Get the date to return (unwrap is safe due to the check above) + // Retrieve and store the current date for output. let current_date = self.next_date_candidate.unwrap(); - // Prepare the *next* candidate for the subsequent call + // Compute and queue the subsequent date for the next call. self.next_date_candidate = Some(find_next_bdate(current_date, self.freq)); - // Decrement the count + // Decrement the remaining period count. self.periods_remaining -= 1; - // Return the stored current date + // Yield the current business date. Some(current_date) } } -// --- Internal helper functions (not part of the public API) --- +// Internal helper functions (private implementation) /// Generates the flat list of business dates for the given range and frequency. /// @@ -589,70 +568,58 @@ fn get_bdates_list_with_freq( end_date_str: &str, freq: BDateFreq, ) -> Result, Box> { - // Parse the start and end dates, returning error if parsing fails. + // Parse input date strings; propagate parsing errors. let start_date = NaiveDate::parse_from_str(start_date_str, "%Y-%m-%d")?; let end_date = NaiveDate::parse_from_str(end_date_str, "%Y-%m-%d")?; - // Handle edge case where end date is before start date. + // Return empty list immediately if the date range is invalid. if start_date > end_date { return Ok(Vec::new()); } - // Collect dates based on the specified frequency. + // Generate dates according to the requested frequency. let mut dates = match freq { BDateFreq::Daily => collect_daily(start_date, end_date), BDateFreq::WeeklyMonday => collect_weekly(start_date, end_date, Weekday::Mon), BDateFreq::WeeklyFriday => collect_weekly(start_date, end_date, Weekday::Fri), - BDateFreq::MonthStart => { - collect_monthly(start_date, end_date, /*want_first_day=*/ true) - } - BDateFreq::MonthEnd => { - collect_monthly(start_date, end_date, /*want_first_day=*/ false) - } - BDateFreq::QuarterStart => { - collect_quarterly(start_date, end_date, /*want_first_day=*/ true) - } - BDateFreq::QuarterEnd => { - collect_quarterly(start_date, end_date, /*want_first_day=*/ false) - } - BDateFreq::YearStart => collect_yearly(start_date, end_date, /*want_first_day=*/ true), - BDateFreq::YearEnd => collect_yearly(start_date, end_date, /*want_first_day=*/ false), + BDateFreq::MonthStart => collect_monthly(start_date, end_date, true), + BDateFreq::MonthEnd => collect_monthly(start_date, end_date, false), + BDateFreq::QuarterStart => collect_quarterly(start_date, end_date, true), + BDateFreq::QuarterEnd => collect_quarterly(start_date, end_date, false), + BDateFreq::YearStart => collect_yearly(start_date, end_date, true), + BDateFreq::YearEnd => collect_yearly(start_date, end_date, false), }; - // Filter out any weekend days. While the core logic aims for business days, - // this ensures robustness against edge cases where computed dates might fall - // on a weekend (e.g., first day of month being Saturday). - // Note: This retain is redundant if collect_* functions are correct, but adds safety. - // It's essential for Daily, less so for others if they always return bdays. + // Exclude weekends to ensure only business days remain. dates.retain(|d| is_weekday(*d)); - // Ensure the final list is sorted. The `collect_*` functions generally - // produce sorted output, but an explicit sort guarantees it. + // Guarantee chronological order of the result. dates.sort(); Ok(dates) } -/* ---------------------- Low-Level Date Collection Functions (Internal) ---------------------- */ -// These functions generate dates within a *range* [start_date, end_date] +/* Low-level date collection routines (private) */ -/// Returns all business days (Mon-Fri) day-by-day within the range. +/// Returns all weekdays from `start_date` through `end_date`, inclusive. fn collect_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec { let mut result = Vec::new(); let mut current = start_date; + + // Iterate one day at a time. while current <= end_date { if is_weekday(current) { result.push(current); } - // Use succ_opt() and expect(), assuming valid date range and no overflow in practical scenarios current = current .succ_opt() - .expect("date overflow near end of supported range"); + .expect("Date overflow near end of supported range"); } + result } -/// Returns the specified `target_weekday` in each week within the range. +/// Returns each occurrence of `target_weekday` within the date range. fn collect_weekly( start_date: NaiveDate, end_date: NaiveDate, @@ -660,79 +627,70 @@ fn collect_weekly( ) -> Vec { let mut result = Vec::new(); - // Find the first target_weekday on or after the start date. + // Find the first matching weekday on or after the start date. let mut current = move_to_weekday_on_or_after(start_date, target_weekday); - // Step through the range in 7-day increments. + // Step through each week until exceeding the end date. while current <= end_date { - // Ensure the found date is actually a weekday (should be Mon/Fri but belt-and-suspenders) + // Only include if still a weekday. if is_weekday(current) { result.push(current); } - // Use checked_add_signed for safety, though basic addition is likely fine for 7 days. current = current .checked_add_signed(Duration::days(7)) - .expect("date overflow adding 7 days"); + .expect("Date overflow when advancing by one week"); } + result } -/// Returns either the first or last business day in each month of the range. +/// Returns either the first or last business day of each month in the range. fn collect_monthly( start_date: NaiveDate, end_date: NaiveDate, want_first_day: bool, ) -> Vec { let mut result = Vec::new(); - let mut year = start_date.year(); let mut month = start_date.month(); - // Helper closure to advance year and month by one month. - let next_month = - |(yr, mo): (i32, u32)| -> (i32, u32) { if mo == 12 { (yr + 1, 1) } else { (yr, mo + 1) } }; + // Advance (year, month) by one month. + let next_month = |(yr, mo): (i32, u32)| { + if mo == 12 { (yr + 1, 1) } else { (yr, mo + 1) } + }; - // Iterate month by month from the start date's month up to or past the end date's month. + // Iterate months from the start date until past the end date. loop { - // Compute the candidate date (first or last business day) for the current month. - // Use _opt and expect(), expecting valid month/year combinations within realistic ranges. + // Determine the candidate business date for this month. let candidate = if want_first_day { first_business_day_of_month(year, month) } else { last_business_day_of_month(year, month) }; - // If the candidate is after the end date, we've gone past the range, so stop. + // Stop if the candidate is beyond the allowed range. if candidate > end_date { break; } - // If the candidate is within the specified range [start_date, end_date], add it. - if candidate >= start_date { - // Ensure it's actually a weekday (should be, but adds safety) - if is_weekday(candidate) { - result.push(candidate); - } + // Include candidate if it falls within [start_date, end_date]. + if candidate >= start_date && is_weekday(candidate) { + result.push(candidate); } - // Note: We don't break if candidate < start_date because a later month's candidate - // might be within the range. - // Check if the current month is the last month we should process + // If we've processed the end date's month, terminate. if year > end_date.year() || (year == end_date.year() && month >= end_date.month()) { - // If we just processed the end_date's month, stop. - // Need >= because we need to include the end date's month itself if its candidate is valid. break; } - // Advance to the next month. + // Move to the next month. let (ny, nm) = next_month((year, month)); year = ny; month = nm; - // Safety break: Stop if we have moved clearly past the end date's year. - // This check is technically redundant given the loop condition above, but harmless. + // Safety guard against unexpected infinite loops. if year > end_date.year() + 1 { - break; // Avoid potential infinite loops in unexpected scenarios + break; } } @@ -864,7 +822,7 @@ fn move_to_weekday_on_or_after(date: NaiveDate, target: Weekday) -> NaiveDate { fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate { // Start with the 1st of the month. Use _opt and expect(), assuming valid Y/M. let mut d = NaiveDate::from_ymd_opt(year, month, 1).expect("invalid year-month combination"); - // If it’s Sat/Sun, move forward until we find a weekday. + // If it's Sat/Sun, move forward until we find a weekday. while !is_weekday(d) { // Use succ_opt() and expect(), assuming valid date and no overflow. d = d.succ_opt().expect("date overflow finding first bday"); @@ -872,112 +830,128 @@ fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate { d } -/// Return the latest business day of the given (year, month). +/// Returns the last business day (Monday-Friday) of the specified month. +/// +/// Calculates the number of days in the month, then steps backward from the +/// month's last day until a weekday is found. +/// +/// # Panics +/// Panics if date construction or predecessor operations underflow. fn last_business_day_of_month(year: i32, month: u32) -> NaiveDate { let last_dom = days_in_month(year, month); - // Use _opt and expect(), assuming valid Y/M/D combination. let mut d = - NaiveDate::from_ymd_opt(year, month, last_dom).expect("invalid year-month-day combination"); - // If it’s Sat/Sun, move backward until we find a weekday. + NaiveDate::from_ymd_opt(year, month, last_dom).expect("Invalid year-month-day combination"); while !is_weekday(d) { - // Use pred_opt() and expect(), assuming valid date and no underflow. - d = d.pred_opt().expect("date underflow finding last bday"); + d = d + .pred_opt() + .expect("Date underflow finding last business day"); } d } -/// Returns the number of days in a given month and year. +/// Returns the number of days in the specified month, correctly handling leap years. +/// +/// Determines the first day of the next month and subtracts one day. +/// +/// # Panics +/// Panics if date construction or predecessor operations underflow. fn days_in_month(year: i32, month: u32) -> u32 { - // A common trick: find the first day of the *next* month, then subtract one day. - // This correctly handles leap years. let (ny, nm) = if month == 12 { (year + 1, 1) } else { (year, month + 1) }; - // Use _opt and expect(), assuming valid Y/M combination (start of next month). - let first_of_next = NaiveDate::from_ymd_opt(ny, nm, 1).expect("invalid next year-month"); - // Use pred_opt() and expect(), assuming valid date and no underflow (first of month - 1). + let first_of_next = + NaiveDate::from_ymd_opt(ny, nm, 1).expect("Invalid next year-month combination"); let last_of_this = first_of_next .pred_opt() - .expect("invalid date before first of month"); + .expect("Date underflow computing last day of month"); last_of_this.day() } -/// Converts a month number (1-12) to a quarter number (1-4). +/// Maps a month (1-12) to its corresponding quarter (1-4). +/// +/// # Panics +/// Panics if `m` is outside the range 1-12. fn month_to_quarter(m: u32) -> u32 { match m { 1..=3 => 1, 4..=6 => 2, 7..=9 => 3, 10..=12 => 4, - _ => panic!("Invalid month: {}", m), // Should not happen with valid dates + _ => panic!("Invalid month: {}", m), } } -/// Returns the 1st day of the month that starts a given (year, quarter). +/// Returns the starting month (1, 4, 7, or 10) for the given quarter (1-4). +/// +/// # Panics +/// Panics if `quarter` is not in 1-4. fn quarter_start_month(quarter: u32) -> u32 { match quarter { 1 => 1, 2 => 4, 3 => 7, 4 => 10, - _ => panic!("invalid quarter: {}", quarter), // This function expects quarter 1-4 + _ => panic!("Invalid quarter: {}", quarter), } } -/// Return the earliest business day in the given (year, quarter). +/// Returns the first business day (Monday-Friday) of the specified quarter. +/// +/// Delegates to `first_business_day_of_month` using the quarter's start month. fn first_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate { let month = quarter_start_month(quarter); first_business_day_of_month(year, month) } -/// Return the last business day in the given (year, quarter). +/// Returns the last business day (Monday-Friday) of the specified quarter. +/// +/// Determines the quarter's final month and delegates to +/// `last_business_day_of_month`. fn last_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate { - // The last month of a quarter is the start month + 2. - let last_month_in_quarter = match quarter { - 1 => 3, - 2 => 6, - 3 => 9, - 4 => 12, - _ => panic!("invalid quarter: {}", quarter), - }; - last_business_day_of_month(year, last_month_in_quarter) + let last_month = quarter * 3; + last_business_day_of_month(year, last_month) } -/// Returns the earliest business day (Mon-Fri) of the given year. +/// Returns the first business day (Monday-Friday) of the specified year. +/// +/// Starts at January 1st and advances to the next weekday if needed. +/// +/// # Panics +/// Panics if date construction or successor operations overflow. fn first_business_day_of_year(year: i32) -> NaiveDate { - // Start with Jan 1st. Use _opt and expect(), assuming valid Y/M/D combination. - let mut d = NaiveDate::from_ymd_opt(year, 1, 1).expect("invalid year for Jan 1st"); - // If Jan 1st is a weekend, move forward to the next weekday. + let mut d = NaiveDate::from_ymd_opt(year, 1, 1).expect("Invalid year for January 1st"); while !is_weekday(d) { - // Use succ_opt() and expect(), assuming valid date and no overflow. d = d .succ_opt() - .expect("date overflow finding first bday of year"); + .expect("Date overflow finding first business day of year"); } d } -/// Returns the last business day (Mon-Fri) of the given year. +/// Returns the last business day (Monday-Friday) of the specified year. +/// +/// Starts at December 31st and moves backward to the previous weekday if needed. +/// +/// # Panics +/// Panics if date construction or predecessor operations underflow. fn last_business_day_of_year(year: i32) -> NaiveDate { - // Start with Dec 31st. Use _opt and expect(), assuming valid Y/M/D combination. - let mut d = NaiveDate::from_ymd_opt(year, 12, 31).expect("invalid year for Dec 31st"); - // If Dec 31st is a weekend, move backward to the previous weekday. + let mut d = NaiveDate::from_ymd_opt(year, 12, 31).expect("Invalid year for December 31st"); while !is_weekday(d) { - // Use pred_opt() and expect(), assuming valid date and no underflow. d = d .pred_opt() - .expect("date underflow finding last bday of year"); + .expect("Date underflow finding last business day of year"); } d } -// --- Generator Helper Functions --- - -/// Finds the *first* valid business date according to the frequency, -/// starting the search *on or after* the given `start_date`. -/// Panics on date overflow/underflow in extreme cases, but generally safe. +/// Finds the first valid business date on or after `start_date` according to `freq`. +/// +/// This may advance across days, weeks, months, quarters, or years depending on `freq`. +/// +/// # Panics +/// Panics on extreme date overflows or underflows. fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> NaiveDate { match freq { BDateFreq::Daily => { @@ -994,28 +968,24 @@ fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> Naive BDateFreq::MonthStart => { let mut candidate = first_business_day_of_month(start_date.year(), start_date.month()); if candidate < start_date { - // If the first bday of the current month is before start_date, - // we need the first bday of the *next* month. - let (next_y, next_m) = if start_date.month() == 12 { + let (ny, nm) = if start_date.month() == 12 { (start_date.year() + 1, 1) } else { (start_date.year(), start_date.month() + 1) }; - candidate = first_business_day_of_month(next_y, next_m); + candidate = first_business_day_of_month(ny, nm); } candidate } BDateFreq::MonthEnd => { let mut candidate = last_business_day_of_month(start_date.year(), start_date.month()); if candidate < start_date { - // If the last bday of current month is before start_date, - // we need the last bday of the *next* month. - let (next_y, next_m) = if start_date.month() == 12 { + let (ny, nm) = if start_date.month() == 12 { (start_date.year() + 1, 1) } else { (start_date.year(), start_date.month() + 1) }; - candidate = last_business_day_of_month(next_y, next_m); + candidate = last_business_day_of_month(ny, nm); } candidate } @@ -1023,14 +993,12 @@ fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> Naive let current_q = month_to_quarter(start_date.month()); let mut candidate = first_business_day_of_quarter(start_date.year(), current_q); if candidate < start_date { - // If the first bday of the current quarter is before start_date, - // we need the first bday of the *next* quarter. - let (next_y, next_q) = if current_q == 4 { + let (ny, nq) = if current_q == 4 { (start_date.year() + 1, 1) } else { (start_date.year(), current_q + 1) }; - candidate = first_business_day_of_quarter(next_y, next_q); + candidate = first_business_day_of_quarter(ny, nq); } candidate } @@ -1038,22 +1006,18 @@ fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> Naive let current_q = month_to_quarter(start_date.month()); let mut candidate = last_business_day_of_quarter(start_date.year(), current_q); if candidate < start_date { - // If the last bday of the current quarter is before start_date, - // we need the last bday of the *next* quarter. - let (next_y, next_q) = if current_q == 4 { + let (ny, nq) = if current_q == 4 { (start_date.year() + 1, 1) } else { (start_date.year(), current_q + 1) }; - candidate = last_business_day_of_quarter(next_y, next_q); + candidate = last_business_day_of_quarter(ny, nq); } candidate } BDateFreq::YearStart => { let mut candidate = first_business_day_of_year(start_date.year()); if candidate < start_date { - // If the first bday of the current year is before start_date, - // we need the first bday of the *next* year. candidate = first_business_day_of_year(start_date.year() + 1); } candidate @@ -1061,8 +1025,6 @@ fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> Naive BDateFreq::YearEnd => { let mut candidate = last_business_day_of_year(start_date.year()); if candidate < start_date { - // If the last bday of the current year is before start_date, - // we need the last bday of the *next* year. candidate = last_business_day_of_year(start_date.year() + 1); } candidate @@ -1070,15 +1032,19 @@ fn find_first_bdate_on_or_after(start_date: NaiveDate, freq: BDateFreq) -> Naive } } -/// Finds the *next* valid business date according to the frequency, -/// given the `current_date` (which is assumed to be a valid date previously generated). -/// Panics on date overflow/underflow in extreme cases, but generally safe. +/// Finds the next business date after `current_date` according to `freq`. +/// +/// Assumes `current_date` was previously generated. Advances by days, weeks, +/// months, quarters, or years as specified. +/// +/// # Panics +/// Panics on extreme date overflows or underflows. fn find_next_bdate(current_date: NaiveDate, freq: BDateFreq) -> NaiveDate { match freq { BDateFreq::Daily => { let mut next_day = current_date .succ_opt() - .expect("Date overflow finding next daily"); + .expect("Date overflow finding next daily date"); while !is_weekday(next_day) { next_day = next_day .succ_opt() @@ -1086,45 +1052,42 @@ fn find_next_bdate(current_date: NaiveDate, freq: BDateFreq) -> NaiveDate { } next_day } - BDateFreq::WeeklyMonday | BDateFreq::WeeklyFriday => { - // Assuming current_date is already a Mon/Fri, the next one is 7 days later. - current_date - .checked_add_signed(Duration::days(7)) - .expect("Date overflow adding 7 days") - } + BDateFreq::WeeklyMonday | BDateFreq::WeeklyFriday => current_date + .checked_add_signed(Duration::days(7)) + .expect("Date overflow adding one week"), BDateFreq::MonthStart => { - let (next_y, next_m) = if current_date.month() == 12 { + let (ny, nm) = if current_date.month() == 12 { (current_date.year() + 1, 1) } else { (current_date.year(), current_date.month() + 1) }; - first_business_day_of_month(next_y, next_m) + first_business_day_of_month(ny, nm) } BDateFreq::MonthEnd => { - let (next_y, next_m) = if current_date.month() == 12 { + let (ny, nm) = if current_date.month() == 12 { (current_date.year() + 1, 1) } else { (current_date.year(), current_date.month() + 1) }; - last_business_day_of_month(next_y, next_m) + last_business_day_of_month(ny, nm) } BDateFreq::QuarterStart => { let current_q = month_to_quarter(current_date.month()); - let (next_y, next_q) = if current_q == 4 { + let (ny, nq) = if current_q == 4 { (current_date.year() + 1, 1) } else { (current_date.year(), current_q + 1) }; - first_business_day_of_quarter(next_y, next_q) + first_business_day_of_quarter(ny, nq) } BDateFreq::QuarterEnd => { let current_q = month_to_quarter(current_date.month()); - let (next_y, next_q) = if current_q == 4 { + let (ny, nq) = if current_q == 4 { (current_date.year() + 1, 1) } else { (current_date.year(), current_q + 1) }; - last_business_day_of_quarter(next_y, next_q) + last_business_day_of_quarter(ny, nq) } BDateFreq::YearStart => first_business_day_of_year(current_date.year() + 1), BDateFreq::YearEnd => last_business_day_of_year(current_date.year() + 1),