ml-kem: Don't use FieldElement for compressed values

[bifurcation]

FieldElement should specifically represent a member of $Z_q$ (can't figure out how to get GitHub to do \mathbb). Right now, it is also used as the output of the compress() method as well, even though these are elements of $Z_{2^d}$. We should define a separate type for this latter type of number, something like CompressedValue<D: CompressionFactor>

[pqcfox]

Just started on a branch to implement this! Would love to be assigned for this if that makes sense--hoping to have a full ready-to-review PR up in the next weekend or two.

[bifurcation]

Cool, thanks for picking this up @pqcfox. If you're interested in a slightly bigger project in this zone, the ml-dsa crate has a "module_lattice" module that I isolated with the idea of applying it to ML-KEM as well. You could take a stab at porting ML-KEM to that framework, though we might need help from @tarcieri to actually get the shared crate set up. If you do go that way, let me know; I have a partial starting point I could share.

[pqcfox]

Ooh, that sounds fantastic!

Would it make sense to you for me to go ahead and complete this issue as described, then move the Compress trait and new implementations over into the module_lattice module once the shared crate is ready?

Thanks!

[bifurcation]

Yeah, I think probably. You should take a look at the module_lattice module to make sure I'm not misremembering, but I think its field element representation would also be sad about representing compressed values. So this work would be complementary.

[bifurcation]

Actually, it looks like there are encoding routines in there, so maybe more thought is needed. Let me follow up once I'm back at my computer in a bit.

[bifurcation]

OK, having glanced at the code and thought about this for a moment:

The module_lattice module defines:

• A Field trait that lets you define a representation of a finite field
• A struct Elem<F: Field> for elements of that field, and similar structs for doing linear algebra on polynomials with Elem coefficients.
• An Encode trait for encoding polynomials.

So it's tempting to say that the right order of operations here is:

• Move ml_kem over to use module_lattice, including using its Encode for encoding
• Define a new Field to represent $Z/2^dZ$
• Change the compress() methods in ml_kem so that their range is the new field, not $Z/qZ$
• Verify that the Encode traits produce the expected values with this new scheme

The main thing that gives me pause there is that the Mul implementation will be wrong for fields that do not have prime order (in particular, binary fields). It's not a practical issue, since ML-KEM doesn't actually do anything with compressed values except encode them. If we wanted to be extra safe, we could define a trait PrimeField: Field trait and only enable Mul when F: PrimeField.

@pqcfox - Sorry if that sounds a little more involved than this issue might have appeared at the start. Totally understand if you're not interested in the bigger project! In any case, I went ahead and filed #112 to track the module_lattice transition.

[pqcfox]

Ooh, I like this idea a lot, esp. using a trait to gate the Mul implementation. Would be very happy to give this a shot once the module_lattice repo is ready for PRs!

Would it make sense to rename Field in this case e.g. to Ring to accommodate $\mathbb{Z}/2^d\mathbb{Z}$?

Thanks so much @bifurcation!

[bifurcation]

Hm. I was going to object that $Z / 2^d Z$ is a field, just with a different multiplication algorithm. But actually $Z / 2^d Z$ is not the same as $GF(2^d)$, which is what I had in mind. But this doesn't actually matter in the context of ML-KEM, since the only thing ML-KEM does with these values is encode/decode them.

So maybe the right answer is actually to define another type that just represents "a value that fits in $d$ bits". Something like FixedWidthInt<F, D> where F: Field, D: EncodingSize. Then you would use this as the input and output of the encode/decode methods.

The trouble with this approach will be avoiding copies, since you'll want to be able to convert between Array<Elem<F>, U256> and Array<FixedWidthInt<F, D>, U256>. In principle, this should be doable, since Elem and FixedWidthInt will have the same representation in memory (an F::Int). But it might require transmute or something similar.

[bifurcation]

That is, you would want to be able to do the following:

For ML-KEM:

Array<Elem> ---(compress)--> Array<FixedWidthInt> --(encode)-> Array<u8>
Array<Elem> <-(decompress)-- Array<FixedWidthInt> <-(decode)-- Array<u8>

For ML-DSA

Array<Elem> --(transmute)-> Array<FixedWidthInt> --(encode)-> Array<u8>
Array<Elem> <-(transmute)-- Array<FixedWidthInt> <-(decode)-- Array<u8>

Of course, maybe just having the "transmute" be array.map(|x| x.into()) is cheap enough.

[pqcfox]

This approach seems perfect to me! Thanks for detailing this out :)

I'll go ahead and move forward with this FixedWidthInt strategy once the module_lattice crate is in place, excited to dive in on this!

[tarcieri]

Nit: would probably suggest naming the crate module-lattice

Not sure where it should go but perhaps it can be developed in the KEMs tree for now and extracted somewhere else

[pqcfox]

Sounds like a plan! @tarcieri, would it make sense for me to go ahead and add a module-lattice crate myself to this repo as part of the draft PR for this issue, then before full review all the details of publishing, where it'll end up, etc. can be figured out?

[tarcieri]

Sounds good to me, but double check with @bifurcation

[bifurcation]

Yeah, that sounds fine to me. Just elaborating the plan a little bit, it seems like the following is roughly the right order of operations:

• Locally:
  • Clone signatures and KEMs
  • Move the module-lattice code to a new crate under KEMs
  • Update ml-kem to use module-lattice (incl. any changes to module-lattice needed, hopefully none!)
  • Update ml-dsa to refer to the new module-lattice crate (incl. any changes)
• Then make PRs:
  1. PR on KEMs to add module-lattice crate and update ml-kem
  2. PR on signatures to convert to module-lattice crate

Basically, since we have a cross-repo dependency, seems like a good idea to do cross-repo testing before we start making PRs.

@tarcieri - It seems like maybe we should publish the module-lattice crate in between (1) and (2) above, so that ml-dsa can refer to it via crates.io. Does that sound right to you? And if so, should we agree on its ultimate home before we do that? I think you had mentioned the utils repo before.

[tarcieri]

@bifurcation that all sounds good to me. we could also consider a new repo for module-lattice rather than utils