feat: support fixed-width Nat types as array and blob indices

[q-uint]

Summary

• Allow Nat8, Nat16, Nat32, and Nat64 values as array indices (and Nat64 for blobs), bypassing bignum conversion for direct word-sized index operations
• All fixed-width Nat types compile to direct indexing, reducing IC instruction counts by 30-59% compared to Nat indices
• Changes span the type checker, both codegen backends, IR checker, both interpreters, and include run/trap/fail/bench tests

IC instruction counts (1000 iters x 256 elements)

Type	Instructions	vs Nat
Nat (baseline)	25,687,306
Nat8	11,036,370	-57%
Nat16	11,079,329	-57%
Nat32	10,567,324	-59%
Nat64	17,991,325	-30%
Nat64+toNat()	20,807,325	-19%

[q-uint]

CI is failing because cachix watch-exec requires CACHIX_AUTH_TOKEN, which isn't available to workflows triggered from fork PRs? The build/test themselves should be fine, it's just the Cachix authentication step that errors out.

[ggreif]

Just an idea. Normally a small Nat will be used for indexing, and checking for that should be trivial. I.e. masking the highest and lowest bits, and if both zero shift the number by 1/2 to the left. This is the offset from the array's start. Maybe we should run the benchmarks agains this speculative (low-hanging fruit) attempt.

[q-uint]

Tried this out: mask bit 0 (scalar tag) and bit 63 (sign), and if both zero, shift right to recover the raw index and go straight to Arr.idx, bypassing the idx_bigint -> BigNum.to_word64_with call chain entirely. The catch is that BigNum.to_word64_with (in MakeCompact) already has an inline tagged-scalar fast path: it checks bit 0, branches, and unboxes with a shift. So the speculative check is doing essentially the same work, just one function call closer to the indexing site.

Benchmark results (256K accesses, enhanced backend):

Index type	Before	After
Nat32	19,538,348	19,538,348
Nat64	23,634,342	23,634,342
Nat	36,968,319	36,712,319 (-0.7%)
Nat64->toNat	29,522,342	29,266,342 (-0.9%)

But maybe I interpreted the suggestion wrong?

4299689

[Kamirus]

CI is failing because cachix watch-exec requires CACHIX_AUTH_TOKEN, which isn't available to workflows triggered from fork PRs? The build/test themselves should be fine, it's just the Cachix authentication step that errors out.

Now the CI works with forks as well

[crusso]

Hi Quint,

Thanks for the PR. It looks like you are essentially overloading array indexing on all unsigned types.

However, this seems to fly against the Motoko KISS principle (although maybe that ship sailed long ago), especially not allowing implicit conversions between types.

What is the motivation? If it is just perf, not programmer convenience, I think we could easily enough optimize any
a.[NatXX.toNat(n)] indexing (in compile_array_index) without all these changes to typing etc.

Would that be enough for you or do you actually want to avoid the user writing those casts?

I'm a little worried that flipping from checking against Nat mode to inferring the type of the index might actually break a fair bit of code.

If we really want to support this, maybe just adding explict a.sub8() etc prims would be enough for people that want to micro-optimize their code.

BTW, I heard a rumour that you want to extend subtyping and allow NatXX<: Nat but that's a non-starter because Motoko's subtyping is by subsumption, not coercion. That is subtyping always compiles to a no-op, not a cast that transforms the value to a new one. To support NatXX <: Nat we would need to ensure the representation of these values is always the same or modify every operation on Nat to allow all of the NatXX representations as inputs. A bit like you've done for the special case of array indexing.

[crusso]

(labelling this DO-NOT-MERGE for now)

[q-uint]

Hi Claudio (@crusso), this was just me going on a wild adventure. My motives are just fun and random perf improvements, and I don't like the number casts personally.

Happy to just close this if it goes against the Motoko design principles, same for the rumored issue #5934. I hope I didn't take up too much time.

[crusso]

For perf, I think we could just spot the immediate cast before indexing and do a fastpath.

For convenience, another idea might be to not overload (just) indexing, but allow the compiler to insert NatX->Nat casts (and NatX/IntX -> casts) where appropriate, guided by bi-directional type checking. I can see this getting in the way of tail calls sometimes. Might also be problematic for type argument inference... an application might check with type arguments (and implicit casting), but not without.

[skilesare]

I'll argue for doing something. This is having a material effect on some things that I'm doing that are currently blowing out the instruction limit. By my calculations, I'm doing over 5.1 million iterations over various different arrays...many of which are small 4-item arrays that could really benefit from a Nat8 loop given @q-uint's numbers above.

I get the KISS stuff, but the reality of the language is that right now it is more performant to do

  var i64 : Nat64 = 0;
  while (i64 <= 7) { // 0..8
      let i = Nat64.toNat(i64);
      t += n[i] + other.n[i];
  };

Than

  var i : Nat = 0;
  while (i64 <= 7) { // 0..8
      t += n[i] + other.n[i];
  };
