feat: Add a limited form of arithmetic on generics

[jfecher]

Description

Problem*

Resolves #4784

Summary*

Adds a limited form of arithmetic operators on numeric generic types in Noir. This includes only 5 operations: +, -, *, /, and %.

Going forward I'm going to distinguish between two definitions for clearer communication between devs & users:

• Numeric generics: generics that use numeric type variables like fn foo<let N: u32>()
• Arithmetic generics: when arithmetic operators are used on numeric generics in a type position: fn foo<let N: u32>() -> [Field; N + 1]

Additional Context

It is explicitly a non-goal to always correctly return when two arithmetic generics are equal. For example, in our current type checking we do not check whether applying the distributive law would make two types equal. So you'd get a type error for trying to use 2*(N + 1) where 2*N + 2 was expected. Users should generally consider the type equality here to be largely structural and try to match that as best as possible. The only check performed besides simplification of constant terms is commutativity for + and *. Even then, the check is only one level deep. So N * M * 3 is not equal to 3 * N * M. Update: we canonicalize commutative operations now, so we sort this and it now type checks.

Documentation*

Check one:

• No documentation needed.
• Documentation included in this PR.
• [For Experimental Features] Documentation to be submitted in a separate PR.
  • I'd like to have this in for some time before we start advertising it.

PR Checklist*

• I have tested the changes locally.
• I have formatted the changes with Prettier and/or cargo fmt on default settings.

[jfecher]

Update: added a canonicalization method for commutative operators. We now accept a lot more cases including:

• (1 + N) + 1 = N + 2
• 3 * B * 3 * A = A * B * 9
• others

[michaeljklein]

Explored whether this is sufficient for some techniques that can be used with partial support for dependent types.

In short, yes, but it would look nicer with:

• Trait impl's on numeric generics
• Associated types

Demo here

// KnownConstant
//
// Summary: when we know something is constant, we can know that expressions involving it will resolve to constants.
// - This guarantees that such values will be known at compile-time
// - Note the need for a wrapper struct (`W`) to appease the type checker:
//   we can solve these fine, but reject raw numeric generics.
//
// Note that it's impossible to define an impl as follows, because of overlap:
//
// impl<let N: u32> KnownConstant for W<N> { .. }

// Wrapper type for numeric generics
//
// NOTE: explicit kinds would allow us to replace u32 with a variable
struct W<let N: u32> { }

impl<let N: u32> Eq for W<N> {
    fn eq(self, other: W<N>) -> bool {
        let _ = self;
        let _ = other;
        true
    }
}

trait KnownConstant {
    fn to_constant(self) -> Field;
}

// implementing KnownConstant for a variable would overlap with W<_ + _> and
// other instances
impl KnownConstant for W<0> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        0
    }
}

impl KnownConstant for W<1> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        1
    }
}

impl KnownConstant for W<2> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        2
    }
}

impl KnownConstant for W<3> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        3
    }
}

impl KnownConstant for W<4> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        4
    }
}

impl KnownConstant for W<5> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        5
    }
}

impl KnownConstant for W<6> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        6
    }
}

impl KnownConstant for W<7> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        7
    }
}

impl KnownConstant for W<8> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        8
    }
}

impl KnownConstant for W<9> {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        9
    }
}

impl<let N: u32, let M: u32> KnownConstant for W<N + M> where W<N>: KnownConstant, W<M>: KnownConstant {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        (N as Field) + (M as Field)
    }
}

impl<let N: u32, let M: u32> KnownConstant for W<N * M> where W<N>: KnownConstant, W<M>: KnownConstant {
    fn to_constant(self) -> Field {
        assert(self == self); // TODO to_constant(_self) failed
        (N as Field) * (M as Field)
    }
}

fn add_comm_to<let N: u32, let M: u32>(_x: W<N + M>) -> W<M + N> where W<N>: KnownConstant, W<M>: KnownConstant {
    W {}
}

