Add wasm-encoder API to directly encode instructions

[samestep]

Following up on a Zulip conversation with @alexcrichton, this PR adds a new wasm_encoder::InstructionSink type, which has an encoding method for each variant of wasm_encoder::Instruction:

/// An encoder for Wasm instructions.
#[derive(Debug)]
pub struct InstructionSink<'a> {
    sink: &'a mut Vec<u8>,
}

impl<'a> InstructionSink<'a> {
    /// Create an instruction encoder pointing to the given byte sink.
    pub fn new(sink: &'a mut Vec<u8>) -> Self {
        Self { sink }
    }

    // ...

    /// Encode [`Instruction::I32Load`].
    pub fn i32_load(&mut self, m: MemArg) -> &mut Self {
        self.sink.push(0x28);
        m.encode(self.sink);
        self
    }

    // ...
}

The impl Encode for Instruction is then refactored to just delegate to these methods.

See https://github.com/samestep/wasm-encoder-performance for more details; in particular, the primary motivation for this PR is an observed 30% difference in performance when encoding sequences of instructions that are known at compile time, such as in crates/wit-component/src/gc.rs. Subjectively I also consider the new API to be more ergonomic.

The one tricky part in designing this PR was the fact that 65 Instruction variants have named fields, so if the new API just used u32 for everything like the old API did, it could be more error-prone than before. To address this, I introduced several new types like this:

/// A Wasm _localidx_.
#[derive(Clone, Copy, Debug)]
pub struct LocalIdx(pub u32);

So, the benefit from named fields in nearly all those variants is now provided more directly by the typechecker, with better ergonomics in the case where the programmer already has a value of the more specific u32 wrapper type in hand; and in the other cases, writing out the name of the type to wrap a u32 is no less ergonomic than would be writing the name of the field in the first place, since often the field name is essentially the same as the type name.

Before I switched to using these wrapper types, I considered instead turning each Instruction variant into its own type rather than its own function, and simply implementing Encode for each such type. I decided against this because, unlike the InstructionSink approach, there'd be no way to access all these new instruction encodings from a Function, except maybe by adding a new trait InstructionVariant or something like that, but that seemed like a whole can of worms that I didn't want to open. I find this InstructionSink approach to be far more ergonomic.

In this same PR I've also refactored as many usages as possible to prefer the new API over the old one, because of the resulting improvements in performance and (in my opinion) ergonomics. If it would be preferable for me to undo that and just introduce the new API without refactoring, please let me know and I can undo those changes. Also please let me know if I should make any changes to the new API itself: the repository I linked above contains a script to generate the bulk of the new code, as well as ast-grep rules to do most of the refactoring in this repo, so it's now relatively easy to iterate on the design.

[alexcrichton]

Thanks again so much for your work on this, I like how this turned out and you're definitely not the only one who feels this is more idiomatic when you've got a static list of instructions!

My main thoughts on this are around the newtype u32 wrappers you've added here. I definitely agree that only having u32 is going to be more error-prone now that the named fields of struct variants aren't available, but that being said I also feel that this new addition doesn't really feel "baked in". For example the wasmparser crate's VisitOperator trait doesn't have typed indices and nothing else in this repository currently has typed indices either. This is in contrast to Wasmtime, however, which has typed indices basically everywhere.

So to be it's a bit of a tradeoff where what's earned with typed indices is type annotations of what's what, but what's lost with type annotations is ergonomics where the crate isn't currently really built to work with typed indicies. For example ideally adding a type to a type section would return a TypeIdx (or something like that), but that's not fleshed out entirely just yet.

There's also a bit of a downside where there's some instructions like memory.copy which can copy between two linear memories and while having MemoryIdx helps once we've lost the named fields it's not clear whether the source or destination comes first.

Overall personally I'd lean towards removing the typed indices in favor of a future refactoring to possibly include them later (but that'd likely be quite a big change). What do you think about that though? Do you think that this is worth the tradeoff of "not integrating well" of sorts with the rest of the crate?

[samestep]

Thanks for the review Alex!

