Implement an iterator for universal newlines

[charliermarsh]

Summary

We need to support CR line endings (as opposed to LF and CRLF line endings, which are already supported). They're rare, but they do appear in Python code, and we tend to panic on any file that uses them.

Our Locator abstraction now supports CR line endings. However, Rust's str#lines implementation does not.

This PR adds a UniversalNewlineIterator implementation that respects all of CR, LF, and CRLF line endings, and plugs it into most of the .lines() call sites.

As an alternative design, it could be nice if we could leverage Locator for this. We've already computed all of the line endings, so we could probably iterate much more efficiently?

Test Plan

Largely relying on automated testing, however, also ran over some known failure cases, like #3404.

[github-actions]

✅ ecosystem check detected no changes.

[charliermarsh]

I have a feeling that I'm breaking some idioms and that this code could be far more efficient, so feel free to give it a harsh review if it deserves one :)

[MichaReiser]

I tried it myself and haven't been able to come up with something much more concise. It's surprisingly complicated

This implementation mutates the underlying string slices rather than tracking start and end positions explicitly (similar to Vec::IntoIter that tracks the start and end pointer). I would need to do some profiling to understand if it is faster than using explicit start and end positions. It might not, because writing a string slice is two words (pointer, length).

/// Like `str#lines`, but accommodates LF, CRLF, and CR line endings,
/// the latter of which are not supported by `str#lines`.
pub struct UniversalNewlineIterator<'a> {
    text: &'a str,
}

impl<'a> UniversalNewlineIterator<'a> {
    pub fn from(text: &'a str) -> UniversalNewlineIterator<'a> {
        UniversalNewlineIterator { text }
    }
}

impl<'a> Iterator for UniversalNewlineIterator<'a> {
    type Item = &'a str;

    #[inline]
    fn next(&mut self) -> Option<&'a str> {
        if self.text.is_empty() {
            return None;
        }

        let line = match self.text.find(['\n', '\r']) {
            // Not the last line
            Some(line_end) => {
                let (line, remainder) = self.text.split_at(line_end);

                self.text = match remainder.as_bytes()[0] {
                    // Explicit branch for `\n` as this is the most likely path
                    b'\n' => &remainder[1..],
                    // '\r\n'
                    b'\r' if remainder.as_bytes().get(1) == Some(&b'\n') => &remainder[2..],
                    // '\r'
                    _ => &remainder[1..],
                };

                line
            }
            // Last line
            None => std::mem::replace(&mut self.text, ""),
        };

        Some(line)
    }
}

impl<'a> DoubleEndedIterator for UniversalNewlineIterator<'_> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.text.is_empty() {
            return None;
        }

        let len = self.text.len();

        // Trim any trailing newlines.
        self.text = match self.text.as_bytes()[len - 1] {
            b'\n' if len > 1 && self.text.as_bytes()[len - 2] == b'\r' => &self.text[..len - 2],
            b'\n' | b'\r' => &self.text[..len - 1],
            _ => &self.text,
        };

        // Find the end of the previous line. The previous line is the text up to, but not including the new line character
        let line = match self.text.rfind(['\n', '\r']) {
            // '\n' or '\r' or '\r\n'
            Some(line_end) => {
                let (remainder, line) = self.text.split_at(line_end + 1);
                self.text = remainder;

                line
            }
            // Last line
            None => std::mem::replace(&mut self.text, ""),
        };

        Some(line)
    }
}

impl FusedIterator for UniversalNewlineIterator<'_> {}

#[cfg(test)]
mod tests {
    use super::UniversalNewlineIterator;

    #[test]
    fn universal_newlines_empty_str() {
        let lines: Vec<_> = UniversalNewlineIterator::from("").collect();
        assert_eq!(lines, Vec::<&str>::default());

        let lines: Vec<_> = UniversalNewlineIterator::from("").rev().collect();
        assert_eq!(lines, Vec::<&str>::default());
    }

