Niche Optimization for `Option<ReferenceId>` and `IndexVec<ReferenceId, T>`

[Boshen]

we have these index types

https://github.com/oxc-project/oxc/blob/712ee0dde8e03d79513b42bad3b6736b566dd60b/crates/oxc_syntax/src/reference.rs#L6-L8

https://github.com/oxc-project/oxc/blob/712ee0dde8e03d79513b42bad3b6736b566dd60b/crates/oxc_syntax/src/symbol.rs#L7-L9

where

https://github.com/oxc-project/oxc/blob/712ee0dde8e03d79513b42bad3b6736b566dd60b/crates/oxc_ast/src/ast/js.rs#L451

https://github.com/oxc-project/oxc/blob/712ee0dde8e03d79513b42bad3b6736b566dd60b/crates/oxc_semantic/src/reference.rs#L22

can be trimmed by niche optimization.

But,

use std::num::NonZeroU32;
define_index_type! {
    pub struct ReferenceId = NonZeroU32;
}

is not supported by the macro define_index_type

error[E0599]: no function or associated item named `max_value` found for struct `NonZero` in the current scope
   --> crates/oxc_syntax/src/reference.rs:7:1
    |
7   | / define_index_type! {
8   | |     pub struct ReferenceId = NonZeroU32;
9   | | }
    | |_^ function or associated item not found in `NonZero<u32>`
    |
note: if you're trying to build a new `NonZero<u32>` consider using one of the following associated functions:
      NonZero::<T>::new
      NonZero::<T>::new_unchecked
   --> /Users/boshen/.rustup/toolchains/1.78.0-aarch64-apple-darwin/lib/rustlib/src/rust/library/core/src/num/nonzero.rs:308:5

but the good news is that we forked it and lives in https://github.com/oxc-project/oxc/tree/main/crates/oxc_index

The task of this issue is to make it work, if possible.

[overlookmotel]

NonMaxU32

Personally, I am not a massive fan of the nonmax crate.

Unlike NonZeroU32, nonmax::NonMaxU32 does not have native compiler support. It's basically a wrapper around NonZeroU32 and stores n ^ u32::MAX internally. So it has a (small) cost to convert the value every time you read or write it.

https://github.com/LPGhatguy/nonmax/blob/645abf30102dd880542865195c36893ee89994b5/src/lib.rs#L132-L137

This differs from Rust's native NonMaxU32 which uses #[rustc_layout_scalar_valid_range_end], but that's not available in userland.

It's annoying that nonmax crate is not "honest" about this, and at first glance looks like it's doing the same thing as the native rustc types.

So unfortunately, there is no NonMaxU32 type available to us which doesn't come with a cost. The cost is small, but in my view best avoided if we can.

What we could do instead

Initialize our indexed Vecs with a dummy entry at index 0, so usable indexes start at 1.

Then 2 options:

1. Use NonZeroU32 as index, and store Option<NonZeroU32> in AST and Semantic etc.
2. Use plain u32 as index and store u32 (no Option). Treat the value 0 as a sentinel meaning "not initialized".

Personally I favour option 2. I assume there must be code all over the place which reads these ID fields and then calls unwrap() on the Option, because we know semantic has run and filled in all the IDs. However, I doubt the compiler can prove that these Options are never None, so it's likely producing unnecessary panicking branches all over the place in the code. Using plain u32 should avoid that.

[overlookmotel]

There's also a much more complicated, but more performant solution.

As far as I'm aware, the Vecs which ReferenceId, SymbolId and ScopeId index into can only grow, never contract.

We could capitalize on this invariant to remove bounds checks. Something like:

pub struct ReferenceVec(Vec<Whatever>);
pub struct ReferenceId(u32);

impl ReferenceId {
  const EMPTY: Self = ReferenceId(0);
}

impl ReferenceVec {
  pub fn new() -> Self {
    // Initialize vec with a dummy at index 0
    let mut inner = Vec::new();
    inner.push(Whatever::default());
    Self(inner)
  }

  pub fn push(&mut self, value: Whatever) -> ReferenceId {
    let index = ReferenceId(self.0.len());
    self.0.push(value);
    index
  }

  // Returns `&Whatever` not `Option<&Whatever>`
  // because `ReferenceId` is always a valid index
  pub fn get(&self, index: ReferenceId) -> &Whatever {
    // `index` must be in bounds, so skip bounds check
    unsafe { self.0.get(index.0).unwrap_unchecked() }
  }

  // NB: No `pop` method
}

i.e. The only way to get a ReferenceId is via push(), and because ReferenceVec never shrinks, once you have a ReferenceId, it always remains valid. So no need for bounds check on get().

Actual implementation would need to be more complex than the above. Would need some mechanism to prevent you using a ReferenceId from one ReferenceVec on another ReferenceVec. But hopefully the above gives the general idea.

NB: Unclear how much of a gain this would be, and whether worth the complexity. In my opinion, we should do the main change (shrink Option<ReferenceId>) and then maybe look into this further optimization later on.

[Boshen]

It seems like you're proposing a new crate for this use case, instead of reusing the crates I found from the community.

Edit: the issue description is changed from NonMax to NonZero.

[overlookmotel]

It seems like you're proposing a new crate for this use case, instead of reusing the crates I found from the community.

I hope not! I thought the reason we swapped index_vec for our own oxc_index was so we could bend it to our use cases. I think @rzvxa had exactly the NonZeroU32 use case in mind when he proposed bringing the indexed vec type "in house".

The append-only-with-no-bounds-checks Vec would I guess have to be a different type. But I don't think we should do that right now - I just mention it as something to have in mind for later.