Seems like we're mostly on the same page; I for sure agree that it feels weird for InstructionSink to be the only type that uses these newtype wrappers. I personally love typed indices so if I could wave a magic wand to make them be used throughout wasm-tools then I'd do so. (I'd also probably be willing to work on such a refactor if that'd be desirable.)

I went back and forth on the design for this PR a lot while I was prototyping it last week, and despite the mismatch, I still think that the right thing to do is to adopt this new InstructionSink API with the newtype u32 wrappers, rather than without. My reasoning is that, for most instructions (e.g. struct_get), typing the names of the newtype wrappers is pretty much the same (ergonomically) as typing the field names. E.g., before:

newfunc.instruction(&Instruction::MemoryInit {
    mem: init.memory,
    data_index: init.segment,
});

vs. after:

insn.memory_init(MemIdx(init.memory), DataIdx(init.segment))

So I still see this as a strict improvement because it's net-neutral ergonomically and net-positive for performance.

Another angle: in constructing this PR, I constrained the design to be strictly backward-compatible. Using the newtype wrappers in other types in the public API would be a backward-incompatible change and I don't have a great intuition for how much the wasm-tools team wants to avoid those, so I just didn't do that in this PR. But if you feel that this change would make sense only after having incorporated these newtype wrappers into the rest of the API (entailing breaking changes) then I'd be fine with that.

Regarding copy instructions, I agree that it's a bit unfortunate. I have a couple thoughts:

1. For all the multi-parameter instructions in this PR, I manually went and checked the Structure section of the Wasm spec to make sure that I was putting the parameters in the same order as the spec. This resulted in a couple changes from the order of fields in Instruction: for instance, Instruction::MemoryCopy has src_mem followed by dst_mem, but the spec lists the dst_mem first. (An interesting example is Instruction::TableInit, which has elem_index before table, just like the Binary Format, but unlike the Structure.)

   Anyways, all this is to say that in the spec, for copy instructions (array.copy, memory.copy, and table.copy), the destination is always listed first. So at least they're all consistent with each other.

2. We've already factored out the MemArg type which is shared among many instruction variants; if we're worried about the ergonomics of copy instructions, we could also create a type to represent the args for copy instructions, and share it among all three?

   struct CopyArg<T> {
       dst: T,
       src: T,
   }

   Not sure what the name should be; unfortunately Copy is already taken.

The only other instructions that have this issue are Instruction::BrOnCast and Instruction::BrOnCastFail; for those as well, we could create a struct to hold the from_ref_type and to_ref_type parameters if we wanted.

Sorry for the long response. Overall, I originally started working on this PR for the performance and you're the maintainer so I'll go with whatever you prefer on any of these ergonomics issues. Personally, though, I really want this new API to be strictly better than the old one whenever the instruction variant is static, and I don't see how to do that without these newtype u32 wrappers. And then hopefully we can make the API more consistent going forward?

[samestep]

By the way, if such a refactoring were to take place to use typed indices more consistently throughout wasm-tools, in what crate would those type definitions live? (I think part of the answer to this question would also include why wasmparser::ValType and wasm_encoder::ValType are distinct in the status quo.)

[alexcrichton]

You bring up all good points! I think personally though I'm not entirely convinced, but I'd also like to cc @fitzgen on this as well. Nick do you have thoughts on the use of newtypes here in wasm-encoder?

My reasoning is that, for most instructions (e.g. struct_get), typing the names of the newtype wrappers is pretty much the same (ergonomically) as typing the field names.

I don't doubt your conclusion here but I also personally place less weight on the transitionary nature of changes like this. I instead try to evaluate the end state in its own right, so comparing:

insn.memory_init(MemIdx(init.memory), DataIdx(init.segment))

vs

insn.memory_init(init.memory, init.segment)

the latter is more desirable to me. Another space that I can fill in some context on:

Using the newtype wrappers in other types in the public API would be a backward-incompatible change and I don't have a great intuition for how much the wasm-tools team wants to avoid those, so I just didn't do that in this PR.