```

...which seems kind of not simple. At least from the user's perspective.

[q-uint]

Claudio's suggestion targets a different scenario than what the PR handles.

My changes make Nat64 (and other NatX) first-class index types. When you write arr[myNat64], codegen sees e2 : Nat64 and emits the dedicated idx_nat64 path directly: unbox, bounds-check, done.

But if someone writes arr[Nat64.toNat(x)], the index expression has type Nat, so codegen hits the plain-Nat branch with the speculative scalar check and potential idx_bigint fallback: even though x was just a Nat64.

The suggestion: pattern-match in codegen to recognize a toNat cast wrapping a NatX, peel it off, and emit the idx_nat64 path using the original x directly. A codegen peephole, no type system changes.

	My PR	Suggestion
Where	Type system + codegen	Codegen only (peephole)
Handles	arr[myNat64]	arr[Nat64.toNat(x)]
How	New index types accepted	Recognize cast in AST, elide it
Benefit	Best perf (-30-59%), clean types	-19%, optimizes legacy code patterns

[crusso]

Benefit Best perf (-30-59%), clean types -19%, optimizes legacy code patterns

@Quint - did you actually try my peephole suggestion? Any idea why it gets less perf back than yours? Did you also apply the peephole for indexing on left of an assignment?

[q-uint]

Seems like I didn't test it properly.

Benchmark (1000 x 256 accesses, enhanced)

Index type	Before	After	Change
Nat	25,687,306	25,687,306	--
Nat8	11,036,370	11,036,370	--
Nat16	11,079,329	11,079,329	--
Nat32	10,567,324	10,567,324	--
Nat64	17,991,325	17,991,325	--
Nat64.toNat()	20,807,325	17,991,325	-13.5%

arr[nat64ToNat(n)] now compiles to the same code as arr[n] with a native Nat64 index.

[q-uint]

The peephole produces identical instruction counts to the native type-change approach:

Index path	Type change	Peephole (toNat)
Nat8	11,036,370	11,036,370
Nat16	11,079,329	11,079,329
Nat32	10,567,324	10,567,324
Nat64	17,991,325	17,991,325

Zero overhead from the toNat() call, it's completely elided at compile time.

[crusso]

So maybe that would be good enough, without overloading indexing. Funnly enough, I just started on a PR to do that, but maybe you can extract one that just does the peephole optimization (or I steal some of your code as well as your test ;-))

If so, it's probably worth double checking the optimization still kicks in if we call

arr[Nat64.toNat(n))]
arr[n.toNat()]

[crusso]

I wonder if that's actually true or just a limitation of our E.if_ helper that doesn't let you supply a operand stack type for each branch.

Wasm 1.0 didn't allow the branches to assume the stack is non-empty. But Wasm 2.0 does.
Unfortunately, I'm not entirely sure we are allowed to use Wasm 2.0 features.

@ggreif are we good with Wasm 2.0 instructions?

If so, we don't need new locals and can avoid the stores and loads of get_va and get_vi.

[crusso]

I guess one could even iterate this for deeper projection paths, but probably not worth it.

[crusso]

I have hunch this approach still won't catch compound array assignments that combine an operator with update, like
a[i] += 1 . Do we care? I'm not sure I do.

[q-uint]

No, a[i] += 1 won't hit the peephole, but it doesn't matter I think? The parser desugars it into let tmp = i; a[tmp] = a[tmp] + 1, so the toNat call binds once into a plain Nat local and both array accesses reuse it.

[crusso]

No, a[i] += 1 won't hit the peephole, but it doesn't matter I think? The parser desugars it into let tmp = i; a[tmp] = a[tmp] + 1, so the toNat call binds once into a plain Nat local and both array accesses reuse it.

Also, our desugaring is pretty inefficient duplicating the field/index computation (and bounds check). So we could optimize this way more anyway (but not in this PR). Opened an issue for that.

#5972

[crusso]

(if the bench mark fails, I think you need to run make accept to update the numbers to use eop. If you just run ../run.sh -a -p it will run the benchmarks with 32-bit.

Or use

EXTRA_MOC_ARGS=--enhanced-orthogonal-persistence ../run.sh -a -p <test>.mo

[q-uint]

The Create GitHub App Token step in test.yml (line 71-76) runs unconditionally, but fork PRs don't have access to vars.GENERIC_CI_RW_APP_ID or secrets.GENERIC_CI_RW_APP_PRIVATE_KEY, causing this failure:

Error: [@octokit/auth-app] appId option is required

The perf step already guards on is_fork (line 86), but the token creation step itself needs a condition too, e.g.:

- name: Create GitHub App Token
  if: github.event.pull_request.head.repo.full_name == github.repository
  uses: actions/create-github-app-token@v3

Without this, every fork PR fails at this step.

@Kamirus

[skilesare]

I kind of ran into this again with Array.tabulate...looks like there is a Prim function that does this that has Nat hardcoded as the func paramater...I wonder if there are savings there as well?