Try another representation

crates/ruff_python_parser/src/lexer.rs

@@ -266,7 +266,7 @@ impl<'src> Lexer<'src> {
     }
 
     fn handle_indentation(&mut self, indentation: Indentation) -> Option<TokenKind> {
-        let token = match self.indentations.current().try_compare(indentation) {
+        match self.indentations.current().try_compare(indentation) {
             // Dedent
             Ok(Ordering::Greater) => {
                 self.pending_indentation = Some(indentation);
@@ -303,15 +303,12 @@ impl<'src> Lexer<'src> {
                 self.indentations.indent(indentation);
                 Some(TokenKind::Indent)
             }
-            Err(_) => {
-                return Some(self.push_error(LexicalError::new(
-                    LexicalErrorType::IndentationError,
-                    self.token_range(),
-                )));
-            }
-        };
 
-        token
+            Err(_) => Some(self.push_error(LexicalError::new(
+                LexicalErrorType::IndentationError,
+                self.token_range(),
+            ))),
+        }
     }
 
     fn skip_whitespace(&mut self) -> Result<(), LexicalError> {

crates/ruff_python_parser/src/lexer/indentation.rs

@@ -1,7 +1,8 @@
-use static_assertions::assert_eq_size;
 use std::cmp::Ordering;
 use std::fmt::Debug;
 
+use static_assertions::assert_eq_size;
+
 /// The column index of an indentation.
 ///
 /// A space increments the column by one. A tab adds up to 2 (if tab size is 2) indices, but just one
@@ -9,22 +10,10 @@ use std::fmt::Debug;
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Default)]
 pub(super) struct Column(u32);
 
-impl Column {
-    pub(super) const fn new(column: u32) -> Self {
-        Self(column)
-    }
-}
-
 /// The number of characters in an indentation. Each character accounts for 1.
 #[derive(Debug, Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Default)]
 pub(super) struct Character(u32);
 
