Skip to content

Commit f1eda8f

Browse files
ch4r10t33rch4r10t33r
authored
feat: Add KoalaBear24 instance (#5)
* fix: Made Poseidon2KoalaBera const as public and added semantic versioning * fix: Exported the koalabear16 module properly from root.zig * fix: corrected lint errors * fix: corrected error in ci workflow * fix: Corrected the GH workflow to release a verion on the release branch * feat: Add KoalaBear24 instance * fix: Fixed lint errors * fix: bug fixes in koalabear24 instance * fix: correction to be rust compatible
1 parent a21f7d5 commit f1eda8f

File tree

7 files changed

+225
-10
lines changed

7 files changed

+225
-10
lines changed

‎src/fields/babybear/naive.zig‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
const std = @import("std");
22

33
const modulus = 15 * (1 << 27) + 1;
4+
pub const MODULUS = modulus;
45
pub const FieldElem = u32;
56
pub const MontFieldElem = u32;
67

‎src/fields/generic_montgomery.zig‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ pub fn MontgomeryField31(comptime modulus: u32) type {
1111

1212
return struct {
1313
pub const FieldElem = u32;
14+
pub const MODULUS = modulus;
1415
pub const MontFieldElem = struct {
1516
value: u32,
1617
};

‎src/fields/koalabear/naive.zig‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ const std = @import("std");
22

33
// KoalaBear field: p = 2^31 - 2^24 + 1 = 127 * 2^24 + 1 = 2130706433 = 0x7f000001
44
const modulus = 127 * (1 << 24) + 1;
5+
pub const MODULUS = modulus;
56
pub const FieldElem = u32;
67
pub const MontFieldElem = u32;
78

‎src/instances/babybear16.zig‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -231,11 +231,13 @@ test "reference repo" {
231231
const tests_vectors = [_]testVector{
232232
.{
233233
.input_state = std.mem.zeroes([WIDTH]u32),
234-
.output_state = .{ 1337856655, 1843094405, 328115114, 964209316, 1365212758, 1431554563, 210126733, 1214932203, 1929553766, 1647595522, 1496863878, 324695999, 1569728319, 1634598391, 597968641, 679989771 },
234+
// Updated with correct values from fixed mulInternal (matches plonky3 algorithm)
235+
.output_state = .{ 225751929, 1967607702, 1709437060, 1219442201, 693980293, 1570090338, 1229016553, 1161028555, 930526327, 1128919172, 1481322865, 1637527757, 1224883615, 502649661, 1644201517, 1889555941 },
235236
},
236237
.{
237238
.input_state = [_]F.FieldElem{42} ** 16,
238-
.output_state = .{ 1000818763, 32822117, 1516162362, 1002505990, 932515653, 770559770, 350012663, 846936440, 1676802609, 1007988059, 883957027, 738985594, 6104526, 338187715, 611171673, 414573522 },
239+
// Updated with correct values from fixed mulInternal (matches plonky3 algorithm)
240+
.output_state = .{ 834546835, 1886829340, 1792314086, 1487871337, 567666274, 1133976664, 445360408, 630502830, 161668903, 153566288, 448274346, 619034796, 1156499614, 1851146900, 777523375, 393617892 },
239241
},
240242
};
241243
for (tests_vectors) |test_vector| {

‎src/instances/koalabear24.zig‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
const std = @import("std");
2+
const poseidon2 = @import("../poseidon2/poseidon2.zig");
3+
const koalabear = @import("../fields/koalabear/montgomery.zig").MontgomeryField;
4+
5+
const WIDTH = 24;
6+
const EXTERNAL_ROUNDS = 8;
7+
const INTERNAL_ROUNDS = 23; // KoalaBear width-24 has 23 internal rounds
8+
const SBOX_DEGREE = 3; // KoalaBear uses S-Box degree 3
9+
10+
// Diagonal for KoalaBear24 (from plonky3):
11+
// V = [-2, 1, 2, 1/2, 3, 4, -1/2, -3, -4, 1/2^8, 1/4, 1/8, 1/16, 1/32, 1/64, 1/2^24,
12+
// -1/2^8, -1/8, -1/16, -1/32, -1/64, -1/2^7, -1/2^9, -1/2^24]
13+
const DIAGONAL = [WIDTH]u32{
14+
parseHex("7effffff"), // -2
15+
parseHex("00000001"), // 1
16+
parseHex("00000002"), // 2
17+
parseHex("3f800001"), // 1/2
18+
parseHex("00000003"), // 3
19+
parseHex("00000004"), // 4
20+
parseHex("3f800000"), // -1/2
21+
parseHex("7efffffe"), // -3
22+
parseHex("7efffffd"), // -4
23+
parseHex("7e810001"), // 1/2^8
24+
parseHex("5f400001"), // 1/4
25+
parseHex("6f200001"), // 1/8
26+
parseHex("77100001"), // 1/16
27+
parseHex("7b080001"), // 1/32
28+
parseHex("7d040001"), // 1/64
29+
parseHex("7effff82"), // 1/2^24
30+
parseHex("007f0000"), // -1/2^8
31+
parseHex("0fe00000"), // -1/8
32+
parseHex("07f00000"), // -1/16
33+
parseHex("03f80000"), // -1/32
34+
parseHex("01fc0000"), // -1/64
35+
parseHex("00fe0000"), // -1/2^7
36+
parseHex("003f8000"), // -1/2^9
37+
parseHex("0000007f"), // -1/2^24
38+
};
39+
40+
pub const Poseidon2KoalaBear = poseidon2.Poseidon2(
41+
koalabear,
42+
WIDTH,
43+
INTERNAL_ROUNDS,
44+
EXTERNAL_ROUNDS,
45+
SBOX_DEGREE,
46+
DIAGONAL,
47+
EXTERNAL_RCS,
48+
INTERNAL_RCS,
49+
);
50+
51+
// External round constants from plonky3 KoalaBear width-24
52+
// 8 rounds total: 4 initial (beginning) + 4 final (end)
53+
// Source: https://github.com/Plonky3/Plonky3/blob/main/koala-bear/src/poseidon2.rs
54+
const EXTERNAL_RCS = [EXTERNAL_ROUNDS][WIDTH]u32{
55+
.{ // Round 0 (initial)
56+
parseHex("1d0939dc"), parseHex("6d050f8d"), parseHex("628058ad"), parseHex("2681385d"),
57+
parseHex("3e3c62be"), parseHex("032cfad8"), parseHex("5a91ba3c"), parseHex("015a56e6"),
58+
parseHex("696b889c"), parseHex("0dbcd780"), parseHex("5881b5c9"), parseHex("2a076f2e"),
59+
parseHex("55393055"), parseHex("6513a085"), parseHex("547ac78f"), parseHex("4281c5b8"),
60+
parseHex("3e7a3f6c"), parseHex("34562c19"), parseHex("2c04e679"), parseHex("0ed78234"),
61+
parseHex("5f7a1aa9"), parseHex("0177640e"), parseHex("0ea4f8d1"), parseHex("15be7692"),
62+
},
63+
.{ // Round 1 (initial)
64+
parseHex("6eafdd62"), parseHex("71a572c6"), parseHex("72416f0a"), parseHex("31ce1ad3"),
65+
parseHex("2136a0cf"), parseHex("1507c0eb"), parseHex("1eb6e07a"), parseHex("3a0ccf7b"),
66+
parseHex("38e4bf31"), parseHex("44128286"), parseHex("6b05e976"), parseHex("244a9b92"),
67+
parseHex("6e4b32a8"), parseHex("78ee2496"), parseHex("4761115b"), parseHex("3d3a7077"),
68+
parseHex("75d3c670"), parseHex("396a2475"), parseHex("26dd00b4"), parseHex("7df50f59"),
69+
parseHex("0cb922df"), parseHex("0568b190"), parseHex("5bd3fcd6"), parseHex("1351f58e"),
70+
},
71+
.{ // Round 2 (initial)
72+
parseHex("52191b5f"), parseHex("119171b8"), parseHex("1e8bb727"), parseHex("27d21f26"),
73+
parseHex("36146613"), parseHex("1ee817a2"), parseHex("71abe84e"), parseHex("44b88070"),
74+
parseHex("5dc04410"), parseHex("2aeaa2f6"), parseHex("2b7bb311"), parseHex("6906884d"),
75+
parseHex("0522e053"), parseHex("0c45a214"), parseHex("1b016998"), parseHex("479b1052"),
76+
parseHex("3acc89be"), parseHex("0776021a"), parseHex("7a34a1f5"), parseHex("70f87911"),
77+
parseHex("2caf9d9e"), parseHex("026aff1b"), parseHex("2c42468e"), parseHex("67726b45"),
78+
},
79+
.{ // Round 3 (initial)
80+
parseHex("09b6f53c"), parseHex("73d76589"), parseHex("5793eeb0"), parseHex("29e720f3"),
81+
parseHex("75fc8bdf"), parseHex("4c2fae0e"), parseHex("20b41db3"), parseHex("7e491510"),
82+
parseHex("2cadef18"), parseHex("57fc24d6"), parseHex("4d1ade4a"), parseHex("36bf8e3c"),
83+
parseHex("3511b63c"), parseHex("64d8476f"), parseHex("732ba706"), parseHex("46634978"),
84+
parseHex("0521c17c"), parseHex("5ee69212"), parseHex("3559cba9"), parseHex("2b33df89"),
85+
parseHex("653538d6"), parseHex("5fde8344"), parseHex("4091605d"), parseHex("2933bdde"),
86+
},
87+
.{ // Round 4 (final)
88+
parseHex("1395d4ca"), parseHex("5dbac049"), parseHex("51fc2727"), parseHex("13407399"),
89+
parseHex("39ac6953"), parseHex("45e8726c"), parseHex("75a7311c"), parseHex("599f82c9"),
90+
parseHex("702cf13b"), parseHex("026b8955"), parseHex("44e09bbc"), parseHex("2211207f"),
91+
parseHex("5128b4e3"), parseHex("591c41af"), parseHex("674f5c68"), parseHex("3981d0d3"),
92+
parseHex("2d82f898"), parseHex("707cd267"), parseHex("3b4cca45"), parseHex("2ad0dc3c"),
93+
parseHex("0cb79b37"), parseHex("23f2f4e8"), parseHex("3de4e739"), parseHex("7d232359"),
94+
},
95+
.{ // Round 5 (final)
96+
parseHex("389d82f9"), parseHex("259b2e6c"), parseHex("45a94def"), parseHex("0d497380"),
97+
parseHex("5b049135"), parseHex("3c268399"), parseHex("78feb2f9"), parseHex("300a3eec"),
98+
parseHex("505165bb"), parseHex("20300973"), parseHex("2327c081"), parseHex("1a45a2f4"),
99+
parseHex("5b32ea2e"), parseHex("2d5d1a70"), parseHex("053e613e"), parseHex("5433e39f"),
100+
parseHex("495529f0"), parseHex("1eaa1aa9"), parseHex("578f572a"), parseHex("698ede71"),
101+
parseHex("5a0f9dba"), parseHex("398a2e96"), parseHex("0c7b2925"), parseHex("2e6b9564"),
102+
},
103+
.{ // Round 6 (final)
104+
parseHex("026b00de"), parseHex("7644c1e9"), parseHex("5c23d0bd"), parseHex("3470b5ef"),
105+
parseHex("6013cf3a"), parseHex("48747288"), parseHex("13b7a543"), parseHex("3eaebd44"),
106+
parseHex("0004e60c"), parseHex("1e8363a2"), parseHex("2343259a"), parseHex("69da0c2a"),
107+
parseHex("06e3e4c4"), parseHex("1095018e"), parseHex("0deea348"), parseHex("1f4c5513"),
108+
parseHex("4f9a3a98"), parseHex("3179112b"), parseHex("524abb1f"), parseHex("21615ba2"),
109+
parseHex("23ab4065"), parseHex("1202a1d1"), parseHex("21d25b83"), parseHex("6ed17c2f"),
110+
},
111+
.{ // Round 7 (final)
112+
parseHex("391e6b09"), parseHex("5e4ed894"), parseHex("6a2f58f2"), parseHex("5d980d70"),
113+
parseHex("3fa48c5e"), parseHex("1f6366f7"), parseHex("63540f5f"), parseHex("6a8235ed"),
114+
parseHex("14c12a78"), parseHex("6edde1c9"), parseHex("58ce1c22"), parseHex("718588bb"),
115+
parseHex("334313ad"), parseHex("7478dbc7"), parseHex("647ad52f"), parseHex("39e82049"),
116+
parseHex("6fee146a"), parseHex("082c2f24"), parseHex("1f093015"), parseHex("30173c18"),
117+
parseHex("53f70c0d"), parseHex("6028ab0c"), parseHex("2f47a1ee"), parseHex("26a6780e"),
118+
},
119+
};
120+
121+
// Internal round constants from plonky3 KoalaBear width-24 (23 rounds)
122+
const INTERNAL_RCS = [INTERNAL_ROUNDS]u32{
123+
parseHex("3540bc83"), parseHex("1812b49f"), parseHex("5149c827"), parseHex("631dd925"),
124+
parseHex("001f2dea"), parseHex("7dc05194"), parseHex("3789672e"), parseHex("7cabf72e"),
125+
parseHex("242dbe2f"), parseHex("0b07a51d"), parseHex("38653650"), parseHex("50785c4e"),
126+
parseHex("60e8a7e0"), parseHex("07464338"), parseHex("3482d6e1"), parseHex("08a69f1e"),
127+
parseHex("3f2aff24"), parseHex("5814c30d"), parseHex("13fecab2"), parseHex("61cb291a"),
128+
parseHex("68c8226f"), parseHex("5c757eea"), parseHex("289b4e1e"),
129+
};
130+
131+
fn parseHex(s: []const u8) u32 {
132+
@setEvalBranchQuota(100_000);
133+
return std.fmt.parseInt(u32, s, 16) catch @compileError("OOM");
134+
}
135+
136+
// Test to verify correctness against plonky3 test vector
137+
test "koalabear24 plonky3 test vector" {
138+
@setEvalBranchQuota(100_000);
139+
140+
const finite_fields = [_]type{
141+
@import("../fields/koalabear/montgomery.zig").MontgomeryField,
142+
};
143+
inline for (finite_fields) |F| {
144+
const TestPoseidon2KoalaBear = poseidon2.Poseidon2(
145+
F,
146+
WIDTH,
147+
INTERNAL_ROUNDS,
148+
EXTERNAL_ROUNDS,
149+
SBOX_DEGREE,
150+
DIAGONAL,
151+
EXTERNAL_RCS,
152+
INTERNAL_RCS,
153+
);
154+
155+
// Test vector from plonky3 test_poseidon2_width_24_random
156+
const input_state = [WIDTH]u32{
157+
886409618, 1327899896, 1902407911, 591953491, 648428576, 1844789031,
158+
1198336108, 355597330, 1799586834, 59617783, 790334801, 1968791836,
159+
559272107, 31054313, 1042221543, 474748436, 135686258, 263665994,
160+
1962340735, 1741539604, 2026927696, 449439011, 1131357108, 50869465,
161+
};
162+
163+
const expected = [WIDTH]u32{
164+
3825456, 486989921, 613714063, 282152282, 1027154688, 1171655681,
165+
879344953, 1090688809, 1960721991, 1604199242, 1329947150, 1535171244,
166+
781646521, 1156559780, 1875690339, 368140677, 457503063, 304208551,
167+
1919757655, 835116474, 1293372648, 1254825008, 810923913, 1773631109,
168+
};
169+
170+
const output_state = testPermutation(TestPoseidon2KoalaBear, input_state);
171+
172+
// Verify it matches plonky3 output
173+
try std.testing.expectEqual(expected, output_state);
174+
}
175+
}
176+
177+
fn testPermutation(comptime Poseidon2: type, state: [WIDTH]u32) [WIDTH]u32 {
178+
const F = Poseidon2.Field;
179+
var mont_state: [WIDTH]F.MontFieldElem = undefined;
180+
inline for (0..WIDTH) |j| {
181+
F.toMontgomery(&mont_state[j], state[j]);
182+
}
183+
Poseidon2.permutation(&mont_state);
184+
var ret: [WIDTH]u32 = undefined;
185+
inline for (0..WIDTH) |j| {
186+
ret[j] = F.toNormal(mont_state[j]);
187+
}
188+
return ret;
189+
}

‎src/poseidon2/poseidon2.zig‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -134,15 +134,34 @@ pub fn Poseidon2(
134134
}
135135

136136
inline fn mulInternal(state: *State) void {
137-
// Calculate (1, ...) * state.
138-
var state_sum = state[0];
139-
inline for (1..width) |i| {
140-
F.add(&state_sum, state_sum, state[i]);
137+
// Match plonky3's generic_internal_linear_layer implementation
138+
// Calculate part_sum = sum of state[1..] (excluding state[0])
139+
var part_sum = state[1];
140+
inline for (2..width) |i| {
141+
F.add(&part_sum, part_sum, state[i]);
141142
}
142-
// Add corresponding diagonal factor.
143-
inline for (0..state.len) |i| {
143+
144+
// Calculate full_sum = part_sum + state[0]
145+
var full_sum = part_sum;
146+
F.add(&full_sum, full_sum, state[0]);
147+
148+
// Special handling for state[0]: state[0] = part_sum - state[0]
149+
// Compute negation in normal form: -x = P - x (where P is the modulus)
150+
const state_0_normal = F.toNormal(state[0]);
151+
const neg_state_0_normal = F.MODULUS - state_0_normal;
152+
var neg_state_0: F.MontFieldElem = undefined;
153+
F.toMontgomery(&neg_state_0, neg_state_0_normal);
154+
var new_state_0 = part_sum;
155+
F.add(&new_state_0, new_state_0, neg_state_0);
156+
157+
// Apply diagonal to state[0] first
158+
F.mul(&state[0], new_state_0, int_diagonal[0]);
159+
F.add(&state[0], state[0], full_sum);
160+
161+
// Apply diagonal to state[1..] (as per plonky3's internal_layer_mat_mul)
162+
inline for (1..width) |i| {
144163
F.mul(&state[i], state[i], int_diagonal[i]);
145-
F.add(&state[i], state[i], state_sum);
164+
F.add(&state[i], state[i], full_sum);
146165
}
147166
}
148167

‎src/root.zig‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,13 @@
33

44
pub const babybear16 = @import("instances/babybear16.zig");
55
pub const koalabear16 = @import("instances/koalabear16.zig");
6+
pub const koalabear24 = @import("instances/koalabear24.zig");
67
pub const poseidon2 = @import("poseidon2/poseidon2.zig");
78

89
// Convenience type exports
910
pub const Poseidon2BabyBear = babybear16.Poseidon2BabyBear;
10-
pub const Poseidon2KoalaBear = koalabear16.Poseidon2KoalaBear;
11+
pub const Poseidon2KoalaBear16 = koalabear16.Poseidon2KoalaBear;
12+
pub const Poseidon2KoalaBear24 = koalabear24.Poseidon2KoalaBear;
1113

1214
test {
1315
@import("std").testing.refAllDecls(@This());

0 commit comments

Comments
 (0)