Printer: Slice based queue and stack

crates/ruff_formatter/src/printer/mod.rs

@@ -60,11 +60,13 @@ impl<'a> Printer<'a> {
             let mut stack = PrintCallStack::new(PrintElementArgs::new(Indention::Level(indent)));
             let mut queue: PrintQueue<'a> = PrintQueue::new(document.as_ref());
 
-            while let Some(element) = queue.pop() {
-                self.print_element(&mut stack, &mut queue, element)?;
-
-                if queue.is_empty() {
-                    self.flush_line_suffixes(&mut queue, &mut stack, None);
+            loop {
+                if let Some(element) = queue.pop() {
+                    self.print_element(&mut stack, &mut queue, element)?;
+                } else {
+                    if !self.flush_line_suffixes(&mut queue, &mut stack, None) {
+                        break;
+                    }
                 }
             }
 
@@ -413,7 +415,7 @@ impl<'a> Printer<'a> {
         queue: &mut PrintQueue<'a>,
         stack: &mut PrintCallStack,
         line_break: Option<&'a FormatElement>,
-    ) {
+    ) -> bool {
         let suffixes = self.state.line_suffixes.take_pending();
 
         if suffixes.len() > 0 {
@@ -437,6 +439,10 @@ impl<'a> Printer<'a> {
                     }
                 }
             }
+
+            true
+        } else {
+            false
         }
     }
 
@@ -771,7 +777,7 @@ struct PrinterState<'a> {
     // Re-used queue to measure if a group fits. Optimisation to avoid re-allocating a new
     // vec every time a group gets measured
     fits_stack: Vec<StackFrame>,
-    fits_queue: Vec<&'a [FormatElement]>,
+    fits_queue: Vec<std::slice::Iter<'a, FormatElement>>,
 }
 
 impl<'a> PrinterState<'a> {

crates/ruff_formatter/src/printer/queue.rs

@@ -1,6 +1,5 @@
 use crate::format_element::tag::TagKind;
 use crate::prelude::Tag;
-use crate::printer::stack::{Stack, StackedStack};
 use crate::printer::{invalid_end_tag, invalid_start_tag};
 use crate::{FormatElement, PrintResult};
 use std::fmt::Debug;
@@ -9,43 +8,11 @@ use std::marker::PhantomData;
 
 /// Queue of [`FormatElement`]s.
 pub(super) trait Queue<'a> {
-    type Stack: Stack<&'a [FormatElement]>;
-
-    fn stack(&self) -> &Self::Stack;
-
-    fn stack_mut(&mut self) -> &mut Self::Stack;
-
-    fn next_index(&self) -> usize;
-
-    fn set_next_index(&mut self, index: usize);
-
     /// Pops the element at the end of the queue.
-    fn pop(&mut self) -> Option<&'a FormatElement> {
-        match self.stack().top() {
-            Some(top_slice) => {
-                // SAFETY: Safe because queue ensures that slices inside `slices` are never empty.
-                let next_index = self.next_index();
-                let element = &top_slice[next_index];
-
-                if next_index + 1 == top_slice.len() {
-                    self.stack_mut().pop().unwrap();
-                    self.set_next_index(0);
-                } else {
-                    self.set_next_index(next_index + 1);
-                }
-
-                Some(element)
-            }
-            None => None,
-        }
-    }
+    fn pop(&mut self) -> Option<&'a FormatElement>;
 
     /// Returns the next element, not traversing into [`FormatElement::Interned`].
-    fn top_with_interned(&self) -> Option<&'a FormatElement> {
-        self.stack()
-            .top()
-            .map(|top_slice| &top_slice[self.next_index()])
-    }
+    fn top_with_interned(&self) -> Option<&'a FormatElement>;
 
     /// Returns the next element, recursively resolving the first element of [`FormatElement::Interned`].
     fn top(&self) -> Option<&'a FormatElement> {
@@ -64,29 +31,10 @@ pub(super) trait Queue<'a> {
     }
 
     /// Queues a slice of elements to process before the other elements in this queue.