    #[test]
    fn universal_newlines_forward() {
        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop\n").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop\n\n").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop", ""]);
    }

    #[test]
    fn universal_newlines_backwards() {
        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop")
            .rev()
            .collect();
        assert_eq!(lines, vec!["bop", "baz", "", "bar", "foo"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\nbaz\rbop\n")
            .rev()
            .collect();

        assert_eq!(lines, vec!["bop", "baz", "", "bar", "foo"]);
    }

    #[test]
    fn universal_newlines_mixed() {
        let mut lines = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop");

        assert_eq!(lines.next_back(), Some("bop"));
        assert_eq!(lines.next(), Some("foo"));
        assert_eq!(lines.next_back(), Some("baz"));
        assert_eq!(lines.next(), Some("bar"));
        assert_eq!(lines.next_back(), Some(""));
        assert_eq!(lines.next(), None);
    }
}

I also added a few more tests.

This is a more "traditional" implementation that uses start and end offsets (similar to how SplitIterator works)

/// Like `str#lines`, but accommodates LF, CRLF, and CR line endings,
/// the latter of which are not supported by `str#lines`.
pub struct UniversalNewlineIterator<'a> {
    text: &'a str,
    start_offset: usize,
    end_offset: usize,
}

impl<'a> UniversalNewlineIterator<'a> {
    pub fn from(text: &'a str) -> UniversalNewlineIterator<'a> {
        UniversalNewlineIterator {
            text,
            start_offset: 0,
            end_offset: text.len(),
        }
    }

    fn as_str(&self) -> &'a str {
        // SAFETY: It's guaranteed that start and end offsets are always positioned at character boundaries.
        unsafe { &self.text.get_unchecked(self.start_offset..self.end_offset) }
    }
}

impl<'a> Iterator for UniversalNewlineIterator<'a> {
    type Item = &'a str;

    #[inline]
    fn next(&mut self) -> Option<&'a str> {
        let text = self.as_str();

        if text.is_empty() {
            return None;
        }

        let line = match text.find(['\n', '\r']) {
            // Not the last line
            Some(line_end) => {
                self.start_offset += line_end
                    + match text.as_bytes()[line_end] {
                        // Explicit branch for `\n` as this is the most likely path
                        b'\n' => 1,
                        // '\r\n'
                        b'\r' if text.as_bytes().get(line_end + 1) == Some(&b'\n') => 2,
                        // '\r'
                        _ => 1,
                    };

                &text[..line_end]
            }
            // Last line
            None => {
                self.start_offset = self.end_offset;
                text
            }
        };

        Some(line)
    }
}

impl<'a> DoubleEndedIterator for UniversalNewlineIterator<'_> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        let mut text = self.as_str();

        if text.is_empty() {
            return None;
        }

        let len = text.len();

        // Trim any trailing newlines.
        text = match text.as_bytes()[len - 1] {
            b'\n' if len > 1 && text.as_bytes()[len - 2] == b'\r' => &text[..len - 2],
            b'\n' | b'\r' => &text[..len - 1],
            _ => &text,
        };

        // Find the end of the previous line. The previous line is the text up to, but not including the new line character
        let line = match text.rfind(['\n', '\r']) {
            // '\n' or '\r' or '\r\n'
            Some(newline_pos) => {
                let line_start = newline_pos + 1;
                let line = &text[line_start..];
                self.end_offset = self.start_offset + line_start;

                line
            }
            // Last line
            None => {
                self.end_offset = self.start_offset;
                text
            }
        };

        Some(line)
    }
}

impl FusedIterator for UniversalNewlineIterator<'_> {}

[charliermarsh]

I took the first suggestion, since the latter requires unsafe, and it'd be our only unsafe block right now.

[charliermarsh]

