Update string nodes for implicit concatenation

[dhruvmanila]

Summary

This PR updates the string nodes (ExprStringLiteral, ExprBytesLiteral, and ExprFString) to account for implicit string concatenation.

Motivation

In Python, implicit string concatenation are joined while parsing because the interpreter doesn't require the information for each part. While that's feasible for an interpreter, it falls short for a static analysis tool where having such information is more useful. Currently, various parts of the code uses the lexer to get the individual string parts.

One of the main challenge this solves is that of string formatting. Currently, the formatter relies on the lexer to get the individual string parts, and formats them including the comments accordingly. But, with PEP 701, f-string can also contain comments. Without this change, it becomes very difficult to add support for f-string formatting.

Implementation

The initial proposal was made in this discussion: #6183 (comment). There were various AST designs which were explored for this task which are available in the linked internal document¹.

The selected variant was the one where the nodes were kept as it is except that the implicit_concatenated field was removed and instead a new struct was added to the Expr* struct. This would be a private struct would contain the actual implementation of how the AST is designed for both single and implicitly concatenated strings.

This implementation is achieved through an enum with two variants: Single and Concatenated to avoid allocating a vector even for single strings. There are various public methods available on the value struct to query certain information regarding the node.

The nodes are structured in the following way:

ExprStringLiteral - "foo" "bar"
|- StringLiteral - "foo"
|- StringLiteral - "bar"

ExprBytesLiteral - b"foo" b"bar"
|- BytesLiteral - b"foo"
|- BytesLiteral - b"bar"

ExprFString - "foo" f"bar {x}"
|- FStringPart::Literal - "foo"
|- FStringPart::FString - f"bar {x}"
  |- StringLiteral - "bar "
  |- FormattedValue - "x"

Visitor

The way the nodes are structured is that the entire string, including all the parts that are implicitly concatenation, is a single node containing individual nodes for the parts. The previous section has a representation of that tree for all the string nodes. This means that new visitor methods are added to visit the individual parts of string, bytes, and f-strings for Visitor, PreorderVisitor, and Transformer.

Test Plan

• cargo insta test --workspace --all-features --unreferenced reject
• Verify that the ecosystem results are unchanged

Footnotes

1. Internal document: https://www.notion.so/astral-sh/Implicit-String-Concatenation-e036345dc48943f89e416c087bf6f6d9?pvs=4 ↩

[dhruvmanila]

Current dependencies on/for this PR:

• main
  • PR Remove leftover constant tuple reference #8062
    • PR New Singleton enum for PatternMatchSingleton node #8063
      • PR Split Constant to individual literal nodes #8064
        • PR Update string nodes for implicit concatenation #7927 👈
          • PR Explicit as_str (no deref), add no allocation methods #8826
            • PR Remove #[allow(unused_variables)] from visitor methods #8828
              • PR New AST nodes for f-string elements #8835
                • PR Introduce StringLike enum #9016

This stack of pull requests is managed by Graphite.

[BurntSushi]

Ah this is much better! Thank you for breaking it up. It made it oodles easier to review. :-)

Overall great work. I especially found the new fstring/string/byte-string types pretty easy to follow and understand. Their structure makes sense to me.

[BurntSushi]

My first thought here is to wonder why, based on the comment, this isn't Single(FString) in order to better encode the invariant.

[BurntSushi]

Hmmm yes, I see, the FStringValue::parts method makes it rather annoying to do that without extra type machinery that probably isn't worth doing.

[dhruvmanila]

Yeah, I tried adding a FStringPartRef enum but then we can't get the mutable reference for it which is required in the Transformer.

enum FStringPartRef<'a> {
	Literal(&'a StringLiteral),
	FString(&'a FString),
}

[charliermarsh]

Interesting, how did this get so much simpler?

[dhruvmanila]

Ah, that's the magic of this refactor 😉

Jokes aside, this basically reverts this commit ac4a4da which fixed the bug where this rule wasn't producing the correct fix for an implicitly concatenated string. But, now that we work on individual parts, that isn't required.

[charliermarsh]

This is impressive work. There's a huge surface area here. You did a great job of breaking down the individual PRs and walking us through it.

I left a few questions, which I'm happy to continued discussing, but erring on the side of approving.

[charliermarsh]

I wonder if these should be smallvec... It's probably worth benchmarking in a separate PR. In practice, I'd bet the vast majority of concatenations are <= 5 elements.

[charliermarsh]

This is an example where... we can probably do this without the concatenation? Imagine if we had a custom matcher that worked on implicit concatenations (i.e., a vector of strings). \cc @BurntSushi in case you have obvious ideas here.

[BurntSushi]

I think the closest "out of the box" experience you can get is stream searching with aho-corasick, but you need to provide a std::io::Read impl. Mildly annoying but doable in this context.

But I would say to write the simple/obvious code for now, and if this ends up popping up on a profile then we can revisit it and do something bespoke here or use aho-corasick.

[dhruvmanila]

I think I would prefer to avoid the complexity for such simple cases but I do understand that this could be useful in similar cases where the strings are long.