chore(mlkem): refactoring, reduce alloc times

This commit is contained in:
Sun Yimin 2025-09-10 11:08:42 +08:00 committed by GitHub
parent 85f4205314
commit 18783be107
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 17 deletions

View File

@ -307,7 +307,7 @@ func ringDecodeAndDecompress10(bb *[encodingSize10]byte) ringElement {
func ringCompressAndEncode(s []byte, f ringElement, d uint8) []byte { func ringCompressAndEncode(s []byte, f ringElement, d uint8) []byte {
var b byte var b byte
var bIdx uint8 var bIdx uint8
for i := 0; i < n; i++ { for i := range n {
c := compress(f[i], d) c := compress(f[i], d)
var cIdx uint8 var cIdx uint8
for cIdx < d { for cIdx < d {
@ -401,8 +401,8 @@ func samplePolyCBD(s []byte, b, η byte) ringElement {
prf := sha3.NewSHAKE256() prf := sha3.NewSHAKE256()
prf.Write(s) prf.Write(s)
prf.Write([]byte{b}) prf.Write([]byte{b})
B := make([]byte, 64*η) var B [maxBytesOf64Mulη]byte
prf.Read(B) prf.Read(B[:64*η])
// SamplePolyCBD simply draws four (2η) bits for each coefficient, and adds // SamplePolyCBD simply draws four (2η) bits for each coefficient, and adds
// the first two and subtracts the last two. // the first two and subtracts the last two.
@ -419,20 +419,15 @@ func samplePolyCBD(s []byte, b, η byte) ringElement {
} }
case 3: case 3:
for i := 0; i < n; i += 4 { for i := 0; i < n; i += 4 {
j := (i / 4) * 3 j := (i >> 2) * 3
b := B[j]
b_7, b_6, b_5, b_4 := b>>7, b>>6&1, b>>5&1, b>>4&1 bits := uint32(B[j]) | uint32(B[j+1])<<8 | uint32(B[j+2])<<16
b_3, b_2, b_1, b_0 := b>>3&1, b>>2&1, b>>1&1, b&1 for k := range 4 {
b = B[j+1] off := 6 * k
b_15, b_14, b_13, b_12 := b>>7, b>>6&1, b>>5&1, b>>4&1 sum := ((bits >> off) & 1) + ((bits >> (off + 1)) & 1) + ((bits >> (off + 2)) & 1)
b_11, b_10, b_9, b_8 := b>>3&1, b>>2&1, b>>1&1, b&1 sub := ((bits >> (off + 3)) & 1) + ((bits >> (off + 4)) & 1) + ((bits >> (off + 5)) & 1)
b = B[j+2] f[i+k] = fieldSub(fieldElement(sum), fieldElement(sub))
b_23, b_22, b_21, b_20 := b>>7, b>>6&1, b>>5&1, b>>4&1 }
b_19, b_18, b_17, b_16 := b>>3&1, b>>2&1, b>>1&1, b&1
f[i] = fieldSub(fieldElement(b_0+b_1+b_2), fieldElement(b_3+b_4+b_5))
f[i+1] = fieldSub(fieldElement(b_6+b_7+b_8), fieldElement(b_9+b_10+b_11))
f[i+2] = fieldSub(fieldElement(b_12+b_13+b_14), fieldElement(b_15+b_16+b_17))
f[i+3] = fieldSub(fieldElement(b_18+b_19+b_20), fieldElement(b_21+b_22+b_23))
} }
} }
return f return f

View File

@ -37,6 +37,7 @@ const (
// ML-KEM global constants. // ML-KEM global constants.
n = 256 n = 256
q = 3329 q = 3329
maxBytesOf64Mulη = 192
// encodingSizeX is the byte size of a ringElement or nttElement encoded // encodingSizeX is the byte size of a ringElement or nttElement encoded
// by ByteEncode_X (FIPS 203, Algorithm 5). // by ByteEncode_X (FIPS 203, Algorithm 5).