(I did some cursory benchmarking, and what you had here minorly outperformed the current version in the PR -- but I'd need to run an actual microbenchmark.)

[MichaReiser]

There's an issue with the iterator implementation that it allows next_back and next to pass each other, which isn't allowed according to the Iterator protocol. Other than that, looks good to me (I proposed an alternative and suggested some more tests but it leads to about the same amount of code)

It would be good to add a few tests we knew were failing before to verify that the new logic handles \r correctly.

I'm undecided if we should create an extension trait to implement universal_lines on &str. It could help discoverability (over UniversalNewlinesIterator::from(str).

[MichaReiser]

I tried it myself and haven't been able to come up with something much more concise. It's surprisingly complicated

This implementation mutates the underlying string slices rather than tracking start and end positions explicitly (similar to Vec::IntoIter that tracks the start and end pointer). I would need to do some profiling to understand if it is faster than using explicit start and end positions. It might not, because writing a string slice is two words (pointer, length).

/// Like `str#lines`, but accommodates LF, CRLF, and CR line endings,
/// the latter of which are not supported by `str#lines`.
pub struct UniversalNewlineIterator<'a> {
    text: &'a str,
}

impl<'a> UniversalNewlineIterator<'a> {
    pub fn from(text: &'a str) -> UniversalNewlineIterator<'a> {
        UniversalNewlineIterator { text }
    }
}

impl<'a> Iterator for UniversalNewlineIterator<'a> {
    type Item = &'a str;

    #[inline]
    fn next(&mut self) -> Option<&'a str> {
        if self.text.is_empty() {
            return None;
        }

        let line = match self.text.find(['\n', '\r']) {
            // Not the last line
            Some(line_end) => {
                let (line, remainder) = self.text.split_at(line_end);

                self.text = match remainder.as_bytes()[0] {
                    // Explicit branch for `\n` as this is the most likely path
                    b'\n' => &remainder[1..],
                    // '\r\n'
                    b'\r' if remainder.as_bytes().get(1) == Some(&b'\n') => &remainder[2..],
                    // '\r'
                    _ => &remainder[1..],
                };

                line
            }
            // Last line
            None => std::mem::replace(&mut self.text, ""),
        };

        Some(line)
    }
}

impl<'a> DoubleEndedIterator for UniversalNewlineIterator<'_> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.text.is_empty() {
            return None;
        }

        let len = self.text.len();

        // Trim any trailing newlines.
        self.text = match self.text.as_bytes()[len - 1] {
            b'\n' if len > 1 && self.text.as_bytes()[len - 2] == b'\r' => &self.text[..len - 2],
            b'\n' | b'\r' => &self.text[..len - 1],
            _ => &self.text,
        };

        // Find the end of the previous line. The previous line is the text up to, but not including the new line character
        let line = match self.text.rfind(['\n', '\r']) {
            // '\n' or '\r' or '\r\n'
            Some(line_end) => {
                let (remainder, line) = self.text.split_at(line_end + 1);
                self.text = remainder;

                line
            }
            // Last line
            None => std::mem::replace(&mut self.text, ""),
        };

        Some(line)
    }
}

impl FusedIterator for UniversalNewlineIterator<'_> {}

#[cfg(test)]
mod tests {
    use super::UniversalNewlineIterator;

    #[test]
    fn universal_newlines_empty_str() {
        let lines: Vec<_> = UniversalNewlineIterator::from("").collect();
        assert_eq!(lines, Vec::<&str>::default());

        let lines: Vec<_> = UniversalNewlineIterator::from("").rev().collect();
        assert_eq!(lines, Vec::<&str>::default());
    }

