Add CellOffsets abstraction (#8814)

Refactor `Notebook::cell_offsets` to use an abstract struct for storing
the cell offsets. This will allow us to add useful methods on it.

crates/ruff_notebook/src/cell.rs

@@ -1,4 +1,7 @@
 use std::fmt;
+use std::ops::{Deref, DerefMut};
+
+use ruff_text_size::TextSize;
 
 use crate::schema::{Cell, SourceValue};
 
@@ -168,3 +171,34 @@ impl Cell {
         lines.any(|line| line.trim_start().starts_with("%%"))
     }
 }
+
+/// Cell offsets are used to keep track of the start and end offsets of each
+/// cell in the concatenated source code.
+#[derive(Clone, Debug, Default, PartialEq)]
+pub struct CellOffsets(Vec<TextSize>);
+
+impl CellOffsets {
+    /// Create a new [`CellOffsets`] with the given capacity.
+    pub(crate) fn with_capacity(capacity: usize) -> Self {
+        Self(Vec::with_capacity(capacity))
+    }
+
+    /// Push a new offset to the end of the [`CellOffsets`].
+    pub(crate) fn push(&mut self, offset: TextSize) {
+        self.0.push(offset);
+    }
+}
+
+impl Deref for CellOffsets {
+    type Target = [TextSize];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for CellOffsets {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}

crates/ruff_notebook/src/lib.rs

@@ -1,5 +1,6 @@
 //! Utils for reading and writing jupyter notebooks
 
+pub use cell::*;
 pub use index::*;
 pub use notebook::*;
 pub use schema::*;

crates/ruff_notebook/src/notebook.rs

@@ -15,6 +15,7 @@ use ruff_diagnostics::{SourceMap, SourceMarker};
 use ruff_source_file::{NewlineWithTrailingNewline, OneIndexed, UniversalNewlineIterator};
 use ruff_text_size::TextSize;
 
+use crate::cell::CellOffsets;
 use crate::index::NotebookIndex;
 use crate::schema::{Cell, RawNotebook, SortAlphabetically, SourceValue};
 
@@ -65,7 +66,7 @@ pub struct Notebook {
     raw: RawNotebook,
     /// The offsets of each cell in the concatenated source code. This includes
     /// the first and last character offsets as well.
-    cell_offsets: Vec<TextSize>,
+    cell_offsets: CellOffsets,
     /// The cell index of all valid code cells in the notebook.
     valid_code_cells: Vec<u32>,
     /// Flag to indicate if the JSON string of the notebook has a trailing newline.
@@ -128,7 +129,7 @@ impl Notebook {
 
         let mut contents = Vec::with_capacity(valid_code_cells.len());
         let mut current_offset = TextSize::from(0);
-        let mut cell_offsets = Vec::with_capacity(valid_code_cells.len());
+        let mut cell_offsets = CellOffsets::with_capacity(valid_code_cells.len());
         cell_offsets.push(TextSize::from(0));
 
         for &idx in &valid_code_cells {
@@ -324,9 +325,9 @@ impl Notebook {
         self.index.take().unwrap_or_else(|| self.build_index())
     }
 
-    /// Return the cell offsets for the concatenated source code corresponding
+    /// Return the [`CellOffsets`] for the concatenated source code corresponding
     /// the Jupyter notebook.
-    pub fn cell_offsets(&self) -> &[TextSize] {
+    pub fn cell_offsets(&self) -> &CellOffsets {
         &self.cell_offsets
     }
 
@@ -503,7 +504,7 @@ print("after empty cells")
             }
         );
         assert_eq!(
-            notebook.cell_offsets(),
+            notebook.cell_offsets().as_ref(),
             &[
                 0.into(),
                 90.into(),