Skip to content

Commit

Permalink
convert token_infos from vector to iterator
Browse files Browse the repository at this point in the history
  • Loading branch information
augustelalande committed May 19, 2024
1 parent 50ba417 commit c9c5e0c
Showing 1 changed file with 78 additions and 59 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -47,42 +47,69 @@ struct TokenInfo {
token_end_within_physical_line: i64,
}

/// Compute the `TokenInfo` of each token.
fn get_token_infos<'a>(
logical_line: &LogicalLine,
locator: &'a Locator,
/// An iterator over `TokenInfo`.
struct TokenInfoIter<'a> {
locator: &'a Locator<'a>,
indexer: &'a Indexer,
) -> Vec<TokenInfo> {
let mut token_infos = Vec::with_capacity(logical_line.tokens().len());
let mut current_line_idx: usize = 0;
// The first physical line occupied by the token, since a token can span multiple physical lines.
let mut first_physical_line_start: usize;
let mut next_continuation;
if let Some(first_token) = logical_line.first_token() {
first_physical_line_start = first_token.range.start().into();
next_continuation = continuation_line_end(first_token, locator, indexer);
} else {
return token_infos;
};

let mut current_physical_line_start: usize;
let mut prev_token: Option<&LogicalLineToken> = None;
for token in logical_line.tokens() {
let mut start_physical_line_idx = current_line_idx;
current_physical_line_start = first_physical_line_start;
tokens: std::slice::Iter<'a, LogicalLineToken>,
current_line_idx: usize,
first_physical_line_start: usize,
next_continuation: Option<TextSize>,
current_physical_line_start: usize,
prev_token: Option<&'a LogicalLineToken>,
}

impl<'a> TokenInfoIter<'a> {
fn new(logical_line: &'a LogicalLine, locator: &'a Locator, indexer: &'a Indexer) -> Self {
let (first_physical_line_start, next_continuation) =
if let Some(first_token) = logical_line.first_token() {
(
first_token.range.start().into(),
continuation_line_end(first_token, locator, indexer),
)
} else {
(0, None)
};

TokenInfoIter {
locator,
indexer,
tokens: logical_line.tokens().iter(),
current_line_idx: 0,
first_physical_line_start,
next_continuation,
current_physical_line_start: 0,
prev_token: None,
}
}
}

impl<'a> Iterator for TokenInfoIter<'a> {
type Item = TokenInfo;

fn next(&mut self) -> Option<Self::Item> {
let token = self.tokens.next()?;

let mut start_physical_line_idx = self.current_line_idx;
self.current_physical_line_start = self.first_physical_line_start;

// Check for escaped ('\') continuation lines between the previous and current tokens.
if let Some(prev_token) = prev_token {
if next_continuation.is_some() && token.start() >= next_continuation.unwrap() {
next_continuation = continuation_line_end(token, locator, indexer);

let trivia =
locator.slice(TextRange::new(prev_token.range.end(), token.range.start()));
for (index, _text) in trivia.match_indices('\n') {
start_physical_line_idx += 1;
current_line_idx = start_physical_line_idx;
first_physical_line_start = usize::from(prev_token.range.end()) + index + 1;
current_physical_line_start = first_physical_line_start;
if let Some(prev_token) = self.prev_token {
if let Some(next_continuation) = self.next_continuation {
if token.start() >= next_continuation {
self.next_continuation =
continuation_line_end(token, self.locator, self.indexer);

let trivia = self
.locator
.slice(TextRange::new(prev_token.range.end(), token.range.start()));
for (index, _text) in trivia.match_indices('\n') {
start_physical_line_idx += 1;
self.current_line_idx = start_physical_line_idx;
self.first_physical_line_start =
usize::from(prev_token.range.end()) + index + 1;
self.current_physical_line_start = self.first_physical_line_start;
}
}
}
}
Expand All @@ -95,39 +122,39 @@ fn get_token_infos<'a>(
| TokenKind::FStringEnd
) {
// Look for newlines within strings.
let trivia = locator.slice(token.range);
let trivia = self.locator.slice(token.range);
for (index, _text) in trivia.match_indices('\n') {
current_line_idx += 1;
current_physical_line_start = usize::from(token.range.start()) + index + 1;
self.current_line_idx += 1;
self.current_physical_line_start = usize::from(token.range.start()) + index + 1;
}
}

token_infos.push(TokenInfo {
let token_info = TokenInfo {
start_physical_line_idx,
end_physical_line_idx: current_line_idx,
end_physical_line_idx: self.current_line_idx,
token_start_within_physical_line: i64::try_from(
usize::from(token.range.start()) - first_physical_line_start,
usize::from(token.range.start()) - self.first_physical_line_start,
)
.expect("Lines are expected to be relatively short."),
token_end_within_physical_line: i64::try_from(
usize::from(token.range.end()) - current_physical_line_start,
usize::from(token.range.end()) - self.current_physical_line_start,
)
.expect("Lines are expected to be relatively short."),
});
};

if matches!(
token.kind,
TokenKind::NonLogicalNewline | TokenKind::Newline
) {
current_line_idx += 1;
first_physical_line_start = token.range.end().into();
self.current_line_idx += 1;
self.first_physical_line_start = token.range.end().into();
} else {
first_physical_line_start = current_physical_line_start;
self.first_physical_line_start = self.current_physical_line_start;
}
prev_token = Some(token);
}
self.prev_token = Some(token);

token_infos
Some(token_info)
}
}

fn continuation_line_end(
Expand Down Expand Up @@ -199,16 +226,6 @@ pub(crate) fn continuation_lines(
// The pycodestyle implementation makes use of negative values,
// converting the indent_size type at the start avoids converting it multiple times later.
let indent_size = i64::try_from(indent_size).expect("Indent size to be relatively small.");
let token_infos = get_token_infos(logical_line, locator, indexer);
let nb_physical_lines = if let Some(last_token_info) = token_infos.last() {
1 + last_token_info.start_physical_line_idx
} else {
1
};

if nb_physical_lines == 1 {
return;
}

// Indent of the first physical line.
let start_indent_level = expand_indent(
Expand All @@ -232,7 +249,7 @@ pub(crate) fn continuation_lines(
// Brackets opened on a line.
let mut brackets_opened = 0u32;
// Relative indents of physical lines.
let mut rel_indent: Vec<i64> = vec![0; nb_physical_lines];
let mut rel_indent: Vec<i64> = vec![0];
// For each depth, collect a list of opening rows.
let mut open_rows: Vec<Vec<usize>> = Vec::with_capacity(max_depth + 1);
open_rows.push(vec![0]);
Expand All @@ -248,7 +265,9 @@ pub(crate) fn continuation_lines(
let mut indent = Vec::with_capacity(max_depth + 1);
indent.push(start_indent_level);

for (token, token_info) in zip(logical_line.tokens(), &token_infos) {
let token_infos_iter = TokenInfoIter::new(logical_line, locator, indexer);
for (token, token_info) in zip(logical_line.tokens(), token_infos_iter) {
rel_indent.resize(token_info.end_physical_line_idx + 1, 0);
let mut is_newline = row < token_info.start_physical_line_idx;
if is_newline {
row = token_info.start_physical_line_idx;
Expand Down

0 comments on commit c9c5e0c

Please sign in to comment.