Fix argument index remapping in format_args flattening.

This commit is contained in:
Mara Bos 2023-01-13 19:01:19 +01:00
parent b7678d48b8
commit b6c988b041

View File

@ -68,33 +68,21 @@ fn flatten_format_args(mut fmt: Cow<'_, FormatArgs>) -> Cow<'_, FormatArgs> {
args.extend(remaining_args); args.extend(remaining_args);
// Correct the indexes that refer to the arguments after the newly inserted arguments. // Correct the indexes that refer to the arguments after the newly inserted arguments.
for piece in &mut fmt.template { for_all_argument_indexes(&mut fmt.template, |index| {
if let FormatArgsPiece::Placeholder(placeholder) = piece if *index >= old_arg_offset {
&& let Ok(index) = &mut placeholder.argument.index
&& *index >= old_arg_offset
{
*index -= old_arg_offset; *index -= old_arg_offset;
*index += new_arg_offset; *index += new_arg_offset;
} }
} });
// Now merge the placeholders: // Now merge the placeholders:
let rest = fmt.template.split_off(i + 1); let rest = fmt.template.split_off(i + 1);
fmt.template.pop(); // remove the placeholder for the nested fmt args. fmt.template.pop(); // remove the placeholder for the nested fmt args.
// Insert the pieces from the nested format args, but correct any
for piece in fmt2.template.drain(..) { // placeholders to point to the correct argument index.
match piece { for_all_argument_indexes(&mut fmt2.template, |index| *index += arg_index);
FormatArgsPiece::Literal(s) => fmt.template.push(FormatArgsPiece::Literal(s)), fmt.template.append(&mut fmt2.template);
FormatArgsPiece::Placeholder(mut p) => {
// Correct the index to refer to the right place into the outer argument list.
if let Ok(n) = &mut p.argument.index {
*n += arg_index;
}
fmt.template.push(FormatArgsPiece::Placeholder(p));
}
}
}
fmt.template.extend(rest); fmt.template.extend(rest);
// Don't increment `i` here, so we recurse into the newly added pieces. // Don't increment `i` here, so we recurse into the newly added pieces.
@ -150,16 +138,17 @@ fn inline_literals(mut fmt: Cow<'_, FormatArgs>) -> Cow<'_, FormatArgs> {
// Drop all the arguments that are marked for removal. // Drop all the arguments that are marked for removal.
let mut remove_it = remove.iter(); let mut remove_it = remove.iter();
fmt.arguments.all_args_mut().retain(|_| remove_it.next() != Some(&Some(true))); fmt.arguments.all_args_mut().retain(|_| remove_it.next() != Some(&Some(true)));
// Calculate the mapping of old to new indexes for the remaining arguments.
let index_map: Vec<usize> = remove
.into_iter()
.scan(0, |i, remove| {
let mapped = *i;
*i += (remove != Some(true)) as usize;
Some(mapped)
})
.collect();
// Correct the indexes that refer to arguments that have shifted position. // Correct the indexes that refer to arguments that have shifted position.
for piece in &mut fmt.template { for_all_argument_indexes(&mut fmt.template, |index| *index = index_map[*index]);
let FormatArgsPiece::Placeholder(placeholder) = piece else { continue };
let Ok(arg_index) = &mut placeholder.argument.index else { continue };
for i in 0..*arg_index {
if remove[i] == Some(true) {
*arg_index -= 1;
}
}
}
} }
fmt fmt
@ -572,3 +561,22 @@ fn may_contain_yield_point(e: &ast::Expr) -> bool {
visitor.visit_expr(e); visitor.visit_expr(e);
visitor.0 visitor.0
} }
fn for_all_argument_indexes(template: &mut [FormatArgsPiece], mut f: impl FnMut(&mut usize)) {
for piece in template {
let FormatArgsPiece::Placeholder(placeholder) = piece else { continue };
if let Ok(index) = &mut placeholder.argument.index {
f(index);
}
if let Some(FormatCount::Argument(FormatArgPosition { index: Ok(index), .. })) =
&mut placeholder.format_options.width
{
f(index);
}
if let Some(FormatCount::Argument(FormatArgPosition { index: Ok(index), .. })) =
&mut placeholder.format_options.precision
{
f(index);
}
}
}