fn polynomial_input<let N: u32, let M: u32>(_x: W<2 * N * N + 1 * M + 2>)
    -> W<2 + N * N * 2 + M> where W<N>: KnownConstant, W<M>: KnownConstant {
    W {}
}


////////////////////////////////////////////////////////////////////////////////////////////////////////////
// Equiv
////////////////////////////////////////////////////////////////////////////////////////////////////////////
//
// Summary: this is a simplification of isomorphism-with-coherence-conditions-style equivalence
// - EquivT is an attempt to use traits directly
// - Equiv below is a successful attempt to use structs instead
//   + Trait associated types would make it much cleaner

trait EquivT<T> {
    fn equivT_to(x: Self) -> T;
    fn equivT_from(x: T) -> Self;
}

impl<let N: u32, let M: u32> EquivT<W<M + N>> for W<N + M> {
    fn equivT_to(_x: Self) -> W<M + N> { W {} }
    fn equivT_from(_x: W<M + N>) -> Self { W {} }
}

impl<let N: u32, let M: u32> EquivT<W<M * N>> for W<N * M> {
    fn equivT_to(_x: Self) -> W<M * N> { W {} }
    fn equivT_from(_x: W<M * N>) -> Self { W {} }
}

// // NOTE: overlapping impl blocks this approach
// // - Could be resolved by allowing trait impl naming and/or selection
//
// impl<let N: u32, let M: u32, let P: u32> EquivT<W<M + (N + P)>> for W<(N + M) + P> {
//     fn equivT_to(_x: Self) -> W<M + (N + P)> { W {} }
//     fn equivT_from(_x: W<M + (N + P)>) -> Self { W {} }
// }

// NOTE: this version w/o ENV's can't be composed
//
// struct Equiv<T, U> {
//     to_: fn(T) -> U,
//     from_: fn(U) -> T,
//     // .. other coherence conditions
// }
struct Equiv<T, TU, U, UT> {
    // TODO(https://github.com/noir-lang/noir/issues/5644):
    // Bug with struct_obj.field_thats_a_fn(x)

    to_: fn[TU](T) -> U,
    fro_: fn[UT](U) -> T,
    // .. other coherence conditions
}

impl<T, TU, U, UT> Equiv<T, TU, U, UT> {
    fn to(self, x: T) -> U {
        let f = self.to_;
        f(x)
    }

    // TODO: bad error span and/or wrong error (appears to overlap w/ stdlib when named 'from')
    //
    // error: duplicate definitions of from found
    //     ┌─ /Users/michaelklein/Coding/rust/noir/tmp/src/main.nr:7:29
    //     │
    //   7 │ // Needed for "raw equivalence" using traits:
    //     │                             ---- first definition found here
    //     ·
    // 175 │     fn from(self, x: U) -> T {
    //     │        ---- second definition found here
    //     │
    fn fro(self, x: U) -> T {
        let f = self.fro_;
        f(x)
    }
}

// NOTE: it would be nice to hide these ENV's using associated types
fn equiv_trans<T, TU, U, UT, UV, V, VU>(x: Equiv<T, TU, U, UT>, y: Equiv<U, UV, V, VU>) -> Equiv<T, (Equiv<U, UV, V, VU>, Equiv<T, TU, U, UT>), V, (Equiv<T, TU, U, UT>, Equiv<U, UV, V, VU>)> {
    Equiv {
        to_: |z| { y.to(x.to(z)) },
        fro_: |z| { x.fro(y.fro(z)) },
    }
}