-    fn extend_back(&mut self, elements: &'a [FormatElement]) {
-        match elements {
-            [] => {
-                // Don't push empty slices
-            }
-            slice => {
-                let next_index = self.next_index();
-                let stack = self.stack_mut();
-                if let Some(top) = stack.pop() {
-                    stack.push(&top[next_index..]);
-                }
-
-                stack.push(slice);
-                self.set_next_index(0);
-            }
-        }
-    }
+    fn extend_back(&mut self, elements: &'a [FormatElement]);
 
     /// Removes top slice.
-    fn pop_slice(&mut self) -> Option<&'a [FormatElement]> {
-        self.set_next_index(0);
-        self.stack_mut().pop()
-    }
+    fn pop_slice(&mut self) -> Option<&'a [FormatElement]>;
 
     /// Skips all content until it finds the corresponding end tag with the given kind.
     fn skip_content(&mut self, kind: TagKind)
@@ -112,45 +60,58 @@ pub(super) trait Queue<'a> {
 /// Queue with the elements to print.
 #[derive(Debug, Default, Clone)]
 pub(super) struct PrintQueue<'a> {
-    slices: Vec<&'a [FormatElement]>,
-    next_index: usize,
+    element_slices: Vec<std::slice::Iter<'a, FormatElement>>,
 }
 
 impl<'a> PrintQueue<'a> {
     pub(super) fn new(slice: &'a [FormatElement]) -> Self {
-        let slices = match slice {
-            [] => Vec::default(),
-            slice => vec![slice],
-        };
-
         Self {
-            slices,
-            next_index: 0,
+            element_slices: if slice.is_empty() {
+                Vec::new()
+            } else {
+                vec![slice.iter()]
+            },
         }
     }
-
-    pub(super) fn is_empty(&self) -> bool {
-        self.slices.is_empty()
-    }
 }
 
 impl<'a> Queue<'a> for PrintQueue<'a> {
-    type Stack = Vec<&'a [FormatElement]>;
-
-    fn stack(&self) -> &Self::Stack {
-        &self.slices
+    fn pop(&mut self) -> Option<&'a FormatElement> {
+        let elements = self.element_slices.last_mut()?;
+        elements.next().or_else(|| {
+            self.element_slices.pop();
+            let elements = self.element_slices.last_mut()?;
+            elements.next()
+        })
     }
 
-    fn stack_mut(&mut self) -> &mut Self::Stack {
-        &mut self.slices
+    fn top_with_interned(&self) -> Option<&'a FormatElement> {
+        let mut slices = self.element_slices.iter().rev();
+        let slice = slices.next()?;
+
+        match slice.as_slice().first() {
+            Some(element) => Some(element),
+            None => {
+                if let Some(next_elements) = slices.next() {
+                    next_elements.as_slice().first()
+                } else {
+                    None
+                }
+            }
+        }
     }
 
-    fn next_index(&self) -> usize {
-        self.next_index
+    fn extend_back(&mut self, elements: &'a [FormatElement]) {
+        if !elements.is_empty() {
+            self.element_slices.push(elements.iter());
+        }
     }
 
-    fn set_next_index(&mut self, index: usize) {
-        self.next_index = index;
+    /// Removes top slice.
+    fn pop_slice(&mut self) -> Option<&'a [FormatElement]> {
+        self.element_slices
+            .pop()
+            .map(|elements| elements.as_slice())
     }
 }
 
