Annotating a compile-time value to a rust struct

✍️ Written on 2022-02-16 in 3630 words.
Part of cs software-development programming-languages rustlang

Update 2022-03-29: Andreas Kage pointed me to an article “Generating static arrays during compile time in Rust” by rustyoctopus. It is wonderfully elaborative and shows a similar thought process like mine. An approach with procedural macros can solve my problem. Thank you!

Motivation

I am eagerly preparing PQC implementations in rust for publication. I was thinking about the design of the public API when I came across a limitation of const generics in rust. The question is: How can I annotate a compile-time value to a rust struct? It turns out, I cannot.

What I mean by annotated value

Consider a key encapsulation mechanism. It features three operations, which we want to design as methods:

type R = Result<(), Box<dyn std::error::Error>>;

struct KEM {}

impl KEM {
  fn crypto_kem_keypair(&self, pk: &mut [u8], sk: &mut [u8]) -> R {
    Ok(())
  }

  fn crypto_kem_enc(&self, c: &mut [u8], key: &mut [u8], pk: &[u8]) -> R {
    Ok(())
  }

  fn crypto_kem_dec(&self, key: &mut [u8], c: &[u8], sk: &[u8]) -> R {
    Ok(())
  }
}

However, there are different variants of this algorithm. Classic McEliece has 10 different parameter sets, NTRU has 4 different parameter sets, and Saber has 3 different parameter sets. Based on the parameter sets, some constants are defined (size of matrices, module/lattice dimension, …).

Some of them even need to be exposed since the parameters affect key sizes, ciphertext size, and shared secret size. And we want to ask the user to allocate the corresponding buffers themselves, right? Allocation flexibility is a requirement to make your implementation work on a wide variety of platforms. For example, if you provide a method returning a pointer to an allocated heap object, you make assumptions about the size of heap.

To summarize: Our KEM struct above models an algorithm. It exists in various variants which have annotated values.

Annotated runtime value

To revise, the following declaration defines an annotated value n:

struct KEM { n: usize }

However, this is a runtime value. We need it at compile-time, because allocation on the stack requires a compile-time value. One obvious way to define compile-time variables are const generics.

A usecase for const generics

type R = Result<(), Box<dyn std::error::Error>>;

struct KEM<const N: usize> {}

impl<const N: usize> KEM<N> {
  fn crypto_kem_keypair(&self, pk: &mut [u8], sk: &mut [u8]) -> R {
    Ok(())
  }

  fn crypto_kem_enc(&self, c: &mut [u8], key: &mut [u8], pk: &[u8]) -> R {
    Ok(())
  }

  fn crypto_kem_dec(&self, key: &mut [u8], c: &[u8], sk: &[u8]) -> R {
    Ok(())
  }
}

I think it is neat that in rust, const always means compile-time variable. In C/C++ it gets more complicated with const, constexpr, and preprocessor macros.

The way I instantiate a KEM now involves specifying a value.

fn main() {
  let _kem = KEM::<42>{ };
}

Modelling variants

fn variant1() -> KEM::<21> { KEM::<21>{ } }
fn variant2() -> KEM::<42> { KEM::<42>{ } }

fn main() {
  let _kem = variant1();
}

In this case, I model variant instantiation by simple functions. Why not a constructor/static method? A static method would already require specifying N just to call that method. So this is a bad approach. As for the function, it is interesting to see that the value of N is now part of the type. We don’t return KEM, but (e.g.) KEM::<21>.

Retrieving the const generic value

Approach 1: type of a value

So, the value N is actually part of the type. But how can I retrieve it?

Is there some way in rust to retrieve the type of variable kem and then get the map of const generics? Some kind of reflection API? No, not that I am aware of. Can I somehow use Any or type_name() to do something similar? This is explicitly discouraged and I won’t parse the argument from a string representation.

Approach 2: struct member

What about the idea to create a struct member n and synchronize n with N?

struct KEM<const N: usize> { n: usize }

fn variant1() -> KEM::<21> { KEM::<21>{ n: 21 } }
fn variant2() -> KEM::<42> { KEM::<42>{ n: 42 } }

fn main() {
  let kem = variant1();
  println!("n = {}", kem.n);

  // error[E0435]: attempt to use a non-constant value in a constant
  //let buffer = [0u8; kem.n];
}

This works, but obviously it is a runtime value and cannot be used to specify the size of a buffer on the stack. Thus the last line does not compile.

