feat(mlkem): initialize mlkem from golang standard library

Author: emmansun

mlkem/field.go

@@ -0,0 +1,271 @@
+// Copyright 2023 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.24
+
+package mlkem
+
+import (
+	"bytes"
+	"crypto/rand"
+	"math/big"
+	mathrand "math/rand/v2"
+	"strconv"
+	"testing"
+)
+
+func TestFieldReduce(t *testing.T) {
+	for a := range uint32(2*q*q) {
+		got := fieldReduce(a)
+		exp := fieldElement(a % q)
+		if got != exp {
+			t.Fatalf("reduce(%d) = %d, expected %d", a, got, exp)
+		}
+	}
+}
+
+func TestFieldAdd(t *testing.T) {
+	for a := range fieldElement(q) {
+		for b := range fieldElement(q) {
+			got := fieldAdd(a, b)
+			exp := (a + b) % q
+			if got != exp {
+				t.Fatalf("%d + %d = %d, expected %d", a, b, got, exp)
+			}
+		}
+	}
+}
+
+func TestFieldSub(t *testing.T) {
+	for a := range fieldElement(q) {
+		for b := range fieldElement(q) {
+			got := fieldSub(a, b)
+			exp := (a - b + q) % q
+			if got != exp {
+				t.Fatalf("%d - %d = %d, expected %d", a, b, got, exp)
+			}
+		}
+	}
+}
+
+func TestFieldMul(t *testing.T) {
+	for a := range fieldElement(q) {
+		for b := range fieldElement(q) {
+			got := fieldMul(a, b)
+			exp := fieldElement((uint32(a) * uint32(b)) % q)
+			if got != exp {
+				t.Fatalf("%d * %d = %d, expected %d", a, b, got, exp)
+			}
+		}
+	}
+}
+
+func TestDecompressCompress(t *testing.T) {
+	for _, bits := range []uint8{1, 4, 10} {
+		for a := uint16(0); a < 1<<bits; a++ {
+			f := decompress(a, bits)
+			if f >= q {
+				t.Fatalf("decompress(%d, %d) = %d >= q", a, bits, f)
+			}
+			got := compress(f, bits)
+			if got != a {
+				t.Fatalf("compress(decompress(%d, %d), %d) = %d", a, bits, bits, got)
+			}
+		}
+
+		for a := fieldElement(0); a < q; a++ {
+			c := compress(a, bits)
+			if c >= 1<<bits {
+				t.Fatalf("compress(%d, %d) = %d >= 2^bits", a, bits, c)
+			}
+			got := decompress(c, bits)
+			diff := min(a-got, got-a, a-got+q, got-a+q)
+			ceil := q / (1 << bits)
+			if diff > fieldElement(ceil) {
+				t.Fatalf("decompress(compress(%d, %d), %d) = %d (diff %d, max diff %d)",
+					a, bits, bits, got, diff, ceil)
+			}
+		}
+	}
+}
+
+func CompressRat(x fieldElement, d uint8) uint16 {
+	if x >= q {
+		panic("x out of range")
+	}
+	if d <= 0 || d >= 12 {
+		panic("d out of range")
+	}
+
+	precise := big.NewRat((1<<d)*int64(x), q) // (2ᵈ / q) * x == (2ᵈ * x) / q
+
+	// FloatString rounds halves away from 0, and our result should always be positive,
+	// so it should work as we expect. (There's no direct way to round a Rat.)
+	rounded, err := strconv.ParseInt(precise.FloatString(0), 10, 64)
+	if err != nil {
+		panic(err)
+	}
+
+	// If we rounded up, `rounded` may be equal to 2ᵈ, so we perform a final reduction.
+	return uint16(rounded % (1 << d))
+}
+
+func TestCompress(t *testing.T) {
+	for d := 1; d < 12; d++ {
+		for n := range q {
+			expected := CompressRat(fieldElement(n), uint8(d))
+			result := compress(fieldElement(n), uint8(d))
+			if result != expected {
+				t.Errorf("compress(%d, %d): got %d, expected %d", n, d, result, expected)
+			}
+		}
+	}
+}
+
+func DecompressRat(y uint16, d uint8) fieldElement {
+	if y >= 1<<d {
+		panic("y out of range")
+	}
+	if d <= 0 || d >= 12 {
+		panic("d out of range")
+	}
+
+	precise := big.NewRat(q*int64(y), 1<<d) // (q / 2ᵈ) * y  ==  (q * y) / 2ᵈ
+
+	// FloatString rounds halves away from 0, and our result should always be positive,
+	// so it should work as we expect. (There's no direct way to round a Rat.)
+	rounded, err := strconv.ParseInt(precise.FloatString(0), 10, 64)
+	if err != nil {
+		panic(err)
+	}
+
+	// If we rounded up, `rounded` may be equal to q, so we perform a final reduction.
+	return fieldElement(rounded % q)
+}
+
+func TestDecompress(t *testing.T) {
+	for d := 1; d < 12; d++ {
+		for n := 0; n < (1 << d); n++ {
+			expected := DecompressRat(uint16(n), uint8(d))
+			result := decompress(uint16(n), uint8(d))
+			if result != expected {
+				t.Errorf("decompress(%d, %d): got %d, expected %d", n, d, result, expected)
+			}
+		}
+	}
+}
+
+func randomRingElement() ringElement {
+	var r ringElement
+	for i := range r {
+		r[i] = fieldElement(mathrand.IntN(q))
+	}
+	return r
+}
+
+func TestEncodeDecode(t *testing.T) {
+	f := randomRingElement()
+	b := make([]byte, 12*n/8)
+	rand.Read(b)
+
+	// Compare ringCompressAndEncode to ringCompressAndEncodeN.
+	e1 := ringCompressAndEncode(nil, f, 10)
+	e2 := ringCompressAndEncode10(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode10 = %x", e1, e2)
+	}
+	e1 = ringCompressAndEncode(nil, f, 4)
+	e2 = ringCompressAndEncode4(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode4 = %x", e1, e2)
+	}
+	e1 = ringCompressAndEncode(nil, f, 1)
+	e2 = ringCompressAndEncode1(nil, f)
+	if !bytes.Equal(e1, e2) {
+		t.Errorf("ringCompressAndEncode = %x, ringCompressAndEncode1 = %x", e1, e2)
+	}
+
+	// Compare ringDecodeAndDecompress to ringDecodeAndDecompressN.
+	g1 := ringDecodeAndDecompress(b[:encodingSize10], 10)
+	g2 := ringDecodeAndDecompress10((*[encodingSize10]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress10 = %v", g1, g2)
+	}
+	g1 = ringDecodeAndDecompress(b[:encodingSize4], 4)
+	g2 = ringDecodeAndDecompress4((*[encodingSize4]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress4 = %v", g1, g2)
+	}
+	g1 = ringDecodeAndDecompress(b[:encodingSize1], 1)
+	g2 = ringDecodeAndDecompress1((*[encodingSize1]byte)(b))
+	if g1 != g2 {
+		t.Errorf("ringDecodeAndDecompress = %v, ringDecodeAndDecompress1 = %v", g1, g2)
+	}
+
+	// Round-trip ringCompressAndEncode and ringDecodeAndDecompress.
+	for d := 1; d < 12; d++ {
+		encodingSize := d * n / 8
+		g := ringDecodeAndDecompress(b[:encodingSize], uint8(d))
+		out := ringCompressAndEncode(nil, g, uint8(d))
+		if !bytes.Equal(out, b[:encodingSize]) {
+			t.Errorf("roundtrip failed for d = %d", d)
+		}
+	}
+
+	// Round-trip ringCompressAndEncodeN and ringDecodeAndDecompressN.
+	g := ringDecodeAndDecompress10((*[encodingSize10]byte)(b))
+	out := ringCompressAndEncode10(nil, g)
+	if !bytes.Equal(out, b[:encodingSize10]) {
+		t.Errorf("roundtrip failed for specialized 10")
+	}
+	g = ringDecodeAndDecompress4((*[encodingSize4]byte)(b))
+	out = ringCompressAndEncode4(nil, g)
+	if !bytes.Equal(out, b[:encodingSize4]) {
+		t.Errorf("roundtrip failed for specialized 4")
+	}
+	g = ringDecodeAndDecompress1((*[encodingSize1]byte)(b))
+	out = ringCompressAndEncode1(nil, g)
+	if !bytes.Equal(out, b[:encodingSize1]) {
+		t.Errorf("roundtrip failed for specialized 1")
+	}
+}
+
+func BitRev7(n uint8) uint8 {
+	if n>>7 != 0 {
+		panic("not 7 bits")
+	}
+	var r uint8
+	r |= n >> 6 & 0b0000_0001
+	r |= n >> 4 & 0b0000_0010
+	r |= n >> 2 & 0b0000_0100
+	r |= n /**/ & 0b0000_1000
+	r |= n << 2 & 0b0001_0000
+	r |= n << 4 & 0b0010_0000
+	r |= n << 6 & 0b0100_0000
+	return r
+}
+
+func TestZetas(t *testing.T) {
+	ζ := big.NewInt(17)
+	q := big.NewInt(q)
+	for k, zeta := range zetas {
+		// ζ^BitRev7(k) mod q
+		exp := new(big.Int).Exp(ζ, big.NewInt(int64(BitRev7(uint8(k)))), q)
+		if big.NewInt(int64(zeta)).Cmp(exp) != 0 {
+			t.Errorf("zetas[%d] = %v, expected %v", k, zeta, exp)
+		}
+	}
+}
+
+func TestGammas(t *testing.T) {
+	ζ := big.NewInt(17)
+	q := big.NewInt(q)
+	for k, gamma := range gammas {
+		// ζ^2BitRev7(i)+1
+		exp := new(big.Int).Exp(ζ, big.NewInt(int64(BitRev7(uint8(k)))*2+1), q)
+		if big.NewInt(int64(gamma)).Cmp(exp) != 0 {
+			t.Errorf("gammas[%d] = %v, expected %v", k, gamma, exp)
+		}
+	}
+}