@@ -161,45 +122,63 @@ impl<'a> Queue<'a> for PrintQueue<'a> {
 #[must_use]
 #[derive(Debug)]
 pub(super) struct FitsQueue<'a, 'print> {
-    stack: StackedStack<'print, &'a [FormatElement]>,
-    next_index: usize,
+    queue: PrintQueue<'a>,
+    rest_elements: std::slice::Iter<'print, std::slice::Iter<'a, FormatElement>>,
 }
 
 impl<'a, 'print> FitsQueue<'a, 'print> {
     pub(super) fn new(
-        print_queue: &'print PrintQueue<'a>,
-        saved: Vec<&'a [FormatElement]>,
+        rest_queue: &'print PrintQueue<'a>,
+        queue_vec: Vec<std::slice::Iter<'a, FormatElement>>,
     ) -> Self {
-        let stack = StackedStack::with_vec(&print_queue.slices, saved);
-
         Self {
-            stack,
-            next_index: print_queue.next_index,
+            queue: PrintQueue {
+                element_slices: queue_vec,
+            },
+            rest_elements: rest_queue.element_slices.iter(),
         }
     }
 
-    pub(super) fn finish(self) -> Vec<&'a [FormatElement]> {
-        self.stack.into_vec()
+    pub(super) fn finish(self) -> Vec<std::slice::Iter<'a, FormatElement>> {
+        self.queue.element_slices
     }
 }
 
 impl<'a, 'print> Queue<'a> for FitsQueue<'a, 'print> {
-    type Stack = StackedStack<'print, &'a [FormatElement]>;
-
-    fn stack(&self) -> &Self::Stack {
-        &self.stack
+    fn pop(&mut self) -> Option<&'a FormatElement> {
+        self.queue.pop().or_else(|| {
+            if let Some(next_slice) = self.rest_elements.next_back() {
+                self.queue.extend_back(next_slice.as_slice());
+                self.queue.pop()
+            } else {
+                None
+            }
+        })
     }
 
-    fn stack_mut(&mut self) -> &mut Self::Stack {
-        &mut self.stack
+    fn top_with_interned(&self) -> Option<&'a FormatElement> {
+        self.queue.top_with_interned().or_else(|| {
+            if let Some(next_elements) = self.rest_elements.as_slice().last() {
+                next_elements.as_slice().first()
+            } else {
+                None
+            }
+        })
     }
 
-    fn next_index(&self) -> usize {
-        self.next_index
+    fn extend_back(&mut self, elements: &'a [FormatElement]) {
+        if !elements.is_empty() {
+            self.queue.extend_back(elements);
+        }
     }
 
-    fn set_next_index(&mut self, index: usize) {
-        self.next_index = index;
+    /// Removes top slice.
+    fn pop_slice(&mut self) -> Option<&'a [FormatElement]> {
+        self.queue.pop_slice().or_else(|| {
+            self.rest_elements
+                .next_back()
+                .map(std::slice::Iter::as_slice)
+        })
     }
 }
 

crates/ruff_formatter/src/printer/stack.rs

@@ -35,7 +35,7 @@ impl<T> Stack<T> for Vec<T> {
 #[derive(Debug, Clone)]
 pub(super) struct StackedStack<'a, T> {
     /// The content of the original stack.
-    original: &'a [T],
+    original: std::slice::Iter<'a, T>,
 
     /// Items that have been pushed since the creation of this stack and aren't part of the `original` stack.
     stack: Vec<T>,
@@ -49,7 +49,10 @@ impl<'a, T> StackedStack<'a, T> {
 
     /// Creates a new stack that uses `stack` for storing its elements.
     pub(super) fn with_vec(original: &'a [T], stack: Vec<T>) -> Self {
-        Self { original, stack }
+        Self {
+            original: original.iter(),
+            stack,
+        }
     }
 
     /// Returns the underlying `stack` vector.
@@ -63,25 +66,23 @@ where
     T: Copy,
 {
     fn pop(&mut self) -> Option<T> {
-        self.stack.pop().or_else(|| match self.original {
-            [rest @ .., last] => {
-                self.original = rest;
-                Some(*last)
-            }
-            _ => None,
-        })
+        self.stack
+            .pop()
+            .or_else(|| self.original.next_back().copied())
     }
 
     fn push(&mut self, value: T) {
         self.stack.push(value);
     }
 
     fn top(&self) -> Option<&T> {
-        self.stack.last().or_else(|| self.original.last())
+        self.stack
+            .last()
+            .or_else(|| self.original.as_slice().last())
     }
 
     fn is_empty(&self) -> bool {
-        self.original.is_empty() && self.stack.is_empty()
+        self.stack.is_empty() && self.original.len() == 0
     }
 }