refactor(codegen): add CodeBuffer to fix soundness hole

[DonIsaac]

What This PR Does

Adds CodeBuffer, a simple wrapper over a Vec<u8> with a protective and reduced API for upholding UTF-8 validity guarantees. Closes #6147.

Note that this struct is actually quite small. Most of the added lines are doc comments.

[graphite-app]

Your org has enabled the Graphite merge queue for merging into main

Add the label “0-merge” to the PR and Graphite will automatically add it to the merge queue when it’s ready to merge. Or use the label “hotfix” to add to the merge queue as a hot fix.

You must have a Graphite account and log in to Graphite in order to use the merge queue. Sign up using this link.

[DonIsaac]

• refactor(codegen): add CodeBuffer to fix soundness hole #6148 👈
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @DonIsaac and the rest of your teammates on Graphite

[codspeed-hq]

CodSpeed Performance Report

Merging #6148 will not alter performance

_{Comparing don/09-28-refactor_codegen_add_codebuffer_to_fix_soundness_hole (204bf55) with main (8fe1b0a)}

Summary

✅ 29 untouched benchmarks

[Dunqing]

Overall looks good

[Boshen]

Is this just String? @overlookmotel said it could just be a string but I didn't spend the time converting it.

[overlookmotel]

I really like this. As you say, we were providing a safe API which allowed doing unsafe things. But, beyond that, the str::from_utf8_unchecked calls in codegen gave me the willies, as without reviewing every line of the entire crate, there was no way to know if it was sound or not. Any code anywhere could sneak in a dodgy byte and violate the safety constraints.

CodeBuffer encapsulates the unsafe code and provides a clear contract for what's allowed and what's not. Nice!

Most of my comments below are just nits about naming. I think it's useful to be clear about the difference between bytes and chars.

[overlookmotel]

Is this just String? @overlookmotel said it could just be a string but I didn't spend the time converting it.

Yes, in one sense it's a re-implementation of String. This could be a wrapper around a String with some added unsafe methods to push bytes to it. But I think the implementation in this PR wrapping a Vec<u8> is fine too.

[overlookmotel]

If you want to remove some of the unsafe calls, you could also copy this: https://doc.rust-lang.org/std/ascii/enum.Char.html

Then to push a space, you'd do:

self.code.print_ascii(AsciiByte::Space);

AsciiByte can only be an ASCII byte, so no need for an assertion to check that it is. i.e. Implementation of print_ascii would be just self.buf.push(ascii_byte) (same as unsafe method print_byte), but it can be a safe method.

This will likely optimize better when it's not trivial to deduce that the byte being pushed is definitely ASCII. e.g. if there's code like this in the codegen:

fn print_delimited_list(&mut self, parts: &[&str], delimiter: u8) {
    for (index, s) in parts.iter().enumerate() {
        if index > 0 {
            // Does compiler know for sure that `delimiter` is ASCII
            // to optimize out the assertion in `print_ascii_byte`?
            self.code.print_ascii_byte(delimiter);
        }
        self.code.push_str(s);
    }
}

Also:

pub(crate) fn print_asciis<I>(&mut self, bytes: I)
    where I: IntoIterator<Item = AsciiByte>,
{
    self.buf.extend(bytes.map(|b| b as u8));
}

// Previously
unsafe { self.code.print_unchecked(std::iter::repeat(b'\t').take(self.indent as usize)); }
// Now
self.code_print_asciis(std::iter::repeat(AsciiByte::Tab).take(self.indent as usize));

If AsciiByte::Space etc is too verbose, could add a bit of macro sugar:

// Without sugar
code.print_ascii(AsciiByte::Space);
// With sugar
code.print_ascii(ascii!(' '));

If you did want to do any of this, I'd suggest doing it in a follow-up PR.

[overlookmotel]

One last thing... I was looking at the assembly for the various print methods, and noticed something odd:

print_char(' ') and print_str(" ") seem to produce tighter assembly than print_ascii_byte(b' ').

https://godbolt.org/z/Kv8sEoEed (ignore all the alloc::raw_vec::* stuff at the top - that's all handling Vec growth)

print_ascii_byte has the "this Vec must grow" branch as the default. That's not ideal because growth strategy is doubling, so e.g. when the Vec has just grown from 1024 bytes to 2048 bytes, it won't need to grow again until another 1024 bytes have been pushed. "Needs to grow" is a very rare occurrence.

Also, it pushes 3 registers to the stack before it does anything else, and it's quite hefty overall. Not good!

print_ascii_byte_v2 is my own version which moves all the "Vec needs to grow" logic into a cold path, and only uses 2 registers.

The question is why does std::vec::Vec have this "bad" behavior? This seems odd, as std is highly optimized. It's implausible it wouldn't be the way it is without good reason.

My best guess is that Vec is optimized for a general use case, and it's not uncommon for a Vec to be empty and then have a single item pushed. In which case it only hits the "needs to grow" path, and never hits "no growth necessary". So my print_ascii_byte_v2 would penalize this common case.

But for our use case here, we are always pushing a ton of bytes. So it may be that print_ascii_byte_v2 would work better for this specific use case.

We may also be able to replicate what Stack does to make all these operations even cheaper still. This push is 1 operation shorter than print_ascii_byte_v2, and (perhaps more importantly) uses only 1 register: https://godbolt.org/z/sna4K84G1

This is all micro-optimization stuff, but as pushing bytes to output is such a large component of what oxc_codegen does, and its such a hot path, it might be worth giving it a go.

(again, that'd be best done in a follow-up PR)

[overlookmotel]

@DonIsaac Just to check - are any of my comments above blockers for you?

I think this PR is a great improvement, and would very much like to get it merged. Removing potential UB is always a high priority, in my opinion.

If you're busy and would like me to finish it off, please say so.

[DonIsaac]

@DonIsaac Just to check - are any of my comments above blockers for you?

I think this PR is a great improvement, and would very much like to get it merged. Removing potential UB is always a high priority, in my opinion.

If you're busy and would like me to finish it off, please say so.

None are blockers, I'm just working towards it slowly. I've had a lot going on these past few days 😄

dismiss

[Boshen]

Merge activity

• Oct 13, 5:14 AM EDT: The merge label '0-merge' was detected. This PR will be added to the Graphite merge queue once it meets the requirements.
• Oct 13, 5:14 AM EDT: A user added this pull request to the Graphite merge queue.
• Oct 13, 5:19 AM EDT: A user merged this pull request with the Graphite merge queue.