Arithmetic on numeric generics

[nventuro]

Problem

There's multiple design patterns that result in the need to declare generic array lenghts that are calculated on some other generic value.

Examples include:

impl<T, const N: usize> [T; N] {
    fn split_first(self) -> (T, [T; N - 1]) {
        ...
    }
}

or

trait Serialize<N> {
    fn serialize(self) -> [Field; N];
}

struct Foo<T> {
    a: T,
    b: T,
}

impl<T, N> Serialize<N * 2> for Foo<T> {
    fn serialize(self) -> [Field; N * 2] where T: Serialize<N> {
      ...
    }
}

Happy Case

I'd want to be able to both have numeric generics without having to resort to the workaround from #4633, and to be able to perform basic arithmetic on them. My current use case is to do addition to both other generics and constants (e.g. N + N + 1), but it seems multiplication and subtraction are common patterns as well.

Project Impact

Blocker

Impact Context

Lack of this feature blocks multiple things, such as flatten as requested by @LogvinovLeon here, or AztecProtocol/aztec-packages#5492 which is required for non-trivial aztec-nr data structures.

This is also related to #4633, since many of the above fall into scenarios where the current workaround to force a generic value to be numeric (i.e. make it an array length) don't work.

Workaround

None

[Thunkar]

Would be awesome to also have the modulo operation (mainly for AES)

pub fn compute_ciphertext<N>(
       self,
       secret: GrumpkinPrivateKey,
       point: GrumpkinPoint
   ) -> [u8; N + 16 - N % 16] where Note: NoteInterface<N> {
...

[jfecher]

The Noir team has been thinking about this issue a lot and I have a few thoughts here as well.

• It seems like slices could be usable for Serialize here now. We've avoided using them in the past but they should be bug free now that nested slices are not allowed. Since nested slices shouldn't be needed here this could be an option, although there could also be some performance implications with appending slices.
• Another option could be changing the Serialize interface so that it accepts a large enough buffer and the generic is no longer needed:

trait Serialize {
    fn serialize<N>(self, state: &mut SerializeState<N>);
}

struct SerializeState<N> {
    buffer: [Field; N],
    index: u64,
}

impl<T, M> Serialize for [T; M] where T: Serialize {
    fn serialize<N>(self, state: &mut SerializeState<N>) {
        for i in 0 .. M {
            self[i].serialize(state);
        }
    }
}

impl Serialize for Field {
    fn serialize<N>(self, state: &mut SerializeState<N>) {
        state.buffer[state.index] = self;
        state.index += 1;
    }
}

struct Foo { x: [u64; 3], y: Field }

impl Serialize for Foo {
    fn serialize<N>(self, state: &mut SerializeState<N>) {
        self.x.serialize(state);
        self.y.serialize(state);
    }
}

The caveat of course being that main would need to supply a buffer that is large enough:

fn serialize_helper(x: impl Serialize, buffer: [Field; N]) -> [Field; N]) {
    let mut state = SerializeState { buffer, index: 0 };
    x.serialize(&mut state);
    state.buffer
}

fn main(foo: Foo) {
    let serialized = serialize_helper(foo, [0; 10]);
}

Once we have metaprogramming we'd be able to write a function which can give us the size of a type:

// Assuming we have
comptime fn size_of_type(typ: Type) -> u64 { ... }

// We could write a macro to make a large enough array
comptime fn make_buffer(typ: Type) -> Expr {
    let mut vec = Vec::new();
    for i in 0 .. size_of_type(typ) {
        vec.push(0);
    }
    x.into_array() // Returns a code snippet of an array of zeroes
}

/// This needs to be a comptime function now so that we can get the type
/// of x without it returning an opaque generic type like `T`
comptime fn serialize_helper(x: TypedExpr, buffer: [Field; N]) -> Expr {
    let buffer = make_buffer!(x.get_type());

    quote {
        let mut state = SerializeState { buffer: $buffer, index: 0 };
        x.serialize(&mut state);
        state.buffer
    }
}

fn main(foo: Foo) {
    let serialized = serialize_helper(foo);
}

The downside of the approach is that you'd still need to pass in the buffer length explicitly if serialize_helper is operating on a generic argument, e.g:

fn bar<T>(foo: T) where T: Serialize {
    // Error here, size_of_type(foo) is unknown
    let serialized = serialize_helper(foo);
}

This could possibly be resolved by a T: Sized constraint which returns a compile-time size however. Another caveat is that the size of each type here is separated from the Serialize function so we may not want a [bool; 4] to be 4 fields large if we could pack them into one field.