fn equiv_refl<let N: u32>() -> Equiv<W<N>, (), W<N>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn add_zero_r<let N: u32>() -> Equiv<W<N + 0>, (), W<N>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn mul_one_r<let N: u32>() -> Equiv<W<N * 1>, (), W<N>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn add_equiv_r<let C: u32, let N: u32, let M: u32, EN, EM>(f: Equiv<W<N>, EN, W<M>, EM>) -> Equiv<W<C + N>, (), W<C + M>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn add_comm<let N: u32, let M: u32>() -> Equiv<W<N + M>, (), W<M + N>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn mul_comm<let N: u32, let M: u32>() -> Equiv<W<N * M>, (), W<M * N>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn add_assoc<let N: u32, let M: u32, let P: u32>() -> Equiv<W<N + (M + P)>, (), W<(N + M) + P>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn mul_assoc<let N: u32, let M: u32, let P: u32>() -> Equiv<W<N * (M * P)>, (), W<(N * M) * P>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

fn mul_add<let N: u32, let M: u32, let P: u32>() -> Equiv<W<N * (M + P)>, (), W<N * M + N * P>, ()> {
    Equiv { to_: |_x| { W {} }, fro_: |_x| { W {} } }
}

// (N + 1) * N == N * N + N
//
// (550 character line, due to ENV's)
fn demo_proof<let N: u32>() -> Equiv<W<(N * (N + 1))>, (Equiv<W<N>, (), W<N>, ()>, Equiv<W<(N * (N + 1))>, (Equiv<W<(N * (N + 1))>, (), W<N>, ()>, Equiv<W<(N * (N + 1))>, (), W<(N * (N + 1))>, ()>), W<N>, (Equiv<W<(N * (N + 1))>, (), W<(N * (N + 1))>, ()>, Equiv<W<(N * (N + 1))>, (), W<N>, ()>)>), W<N>, (Equiv<W<(N * (N + 1))>, (Equiv<W<(N * (N + 1))>, (), W<N>, ()>, Equiv<W<(N * (N + 1))>, (), W<(N * (N + 1))>, ()>), W<N>, (Equiv<W<(N * (N + 1))>, (), W<(N * (N + 1))>, ()>, Equiv<W<(N * (N + 1))>, (), W<N>, ()>)>, Equiv<W<N>, (), W<N>, ()>)> {
    let p1: Equiv<W<(N + 1) * N>, (), W<N * (N + 1)>, ()> = mul_comm();
    let p2: Equiv<W<N * (N + 1)>, (), W<N * N + N * 1>, ()> = mul_add::<N, N, 1>();
    let p3_sub: Equiv<W<N * 1>, (), W<N>, ()> = mul_one_r();
    let p3: Equiv<W<N * N + N * 1>, (), W<N * N + N>, ()> = add_equiv_r::<N * N, N * 1, N, _, _>(p3_sub);
    equiv_trans(equiv_trans(p1, p2), p3)
}

fn main() {
    let quadratic: W<2 * 2 * 2 + 1 * 1 + 2> = W {};

    // Bug: both should evaluate to W<11>
    //
    // error: Expected type W<7>, found type W<11>
    //     ┌─ /Users/michaelklein/Coding/rust/noir/tmp/src/main.nr:103:56
    //     │
    // 103 │     let alternate_quadratic = polynomial_input::<2, 1>(quadratic);
    //     │                                                        ---------
    //     │
    // let alternate_quadratic = polynomial_input::<2, 1>(quadratic);

    let _ = add_comm_to::<2 * 2 * 2 + 1 * 1, 2>(quadratic);
}

[jfecher]

let quadratic: W<2 * 2 * 2 + 1 * 1 + 2> = W {};

Why the 2 * 2 ... expression there? We should constant fold that before we get to testing arithmetic generics.

I'd think a lot of the impls should be unneeded as well. Since we already have N + M = M + N in unification we shouldn't also need a commutative impl for add.

[michaeljklein]

@jfecher I was checking that everything resolves properly, but will also typecheck against add_comm for known inputs. E.g. I thought it would be worthwhile to see how N + 1 where N = 1 vs 2 unifies.

demo_proof is more representative of what can currently be done with unification.

[vezenovm]

Although limited these already seem like a big UX improvement for most cases people want and they are under a flag, so I am approving.

[michaeljklein]

LGTM