-impl Character {
-    pub(super) const fn new(characters: u32) -> Self {
-        Self(characters)
-    }
-}
-
 /// The [Indentation](https://docs.python.org/3/reference/lexical_analysis.html#indentation) of a logical line.
 #[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
 pub(super) struct Indentation {
@@ -35,10 +24,15 @@ pub(super) struct Indentation {
 impl Indentation {
     const TAB_SIZE: u32 = 2;
 
+    const ROOT: Indentation = Indentation {
+        column: Column(0),
+        character: Character(0),
+    };
+
     pub(super) const fn root() -> Self {
         Self {
-            column: Column::new(0),
-            character: Character::new(0),
+            column: Column(0),
+            character: Character(0),
         }
     }
 
@@ -77,6 +71,14 @@ impl Indentation {
             Err(UnexpectedIndentation)
         }
     }
+
+    /// Computes the indentation at the given level based on the current indentation.
+    const fn at(self, level: u32) -> Self {
+        Self {
+            character: Character(self.character.0 * level),
+            column: Column(self.column.0 * level),
+        }
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq)]
@@ -84,16 +86,38 @@ pub(super) struct UnexpectedIndentation;
 
 /// The indentations stack is used to keep track of the current indentation level
 /// [See Indentation](docs.python.org/3/reference/lexical_analysis.html#indentation).
-#[derive(Debug, Clone, Default)]
+#[derive(Debug)]
 pub(super) struct Indentations {
-    stack: Vec<Indentation>,
+    inner: IndentationsInner,
+}
+
+#[derive(Debug, Clone)]
+enum IndentationsInner {
+    Stack(Vec<Indentation>),
+    Counter(IndentationCounter),
+}
+
+impl Default for Indentations {
+    fn default() -> Self {
+        Indentations {
+            inner: IndentationsInner::Counter(IndentationCounter::default()),
+        }
+    }
 }
 
 impl Indentations {
     pub(super) fn indent(&mut self, indent: Indentation) {
         debug_assert_eq!(self.current().try_compare(indent), Ok(Ordering::Less));
 
-        self.stack.push(indent);
+        match &mut self.inner {
+            IndentationsInner::Stack(indentations) => indentations.push(indent),
+            IndentationsInner::Counter(indentations) => {
+                if indentations.indent(indent) {
+                    return;
+                }
+                self.make_stack().push(indent);
+            }
+        }
     }
 
     /// Dedent one level to eventually reach `new_indentation`.
@@ -105,7 +129,7 @@ impl Indentations {
     ) -> Result<Option<Indentation>, UnexpectedIndentation> {
         let previous = self.dedent();
 
-        match new_indentation.try_compare(*self.current())? {
+        match new_indentation.try_compare(self.current())? {
             Ordering::Less | Ordering::Equal => Ok(previous),
             // ```python
             // if True:
@@ -117,40 +141,113 @@ impl Indentations {
     }
 
     pub(super) fn dedent(&mut self) -> Option<Indentation> {
-        self.stack.pop()
+        match &mut self.inner {
+            IndentationsInner::Stack(indentations) => indentations.pop(),
+            IndentationsInner::Counter(indentations) => indentations.dedent(),
+        }
     }
 
-    pub(super) fn current(&self) -> &Indentation {
-        static ROOT: Indentation = Indentation::root();
-        self.stack.last().unwrap_or(&ROOT)
+    pub(super) fn current(&self) -> Indentation {
+        match &self.inner {
+            IndentationsInner::Stack(indentations) => {
+                *indentations.last().unwrap_or(&Indentation::ROOT)
+            }
+            IndentationsInner::Counter(indentations) => indentations.current(),
+        }
     }
 
     pub(crate) fn checkpoint(&self) -> IndentationsCheckpoint {
-        IndentationsCheckpoint(self.stack.clone())
+        IndentationsCheckpoint(self.inner.clone())
     }
 
     pub(crate) fn rewind(&mut self, checkpoint: IndentationsCheckpoint) {
-        self.stack = checkpoint.0;
+        self.inner = checkpoint.0;
+    }
+
+    fn make_stack(&mut self) -> &mut Vec<Indentation> {
+        if let IndentationsInner::Counter(IndentationCounter { first, level, .. }) = self.inner {
+            if level == 0 {
+                *self = Indentations {
+                    inner: IndentationsInner::Stack(vec![]),
+                };
+            } else {
+                *self = Indentations {
+                    inner: IndentationsInner::Stack(first.map_or_else(Vec::new, |first| {
+                        (1..=level).map(|level| first.at(level)).collect()
+                    })),
+                };
+            }
+        }
+        match &mut self.inner {
+            IndentationsInner::Stack(stack) => stack,
+            IndentationsInner::Counter(_) => unreachable!(),
+        }
     }
 }
 
-#[derive(Debug, Clone)]
-pub(crate) struct IndentationsCheckpoint(Vec<Indentation>);
+#[derive(Debug, Default, Clone)]
+struct IndentationCounter {
+    /// The current indentation.
+    current: Indentation,
+    /// The first indentation in the source code.
+    first: Option<Indentation>,
+    /// The current indentation level.
+    level: u32,
+}
+
+impl IndentationCounter {
+    fn indent(&mut self, indent: Indentation) -> bool {
+        if let Some(first) = self.first {
+            if first.at(self.level + 1) == indent {
+                self.current = indent;
+                self.level += 1;
+                true
+            } else {
+                false
+            }
+        } else {
+            self.first = Some(indent);
+            self.current = indent;
+            self.level = 1;
+            true
+        }
+    }
+
+    fn dedent(&mut self) -> Option<Indentation> {
+        if self.level == 0 {
+            None
+        } else if let Some(first) = self.first {
+            let current = self.current;
+            self.level -= 1;
+            self.current = first.at(self.level);
+            Some(current)
+        } else {
+            unreachable!()
+        }
+    }
+
+    fn current(&self) -> Indentation {
+        self.current
+    }
+}
+
+pub(super) struct IndentationsCheckpoint(IndentationsInner);
 
 assert_eq_size!(Indentation, u64);
 
 #[cfg(test)]
 mod tests {
-    use super::{Character, Column, Indentation};
     use std::cmp::Ordering;
 
+    use super::{Character, Column, Indentation};
+
     #[test]
     fn indentation_try_compare() {
-        let tab = Indentation::new(Column::new(8), Character::new(1));
+        let tab = Indentation::new(Column(8), Character(1));
 
         assert_eq!(tab.try_compare(tab), Ok(Ordering::Equal));
 
-        let two_tabs = Indentation::new(Column::new(16), Character::new(2));
+        let two_tabs = Indentation::new(Column(16), Character(2));
         assert_eq!(two_tabs.try_compare(tab), Ok(Ordering::Greater));
         assert_eq!(tab.try_compare(two_tabs), Ok(Ordering::Less));
     }