Currently wasm-encoder, and almost all other crates in this repository, provide no strict guarantees about backwards compatibilty. Instead the release process we have always issues a new major version bump. In that sense we're technically free to update APIs whenever we want, but naturally we try to reduce churn where possible to not unnecessarily cause headaches. Basically breaking changes are fine, but ideally done at-once instead of spread out over time.

In that sense I'm not worried about changing existing APIs. For example I think it would be reasonable to change the APIs in CoreTypeEncoder to all return TypeIdx. The MemorySection::memory function could change from returning &mut Self to returning MemoryIdx as well for example.

To me the main benefit of typed indices though is that you worry about construction of the index in one location and otherwise generally don't worry about it. You get compiler errors when the wires are crossed by accident but for the most part you aren't responsible for constructing indices manually. That to me is the main drawback of this current state of affairs which is that the typed index has to be created manually every time. Ideally the indices are already stored at rest in their typed form instead.

Regarding copy instructions, I agree that it's a bit unfortunate. I have a couple thoughts:

Thinking more on this with what you've brought up as well, I think it's probably best to basically ignore the concern I'm thinking of. While it's possible to do something like CopyArg that feels like sort of overkill for a pretty niche problem. I think it's reasonable to rely on documentation/IDEs with parameter name hints and such to cover this use case.

By the way, if such a refactoring were to take place to use typed indices more consistently throughout wasm-tools, in what crate would those type definitions live?

Heh this is an interesting question :)

The requirements of wasm-encoder and wasmparser are different enough that I think it's a good reason to not merge the two type hierarchies into one crate at this time. (historically they both were initially developed independently which is why they didn't start merged). There's subtle requirements about creating things vs parsing things which leads to domain-specific problems which is what leads me to the feeling that we shouldn't entirely merge everything.

Now that being said typed indices are far simpler than the representation of a function type as it's unconditionally a typed wrapper around a 32-bit integer. There's still questions about what exactly does the API look like but that's probably going to be the same for wasmparser and wasm-encoder. I think it's probably best though to start in one crate and maybe move out from there as opposed to trying to unify everything initially.

Basically tl;dr; I think it's reasonable to live in wasm-encoder to start and we can evaluate moving it elsewhere later.

---

Ok well that's sort of an equally long reply, en garde! In any case though personally I still feel that I'd like to move in the direction of typed indices in the future but it's best done as a follow-up. The main benefit would be to avoid having end users manually construct indexes themselves and instead they're naturally created as part of the rest of the wasm-encoder API.

[samestep]

This all makes sense, thanks Alex! If you and Nick want me to remove the newtypes from this PR and then move forward to hopefully merge it, I'm happy to do that.

Would you be open to me starting on a separate PR to use typed indices in wasm-encoder?

Also, as an aside: even after weaving typed indices throughout the whole API, that still won't mean users never have to construct them themselves. For instance, with literal constants, they'd still have to write .local_get(LocalIdx(0)) rather than .local_get(0). But for most other cases hopefully it should "just work."

[fitzgen]

First: @samestep thanks for taking this PR on! Its great to have both an instruction enum and individual methods. As I think Alex mentioned on Zulip, you end up wanting both in the limit because which one you reach for depends on the exact context of your usage.

I think typed indices in general are great, but agree with Alex that we should probably try to split that out into a separate follow up. That follow up should ideally add support for them to the whole workspace rather than just wasm-encoder, or at least add the infrastructure/groundwork for that, via annotating immediates with a @table-index (or whatever) type in the for-each-instruction macro.

The requirements of wasm-encoder and wasmparser are different enough that I think it's a good reason to not merge the two type hierarchies into one crate at this time. (historically they both were initially developed independently which is why they didn't start merged). There's subtle requirements about creating things vs parsing things which leads to domain-specific problems which is what leads me to the feeling that we shouldn't entirely merge everything.

Agreed that we don't want to have a separate wasm-tools-types crate or whatever.

That said, I do think that wasmparser could benefit from using newtypes for various indices in the same way that wasm-encoder can. (See comment above about adding type annotations to the for-each-instruction macro.)