Approach 3: struct type

In rust, there is a type keyword. And I can use it in an impl block, right? Can it define the buffer type of appropriate size?

impl<const N: usize> KEM<N> {
  type BufN = [u8; N];
}

Dang… no.

error[E0658]: inherent associated types are unstable
 --> test.rs:6:5
  |
6 |     type BufN = [u8; N];
  |     ^^^^^^^^^^^^^^^^^^^^
  |
  = note: see issue #8995 <https://github.com/rust-lang/rust/issues/8995> for more information
  = help: add `#![feature(inherent_associated_types)]` to the crate attributes to enable

So, not yet.

Approach 4: associated type

But then we could use associated types of traits. In the context of traits, type declarations are a stable feature.

struct KEM<const N: usize> { }

trait Data {
  type BufN;
}

impl<const N: usize> Data for KEM<N> {
  type BufN = [u8; N];
}

fn main() {
  let kem = KEM::<42>{}; // or “variant1();”
  // invalid syntax
  //let buf: kem.BufN = std::default::Default::default();
}

The issue is that, we cannot use this type explicitly in any declaration. There is no way to address it. I avoided it on the right-hand side of the assignment with the Default trait, but still need it on the left-hand side.

Approach 5: Constant function

Can we proxy the value through a constant function?

struct KEM<const N: usize> { }

impl<const N: usize> KEM<N> {
  const fn n(&self) -> usize { N }
}

fn main() {
  let kem = KEM::<42>{};
  let buffer = [0u8; kem.n()];
}

The error message gives us some helpful feedback:

error[E0435]: attempt to use a non-constant value in a constant
 --> test.rs:9:22
  |
8 |   let kem = KEM::<42>{};
  |   ------- help: consider using `const` instead of `let`: `const kem`
9 |   let buffer = [0u8; kem.n()];
  |                      ^^^ non-constant value

error: aborting due to previous error

For more information about this error, try `rustc --explain E0435`.

So, can we make it a const variable?

struct KEM<const N: usize> { }

impl<const N: usize> KEM<N> {
  const fn n(&self) -> usize { N }
}

fn main() {
  const kem: KEM::<42> = KEM::<42>{};
  let buffer = [0u8; kem.n()];
}

Yes, this actually works. An issue is that we need to explicitly write down the parameters when declaring the const variable. So if we replace KEM::<42>{} by a call variant1() we would still have to list all parameters which is exactly what I want to avoid. It might be advantageous in your usecase however; maybe not. So, what are we actually aiming for?

The final design

In the end, there does not seem to be a neat approach to annotate a run-time variable to a struct. But if we replace the struct by a module, we can still achieve a decent API design (I added all pub keywords unlike before to make it an executable example, super just refers to the module higher in the hierarchy than the current one):

type R = Result<(), Box<dyn std::error::Error>>;

pub struct Kem<const N: usize> {}

impl<const N: usize> Kem<N> {
  pub fn crypto_kem_keypair(&self, pk: &mut [u8], sk: &mut [u8]) -> R {
    Ok(())
  }

  pub fn crypto_kem_enc(&self, c: &mut [u8], key: &mut [u8], pk: &[u8]) -> R {
    Ok(())
  }

  pub fn crypto_kem_dec(&self, key: &mut [u8], c: &[u8], sk: &[u8]) -> R {
    Ok(())
  }
}

mod variant1 {
  pub const N: usize = 21;
  pub const KEM: super::Kem::<N> = super::Kem::<N>{};
}

mod variant2 {
  pub const N: usize = 42;
  pub const KEM: super::Kem::<N> = super::Kem::<N>{};
}

fn main() -> R {
  use variant1::*;
  let mut pk = [0u8; N];
  let mut sk = [0u8; N];
  KEM.crypto_kem_keypair(&mut pk, &mut sk)
}

This approach has the obvious disadvantage, that I instantiate the struct globally. So since the Kem struct does not have any variable runtime members, I just instantiate the variants I need. But if Kem is customized by runtime variables, this approach fails. But even in this case, I have the cumbersome alternative to extract relevant parameters of the module specifying the desired variant.

Conclusion

I did not find a way to reasonably attach a compile-time constant to a struct. In a limited set of usecases, a const function can be helpful. However, attaching a compile-time constant to a module is a viable alternative for me. And as pointed out in the update text, an approach with procedural macros could solve my problem. But it is more intricate.