    #[test]
    fn universal_newlines_forward() {
        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop\n").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop\n\n").collect();
        assert_eq!(lines, vec!["foo", "bar", "", "baz", "bop", ""]);
    }

    #[test]
    fn universal_newlines_backwards() {
        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop")
            .rev()
            .collect();
        assert_eq!(lines, vec!["bop", "baz", "", "bar", "foo"]);

        let lines: Vec<_> = UniversalNewlineIterator::from("foo\nbar\n\nbaz\rbop\n")
            .rev()
            .collect();

        assert_eq!(lines, vec!["bop", "baz", "", "bar", "foo"]);
    }

    #[test]
    fn universal_newlines_mixed() {
        let mut lines = UniversalNewlineIterator::from("foo\nbar\n\r\nbaz\rbop");

        assert_eq!(lines.next_back(), Some("bop"));
        assert_eq!(lines.next(), Some("foo"));
        assert_eq!(lines.next_back(), Some("baz"));
        assert_eq!(lines.next(), Some("bar"));
        assert_eq!(lines.next_back(), Some(""));
        assert_eq!(lines.next(), None);
    }
}

I also added a few more tests.

This is a more "traditional" implementation that uses start and end offsets (similar to how SplitIterator works)

/// Like `str#lines`, but accommodates LF, CRLF, and CR line endings,
/// the latter of which are not supported by `str#lines`.
pub struct UniversalNewlineIterator<'a> {
    text: &'a str,
    start_offset: usize,
    end_offset: usize,
}

impl<'a> UniversalNewlineIterator<'a> {
    pub fn from(text: &'a str) -> UniversalNewlineIterator<'a> {
        UniversalNewlineIterator {
            text,
            start_offset: 0,
            end_offset: text.len(),
        }
    }

    fn as_str(&self) -> &'a str {
        // SAFETY: It's guaranteed that start and end offsets are always positioned at character boundaries.
        unsafe { &self.text.get_unchecked(self.start_offset..self.end_offset) }
    }
}

impl<'a> Iterator for UniversalNewlineIterator<'a> {
    type Item = &'a str;

    #[inline]
    fn next(&mut self) -> Option<&'a str> {
        let text = self.as_str();

        if text.is_empty() {
            return None;
        }

        let line = match text.find(['\n', '\r']) {
            // Not the last line
            Some(line_end) => {
                self.start_offset += line_end
                    + match text.as_bytes()[line_end] {
                        // Explicit branch for `\n` as this is the most likely path
                        b'\n' => 1,
                        // '\r\n'
                        b'\r' if text.as_bytes().get(line_end + 1) == Some(&b'\n') => 2,
                        // '\r'
                        _ => 1,
                    };

                &text[..line_end]
            }
            // Last line
            None => {
                self.start_offset = self.end_offset;
                text
            }
        };

        Some(line)
    }
}

impl<'a> DoubleEndedIterator for UniversalNewlineIterator<'_> {
    #[inline]
    fn next_back(&mut self) -> Option<Self::Item> {
        let mut text = self.as_str();

        if text.is_empty() {
            return None;
        }

        let len = text.len();

        // Trim any trailing newlines.
        text = match text.as_bytes()[len - 1] {
            b'\n' if len > 1 && text.as_bytes()[len - 2] == b'\r' => &text[..len - 2],
            b'\n' | b'\r' => &text[..len - 1],
            _ => &text,
        };

        // Find the end of the previous line. The previous line is the text up to, but not including the new line character
        let line = match text.rfind(['\n', '\r']) {
            // '\n' or '\r' or '\r\n'
            Some(newline_pos) => {
                let line_start = newline_pos + 1;
                let line = &text[line_start..];
                self.end_offset = self.start_offset + line_start;

                line
            }
            // Last line
            None => {
                self.end_offset = self.start_offset;
                text
            }
        };

        Some(line)
    }
}

impl FusedIterator for UniversalNewlineIterator<'_> {}

[MichaReiser]

Why did we have to add the serde dependency?

[charliermarsh]

Sorry, I should've documented this. I think it's an issue with how I added the serde feature in RustPython. It looks like it doesn't have the derive feature enabled. So we need to use "our" serde, which does have derive, in order for it to compile. (This is based on cursory guess, I didn't look deeply, but I'll create an issue and document.)

[charliermarsh]

I went ahead and added an extension trait, but if you have any feedback @MichaReiser happy to address in a follow-up.