--[=[
	Stateful SHAKE XOF (Extensible Output Function)

	Non-OOP implementation with reusable state buffers.
	For ML-DSA sampling functions that use rejection sampling.

	Example usage:
		local XOF = require(script)

		XOF.Reset128()
		XOF.Absorb128(message)
		local Chunk1 = XOF.Squeeze128(168)
		local Chunk2 = XOF.Squeeze128(168)
--]=]

--!strict
--!optimize 2
--!native

local RATE_128 = 168
local RATE_256 = 136

local LanesLow128 = buffer.create(100)
local LanesHigh128 = buffer.create(100)
local SqueezeOffset128 = 0

local LanesLow256 = buffer.create(100)
local LanesHigh256 = buffer.create(100)
local SqueezeOffset256 = 0

local ZeroBuffer128 = buffer.create(RATE_128)
local ZeroBuffer256 = buffer.create(RATE_256)

local PaddedBuffer128 = buffer.create(RATE_128)
local PaddedBuffer256 = buffer.create(RATE_256)

local SqueezeBuffer128_168 = buffer.create(168)
local SqueezeBuffer256_128 = buffer.create(128)
local SqueezeBuffer256_192 = buffer.create(192)

local LOW_ROUND, HIGH_ROUND = buffer.create(96), buffer.create(96)
do
	local HighFactorKeccak = 0
	local ShiftRegister = 29
	local function GetNextBit(): number
		local Result = ShiftRegister % 2
		ShiftRegister = bit32.bxor((ShiftRegister - Result) // 2, 142 * Result)
		return Result
	end

	for Index = 0, 23 do
		local LowValue: number, Multiplier: number = 0, nil

		for _ = 1, 6 do
			Multiplier = if Multiplier then Multiplier * Multiplier * 2 else 1
			LowValue += GetNextBit() * Multiplier
		end

		local HighValue = GetNextBit() * Multiplier
		buffer.writeu32(HIGH_ROUND, Index * 4, HighValue)
		buffer.writeu32(LOW_ROUND, Index * 4, LowValue + HighValue * HighFactorKeccak)
	end
end

local function Keccak(
	LanesLow: buffer,
	LanesHigh: buffer,
	InputBuffer: buffer,
	Offset: number,
	Size: number,
	BlockSizeInBytes: number
): ()
	local QuadWordsQuantity = BlockSizeInBytes // 8
	local RCHigh, RCLow = HIGH_ROUND, LOW_ROUND
	local XOR = bit32.bxor

	for Position = Offset, Offset + Size - 1, BlockSizeInBytes do
		for Index = 0, (QuadWordsQuantity - 1) * 4, 4 do
			local BufferPos = Position + Index * 2

			buffer.writeu32(
				LanesLow,
				Index,
				XOR(buffer.readu32(LanesLow, Index), buffer.readu32(InputBuffer, BufferPos))
			)

			buffer.writeu32(
				LanesHigh,
				Index,
				XOR(buffer.readu32(LanesHigh, Index), buffer.readu32(InputBuffer, BufferPos + 4))
			)
		end

		local Lane01Low, Lane01High = buffer.readu32(LanesLow, 0), buffer.readu32(LanesHigh, 0)
		local Lane02Low, Lane02High = buffer.readu32(LanesLow, 4), buffer.readu32(LanesHigh, 4)
		local Lane03Low, Lane03High = buffer.readu32(LanesLow, 8), buffer.readu32(LanesHigh, 8)

		local Lane04Low, Lane04High = buffer.readu32(LanesLow, 12), buffer.readu32(LanesHigh, 12)
		local Lane05Low, Lane05High = buffer.readu32(LanesLow, 16), buffer.readu32(LanesHigh, 16)
		local Lane06Low, Lane06High = buffer.readu32(LanesLow, 20), buffer.readu32(LanesHigh, 20)

		local Lane07Low, Lane07High = buffer.readu32(LanesLow, 24), buffer.readu32(LanesHigh, 24)
		local Lane08Low, Lane08High = buffer.readu32(LanesLow, 28), buffer.readu32(LanesHigh, 28)
		local Lane09Low, Lane09High = buffer.readu32(LanesLow, 32), buffer.readu32(LanesHigh, 32)

		local Lane10Low, Lane10High = buffer.readu32(LanesLow, 36), buffer.readu32(LanesHigh, 36)
		local Lane11Low, Lane11High = buffer.readu32(LanesLow, 40), buffer.readu32(LanesHigh, 40)
		local Lane12Low, Lane12High = buffer.readu32(LanesLow, 44), buffer.readu32(LanesHigh, 44)

		local Lane13Low, Lane13High = buffer.readu32(LanesLow, 48), buffer.readu32(LanesHigh, 48)
		local Lane14Low, Lane14High = buffer.readu32(LanesLow, 52), buffer.readu32(LanesHigh, 52)
		local Lane15Low, Lane15High = buffer.readu32(LanesLow, 56), buffer.readu32(LanesHigh, 56)

		local Lane16Low, Lane16High = buffer.readu32(LanesLow, 60), buffer.readu32(LanesHigh, 60)
		local Lane17Low, Lane17High = buffer.readu32(LanesLow, 64), buffer.readu32(LanesHigh, 64)
		local Lane18Low, Lane18High = buffer.readu32(LanesLow, 68), buffer.readu32(LanesHigh, 68)

		local Lane19Low, Lane19High = buffer.readu32(LanesLow, 72), buffer.readu32(LanesHigh, 72)
		local Lane20Low, Lane20High = buffer.readu32(LanesLow, 76), buffer.readu32(LanesHigh, 76)
		local Lane21Low, Lane21High = buffer.readu32(LanesLow, 80), buffer.readu32(LanesHigh, 80)

		local Lane22Low, Lane22High = buffer.readu32(LanesLow, 84), buffer.readu32(LanesHigh, 84)
		local Lane23Low, Lane23High = buffer.readu32(LanesLow, 88), buffer.readu32(LanesHigh, 88)
		local Lane24Low, Lane24High = buffer.readu32(LanesLow, 92), buffer.readu32(LanesHigh, 92)

		local Lane25Low, Lane25High = buffer.readu32(LanesLow, 96), buffer.readu32(LanesHigh, 96)

		for RoundIndex = 0, 92, 4 do
			local Column1Low, Column1High =
				XOR(Lane01Low, Lane06Low, Lane11Low, Lane16Low, Lane21Low),
				XOR(Lane01High, Lane06High, Lane11High, Lane16High, Lane21High)
			local Column2Low, Column2High =
				XOR(Lane02Low, Lane07Low, Lane12Low, Lane17Low, Lane22Low),
				XOR(Lane02High, Lane07High, Lane12High, Lane17High, Lane22High)
			local Column3Low, Column3High =
				XOR(Lane03Low, Lane08Low, Lane13Low, Lane18Low, Lane23Low),
				XOR(Lane03High, Lane08High, Lane13High, Lane18High, Lane23High)
			local Column4Low, Column4High =
				XOR(Lane04Low, Lane09Low, Lane14Low, Lane19Low, Lane24Low),
				XOR(Lane04High, Lane09High, Lane14High, Lane19High, Lane24High)
			local Column5Low, Column5High =
				XOR(Lane05Low, Lane10Low, Lane15Low, Lane20Low, Lane25Low),
				XOR(Lane05High, Lane10High, Lane15High, Lane20High, Lane25High)

			local DeltaLow, DeltaHigh =
				XOR(Column1Low, Column3Low * 2 + Column3High // 2147483648),
				XOR(Column1High, Column3High * 2 + Column3Low // 2147483648)
			local Temp0Low, Temp0High = XOR(DeltaLow, Lane02Low), XOR(DeltaHigh, Lane02High)
			local Temp1Low, Temp1High = XOR(DeltaLow, Lane07Low), XOR(DeltaHigh, Lane07High)
			local Temp2Low, Temp2High = XOR(DeltaLow, Lane12Low), XOR(DeltaHigh, Lane12High)
			local Temp3Low, Temp3High = XOR(DeltaLow, Lane17Low), XOR(DeltaHigh, Lane17High)
			local Temp4Low, Temp4High = XOR(DeltaLow, Lane22Low), XOR(DeltaHigh, Lane22High)

			Lane02Low = Temp1Low // 1048576 + (Temp1High * 4096)
			Lane02High = Temp1High // 1048576 + (Temp1Low * 4096)
			Lane07Low = Temp3Low // 524288 + (Temp3High * 8192)
			Lane07High = Temp3High // 524288 + (Temp3Low * 8192)
			Lane12Low = Temp0Low * 2 + Temp0High // 2147483648
			Lane12High = Temp0High * 2 + Temp0Low // 2147483648
			Lane17Low = Temp2Low * 1024 + Temp2High // 4194304
			Lane17High = Temp2High * 1024 + Temp2Low // 4194304
			Lane22Low = Temp4Low * 4 + Temp4High // 1073741824
			Lane22High = Temp4High * 4 + Temp4Low // 1073741824

			DeltaLow = XOR(Column2Low, Column4Low * 2 + Column4High // 2147483648)
			DeltaHigh = XOR(Column2High, Column4High * 2 + Column4Low // 2147483648)
			Temp0Low = XOR(DeltaLow, Lane03Low)
			Temp0High = XOR(DeltaHigh, Lane03High)
			Temp1Low = XOR(DeltaLow, Lane08Low)
			Temp1High = XOR(DeltaHigh, Lane08High)
			Temp2Low = XOR(DeltaLow, Lane13Low)
			Temp2High = XOR(DeltaHigh, Lane13High)
			Temp3Low = XOR(DeltaLow, Lane18Low)
			Temp3High = XOR(DeltaHigh, Lane18High)
			Temp4Low = XOR(DeltaLow, Lane23Low)
			Temp4High = XOR(DeltaHigh, Lane23High)

			Lane03Low = Temp2Low // 2097152 + (Temp2High * 2048)
			Lane03High = Temp2High // 2097152 + (Temp2Low * 2048)
			Lane08Low = Temp4Low // 8 + bit32.bor(Temp4High * 536870912, 0)
			Lane08High = Temp4High // 8 + bit32.bor(Temp4Low * 536870912, 0)
			Lane13Low = Temp1Low * 64 + Temp1High // 67108864
			Lane13High = Temp1High * 64 + Temp1Low // 67108864
			Lane18Low = (Temp3Low * 32768) + Temp3High // 131072
			Lane18High = (Temp3High * 32768) + Temp3Low // 131072
			Lane23Low = Temp0Low // 4 + bit32.bor(Temp0High * 1073741824, 0)
			Lane23High = Temp0High // 4 + bit32.bor(Temp0Low * 1073741824, 0)

			DeltaLow = XOR(Column3Low, Column5Low * 2 + Column5High // 2147483648)
			DeltaHigh = XOR(Column3High, Column5High * 2 + Column5Low // 2147483648)
			Temp0Low = XOR(DeltaLow, Lane04Low)
			Temp0High = XOR(DeltaHigh, Lane04High)
			Temp1Low = XOR(DeltaLow, Lane09Low)
			Temp1High = XOR(DeltaHigh, Lane09High)
			Temp2Low = XOR(DeltaLow, Lane14Low)
			Temp2High = XOR(DeltaHigh, Lane14High)
			Temp3Low = XOR(DeltaLow, Lane19Low)
			Temp3High = XOR(DeltaHigh, Lane19High)
			Temp4Low = XOR(DeltaLow, Lane24Low)
			Temp4High = XOR(DeltaHigh, Lane24High)

			Lane04Low = bit32.bor(Temp3Low * 2097152, 0) + Temp3High // 2048
			Lane04High = bit32.bor(Temp3High * 2097152, 0) + Temp3Low // 2048
			Lane09Low = bit32.bor(Temp0Low * 268435456, 0) + Temp0High // 16
			Lane09High = bit32.bor(Temp0High * 268435456, 0) + Temp0Low // 16
			Lane14Low = bit32.bor(Temp2Low * 33554432, 0) + Temp2High // 128
			Lane14High = bit32.bor(Temp2High * 33554432, 0) + Temp2Low // 128
			Lane19Low = Temp4Low // 256 + bit32.bor(Temp4High * 16777216, 0)
			Lane19High = Temp4High // 256 + bit32.bor(Temp4Low * 16777216, 0)
			Lane24Low = Temp1Low // 512 + bit32.bor(Temp1High * 8388608, 0)
			Lane24High = Temp1High // 512 + bit32.bor(Temp1Low * 8388608, 0)
			DeltaLow = XOR(Column4Low, Column1Low * 2 + Column1High // 2147483648)
			DeltaHigh = XOR(Column4High, Column1High * 2 + Column1Low // 2147483648)

			Temp0Low = XOR(DeltaLow, Lane05Low)
			Temp0High = XOR(DeltaHigh, Lane05High)
			Temp1Low = XOR(DeltaLow, Lane10Low)
			Temp1High = XOR(DeltaHigh, Lane10High)
			Temp2Low = XOR(DeltaLow, Lane15Low)
			Temp2High = XOR(DeltaHigh, Lane15High)
			Temp3Low = XOR(DeltaLow, Lane20Low)
			Temp3High = XOR(DeltaHigh, Lane20High)
			Temp4Low = XOR(DeltaLow, Lane25Low)
			Temp4High = XOR(DeltaHigh, Lane25High)

			Lane05Low = (Temp4Low * 16384) + Temp4High // 262144
			Lane05High = (Temp4High * 16384) + Temp4Low // 262144
			Lane10Low = bit32.bor(Temp1Low * 1048576, 0) + Temp1High // 4096
			Lane10High = bit32.bor(Temp1High * 1048576, 0) + Temp1Low // 4096
			Lane15Low = Temp3Low * 256 + Temp3High // 16777216
			Lane15High = Temp3High * 256 + Temp3Low // 16777216
			Lane20Low = bit32.bor(Temp0Low * 134217728, 0) + Temp0High // 32
			Lane20High = bit32.bor(Temp0High * 134217728, 0) + Temp0Low // 32
			Lane25Low = Temp2Low // 33554432 + Temp2High * 128
			Lane25High = Temp2High // 33554432 + Temp2Low * 128

			DeltaLow = XOR(Column5Low, Column2Low * 2 + Column2High // 2147483648)
			DeltaHigh = XOR(Column5High, Column2High * 2 + Column2Low // 2147483648)
			Temp1Low = XOR(DeltaLow, Lane06Low)
			Temp1High = XOR(DeltaHigh, Lane06High)
			Temp2Low = XOR(DeltaLow, Lane11Low)
			Temp2High = XOR(DeltaHigh, Lane11High)
			Temp3Low = XOR(DeltaLow, Lane16Low)
			Temp3High = XOR(DeltaHigh, Lane16High)
			Temp4Low = XOR(DeltaLow, Lane21Low)
			Temp4High = XOR(DeltaHigh, Lane21High)
			Lane06Low = Temp2Low * 8 + Temp2High // 536870912
			Lane06High = Temp2High * 8 + Temp2Low // 536870912
			Lane11Low = (Temp4Low * 262144) + Temp4High // 16384
			Lane11High = (Temp4High * 262144) + Temp4Low // 16384
			Lane16Low = Temp1Low // 268435456 + Temp1High * 16
			Lane16High = Temp1High // 268435456 + Temp1Low * 16
			Lane21Low = Temp3Low // 8388608 + Temp3High * 512
			Lane21High = Temp3High // 8388608 + Temp3Low * 512
			Lane01Low = XOR(DeltaLow, Lane01Low)
			Lane01High = XOR(DeltaHigh, Lane01High)

			Lane01Low, Lane02Low, Lane03Low, Lane04Low, Lane05Low =
				XOR(Lane01Low, bit32.band(-1 - Lane02Low, Lane03Low)),
				XOR(Lane02Low, bit32.band(-1 - Lane03Low, Lane04Low)),
				XOR(Lane03Low, bit32.band(-1 - Lane04Low, Lane05Low)),
				XOR(Lane04Low, bit32.band(-1 - Lane05Low, Lane01Low)),
				XOR(Lane05Low, bit32.band(-1 - Lane01Low, Lane02Low))
			Lane01High, Lane02High, Lane03High, Lane04High, Lane05High =
				XOR(Lane01High, bit32.band(-1 - Lane02High, Lane03High)),
				XOR(Lane02High, bit32.band(-1 - Lane03High, Lane04High)),
				XOR(Lane03High, bit32.band(-1 - Lane04High, Lane05High)),
				XOR(Lane04High, bit32.band(-1 - Lane05High, Lane01High)),
				XOR(Lane05High, bit32.band(-1 - Lane01High, Lane02High))
			Lane06Low, Lane07Low, Lane08Low, Lane09Low, Lane10Low =
				XOR(Lane09Low, bit32.band(-1 - Lane10Low, Lane06Low)),
				XOR(Lane10Low, bit32.band(-1 - Lane06Low, Lane07Low)),
				XOR(Lane06Low, bit32.band(-1 - Lane07Low, Lane08Low)),
				XOR(Lane07Low, bit32.band(-1 - Lane08Low, Lane09Low)),
				XOR(Lane08Low, bit32.band(-1 - Lane09Low, Lane10Low))
			Lane06High, Lane07High, Lane08High, Lane09High, Lane10High =
				XOR(Lane09High, bit32.band(-1 - Lane10High, Lane06High)),
				XOR(Lane10High, bit32.band(-1 - Lane06High, Lane07High)),
				XOR(Lane06High, bit32.band(-1 - Lane07High, Lane08High)),
				XOR(Lane07High, bit32.band(-1 - Lane08High, Lane09High)),
				XOR(Lane08High, bit32.band(-1 - Lane09High, Lane10High))
			Lane11Low, Lane12Low, Lane13Low, Lane14Low, Lane15Low =
				XOR(Lane12Low, bit32.band(-1 - Lane13Low, Lane14Low)),
				XOR(Lane13Low, bit32.band(-1 - Lane14Low, Lane15Low)),
				XOR(Lane14Low, bit32.band(-1 - Lane15Low, Lane11Low)),
				XOR(Lane15Low, bit32.band(-1 - Lane11Low, Lane12Low)),
				XOR(Lane11Low, bit32.band(-1 - Lane12Low, Lane13Low))
			Lane11High, Lane12High, Lane13High, Lane14High, Lane15High =
				XOR(Lane12High, bit32.band(-1 - Lane13High, Lane14High)),
				XOR(Lane13High, bit32.band(-1 - Lane14High, Lane15High)),
				XOR(Lane14High, bit32.band(-1 - Lane15High, Lane11High)),
				XOR(Lane15High, bit32.band(-1 - Lane11High, Lane12High)),
				XOR(Lane11High, bit32.band(-1 - Lane12High, Lane13High))
			Lane16Low, Lane17Low, Lane18Low, Lane19Low, Lane20Low =
				XOR(Lane20Low, bit32.band(-1 - Lane16Low, Lane17Low)),
				XOR(Lane16Low, bit32.band(-1 - Lane17Low, Lane18Low)),
				XOR(Lane17Low, bit32.band(-1 - Lane18Low, Lane19Low)),
				XOR(Lane18Low, bit32.band(-1 - Lane19Low, Lane20Low)),
				XOR(Lane19Low, bit32.band(-1 - Lane20Low, Lane16Low))
			Lane16High, Lane17High, Lane18High, Lane19High, Lane20High =
				XOR(Lane20High, bit32.band(-1 - Lane16High, Lane17High)),
				XOR(Lane16High, bit32.band(-1 - Lane17High, Lane18High)),
				XOR(Lane17High, bit32.band(-1 - Lane18High, Lane19High)),
				XOR(Lane18High, bit32.band(-1 - Lane19High, Lane20High)),
				XOR(Lane19High, bit32.band(-1 - Lane20High, Lane16High))
			Lane21Low, Lane22Low, Lane23Low, Lane24Low, Lane25Low =
				XOR(Lane23Low, bit32.band(-1 - Lane24Low, Lane25Low)),
				XOR(Lane24Low, bit32.band(-1 - Lane25Low, Lane21Low)),
				XOR(Lane25Low, bit32.band(-1 - Lane21Low, Lane22Low)),
				XOR(Lane21Low, bit32.band(-1 - Lane22Low, Lane23Low)),
				XOR(Lane22Low, bit32.band(-1 - Lane23Low, Lane24Low))
			Lane21High, Lane22High, Lane23High, Lane24High, Lane25High =
				XOR(Lane23High, bit32.band(-1 - Lane24High, Lane25High)),
				XOR(Lane24High, bit32.band(-1 - Lane25High, Lane21High)),
				XOR(Lane25High, bit32.band(-1 - Lane21High, Lane22High)),
				XOR(Lane21High, bit32.band(-1 - Lane22High, Lane23High)),
				XOR(Lane22High, bit32.band(-1 - Lane23High, Lane24High))

			Lane01Low = XOR(Lane01Low, buffer.readu32(RCLow, RoundIndex))
			Lane01High = XOR(Lane01High, buffer.readu32(RCHigh, RoundIndex))
		end

		buffer.writeu32(LanesLow, 0, Lane01Low)
		buffer.writeu32(LanesHigh, 0, Lane01High)
		buffer.writeu32(LanesLow, 4, Lane02Low)
		buffer.writeu32(LanesHigh, 4, Lane02High)
		buffer.writeu32(LanesLow, 8, Lane03Low)
		buffer.writeu32(LanesHigh, 8, Lane03High)
		buffer.writeu32(LanesLow, 12, Lane04Low)
		buffer.writeu32(LanesHigh, 12, Lane04High)
		buffer.writeu32(LanesLow, 16, Lane05Low)
		buffer.writeu32(LanesHigh, 16, Lane05High)
		buffer.writeu32(LanesLow, 20, Lane06Low)
		buffer.writeu32(LanesHigh, 20, Lane06High)
		buffer.writeu32(LanesLow, 24, Lane07Low)
		buffer.writeu32(LanesHigh, 24, Lane07High)
		buffer.writeu32(LanesLow, 28, Lane08Low)
		buffer.writeu32(LanesHigh, 28, Lane08High)
		buffer.writeu32(LanesLow, 32, Lane09Low)
		buffer.writeu32(LanesHigh, 32, Lane09High)
		buffer.writeu32(LanesLow, 36, Lane10Low)
		buffer.writeu32(LanesHigh, 36, Lane10High)
		buffer.writeu32(LanesLow, 40, Lane11Low)
		buffer.writeu32(LanesHigh, 40, Lane11High)
		buffer.writeu32(LanesLow, 44, Lane12Low)
		buffer.writeu32(LanesHigh, 44, Lane12High)
		buffer.writeu32(LanesLow, 48, Lane13Low)
		buffer.writeu32(LanesHigh, 48, Lane13High)
		buffer.writeu32(LanesLow, 52, Lane14Low)
		buffer.writeu32(LanesHigh, 52, Lane14High)
		buffer.writeu32(LanesLow, 56, Lane15Low)
		buffer.writeu32(LanesHigh, 56, Lane15High)
		buffer.writeu32(LanesLow, 60, Lane16Low)
		buffer.writeu32(LanesHigh, 60, Lane16High)
		buffer.writeu32(LanesLow, 64, Lane17Low)
		buffer.writeu32(LanesHigh, 64, Lane17High)
		buffer.writeu32(LanesLow, 68, Lane18Low)
		buffer.writeu32(LanesHigh, 68, Lane18High)
		buffer.writeu32(LanesLow, 72, Lane19Low)
		buffer.writeu32(LanesHigh, 72, Lane19High)
		buffer.writeu32(LanesLow, 76, Lane20Low)
		buffer.writeu32(LanesHigh, 76, Lane20High)
		buffer.writeu32(LanesLow, 80, Lane21Low)
		buffer.writeu32(LanesHigh, 80, Lane21High)
		buffer.writeu32(LanesLow, 84, Lane22Low)
		buffer.writeu32(LanesHigh, 84, Lane22High)
		buffer.writeu32(LanesLow, 88, Lane23Low)
		buffer.writeu32(LanesHigh, 88, Lane23High)
		buffer.writeu32(LanesLow, 92, Lane24Low)
		buffer.writeu32(LanesHigh, 92, Lane24High)
		buffer.writeu32(LanesLow, 96, Lane25Low)
		buffer.writeu32(LanesHigh, 96, Lane25High)
	end
end

local XOF = {}

function XOF.Reset128()
	buffer.fill(LanesLow128, 0, 0, 100)
	buffer.fill(LanesHigh128, 0, 0, 100)
	SqueezeOffset128 = 0
end

function XOF.Reset256()
	buffer.fill(LanesLow256, 0, 0, 100)
	buffer.fill(LanesHigh256, 0, 0, 100)
	SqueezeOffset256 = 0
end

function XOF.Absorb128(Message: buffer)
	local MessageLength = buffer.len(Message)
	local RateBytes = RATE_128

	local PaddedMessage = PaddedBuffer128
	buffer.fill(PaddedMessage, 0, 0, RateBytes)

	if MessageLength > 0 then
		buffer.copy(PaddedMessage, 0, Message, 0, MessageLength)
	end

	if RateBytes - MessageLength == 1 then
		buffer.writeu8(PaddedMessage, MessageLength, 0x9F)
	else
		buffer.writeu8(PaddedMessage, MessageLength, 0x1F)
		buffer.writeu8(PaddedMessage, RateBytes - 1, 0x80)
	end

	Keccak(LanesLow128, LanesHigh128, PaddedMessage, 0, RateBytes, RateBytes)
end

function XOF.Absorb256(Message: buffer)
	local MessageLength = buffer.len(Message)
	local RateBytes = RATE_256

	local PaddedMessage = PaddedBuffer256
	buffer.fill(PaddedMessage, 0, 0, RateBytes)

	if MessageLength > 0 then
		buffer.copy(PaddedMessage, 0, Message, 0, MessageLength)
	end

	if RateBytes - MessageLength == 1 then
		buffer.writeu8(PaddedMessage, MessageLength, 0x9F)
	else
		buffer.writeu8(PaddedMessage, MessageLength, 0x1F)
		buffer.writeu8(PaddedMessage, RateBytes - 1, 0x80)
	end

	Keccak(LanesLow256, LanesHigh256, PaddedMessage, 0, RateBytes, RateBytes)
end

function XOF.Squeeze128Into(Output: buffer, OutputBytes: number, OutputOffset: number?)
	local OutOffset = OutputOffset or 0
	local RateBytes = RATE_128
	local LanesLow = LanesLow128
	local LanesHigh = LanesHigh128
	local SqueezeOffset = SqueezeOffset128
	local ZeroBuffer = ZeroBuffer128

	local Written = 0
	while Written < OutputBytes do
		if SqueezeOffset >= RateBytes then
			Keccak(LanesLow, LanesHigh, ZeroBuffer, 0, RateBytes, RateBytes)
			SqueezeOffset = 0
		end

		local BytesThisRound = RateBytes - SqueezeOffset
		if BytesThisRound > OutputBytes - Written then
			BytesThisRound = OutputBytes - Written
		end

		local ByteIndex = 0
		while ByteIndex < BytesThisRound do
			local AbsoluteIndex = SqueezeOffset + ByteIndex
			local Lane = bit32.rshift(AbsoluteIndex, 3)
			local ByteInLane = bit32.band(AbsoluteIndex, 7)
			local LaneOffset = bit32.lshift(Lane, 2)

			if ByteInLane == 0 and ByteIndex + 8 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesLow, LaneOffset))
				buffer.writeu32(Output, OutOffset + Written + ByteIndex + 4, buffer.readu32(LanesHigh, LaneOffset))
				ByteIndex += 8
			elseif ByteInLane == 0 and ByteIndex + 4 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesLow, LaneOffset))
				ByteIndex += 4
			elseif ByteInLane == 4 and ByteIndex + 4 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesHigh, LaneOffset))
				ByteIndex += 4
			else
				local Value
				if ByteInLane < 4 then
					Value = bit32.extract(buffer.readu32(LanesLow, LaneOffset), bit32.lshift(ByteInLane, 3), 8)
				else
					Value = bit32.extract(buffer.readu32(LanesHigh, LaneOffset), bit32.lshift(ByteInLane - 4, 3), 8)
				end
				buffer.writeu8(Output, OutOffset + Written + ByteIndex, Value)
				ByteIndex += 1
			end
		end

		Written += BytesThisRound
		SqueezeOffset += BytesThisRound
	end

	SqueezeOffset128 = SqueezeOffset
end

function XOF.Squeeze128(OutputBytes: number): buffer
	local Output = if OutputBytes == 168 then SqueezeBuffer128_168 else buffer.create(OutputBytes)
	XOF.Squeeze128Into(Output, OutputBytes, 0)
	return Output
end

function XOF.Squeeze256Into(Output: buffer, OutputBytes: number, OutputOffset: number?)
	local OutOffset = OutputOffset or 0
	local RateBytes = RATE_256
	local LanesLow = LanesLow256
	local LanesHigh = LanesHigh256
	local SqueezeOffset = SqueezeOffset256
	local ZeroBuffer = ZeroBuffer256

	local Written = 0
	while Written < OutputBytes do
		if SqueezeOffset >= RateBytes then
			Keccak(LanesLow, LanesHigh, ZeroBuffer, 0, RateBytes, RateBytes)
			SqueezeOffset = 0
		end

		local BytesThisRound = RateBytes - SqueezeOffset
		if BytesThisRound > OutputBytes - Written then
			BytesThisRound = OutputBytes - Written
		end

		local ByteIndex = 0
		while ByteIndex < BytesThisRound do
			local AbsoluteIndex = SqueezeOffset + ByteIndex
			local Lane = bit32.rshift(AbsoluteIndex, 3)
			local ByteInLane = bit32.band(AbsoluteIndex, 7)
			local LaneOffset = bit32.lshift(Lane, 2)

			if ByteInLane == 0 and ByteIndex + 8 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesLow, LaneOffset))
				buffer.writeu32(Output, OutOffset + Written + ByteIndex + 4, buffer.readu32(LanesHigh, LaneOffset))
				ByteIndex += 8
			elseif ByteInLane == 0 and ByteIndex + 4 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesLow, LaneOffset))
				ByteIndex += 4
			elseif ByteInLane == 4 and ByteIndex + 4 <= BytesThisRound then
				buffer.writeu32(Output, OutOffset + Written + ByteIndex, buffer.readu32(LanesHigh, LaneOffset))
				ByteIndex += 4
			else
				local Value
				if ByteInLane < 4 then
					Value = bit32.extract(buffer.readu32(LanesLow, LaneOffset), bit32.lshift(ByteInLane, 3), 8)
				else
					Value = bit32.extract(buffer.readu32(LanesHigh, LaneOffset), bit32.lshift(ByteInLane - 4, 3), 8)
				end
				buffer.writeu8(Output, OutOffset + Written + ByteIndex, Value)
				ByteIndex += 1
			end
		end

		Written += BytesThisRound
		SqueezeOffset += BytesThisRound
	end

	SqueezeOffset256 = SqueezeOffset
end

function XOF.Squeeze256(OutputBytes: number): buffer
	local Output = if OutputBytes == 128
		then SqueezeBuffer256_128
		elseif OutputBytes == 192 then SqueezeBuffer256_192
		else buffer.create(OutputBytes)
	XOF.Squeeze256Into(Output, OutputBytes, 0)
	return Output
end

return XOF


--[=[
	ML-DSA (FIPS 204) Digital Signature Algorithm

	Post quantum digital signature scheme based on lattice cryptography.
	Has three security levels: ML-DSA-44, ML-DSA-65, ML-DSA-87.

	Example usage:
		local MLDSA = require("@self/MlDSA")

		local PubKey, SecKey = MLDSA.ML_DSA_44.GenerateKeys()
		local Message = buffer.fromstring("Hello World")
		local Context = buffer.create(0)
		local Signature = buffer.create(MLDSA.ML_DSA_44.SigByteLen)
		local Random = buffer.create(32)

		local Success = MLDSA.ML_DSA_44.Sign(Random, SecKey, Message, Context, Signature)
		local Valid = MLDSA.ML_DSA_44.Verify(PubKey, Message, Context, Signature)
--]=]

--!strict
--!optimize 2
--!native

local SHA3 = require("@self/SHA3")
local NTT = require("@self/NTT")

local PolyVec = require("@self/PolyVec")
local BitPacking = require("@self/Pack")

local Sampling = require("@self/Sampling")
local Params = require("@self/Params")

local Utils = require("@self/Utils")
local CSPRNG = require("@self/CSPRNG")

local KEYGEN_SEED_BYTE_LEN = 32
local RND_BYTE_LEN = 32

local N = 256
local Q = 8380417

local Mldsa = {
	CSPRNG = CSPRNG,
}

local function CompareBufferSlices(A: buffer, AOffset: number, B: buffer, BOffset: number, Length: number): boolean
	local Difference = 0

	if Length == 32 then
		Difference = bit32.bor(
			bit32.bxor(buffer.readu32(A, AOffset), buffer.readu32(B, BOffset)),
			bit32.bxor(buffer.readu32(A, AOffset + 4), buffer.readu32(B, BOffset + 4)),
			bit32.bxor(buffer.readu32(A, AOffset + 8), buffer.readu32(B, BOffset + 8)),
			bit32.bxor(buffer.readu32(A, AOffset + 12), buffer.readu32(B, BOffset + 12)),
			bit32.bxor(buffer.readu32(A, AOffset + 16), buffer.readu32(B, BOffset + 16)),
			bit32.bxor(buffer.readu32(A, AOffset + 20), buffer.readu32(B, BOffset + 20)),
			bit32.bxor(buffer.readu32(A, AOffset + 24), buffer.readu32(B, BOffset + 24)),
			bit32.bxor(buffer.readu32(A, AOffset + 28), buffer.readu32(B, BOffset + 28))
		)
		return Difference == 0
	end

	local AlignedLength = bit32.band(Length, bit32.bnot(3))
	for I = 0, AlignedLength - 4, 4 do
		Difference = bit32.bor(Difference, bit32.bxor(buffer.readu32(A, AOffset + I), buffer.readu32(B, BOffset + I)))
	end

	for I = AlignedLength, Length - 1 do
		Difference = bit32.bor(Difference, bit32.bxor(buffer.readu8(A, AOffset + I), buffer.readu8(B, BOffset + I)))
	end

	return Difference == 0
end

local function Keygen(Seed: buffer, PubKey: buffer, SecKey: buffer, K: number, L: number, D: number, Eta: number)
	if not Params.CheckKeygenParams(K, L, D, Eta) then
		error("Invalid keygen parameters")
	end

	local T1Bw = Utils.BitWidth(Q) - D
	local EtaBw = Utils.BitWidth(2 * Eta)
	local S1Len = L * EtaBw * 32
	local S2Len = K * EtaBw * 32
	local T0Rng = bit32.lshift(1, D - 1)

	local PkOff1 = 32
	local SkOff0 = 0
	local SkOff1 = SkOff0 + 32
	local SkOff2 = SkOff1 + 32
	local SkOff3 = SkOff2 + 64
	local SkOff4 = SkOff3 + S1Len
	local SkOff5 = SkOff4 + S2Len

	local DomainSeparator = buffer.create(2)
	buffer.writeu8(DomainSeparator, 0, K)
	buffer.writeu8(DomainSeparator, 1, L)

	local SeedHashInput = buffer.create(34)
	buffer.copy(SeedHashInput, 0, Seed, 0, 32)
	buffer.copy(SeedHashInput, 32, DomainSeparator, 0, 2)

	local SeedHash = SHA3.SHAKE256(SeedHashInput, 128)

	local Rho = buffer.create(32)
	local RhoPrime = buffer.create(64)
	local Key = buffer.create(32)

	buffer.copy(Rho, 0, SeedHash, 0, 32)
	buffer.copy(RhoPrime, 0, SeedHash, 32, 64)
	buffer.copy(Key, 0, SeedHash, 96, 32)

	local A = buffer.create(K * L * N * 4)
	Sampling.ExpandA(Rho, A, K, L)

	local S1 = buffer.create(L * N * 4)
	local S2 = buffer.create(K * N * 4)

	Sampling.ExpandS(RhoPrime, S1, Eta, L, 0)
	Sampling.ExpandS(RhoPrime, S2, Eta, K, L)

	local S1Prime = buffer.create(L * N * 4)
	PolyVec.Copy(S1, S1Prime, L)
	PolyVec.ForwardNTT(S1Prime, L)

	local T = buffer.create(K * N * 4)
	PolyVec.MatrixMultiply(A, S1Prime, T, K, L, L, 1)
	PolyVec.InverseNTT(T, K)
	PolyVec.AddTo(S2, T, K)

	local T1 = buffer.create(K * N * 4)
	local T0 = buffer.create(K * N * 4)
	PolyVec.Power2Round(T, T1, T0, K, D)

	buffer.copy(PubKey, 0, Rho, 0, 32)
	local T1Encoded = buffer.create(buffer.len(PubKey) - PkOff1)
	PolyVec.Encode(T1, T1Encoded, K, T1Bw)
	buffer.copy(PubKey, PkOff1, T1Encoded, 0, buffer.len(T1Encoded))

	local Tr = SHA3.SHAKE256(PubKey, 64)

	buffer.copy(SecKey, SkOff0, Rho, 0, 32)
	buffer.copy(SecKey, SkOff1, Key, 0, 32)
	buffer.copy(SecKey, SkOff2, Tr, 0, 64)

	PolyVec.SubFromX(S1, L, Eta)
	PolyVec.SubFromX(S2, K, Eta)

	local S1Encoded = buffer.create(S1Len)
	local S2Encoded = buffer.create(S2Len)
	PolyVec.Encode(S1, S1Encoded, L, EtaBw)
	PolyVec.Encode(S2, S2Encoded, K, EtaBw)

	buffer.copy(SecKey, SkOff3, S1Encoded, 0, S1Len)
	buffer.copy(SecKey, SkOff4, S2Encoded, 0, S2Len)

	PolyVec.SubFromX(T0, K, T0Rng)

	local T0Encoded = buffer.create(buffer.len(SecKey) - SkOff5)
	PolyVec.Encode(T0, T0Encoded, K, D)
	buffer.copy(SecKey, SkOff5, T0Encoded, 0, buffer.len(T0Encoded))
end

local function Sign(
	Rnd: buffer,
	SecKey: buffer,
	Msg: buffer,
	Ctx: buffer,
	Sig: buffer,
	K: number,
	L: number,
	D: number,
	Eta: number,
	Gamma1: number,
	Gamma2: number,
	Tau: number,
	Beta: number,
	Omega: number,
	Lambda: number
): boolean
	if not Params.CheckSigningParams(K, L, D, Eta, Gamma1, Gamma2, Tau, Beta, Omega, Lambda) then
		error("Invalid signing parameters")
	end

	if buffer.len(Ctx) > 255 then
		return false
	end

	local T0Rng = bit32.lshift(1, D - 1)
	local EtaBw = Utils.BitWidth(2 * Eta)
	local S1Len = L * EtaBw * 32
	local S2Len = K * EtaBw * 32
	local Alpha = bit32.lshift(Gamma2, 1)
	local M = math.floor((Q - 1) / Alpha)
	local W1Bw = Utils.BitWidth(M - 1)
	local CTildaSize = math.floor((2 * Lambda) / 8)
	local Gamma1Bw = Utils.BitWidth(Gamma1)

	local SkOff0 = 0
	local SkOff1 = SkOff0 + 32
	local SkOff2 = SkOff1 + 32
	local SkOff3 = SkOff2 + 64
	local SkOff4 = SkOff3 + S1Len
	local SkOff5 = SkOff4 + S2Len

	local Rho = buffer.create(32)
	local Key = buffer.create(32)
	local Tr = buffer.create(64)

	buffer.copy(Rho, 0, SecKey, SkOff0, 32)
	buffer.copy(Key, 0, SecKey, SkOff1, 32)
	buffer.copy(Tr, 0, SecKey, SkOff2, 64)

	local A = buffer.create(K * L * N * 4)
	Sampling.ExpandA(Rho, A, K, L)

	local DomainSeparator = buffer.create(2)
	buffer.writeu8(DomainSeparator, 0, 0)
	buffer.writeu8(DomainSeparator, 1, buffer.len(Ctx))

	local MuInput = buffer.create(64 + 2 + buffer.len(Ctx) + buffer.len(Msg))
	local Offset = 0
	buffer.copy(MuInput, Offset, Tr, 0, 64)
	Offset += 64
	buffer.copy(MuInput, Offset, DomainSeparator, 0, 2)
	Offset += 2
	buffer.copy(MuInput, Offset, Ctx, 0, buffer.len(Ctx))
	Offset += buffer.len(Ctx)
	buffer.copy(MuInput, Offset, Msg, 0, buffer.len(Msg))

	local Mu = SHA3.SHAKE256(MuInput, 64)

	local RhoPrimeInput = buffer.create(32 + 32 + 64)
	buffer.copy(RhoPrimeInput, 0, Key, 0, 32)
	buffer.copy(RhoPrimeInput, 32, Rnd, 0, 32)
	buffer.copy(RhoPrimeInput, 64, Mu, 0, 64)

	local RhoPrime = SHA3.SHAKE256(RhoPrimeInput, 64)

	local S1 = buffer.create(L * N * 4)
	local S2 = buffer.create(K * N * 4)
	local T0 = buffer.create(K * N * 4)

	local S1Encoded = buffer.create(S1Len)
	local S2Encoded = buffer.create(S2Len)
	local T0Encoded = buffer.create(buffer.len(SecKey) - SkOff5)

	buffer.copy(S1Encoded, 0, SecKey, SkOff3, S1Len)
	buffer.copy(S2Encoded, 0, SecKey, SkOff4, S2Len)
	buffer.copy(T0Encoded, 0, SecKey, SkOff5, buffer.len(T0Encoded))

	PolyVec.Decode(S1Encoded, S1, L, EtaBw)
	PolyVec.Decode(S2Encoded, S2, K, EtaBw)
	PolyVec.Decode(T0Encoded, T0, K, D)

	PolyVec.SubFromX(S1, L, Eta)
	PolyVec.SubFromX(S2, K, Eta)
	PolyVec.SubFromX(T0, K, T0Rng)

	PolyVec.ForwardNTT(S1, L)
	PolyVec.ForwardNTT(S2, K)
	PolyVec.ForwardNTT(T0, K)

	local HasSigned = false
	local Kappa = 0

	local Y = buffer.create(L * N * 4)
	local YPrime = buffer.create(L * N * 4)
	local W = buffer.create(K * N * 4)
	local W1 = buffer.create(K * N * 4)
	local C = buffer.create(N * 4)
	local Z = buffer.create(L * N * 4)
	local R0 = buffer.create(K * N * 4)
	local R1 = buffer.create(K * N * 4)
	local H0 = buffer.create(K * N * 4)
	local H1 = buffer.create(K * N * 4)
	local H = buffer.create(K * N * 4)
	local CTilda = buffer.create(CTildaSize)
	local W1Encoded = buffer.create(K * W1Bw * 32)

	while not HasSigned do
		Sampling.ExpandMask(RhoPrime, Kappa, Y, Gamma1, L)

		PolyVec.Copy(Y, YPrime, L)
		PolyVec.ForwardNTT(YPrime, L)
		PolyVec.MatrixMultiply(A, YPrime, W, K, L, L, 1)
		PolyVec.InverseNTT(W, K)

		PolyVec.HighBits(W, W1, K, Alpha)
		PolyVec.Encode(W1, W1Encoded, K, W1Bw)

		local ChallengeInput = buffer.create(64 + buffer.len(W1Encoded))
		buffer.copy(ChallengeInput, 0, Mu, 0, 64)
		buffer.copy(ChallengeInput, 64, W1Encoded, 0, buffer.len(W1Encoded))

		local CTildaFull = SHA3.SHAKE256(ChallengeInput, CTildaSize)
		buffer.copy(CTilda, 0, CTildaFull, 0, CTildaSize)

		Sampling.SampleInBall(CTilda, C, Tau, Lambda)
		NTT.ForwardNTT(C)

		PolyVec.MultiplyByPoly(C, S1, Z, L)
		PolyVec.InverseNTT(Z, L)
		PolyVec.AddTo(Y, Z, L)

		PolyVec.MultiplyByPoly(C, S2, R1, K)
		PolyVec.InverseNTT(R1, K)
		PolyVec.Negate(R1, K)
		PolyVec.AddTo(W, R1, K)
		PolyVec.LowBits(R1, R0, K, Alpha)

		local ZNorm = PolyVec.InfinityNorm(Z, L)
		local R0Norm = PolyVec.InfinityNorm(R0, K)

		if ZNorm >= (Gamma1 - Beta) or R0Norm >= (Gamma2 - Beta) then
			HasSigned = false
		else
			PolyVec.MultiplyByPoly(C, T0, H0, K)
			PolyVec.InverseNTT(H0, K)

			PolyVec.Copy(H0, H1, K)
			PolyVec.Negate(H0, K)
			PolyVec.AddTo(H1, R1, K)
			PolyVec.MakeHint(H0, R1, H, K, Alpha)

			local CT0Norm = PolyVec.InfinityNorm(H1, K)
			local Count1s = PolyVec.Count1s(H, K)

			if CT0Norm >= Gamma2 or Count1s > Omega then
				HasSigned = false
			else
				HasSigned = true
			end
		end

		Kappa += L
	end

	local SigOff0 = 0
	local SigOff1 = SigOff0 + CTildaSize
	local SigOff2 = SigOff1 + (32 * L * Gamma1Bw)

	buffer.copy(Sig, SigOff0, CTilda, 0, CTildaSize)

	PolyVec.SubFromX(Z, L, Gamma1)
	local ZEncoded = buffer.create(32 * L * Gamma1Bw)
	PolyVec.Encode(Z, ZEncoded, L, Gamma1Bw)
	buffer.copy(Sig, SigOff1, ZEncoded, 0, buffer.len(ZEncoded))

	local HEncoded = buffer.create(buffer.len(Sig) - SigOff2)
	BitPacking.EncodeHintBits(H, HEncoded, K, Omega)
	buffer.copy(Sig, SigOff2, HEncoded, 0, buffer.len(HEncoded))

	return HasSigned
end

function Verify(
	PubKey: buffer,
	Msg: buffer,
	Ctx: buffer,
	Sig: buffer,
	K: number,
	L: number,
	D: number,
	Gamma1: number,
	Gamma2: number,
	Tau: number,
	Beta: number,
	Omega: number,
	Lambda: number
): boolean
	if not Params.CheckVerifyParams(K, L, D, Gamma1, Gamma2, Tau, Beta, Omega, Lambda) then
		error("Invalid verify parameters")
	end

	if buffer.len(Ctx) > 255 then
		return false
	end

	local T1Bw = Utils.BitWidth(Q) - D
	local Gamma1Bw = Utils.BitWidth(Gamma1)
	local CTildaSize = math.floor((2 * Lambda) / 8)

	local ExpectedPkLen = Utils.PubKeyLen(K, D)
	if buffer.len(PubKey) ~= ExpectedPkLen then
		return false
	end

	local ExpectedSigLen = Utils.SigLen(K, L, Gamma1, Omega, Lambda)
	if buffer.len(Sig) ~= ExpectedSigLen then
		return false
	end

	local Alpha = bit32.lshift(Gamma2, 1)
	local M = math.floor((Q - 1) / Alpha)
	local W1Bw = Utils.BitWidth(M - 1)

	local SigOff0 = 0
	local SigOff1 = SigOff0 + CTildaSize
	local SigOff2 = SigOff1 + (32 * L * Gamma1Bw)
	local PkOff0 = 0
	local PkOff1 = PkOff0 + 32

	local CTilda = buffer.create(CTildaSize)
	local ZEncoded = buffer.create(32 * L * Gamma1Bw)
	local HEncoded = buffer.create(buffer.len(Sig) - SigOff2)

	buffer.copy(CTilda, 0, Sig, SigOff0, CTildaSize)
	buffer.copy(ZEncoded, 0, Sig, SigOff1, buffer.len(ZEncoded))
	buffer.copy(HEncoded, 0, Sig, SigOff2, buffer.len(HEncoded))

	local VerificationFailed = false

	local H = buffer.create(K * N * 4)
	local HintFailed = BitPacking.DecodeHintBits(HEncoded, H, K, Omega)
	VerificationFailed = VerificationFailed or HintFailed

	local Count1s = PolyVec.Count1s(H, K)
	VerificationFailed = VerificationFailed or (Count1s > Omega)

	local C = buffer.create(N * 4)
	Sampling.SampleInBall(CTilda, C, Tau, Lambda)
	NTT.ForwardNTT(C)

	local Z = buffer.create(L * N * 4)
	PolyVec.Decode(ZEncoded, Z, L, Gamma1Bw)
	PolyVec.SubFromX(Z, L, Gamma1)

	local ZNorm = PolyVec.InfinityNorm(Z, L)
	VerificationFailed = VerificationFailed or (ZNorm >= (Gamma1 - Beta))

	local Rho = buffer.create(32)
	local T1Encoded = buffer.create(buffer.len(PubKey) - PkOff1)

	buffer.copy(Rho, 0, PubKey, PkOff0, 32)
	buffer.copy(T1Encoded, 0, PubKey, PkOff1, buffer.len(T1Encoded))

	local A = buffer.create(K * L * N * 4)
	local T1 = buffer.create(K * N * 4)

	Sampling.ExpandA(Rho, A, K, L)
	PolyVec.Decode(T1Encoded, T1, K, T1Bw)

	local Tr = SHA3.SHAKE256(PubKey, 64)

	local DomainSeparator = buffer.create(2)
	buffer.writeu8(DomainSeparator, 0, 0)
	buffer.writeu8(DomainSeparator, 1, buffer.len(Ctx))

	local MuInput = buffer.create(64 + 2 + buffer.len(Ctx) + buffer.len(Msg))
	local Offset = 0
	buffer.copy(MuInput, Offset, Tr, 0, 64)
	Offset += 64
	buffer.copy(MuInput, Offset, DomainSeparator, 0, 2)
	Offset += 2
	buffer.copy(MuInput, Offset, Ctx, 0, buffer.len(Ctx))
	Offset += buffer.len(Ctx)
	buffer.copy(MuInput, Offset, Msg, 0, buffer.len(Msg))

	local Mu = SHA3.SHAKE256(MuInput, 64)

	local W0 = buffer.create(K * N * 4)
	local W1 = buffer.create(K * N * 4)
	local W2 = buffer.create(K * N * 4)

	PolyVec.ForwardNTT(Z, L)
	PolyVec.MatrixMultiply(A, Z, W0, K, L, L, 1)

	PolyVec.LeftShift(T1, K, D)
	PolyVec.ForwardNTT(T1, K)
	PolyVec.MultiplyByPoly(C, T1, W2, K)
	PolyVec.Negate(W2, K)

	PolyVec.AddTo(W0, W2, K)
	PolyVec.InverseNTT(W2, K)

	PolyVec.UseHint(H, W2, W1, K, Alpha)

	local W1Encoded = buffer.create(K * W1Bw * 32)
	PolyVec.Encode(W1, W1Encoded, K, W1Bw)

	local ChallengeInput = buffer.create(64 + buffer.len(W1Encoded))
	buffer.copy(ChallengeInput, 0, Mu, 0, 64)
	buffer.copy(ChallengeInput, 64, W1Encoded, 0, buffer.len(W1Encoded))

	local CTildaPrime = SHA3.SHAKE256(ChallengeInput, CTildaSize)

	local CompareResult = CompareBufferSlices(CTilda, 0, CTildaPrime, 0, CTildaSize)
	VerificationFailed = VerificationFailed or not CompareResult

	return not VerificationFailed
end

Mldsa.ML_DSA_44 = {
	D = 13,
	Tau = 39,
	Gamma1 = bit32.lshift(1, 17),
	Gamma2 = math.floor((Q - 1) / 88),
	K = 4,
	L = 4,
	Eta = 2,
	Beta = 39 * 2,
	Omega = 80,
	Lambda = 128,

	KeygenSeedByteLen = KEYGEN_SEED_BYTE_LEN,
	PubKeyByteLen = Utils.PubKeyLen(4, 13),
	SecKeyByteLen = Utils.SecKeyLen(4, 4, 2, 13),
	SigningSeedByteLen = RND_BYTE_LEN,
	SigByteLen = Utils.SigLen(4, 4, 131072, 80, 128),

	KeyGen = function(Seed: buffer, PubKey: buffer, SecKey: buffer)
		Keygen(Seed, PubKey, SecKey, 4, 4, 13, 2)
	end,

	Sign = function(Rnd: buffer, SecKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Sign(Rnd, SecKey, Msg, Ctx, Sig, 4, 4, 13, 2, 131072, 95232, 39, 78, 80, 128)
	end,

	Verify = function(PubKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Verify(PubKey, Msg, Ctx, Sig, 4, 4, 13, 131072, 95232, 39, 78, 80, 128)
	end,

	GenerateKeys = function(): (buffer, buffer)
		local Seed = CSPRNG.RandomBytes(32)
		local PubKey = buffer.create(1312)
		local SecKey = buffer.create(2560)
		Mldsa.ML_DSA_44.KeyGen(Seed, PubKey, SecKey)
		return PubKey, SecKey
	end,
}

Mldsa.ML_DSA_65 = {
	D = 13,
	Tau = 49,
	Gamma1 = bit32.lshift(1, 19),
	Gamma2 = math.floor((Q - 1) / 32),
	K = 6,
	L = 5,
	Eta = 4,
	Beta = 49 * 4,
	Omega = 55,
	Lambda = 192,

	KeygenSeedByteLen = KEYGEN_SEED_BYTE_LEN,
	PubKeyByteLen = Utils.PubKeyLen(6, 13),
	SecKeyByteLen = Utils.SecKeyLen(6, 5, 4, 13),
	SigningSeedByteLen = RND_BYTE_LEN,
	SigByteLen = Utils.SigLen(6, 5, 524288, 55, 192),

	KeyGen = function(Seed: buffer, PubKey: buffer, SecKey: buffer)
		Keygen(Seed, PubKey, SecKey, 6, 5, 13, 4)
	end,

	Sign = function(Rnd: buffer, SecKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Sign(Rnd, SecKey, Msg, Ctx, Sig, 6, 5, 13, 4, 524288, 261888, 49, 196, 55, 192)
	end,

	Verify = function(PubKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Verify(PubKey, Msg, Ctx, Sig, 6, 5, 13, 524288, 261888, 49, 196, 55, 192)
	end,

	GenerateKeys = function(): (buffer, buffer)
		local Seed = CSPRNG.RandomBytes(32)
		local PubKey = buffer.create(1952)
		local SecKey = buffer.create(4032)
		Mldsa.ML_DSA_65.KeyGen(Seed, PubKey, SecKey)
		return PubKey, SecKey
	end,
}

Mldsa.ML_DSA_87 = {
	D = 13,
	Tau = 60,
	Gamma1 = bit32.lshift(1, 19),
	Gamma2 = math.floor((Q - 1) / 32),
	K = 8,
	L = 7,
	Eta = 2,
	Beta = 60 * 2,
	Omega = 75,
	Lambda = 256,

	KeygenSeedByteLen = KEYGEN_SEED_BYTE_LEN,
	PubKeyByteLen = Utils.PubKeyLen(8, 13),
	SecKeyByteLen = Utils.SecKeyLen(8, 7, 2, 13),
	SigningSeedByteLen = RND_BYTE_LEN,
	SigByteLen = Utils.SigLen(8, 7, 524288, 75, 256),

	KeyGen = function(Seed: buffer, PubKey: buffer, SecKey: buffer)
		Keygen(Seed, PubKey, SecKey, 8, 7, 13, 2)
	end,

	Sign = function(Rnd: buffer, SecKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Sign(Rnd, SecKey, Msg, Ctx, Sig, 8, 7, 13, 2, 524288, 261888, 60, 120, 75, 256)
	end,

	Verify = function(PubKey: buffer, Msg: buffer, Ctx: buffer, Sig: buffer): boolean
		return Verify(PubKey, Msg, Ctx, Sig, 8, 7, 13, 524288, 261888, 60, 120, 75, 256)
	end,

	GenerateKeys = function(): (buffer, buffer)
		local Seed = CSPRNG.RandomBytes(32)
		local PubKey = buffer.create(2592)
		local SecKey = buffer.create(4896)
		Mldsa.ML_DSA_87.KeyGen(Seed, PubKey, SecKey)
		return PubKey, SecKey
	end,
}

Mldsa.PubKeyLen = Utils.PubKeyLen
Mldsa.SecKeyLen = Utils.SecKeyLen
Mldsa.SigLen = Utils.SigLen

return Mldsa


--[=[
	ML-KEM Compression Module

	Polynomial coefficient compression and decompression operations for ML-KEM.
	Maps field elements to and from reduced bit representations.
--]=]

--!strict
--!optimize 2
--!native

local Params = require("./Params")

local N = 256
local Q = 3329

local Compression = {}

function Compression.Compress(Value: number, D: number): number
	local Shifted = Value * (2 ^ D)
	local QHalf = 1664
	local Rounded = Shifted + QHalf
	local Quotient = Rounded // Q
	local Mask = (2 ^ D) - 1

	return bit32.band(Quotient, Mask)
end

function Compression.Decompress(Value: number, D: number): number
	local Product = Value * Q
	local Quotient = bit32.rshift(Product, D)
	local RoundBit = bit32.band(bit32.rshift(Product, D - 1), 1)
	return Quotient + RoundBit
end

function Compression.PolyCompress(Poly: buffer, D: number)
	if not Params.CheckD(D) then
		error("Invalid compression parameter d")
	end

	local Modulus = Q

	if D == 1 then
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Shifted = Value * 2
			local Rounded = Shifted + 1664
			local Quotient = Rounded // Modulus
			buffer.writeu16(Poly, Index * 2, bit32.band(Quotient, 1))
		end
	elseif D == 4 then
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Shifted = Value * 16
			local Rounded = Shifted + 1664
			local Quotient = Rounded // Modulus
			buffer.writeu16(Poly, Index * 2, bit32.band(Quotient, 15))
		end
	elseif D == 5 then
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Shifted = Value * 32
			local Rounded = Shifted + 1664
			local Quotient = Rounded // Modulus
			buffer.writeu16(Poly, Index * 2, bit32.band(Quotient, 31))
		end
	elseif D == 10 then
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Shifted = Value * 1024
			local Rounded = Shifted + 1664
			local Quotient = Rounded // Modulus
			buffer.writeu16(Poly, Index * 2, bit32.band(Quotient, 1023))
		end
	elseif D == 11 then
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Shifted = Value * 2048
			local Rounded = Shifted + 1664
			local Quotient = Rounded // Modulus
			buffer.writeu16(Poly, Index * 2, bit32.band(Quotient, 2047))
		end
	else
		local Compress = Compression.Compress
		for Index = 0, N - 1 do
			local Value = buffer.readu16(Poly, Index * 2)
			local Compressed = Compress(Value, D)
			buffer.writeu16(Poly, Index * 2, Compressed)
		end
	end
end

function Compression.PolyDecompress(Poly: buffer, D: number)
	if not Params.CheckD(D) then
		error("Invalid decompression parameter d")
	end

	local Modulus = Q

	if D == 1 then
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Product = Compressed * Modulus
			local Quotient = bit32.rshift(Product, 1)
			local RoundBit = bit32.band(Product, 1)
			buffer.writeu16(Poly, Index * 2, Quotient + RoundBit)
		end
	elseif D == 4 then
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Product = Compressed * Modulus
			local Quotient = bit32.rshift(Product, 4)
			local RoundBit = bit32.band(bit32.rshift(Product, 3), 1)
			buffer.writeu16(Poly, Index * 2, Quotient + RoundBit)
		end
	elseif D == 5 then
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Product = Compressed * Modulus
			local Quotient = bit32.rshift(Product, 5)
			local RoundBit = bit32.band(bit32.rshift(Product, 4), 1)
			buffer.writeu16(Poly, Index * 2, Quotient + RoundBit)
		end
	elseif D == 10 then
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Product = Compressed * Modulus
			local Quotient = bit32.rshift(Product, 10)
			local RoundBit = bit32.band(bit32.rshift(Product, 9), 1)
			buffer.writeu16(Poly, Index * 2, Quotient + RoundBit)
		end
	elseif D == 11 then
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Product = Compressed * Modulus
			local Quotient = bit32.rshift(Product, 11)
			local RoundBit = bit32.band(bit32.rshift(Product, 10), 1)
			buffer.writeu16(Poly, Index * 2, Quotient + RoundBit)
		end
	else
		local Decompress = Compression.Decompress
		for Index = 0, N - 1 do
			local Compressed = buffer.readu16(Poly, Index * 2)
			local Decompressed = Decompress(Compressed, D)
			buffer.writeu16(Poly, Index * 2, Decompressed)
		end
	end
end

return Compression


--[=[
	ML-KEM Field Arithmetic Module

	Prime field Zq arithmetic operations for ML-KEM.
	Modular arithmetic over the field Z_q where q = 3329.

	Example usage:
		local Field = require(script.Field)

		local Sum = Field.Add(1234, 567)
		local Product = Field.Multiply(1234, 567)
		local Inverse = Field.Invert(1234)

		local Buffer = Field.BufferCreate(256)
		Field.BufferReduce(Buffer, 256)

		local Sum = Field.BufferAdd(BufferA, BufferB, Result, 256)
--]=]

--!strict
--!optimize 2
--!native

local Q = 13 * 256 + 1
local Q_BIT_WIDTH = 12
local WIDTH = (2 ^ (2 * Q_BIT_WIDTH))
local R = math.floor((2 ^ (2 * Q_BIT_WIDTH)) / Q)

local Field = {}

function Field.BarrettReduce(Value: number): number
	return Value % Q
end

local function ReducedAdd(A: number, B: number): number
	local Sum = A + B
	return if Sum >= Q then Sum - Q else Sum
end

local function ReducedSub(A: number, B: number): number
	local Diff = A - B
	return if Diff < 0 then Diff + Q else Diff
end

function Field.Add(A: number, B: number): number
	return ReducedAdd(A, B)
end

function Field.Subtract(A: number, B: number): number
	return ReducedSub(A, B)
end

function Field.Multiply(A: number, B: number): number
	return (A * B) % Q
end

function Field.Negate(A: number): number
	return if A == 0 then 0 else Q - A
end

function Field.Power(Base: number, Exponent: number): number
	local Result = if bit32.band(Exponent, 1) == 1 then Base else 1
	local CurrentBase = Base

	local Exp = Exponent
	while Exp > 1 do
		Exp = bit32.rshift(Exp, 1)
		CurrentBase = Field.Multiply(CurrentBase, CurrentBase)

		if bit32.band(Exp, 1) == 1 then
			Result = Field.Multiply(Result, CurrentBase)
		end
	end

	return Result
end

function Field.Invert(A: number): number
	if A == 0 then
		return 0
	end

	return Field.Power(A, Q - 2)
end

function Field.Divide(A: number, B: number): number
	return Field.Multiply(A, Field.Invert(B))
end

function Field.BufferReduce(Buffer: buffer, Count: number)
	local Modulus = Q
	for Index = 0, Count - 1 do
		local Value = buffer.readu16(Buffer, Index * 2)
		buffer.writeu16(Buffer, Index * 2, (Value % Modulus))
	end
end

function Field.BufferSubtract(A: buffer, B: buffer, Result: buffer, Count: number)
	for Index = 0, Count - 1 do
		local AValue = buffer.readu16(A, Index * 2)
		local BValue = buffer.readu16(B, Index * 2)
		buffer.writeu16(Result, Index * 2, Field.Subtract(AValue, BValue))
	end
end

return Field


--[=[
	ML-KEM Number Theoretic Transform Module

	Number Theoretic Transform (NTT) operations for polynomial multiplication.
	Cooley-Tukey NTT and Gentleman-Sande inverse NTT algorithms

	Example usage:
		local NTT = require(script.NTT)

		local Poly = buffer.create(512) -- 256 coefficients * 2 bytes
		NTT.Ntt(Poly) -- Forward transform
		NTT.Intt(Poly) -- Inverse transform
--]=]

--!strict
--!optimize 2
--!native

local Field = require("./Field")
local Utils = require("./Utils")

local LOG2N = 8
local N = 2 ^ LOG2N

local Q = 3329
local ZETA = 17
local INV_N = Field.Invert(N / 2)

local Ntt = {}

local NTT_ZETA_EXP = buffer.create((N / 2) * 2)
do
	for I = 0, N / 2 - 1 do
		local BitRevIndex = Utils.BitReverse(I, LOG2N - 1)
		local ZetaExp = Field.Power(ZETA, BitRevIndex)
		buffer.writeu16(NTT_ZETA_EXP, I * 2, ZetaExp)
	end
end

local INTT_ZETA_EXP = buffer.create((N / 2) * 2)
do
	for I = 0, N / 2 - 1 do
		local NttZeta = buffer.readu16(NTT_ZETA_EXP, I * 2)
		local NegZeta = if NttZeta == 0 then 0 else Q - NttZeta
		buffer.writeu16(INTT_ZETA_EXP, I * 2, NegZeta)
	end
end

local POLY_MUL_ZETA_EXP = buffer.create((N / 2) * 2)
do
	for I = 0, N / 2 - 1 do
		local BitRevIndex = Utils.BitReverse(I, LOG2N - 1)
		local Exponent = bit32.bxor(bit32.lshift(BitRevIndex, 1), 1)
		local ZetaExp = Field.Power(ZETA, Exponent)
		buffer.writeu16(POLY_MUL_ZETA_EXP, I * 2, ZetaExp)
	end
end

local POWERS_OF_TWO = table.create(LOG2N)
for I = 0, LOG2N - 1 do
	POWERS_OF_TWO[I] = { 2 ^ I, (2 ^ I) * 2 }
end

function Ntt.Ntt(Poly: buffer)
	local Zeta = NTT_ZETA_EXP
	local Powers = POWERS_OF_TWO
	local Modulus = Q

	for L = LOG2N - 1, 1, -1 do
		local Len = Powers[L][1]
		local LenX2 = Powers[L][2]
		local KBeg = bit32.rshift(N, L + 1)

		for Start = 0, N - 1, LenX2 do
			local KNow = KBeg + bit32.rshift(Start, L + 1)
			local ZetaExp = buffer.readu16(Zeta, KNow * 2)

			for I = Start, Start + Len - 1 do
				local OffsetI = I * 2
				local OffsetILen = (I + Len) * 2

				local PolyI = buffer.readu16(Poly, OffsetI)
				local PolyILen = buffer.readu16(Poly, OffsetILen)

				local Tmp = (ZetaExp * PolyILen) % Modulus

				local SubResult = if PolyI >= Tmp then PolyI - Tmp else PolyI - Tmp + Modulus

				local Sum = PolyI + Tmp
				local AddResult = if Sum >= Modulus then Sum - Modulus else Sum

				buffer.writeu16(Poly, OffsetILen, SubResult)
				buffer.writeu16(Poly, OffsetI, AddResult)
			end
		end
	end
end

function Ntt.Intt(Poly: buffer)
	local Zeta = INTT_ZETA_EXP
	local InvN = INV_N
	local Powers = POWERS_OF_TWO
	local Modulus = Q

	for L = 1, LOG2N - 1 do
		local Len = Powers[L][1]
		local LenX2 = Powers[L][2]
		local KBeg = bit32.rshift(N, L) - 1

		for Start = 0, N - 1, LenX2 do
			local KNow = KBeg - bit32.rshift(Start, L + 1)
			local NegZetaExp = buffer.readu16(Zeta, KNow * 2)

			for I = Start, Start + Len - 1 do
				local OffsetI = I * 2
				local OffsetILen = (I + Len) * 2

				local PolyI = buffer.readu16(Poly, OffsetI)
				local PolyILen = buffer.readu16(Poly, OffsetILen)

				local Sum = PolyI + PolyILen

				local AddResult = if Sum >= Modulus then Sum - Modulus else Sum
				local SubResult = if PolyI >= PolyILen then PolyI - PolyILen else PolyI - PolyILen + Modulus

				local MulResult = (SubResult * NegZetaExp) % Modulus

				buffer.writeu16(Poly, OffsetI, AddResult)
				buffer.writeu16(Poly, OffsetILen, MulResult)
			end
		end
	end

	for I = 0, N - 1 do
		local Offset = I * 2
		local Coeff = buffer.readu16(Poly, Offset)

		local Scaled = (Coeff * InvN) % Q
		buffer.writeu16(Poly, Offset, Scaled)
	end
end

function Ntt.NttAt(Poly: buffer, BaseOffset: number)
	local Zeta = NTT_ZETA_EXP
	local Powers = POWERS_OF_TWO
	local Modulus = Q

	for L = LOG2N - 1, 1, -1 do
		local Len = Powers[L][1]
		local LenX2 = Powers[L][2]
		local KBeg = bit32.rshift(N, L + 1)

		for Start = 0, N - 1, LenX2 do
			local KNow = KBeg + bit32.rshift(Start, L + 1)
			local ZetaExp = buffer.readu16(Zeta, KNow * 2)

			for I = Start, Start + Len - 1 do
				local OffsetI = BaseOffset + I * 2
				local OffsetILen = BaseOffset + (I + Len) * 2

				local PolyI = buffer.readu16(Poly, OffsetI)
				local PolyILen = buffer.readu16(Poly, OffsetILen)

				local Tmp = (ZetaExp * PolyILen) % Modulus

				local SubResult = if PolyI >= Tmp then PolyI - Tmp else PolyI - Tmp + Modulus
				local Sum = PolyI + Tmp
				local AddResult = if Sum >= Modulus then Sum - Modulus else Sum

				buffer.writeu16(Poly, OffsetILen, SubResult)
				buffer.writeu16(Poly, OffsetI, AddResult)
			end
		end
	end
end

function Ntt.InttAt(Poly: buffer, BaseOffset: number)
	local Zeta = INTT_ZETA_EXP
	local InvN = INV_N
	local Powers = POWERS_OF_TWO
	local Modulus = Q

	for L = 1, LOG2N - 1 do
		local Len = Powers[L][1]
		local LenX2 = Powers[L][2]
		local KBeg = bit32.rshift(N, L) - 1

		for Start = 0, N - 1, LenX2 do
			local KNow = KBeg - bit32.rshift(Start, L + 1)
			local NegZetaExp = buffer.readu16(Zeta, KNow * 2)

			for I = Start, Start + Len - 1 do
				local OffsetI = BaseOffset + I * 2
				local OffsetILen = BaseOffset + (I + Len) * 2

				local PolyI = buffer.readu16(Poly, OffsetI)
				local PolyILen = buffer.readu16(Poly, OffsetILen)

				local Sum = PolyI + PolyILen
				local AddResult = if Sum >= Modulus then Sum - Modulus else Sum
				local SubResult = if PolyI >= PolyILen then PolyI - PolyILen else PolyI - PolyILen + Modulus

				local MulResult = (SubResult * NegZetaExp) % Modulus

				buffer.writeu16(Poly, OffsetI, AddResult)
				buffer.writeu16(Poly, OffsetILen, MulResult)
			end
		end
	end

	for I = 0, N - 1 do
		local Offset = BaseOffset + I * 2
		local Coeff = buffer.readu16(Poly, Offset)
		buffer.writeu16(Poly, Offset, (Coeff * InvN) % Modulus)
	end
end

function Ntt.PolyAddAt(F: buffer, FOffset: number, G: buffer, GOffset: number, H: buffer, HOffset: number)
	local Modulus = Q
	for Index = 0, 255 do
		local Offset = Index * 2
		local AValue = buffer.readu16(F, FOffset + Offset)
		local BValue = buffer.readu16(G, GOffset + Offset)
		local Sum = AValue + BValue
		buffer.writeu16(H, HOffset + Offset, if Sum >= Modulus then Sum - Modulus else Sum)
	end
end

function Ntt.PolySubAt(F: buffer, FOffset: number, G: buffer, GOffset: number, H: buffer, HOffset: number)
	local Modulus = Q
	for Index = 0, 255 do
		local Offset = Index * 2
		local AValue = buffer.readu16(F, FOffset + Offset)
		local BValue = buffer.readu16(G, GOffset + Offset)
		local Diff = AValue - BValue
		buffer.writeu16(H, HOffset + Offset, if Diff < 0 then Diff + Modulus else Diff)
	end
end

function Ntt.PolyMul(F: buffer, G: buffer, H: buffer)
	local PolyZeta = POLY_MUL_ZETA_EXP
	local Modulus = Q

	for I = 0, N / 2 - 1 do
		local Offset = I * 4
		local ZetaExp = buffer.readu16(PolyZeta, I * 2)

		local F0 = buffer.readu16(F, Offset)
		local F1 = buffer.readu16(F, Offset + 2)
		local G0 = buffer.readu16(G, Offset)
		local G1 = buffer.readu16(G, Offset + 2)

		local Mul1 = (F0 * G0) % Modulus
		local Mul2 = (F1 * G1) % Modulus
		local Mul3 = (Mul2 * ZetaExp) % Modulus

		local Sum0 = Mul1 + Mul3
		local H0 = if Sum0 >= Modulus then Sum0 - Modulus else Sum0

		local Mul4 = (F0 * G1) % Modulus
		local Mul5 = (F1 * G0) % Modulus

		local Sum1 = Mul4 + Mul5
		local H1 = if Sum1 >= Modulus then Sum1 - Modulus else Sum1

		buffer.writeu16(H, Offset, H0)
		buffer.writeu16(H, Offset + 2, H1)
	end
end

function Ntt.PolyMulAt(F: buffer, FOffset: number, G: buffer, GOffset: number, H: buffer, HOffset: number)
	local PolyZeta = POLY_MUL_ZETA_EXP
	local Modulus = Q

	for I = 0, 127 do
		local LocalOffset = I * 4
		local ZetaExp = buffer.readu16(PolyZeta, I * 2)

		local F0 = buffer.readu16(F, FOffset + LocalOffset)
		local F1 = buffer.readu16(F, FOffset + LocalOffset + 2)
		local G0 = buffer.readu16(G, GOffset + LocalOffset)
		local G1 = buffer.readu16(G, GOffset + LocalOffset + 2)

		local Mul1 = (F0 * G0) % Modulus
		local Mul2 = (F1 * G1) % Modulus
		local Mul3 = (Mul2 * ZetaExp) % Modulus

		local Sum0 = Mul1 + Mul3
		local H0 = if Sum0 >= Modulus then Sum0 - Modulus else Sum0

		local Mul4 = (F0 * G1) % Modulus
		local Mul5 = (F1 * G0) % Modulus

		local Sum1 = Mul4 + Mul5
		local H1 = if Sum1 >= Modulus then Sum1 - Modulus else Sum1

		buffer.writeu16(H, HOffset + LocalOffset, H0)
		buffer.writeu16(H, HOffset + LocalOffset + 2, H1)
	end
end

function Ntt.PolyAdd(F: buffer, G: buffer, H: buffer)
	local Modulus = Q
	for Index = 0, N - 1 do
		local Offset = Index * 2
		local AValue = buffer.readu16(F, Offset)
		local BValue = buffer.readu16(G, Offset)

		local Sum = AValue + BValue
		local Result = if Sum >= Modulus then Sum - Modulus else Sum
		buffer.writeu16(H, Offset, Result)
	end
end

function Ntt.PolySub(F: buffer, G: buffer, H: buffer)
	local Modulus = Q
	for Index = 0, N - 1 do
		local Offset = Index * 2
		local AValue = buffer.readu16(F, Offset)
		local BValue = buffer.readu16(G, Offset)

		local Diff = AValue - BValue
		local Result = if Diff < 0 then Diff + Modulus else Diff
		buffer.writeu16(H, Offset, Result)
	end
end

return Ntt


--[=[
	ML-KEM Public Key Encryption Module

	K-PKE (CPA-secure public key encryption) implementation for ML-KEM.
	Provides key generation, encryption, and decryption operations
	as specified in NIST FIPS 203.

	Example usage:
		local PKE = require(script.PKE)

		local Seed = buffer.fromstring("32_byte_seed_here...")
		local PublicKey, SecretKey = PKE.MLKEM_1024.KeyGen(Seed)

		local Message = buffer.create(32)
		local Randomness = buffer.create(32)
		local Ciphertext = PKE.MLKEM_1024.Encrypt(PublicKey, Message, Randomness)

		local Plaintext = PKE.MLKEM_1024.Decrypt(SecretKey, Ciphertext)
--]=]

--!strict
--!optimize 2
--!native

local Ntt = require("./NTT")

local PolyVec = require("./PolyVec")
local Sampling = require("./Sampling")

local Serialize = require("./Serialize")
local Compression = require("./Compression")

local Params = require("./Params")
local Utils = require("./Utils")
local SHA3 = require("./SHA3")

local LOG2N = 8
local N = 2 ^ LOG2N
local POLY_SIZE = N * 2

local TEMP_V = buffer.create(N * 2)
local TEMP_V2 = buffer.create(N * 2)
local TEMP_SINGLE_POLY = buffer.create(N * 2)
local TEMP_R_POLY = buffer.create(N * 2)
local TEMP_T = buffer.create(N * 2)
local TEMP_T2 = buffer.create(N * 2)
local TEMP_U_POLY = buffer.create(N * 2)

local KEYGEN_G_INPUT = buffer.create(33)
local KEYGEN_RHO = buffer.create(32)
local KEYGEN_SIGMA = buffer.create(32)

local Pke = {}

function Pke.KeyGen(K: number, Eta1: number, Seed: buffer): (buffer, buffer)
	if not Params.CheckKeygenParams(K, Eta1) then
		error("Invalid keygen parameters")
	end
	if buffer.len(Seed) ~= 32 then
		error("Seed must be 32 bytes")
	end

	local GInput = KEYGEN_G_INPUT
	buffer.copy(GInput, 0, Seed, 0, 32)
	buffer.writeu8(GInput, 32, K)

	local GOutput = SHA3.SHA3_512(GInput)
	local Rho = KEYGEN_RHO
	local Sigma = KEYGEN_SIGMA
	buffer.copy(Rho, 0, GOutput, 0, 32)
	buffer.copy(Sigma, 0, GOutput, 32, 32)

	local APrime = Sampling.GenerateMatrix(K, Rho, false)

	local Nonce = 0
	local S = Sampling.GenerateVector(K, Eta1, Sigma, Nonce)
	Nonce += K
	local E = Sampling.GenerateVector(K, Eta1, Sigma, Nonce)

	PolyVec.VecNtt(S, K)
	PolyVec.VecNtt(E, K)

	local TPrime = PolyVec.VecCreate(K)
	PolyVec.MatVecMultiply(APrime, S, TPrime, K, K, K)
	PolyVec.VecAddTo(E, TPrime, K)

	local PublicKeyLen = Utils.GetPkePublicKeyLen(K)
	local PublicKey = buffer.create(PublicKeyLen)

	local EncodedTPrime = PolyVec.VecEncode(TPrime, K, 12)
	local RhoOffset = K * 12 * 32
	buffer.copy(PublicKey, 0, EncodedTPrime, 0, RhoOffset)
	buffer.copy(PublicKey, RhoOffset, Rho, 0, 32)

	return PublicKey, PolyVec.VecEncode(S, K, 12)
end

function Pke.Encrypt(
	K: number,
	Eta1: number,
	Eta2: number,
	Du: number,
	Dv: number,
	PublicKey: buffer,
	Message: buffer,
	Randomness: buffer
): buffer
	if not Params.CheckEncryptParams(K, Eta1, Eta2, Du, Dv) then
		error("Invalid encryption parameters")
	end
	if buffer.len(PublicKey) ~= Utils.GetPkePublicKeyLen(K) then
		error("Invalid public key length")
	end
	if buffer.len(Message) ~= 32 then
		error("Message must be 32 bytes")
	end
	if buffer.len(Randomness) ~= 32 then
		error("Randomness must be 32 bytes")
	end

	local RhoOffset = K * 12 * 32
	local EncodedTPrime = buffer.create(RhoOffset)
	local Rho = buffer.create(32)
	buffer.copy(EncodedTPrime, 0, PublicKey, 0, RhoOffset)
	buffer.copy(Rho, 0, PublicKey, RhoOffset, 32)

	local TPrime = PolyVec.VecDecode(EncodedTPrime, K, 12)
	local ReEncodedTPrime = PolyVec.VecEncode(TPrime, K, 12)
	local IsValid = Utils.CtMemcmp(EncodedTPrime, ReEncodedTPrime)
	if IsValid ~= 0xFFFFFFFF then
		error("Key encapsulation verification failed")
	end

	local APrime = Sampling.GenerateMatrix(K, Rho, true)

	local Nonce = 0
	local R = Sampling.GenerateVector(K, Eta1, Randomness, Nonce)
	Nonce += K
	local E1 = Sampling.GenerateVector(K, Eta2, Randomness, Nonce)
	Nonce += K
	local E2 = Sampling.GenerateNoisePoly(Eta2, Randomness, Nonce)

	PolyVec.VecNtt(R, K)

	local U = PolyVec.VecCreate(K)
	PolyVec.MatVecMultiply(APrime, R, U, K, K, K)

	PolyVec.VecIntt(U, K)
	PolyVec.VecAddTo(E1, U, K)

	local V = TEMP_V
	local TempV = TEMP_V2
	local SinglePoly = TEMP_SINGLE_POLY
	local RPoly = TEMP_R_POLY
	buffer.fill(V, 0, 0, N * 2)

	local PolySize = POLY_SIZE
	for I = 0, K - 1 do
		local TPolyOffset = I * PolySize
		local RPolyOffset = I * PolySize
		buffer.copy(SinglePoly, 0, TPrime, TPolyOffset, PolySize)
		buffer.copy(RPoly, 0, R, RPolyOffset, PolySize)
		Ntt.PolyMul(SinglePoly, RPoly, TempV)
		Ntt.PolyAdd(V, TempV, V)
	end

	Ntt.Intt(V)
	Ntt.PolyAdd(V, E2, V)

	local M = Serialize.Decode(Message, 1)
	Compression.PolyDecompress(M, 1)
	Ntt.PolyAdd(V, M, V)

	PolyVec.VecCompress(U, K, Du)
	local EncodedU = PolyVec.VecEncode(U, K, Du)

	Compression.PolyCompress(V, Dv)
	local EncodedV = Serialize.Encode(V, Dv)

	local CiphertextLen = Utils.GetPkeCipherTextLen(K, Du, Dv)
	local Ciphertext = buffer.create(CiphertextLen)
	local ULen = buffer.len(EncodedU)
	buffer.copy(Ciphertext, 0, EncodedU, 0, ULen)
	buffer.copy(Ciphertext, ULen, EncodedV, 0, buffer.len(EncodedV))

	return Ciphertext
end

function Pke.Decrypt(K: number, Du: number, Dv: number, SecretKey: buffer, Ciphertext: buffer): buffer
	if not Params.CheckDecryptParams(K, Du, Dv) then
		error("Invalid decryption parameters")
	end
	if buffer.len(SecretKey) ~= Utils.GetPkeSecretKeyLen(K) then
		error("Invalid secret key length")
	end
	if buffer.len(Ciphertext) ~= Utils.GetPkeCipherTextLen(K, Du, Dv) then
		error("Invalid ciphertext length")
	end

	local ULen = K * Du * 32
	local VLen = Dv * 32

	local EncodedU = buffer.create(ULen)
	local EncodedV = buffer.create(VLen)
	buffer.copy(EncodedU, 0, Ciphertext, 0, ULen)
	buffer.copy(EncodedV, 0, Ciphertext, ULen, VLen)

	local U = PolyVec.VecDecode(EncodedU, K, Du)
	PolyVec.VecDecompress(U, K, Du)

	local V = Serialize.Decode(EncodedV, Dv)
	Compression.PolyDecompress(V, Dv)

	local SPrime = PolyVec.VecDecode(SecretKey, K, 12)

	PolyVec.VecNtt(U, K)

	local T = buffer.create(N * 2)
	local TempT = buffer.create(N * 2)
	local SinglePoly = buffer.create(N * 2)

	local T = TEMP_T
	local TempT = TEMP_T2
	local SinglePoly = TEMP_SINGLE_POLY
	local UPoly = TEMP_U_POLY
	buffer.fill(T, 0, 0, N * 2)

	local PolySize = POLY_SIZE
	for I = 0, K - 1 do
		local SPolyOffset = I * PolySize
		local UPolyOffset = I * PolySize

		buffer.copy(SinglePoly, 0, SPrime, SPolyOffset, PolySize)
		buffer.copy(UPoly, 0, U, UPolyOffset, PolySize)

		Ntt.PolyMul(SinglePoly, UPoly, TempT)
		Ntt.PolyAdd(T, TempT, T)
	end

	Ntt.Intt(T)
	Ntt.PolySub(V, T, V)

	Compression.PolyCompress(V, 1)
	local Plaintext = Serialize.Encode(V, 1)

	return Plaintext
end

Pke.MLKEM_512 = {
	KeyGen = function(Seed: buffer): (buffer, buffer)
		return Pke.KeyGen(2, 3, Seed)
	end,

	Encrypt = function(PublicKey: buffer, Message: buffer, Randomness: buffer): buffer
		return Pke.Encrypt(2, 3, 2, 10, 4, PublicKey, Message, Randomness)
	end,

	Decrypt = function(SecretKey: buffer, Ciphertext: buffer): buffer
		return Pke.Decrypt(2, 10, 4, SecretKey, Ciphertext)
	end,
}

Pke.MLKEM_768 = {
	KeyGen = function(Seed: buffer): (buffer, buffer)
		return Pke.KeyGen(3, 2, Seed)
	end,

	Encrypt = function(PublicKey: buffer, Message: buffer, Randomness: buffer): buffer
		return Pke.Encrypt(3, 2, 2, 10, 4, PublicKey, Message, Randomness)
	end,

	Decrypt = function(SecretKey: buffer, Ciphertext: buffer): buffer
		return Pke.Decrypt(3, 10, 4, SecretKey, Ciphertext)
	end,
}

Pke.MLKEM_1024 = {
	KeyGen = function(Seed: buffer): (buffer, buffer)
		return Pke.KeyGen(4, 2, Seed)
	end,

	Encrypt = function(PublicKey: buffer, Message: buffer, Randomness: buffer): buffer
		return Pke.Encrypt(4, 2, 2, 11, 5, PublicKey, Message, Randomness)
	end,

	Decrypt = function(SecretKey: buffer, Ciphertext: buffer): buffer
		return Pke.Decrypt(4, 11, 5, SecretKey, Ciphertext)
	end,
}

return Pke


--[=[
	ML-KEM Parameters Module

	Parameter validation functions for ML-KEM.
	 NIST FIPS 203 specification.

	Example usage:
		local Params = require(script.Params)
		local IsValid = Params.CheckKeygenParams(4, 2) -- true

		local ValidD = Params.CheckD(11) -- true
		local ValidEncrypt = Params.CheckEncryptParams(4, 2, 2, 11, 5) -- true
--]=]

--!strict
--!optimize 2
--!native

local Q_BIT_WIDTH = 12

local Params = {}

function Params.CheckD(D: number): boolean
	return D < 12
end

function Params.CheckEta(Eta: number): boolean
	return (Eta == 2) or (Eta == 3)
end

function Params.CheckK(K: number): boolean
	return (K == 2) or (K == 3) or (K == 4)
end

function Params.CheckL(L: number): boolean
	return (L == 1) or (L == 4) or (L == 5) or (L == 10) or (L == 11) or (L == 12)
end

function Params.CheckMatrixDim(ACols: number, BRows: number): boolean
	return ACols == BRows
end

function Params.CheckKeygenParams(K: number, Eta1: number): boolean
	local Flag0 = (K == 2) and (Eta1 == 3)
	local Flag1 = (K == 3) and (Eta1 == 2)
	local Flag2 = (K == 4) and (Eta1 == 2)

	return Flag0 or Flag1 or Flag2
end

function Params.CheckEncryptParams(K: number, Eta1: number, Eta2: number, Du: number, Dv: number): boolean
	local Flag0 = (K == 2) and (Eta1 == 3) and (Eta2 == 2) and (Du == 10) and (Dv == 4)
	local Flag1 = (K == 3) and (Eta1 == 2) and (Eta2 == 2) and (Du == 10) and (Dv == 4)
	local Flag2 = (K == 4) and (Eta1 == 2) and (Eta2 == 2) and (Du == 11) and (Dv == 5)

	return Flag0 or Flag1 or Flag2
end

function Params.CheckDecryptParams(K: number, Du: number, Dv: number): boolean
	local Flag0 = (K == 2) and (Du == 10) and (Dv == 4)
	local Flag1 = (K == 3) and (Du == 10) and (Dv == 4)
	local Flag2 = (K == 4) and (Du == 11) and (Dv == 5)

	return Flag0 or Flag1 or Flag2
end

function Params.CheckEncapParams(K: number, Eta1: number, Eta2: number, Du: number, Dv: number): boolean
	return Params.CheckEncryptParams(K, Eta1, Eta2, Du, Dv)
end

function Params.CheckDecapParams(K: number, Eta1: number, Eta2: number, Du: number, Dv: number): boolean
	return Params.CheckEncapParams(K, Eta1, Eta2, Du, Dv)
end

return Params


--[=[
	ML-KEM Polynomial Vector Module

	Polynomial vector and matrix operations for ML-KEM.
	Vector arithmetic, matrix multiplication, and encoding/decoding
	operations on vectors of degree-255 polynomials.
--]=]

--!strict
--!optimize 2
--!native

local Ntt = require("./NTT")
local MlKemParams = require("./Params")
local MlKemSerialize = require("./Serialize")
local Compression = require("./Compression")

local LOG2N = 8
local N = 2 ^ LOG2N
local Q = 3329
local POLY_SIZE = N * 2

local TEMP_POLY_A = buffer.create(N * 2)
local TEMP_POLY_B = buffer.create(N * 2)
local TEMP_POLY_MUL = buffer.create(N * 2)
local TEMP_POLY_ACCUM = buffer.create(N * 2)
local TEMP_POLY_NTT = buffer.create(N * 2)
local TEMP_POLY_ENCODE = buffer.create(N * 2)
local TEMP_POLY_COMPRESS = buffer.create(N * 2)

local PolyVec = {}

function PolyVec.VecCreate(K: number): buffer
	return buffer.create(K * POLY_SIZE)
end

function PolyVec.MatCreate(Rows: number, Cols: number): buffer
	return buffer.create(Rows * Cols * POLY_SIZE)
end

function PolyVec.MatSetPoly(Mat: buffer, Row: number, Col: number, Cols: number, Poly: buffer)
	local Index = Row * Cols + Col
	local Offset = Index * POLY_SIZE

	buffer.copy(Mat, Offset, Poly, 0, POLY_SIZE)
end

function PolyVec.MatVecMultiply(A: buffer, B: buffer, C: buffer, ARows: number, ACols: number, BRows: number)
	local PolySize = POLY_SIZE
	local AccumPoly = TEMP_POLY_ACCUM
	local TempPoly = TEMP_POLY_MUL

	local PolyMul = Ntt.PolyMulAt
	local PolyAdd = Ntt.PolyAddAt

	for I = 0, ARows - 1 do
		local COffset = I * PolySize
		buffer.fill(AccumPoly, 0, 0, PolySize)

		for K = 0, ACols - 1 do
			local AOffset = (I * ACols + K) * PolySize
			local BOffset = K * PolySize

			PolyMul(A, AOffset, B, BOffset, TempPoly, 0)
			PolyAdd(AccumPoly, 0, TempPoly, 0, AccumPoly, 0)
		end

		buffer.copy(C, COffset, AccumPoly, 0, PolySize)
	end
end

function PolyVec.VecNtt(Vec: buffer, K: number)
	local PolySize = POLY_SIZE
	local NttAt = Ntt.NttAt
	for I = 0, K - 1 do
		NttAt(Vec, I * PolySize)
	end
end

function PolyVec.VecIntt(Vec: buffer, K: number)
	local PolySize = POLY_SIZE
	local InttAt = Ntt.InttAt

	for I = 0, K - 1 do
		InttAt(Vec, I * PolySize)
	end
end

function PolyVec.VecAddTo(Src: buffer, Dst: buffer, K: number)
	local Modulus = Q
	local Total = K * N * 2

	for Offset = 0, Total - 2, 2 do
		local SrcValue = buffer.readu16(Src, Offset)
		local DstValue = buffer.readu16(Dst, Offset)
		local Sum = SrcValue + DstValue
		buffer.writeu16(Dst, Offset, if Sum >= Modulus then Sum - Modulus else Sum)
	end
end

function PolyVec.VecAdd(A: buffer, B: buffer, Result: buffer, K: number)
	if not (MlKemParams.CheckK(K) or K == 1) then
		error("Invalid vector dimension K")
	end

	local TotalCoeffs = K * N
	local Modulus = Q
	for Index = 0, TotalCoeffs - 1 do
		local Offset = Index * 2
		local AValue = buffer.readu16(A, Offset)
		local BValue = buffer.readu16(B, Offset)
		local Sum = AValue + BValue
		local AddResult = if Sum >= Modulus then Sum - Modulus else Sum
		buffer.writeu16(Result, Offset, AddResult)
	end
end

function PolyVec.VecEncode(Vec: buffer, K: number, L: number): buffer
	local OutputSize = K * 32 * L
	local Output = buffer.create(OutputSize)
	local PolySize = POLY_SIZE
	local ChunkSize = 32 * L

	local EncodeAt = MlKemSerialize.EncodeAt
	local Encode = MlKemSerialize.Encode

	if L == 12 then
		for I = 0, K - 1 do
			EncodeAt(Vec, I * PolySize, Output, I * ChunkSize, 12)
		end
	else
		local Poly = TEMP_POLY_ENCODE
		for I = 0, K - 1 do
			buffer.copy(Poly, 0, Vec, I * PolySize, PolySize)
			local EncodedPoly = Encode(Poly, L)
			buffer.copy(Output, I * ChunkSize, EncodedPoly, 0, ChunkSize)
		end
	end

	return Output
end

function PolyVec.VecDecode(Data: buffer, K: number, L: number): buffer
	local Vec = buffer.create(K * POLY_SIZE)
	local PolySize = POLY_SIZE
	local ChunkSize = 32 * L

	local DecodeAt = MlKemSerialize.DecodeAt
	local Decode = MlKemSerialize.Decode

	if L == 12 then
		for I = 0, K - 1 do
			DecodeAt(Data, I * ChunkSize, Vec, I * PolySize, 12)
		end
	else
		local PolyData = buffer.create(ChunkSize)
		for I = 0, K - 1 do
			buffer.copy(PolyData, 0, Data, I * ChunkSize, ChunkSize)
			local DecodedPoly = Decode(PolyData, L)
			buffer.copy(Vec, I * PolySize, DecodedPoly, 0, PolySize)
		end
	end

	return Vec
end

function PolyVec.VecCompress(Vec: buffer, K: number, D: number)
	local Poly = TEMP_POLY_COMPRESS
	local PolySize = POLY_SIZE

	for I = 0, K - 1 do
		local Offset = I * PolySize
		buffer.copy(Poly, 0, Vec, Offset, PolySize)
		Compression.PolyCompress(Poly, D)
		buffer.copy(Vec, Offset, Poly, 0, PolySize)
	end
end

function PolyVec.VecDecompress(Vec: buffer, K: number, D: number)
	local Poly = TEMP_POLY_COMPRESS
	local PolySize = POLY_SIZE

	for I = 0, K - 1 do
		local Offset = I * PolySize
		buffer.copy(Poly, 0, Vec, Offset, PolySize)
		Compression.PolyDecompress(Poly, D)
		buffer.copy(Vec, Offset, Poly, 0, PolySize)
	end
end

return PolyVec


--[=[
	ML-KEM Sampling Module

	Polynomial sampling operations for ML-KEM.
	 Uniform sampling from Rq and centered binomial distribution sampling.

	Example usage:
		local Sampling = require(script.Sampling)

		local Rho = buffer.create(32)
		local Matrix = Sampling.GenerateMatrix(3, Rho, false)

		local Sigma = buffer.create(32)
		local Vector = Sampling.GenerateVector(3, 2, Sigma, 0)

		local PrfOutput = buffer.create(128)
		local NoisePoly = Sampling.SamplePolyCbd(PrfOutput, 2)
--]=]

--!strict
--!optimize 2
--!native

local PolyVec = require("./PolyVec")
local MlKemParams = require("./Params")
local SHA3 = require("./SHA3")
local XOF = require("./XOF")

local LOG2N = 8
local N = 2 ^ LOG2N
local POLY_SIZE = N * 2

local Q = 13 * 256 + 1
local Q_BIT_WIDTH = 12
local WIDTH = (2 ^ (2 * Q_BIT_WIDTH))
local R = math.floor((2 ^ (2 * Q_BIT_WIDTH)) / Q)

local SAMPLE_NTT_CHUNK = buffer.create(168)
local SAMPLE_NTT_POLY = buffer.create(N * 2)
local GEN_VECTOR_PRF_OUTPUT = buffer.create(192)
local GEN_VECTOR_PRF_INPUT = buffer.create(33)

local Sampling = {}

local function SampleNtt(XofInput: buffer, Poly: buffer)
	local CoeffIndex = 0
	local Chunk = SAMPLE_NTT_CHUNK

	XOF.Reset128()
	XOF.Absorb128(XofInput)

	while CoeffIndex < N do
		XOF.Squeeze128Into(Chunk, 168, 0)
		local Offset = 0

		while Offset + 2 < 168 and CoeffIndex < N do
			local B0 = buffer.readu8(Chunk, Offset)
			local B1 = buffer.readu8(Chunk, Offset + 1)
			local B2 = buffer.readu8(Chunk, Offset + 2)

			local D1 = bit32.bor(B0, bit32.lshift(bit32.band(B1, 0x0F), 8))
			local D2 = bit32.bor(bit32.rshift(B1, 4), bit32.lshift(B2, 4))

			if D1 < Q then
				buffer.writeu16(Poly, CoeffIndex * 2, D1)
				CoeffIndex += 1
			end

			if D2 < Q and CoeffIndex < N then
				buffer.writeu16(Poly, CoeffIndex * 2, D2)
				CoeffIndex += 1
			end

			Offset += 3
		end
	end
end

function Sampling.GenerateMatrix(K: number, Rho: buffer, Transpose: boolean): buffer
	local Matrix = PolyVec.MatCreate(K, K)
	local XofInput = buffer.create(34)
	buffer.copy(XofInput, 0, Rho, 0, 32)

	local PolySize = POLY_SIZE
	local Poly = SAMPLE_NTT_POLY
	local Sample = SampleNtt

	for I = 0, K - 1 do
		for J = 0, K - 1 do
			if Transpose then
				buffer.writeu8(XofInput, 32, I)
				buffer.writeu8(XofInput, 33, J)
			else
				buffer.writeu8(XofInput, 32, J)
				buffer.writeu8(XofInput, 33, I)
			end

			Sample(XofInput, Poly)
			local Index = (I * K + J) * PolySize
			buffer.copy(Matrix, Index, Poly, 0, PolySize)
		end
	end

	return Matrix
end

local function SamplePolyCbd2(Prf: buffer, Poly: buffer, Offset: number?)
	local BaseOffset = Offset or 0
	local Mask2 = 0x03
	local Modulus = Q

	for I = 0, 127 do
		local PolyOffset = I * 2
		local Word = buffer.readu8(Prf, I)

		local T0 = bit32.band(Word, 0x55)
		local T1 = bit32.band(bit32.rshift(Word, 1), 0x55)
		local T2 = T0 + T1

		local A0 = bit32.band(T2, Mask2)
		local B0 = bit32.band(bit32.rshift(T2, 2), Mask2)
		local Diff0 = A0 - B0
		local Coeff0 = if Diff0 < 0 then Diff0 + Modulus else Diff0

		local A1 = bit32.band(bit32.rshift(T2, 4), Mask2)
		local B1 = bit32.band(bit32.rshift(T2, 6), Mask2)
		local Diff1 = A1 - B1
		local Coeff1 = if Diff1 < 0 then Diff1 + Modulus else Diff1

		buffer.writeu16(Poly, BaseOffset + PolyOffset * 2, Coeff0)
		buffer.writeu16(Poly, BaseOffset + (PolyOffset + 1) * 2, Coeff1)
	end
end

local function SamplePolyCbd3(Prf: buffer, Poly: buffer, Offset: number?)
	local BaseOffset = Offset or 0
	local Mask24 = 0x249249
	local Mask3 = 0x07
	local Modulus = Q

	for I = 0, 63 do
		local ByteOffset = I * 3
		local PolyOffset = I * 4

		local B0 = buffer.readu8(Prf, ByteOffset)
		local B1 = buffer.readu8(Prf, ByteOffset + 1)
		local B2 = buffer.readu8(Prf, ByteOffset + 2)
		local Word = bit32.bor(B0, bit32.lshift(B1, 8), bit32.lshift(B2, 16))

		local T0 = bit32.band(Word, Mask24)
		local T1 = bit32.band(bit32.rshift(Word, 1), Mask24)
		local T2 = bit32.band(bit32.rshift(Word, 2), Mask24)
		local T3 = T0 + T1 + T2

		for J = 0, 3 do
			local JOffset = J * 6
			local A = bit32.band(bit32.rshift(T3, JOffset), Mask3)
			local B = bit32.band(bit32.rshift(T3, JOffset + 3), Mask3)
			local Diff = A - B
			local Coeff = if Diff < 0 then Diff + Modulus else Diff

			buffer.writeu16(Poly, BaseOffset + (PolyOffset + J) * 2, Coeff)
		end
	end
end

function Sampling.SamplePolyCbd(Prf: buffer, Eta: number): buffer
	local Poly = buffer.create(N * 2)

	if Eta == 2 then
		SamplePolyCbd2(Prf, Poly, 0)
	else
		SamplePolyCbd3(Prf, Poly, 0)
	end

	return Poly
end

function Sampling.GenerateVector(K: number, Eta: number, Sigma: buffer, Nonce: number): buffer
	local Vec = PolyVec.VecCreate(K)
	local PrfInput = GEN_VECTOR_PRF_INPUT
	local PrfOutput = GEN_VECTOR_PRF_OUTPUT
	buffer.copy(PrfInput, 0, Sigma, 0, 32)

	local PolySize = POLY_SIZE
	local PrfLen = 64 * Eta

	for I = 0, K - 1 do
		buffer.writeu8(PrfInput, 32, Nonce + I)

		XOF.Reset256()
		XOF.Absorb256(PrfInput)
		XOF.Squeeze256Into(PrfOutput, PrfLen, 0)

		local Offset = I * PolySize
		if Eta == 2 then
			SamplePolyCbd2(PrfOutput, Vec, Offset)
		else
			SamplePolyCbd3(PrfOutput, Vec, Offset)
		end
	end

	return Vec
end

function Sampling.GenerateNoisePoly(Eta: number, Sigma: buffer, Nonce: number): buffer
	if not MlKemParams.CheckEta(Eta) then
		error("Invalid eta parameter")
	end
	if buffer.len(Sigma) ~= 32 then
		error("Sigma must be 32 bytes")
	end

	local PrfInput = buffer.create(33)
	buffer.copy(PrfInput, 0, Sigma, 0, 32)
	buffer.writeu8(PrfInput, 32, Nonce)

	local PrfOutput = SHA3.SHAKE256(PrfInput, 64 * Eta)
	return Sampling.SamplePolyCbd(PrfOutput, Eta)
end

return Sampling


--[=[
	ML-KEM Serialization Module

	Polynomial encoding and decoding operations for ML-KEM.
	Converts polynomials to/from byte arrays with different bit lengths.

	Example usage:
		local Serialize = require(script.Serialize)

		local Poly = buffer.create(512) -- 256 coefficients
		local Encoded = Serialize.Encode(Poly, 12)
		local Decoded = Serialize.Decode(Encoded, 12)
--]=]

--!strict
--!optimize 2
--!native

local Field = require("./Field")
local Params = require("./Params")

local LOG2N = 8
local N = 2 ^ LOG2N
local POLY_SIZE = N * 2

local Q = 13 * 256 + 1
local Q_BIT_WIDTH = 12
local WIDTH = (2 ^ (2 * Q_BIT_WIDTH))
local R = math.floor((2 ^ (2 * Q_BIT_WIDTH)) / Q)

local Serialize = {}

function Serialize.Encode(Poly: buffer, L: number): buffer
	if not Params.CheckL(L) then
		error("Invalid encoding parameter l")
	end

	local OutputSize = 32 * L
	local Output = buffer.create(OutputSize)

	if L == 1 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local Offset = I * 8
			local Byte = 0

			for J = 0, 7 do
				local Coeff = buffer.readu16(Poly, (Offset + J) * 2)
				local Bit = bit32.band(Coeff, 1)
				Byte = bit32.bor(Byte, bit32.lshift(Bit, J))
			end

			buffer.writeu8(Output, I, Byte)
		end
	elseif L == 4 then
		local IterCount = N / 2

		for I = 0, IterCount - 1 do
			local Offset = I * 2
			local C0 = bit32.band(buffer.readu16(Poly, Offset * 2), 0xF)
			local C1 = bit32.band(buffer.readu16(Poly, (Offset + 1) * 2), 0xF)

			local Byte = bit32.bor(C0, bit32.lshift(C1, 4))
			buffer.writeu8(Output, I, Byte)
		end
	elseif L == 5 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 8
			local ByteOffset = I * 5

			local T0 = buffer.readu16(Poly, (PolyOffset + 0) * 2)
			local T1 = buffer.readu16(Poly, (PolyOffset + 1) * 2)
			local T2 = buffer.readu16(Poly, (PolyOffset + 2) * 2)
			local T3 = buffer.readu16(Poly, (PolyOffset + 3) * 2)
			local T4 = buffer.readu16(Poly, (PolyOffset + 4) * 2)
			local T5 = buffer.readu16(Poly, (PolyOffset + 5) * 2)
			local T6 = buffer.readu16(Poly, (PolyOffset + 6) * 2)
			local T7 = buffer.readu16(Poly, (PolyOffset + 7) * 2)

			buffer.writeu8(
				Output,
				ByteOffset + 0,
				bit32.bor(bit32.band(T0, 0x1F), bit32.lshift(bit32.band(T1, 0x07), 5))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 1,
				bit32.bor(
					bit32.rshift(T1, 3),
					bit32.lshift(bit32.band(T2, 0x1F), 2),
					bit32.lshift(bit32.band(T3, 0x01), 7)
				)
			)
			buffer.writeu8(
				Output,
				ByteOffset + 2,
				bit32.bor(bit32.rshift(T3, 1), bit32.lshift(bit32.band(T4, 0x0F), 4))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 3,
				bit32.bor(
					bit32.rshift(T4, 4),
					bit32.lshift(bit32.band(T5, 0x1F), 1),
					bit32.lshift(bit32.band(T6, 0x03), 6)
				)
			)
			buffer.writeu8(
				Output,
				ByteOffset + 4,
				bit32.bor(bit32.rshift(T6, 2), bit32.lshift(bit32.band(T7, 0x1F), 3))
			)
		end
	elseif L == 10 then
		local IterCount = N / 4

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 4
			local ByteOffset = I * 5

			local T0 = buffer.readu16(Poly, (PolyOffset + 0) * 2)
			local T1 = buffer.readu16(Poly, (PolyOffset + 1) * 2)
			local T2 = buffer.readu16(Poly, (PolyOffset + 2) * 2)
			local T3 = buffer.readu16(Poly, (PolyOffset + 3) * 2)

			buffer.writeu8(Output, ByteOffset + 0, bit32.band(T0, 0xFF))
			buffer.writeu8(
				Output,
				ByteOffset + 1,
				bit32.bor(bit32.rshift(T0, 8), bit32.lshift(bit32.band(T1, 0x3F), 2))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 2,
				bit32.bor(bit32.rshift(T1, 6), bit32.lshift(bit32.band(T2, 0x0F), 4))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 3,
				bit32.bor(bit32.rshift(T2, 4), bit32.lshift(bit32.band(T3, 0x03), 6))
			)
			buffer.writeu8(Output, ByteOffset + 4, bit32.rshift(T3, 2))
		end
	elseif L == 11 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 8
			local ByteOffset = I * 11

			local T0 = buffer.readu16(Poly, (PolyOffset + 0) * 2)
			local T1 = buffer.readu16(Poly, (PolyOffset + 1) * 2)
			local T2 = buffer.readu16(Poly, (PolyOffset + 2) * 2)
			local T3 = buffer.readu16(Poly, (PolyOffset + 3) * 2)
			local T4 = buffer.readu16(Poly, (PolyOffset + 4) * 2)
			local T5 = buffer.readu16(Poly, (PolyOffset + 5) * 2)
			local T6 = buffer.readu16(Poly, (PolyOffset + 6) * 2)
			local T7 = buffer.readu16(Poly, (PolyOffset + 7) * 2)

			buffer.writeu8(Output, ByteOffset + 0, bit32.band(T0, 0xFF))
			buffer.writeu8(
				Output,
				ByteOffset + 1,
				bit32.bor(bit32.rshift(T0, 8), bit32.lshift(bit32.band(T1, 0x1F), 3))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 2,
				bit32.bor(bit32.rshift(T1, 5), bit32.lshift(bit32.band(T2, 0x03), 6))
			)
			buffer.writeu8(Output, ByteOffset + 3, bit32.rshift(T2, 2))
			buffer.writeu8(
				Output,
				ByteOffset + 4,
				bit32.bor(bit32.rshift(T2, 10), bit32.lshift(bit32.band(T3, 0x7F), 1))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 5,
				bit32.bor(bit32.rshift(T3, 7), bit32.lshift(bit32.band(T4, 0x0F), 4))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 6,
				bit32.bor(bit32.rshift(T4, 4), bit32.lshift(bit32.band(T5, 0x01), 7))
			)
			buffer.writeu8(Output, ByteOffset + 7, bit32.rshift(T5, 1))
			buffer.writeu8(
				Output,
				ByteOffset + 8,
				bit32.bor(bit32.rshift(T5, 9), bit32.lshift(bit32.band(T6, 0x3F), 2))
			)
			buffer.writeu8(
				Output,
				ByteOffset + 9,
				bit32.bor(bit32.rshift(T6, 6), bit32.lshift(bit32.band(T7, 0x07), 5))
			)
			buffer.writeu8(Output, ByteOffset + 10, bit32.rshift(T7, 3))
		end
	else
		local IterCount = N / 2

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 2
			local ByteOffset = I * 3

			local T0 = buffer.readu16(Poly, (PolyOffset + 0) * 2)
			local T1 = buffer.readu16(Poly, (PolyOffset + 1) * 2)

			buffer.writeu8(Output, ByteOffset + 0, bit32.band(T0, 0xFF))
			buffer.writeu8(
				Output,
				ByteOffset + 1,
				bit32.bor(bit32.rshift(T0, 8), bit32.lshift(bit32.band(T1, 0x0F), 4))
			)
			buffer.writeu8(Output, ByteOffset + 2, bit32.rshift(T1, 4))
		end
	end

	return Output
end

function Serialize.EncodeAt(Poly: buffer, PolyOffset: number, Output: buffer, OutOffset: number, L: number)
	if L == 12 then
		for I = 0, 127 do
			local PolyIdx = PolyOffset + I * 4
			local ByteIdx = OutOffset + I * 3

			local T0 = buffer.readu16(Poly, PolyIdx)
			local T1 = buffer.readu16(Poly, PolyIdx + 2)

			buffer.writeu8(Output, ByteIdx, bit32.band(T0, 0xFF))
			buffer.writeu8(Output, ByteIdx + 1, bit32.bor(bit32.rshift(T0, 8), bit32.lshift(bit32.band(T1, 0x0F), 4)))
			buffer.writeu8(Output, ByteIdx + 2, bit32.rshift(T1, 4))
		end
	end
end

function Serialize.Decode(Data: buffer, L: number): buffer
	if not Params.CheckL(L) then
		error("Invalid encoding parameter l")
	end

	local Poly = buffer.create(N * 2)
	if L == 1 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local Byte = buffer.readu8(Data, I)
			local Offset = I * 8

			for J = 0, 7 do
				local Bit = bit32.band(bit32.rshift(Byte, J), 1)
				buffer.writeu16(Poly, (Offset + J) * 2, Bit)
			end
		end
	elseif L == 4 then
		local IterCount = N / 2

		for I = 0, IterCount - 1 do
			local Byte = buffer.readu8(Data, I)
			local Offset = I * 2

			buffer.writeu16(Poly, (Offset + 0) * 2, bit32.band(Byte, 0x0F))
			buffer.writeu16(Poly, (Offset + 1) * 2, bit32.rshift(Byte, 4))
		end
	elseif L == 5 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 8
			local ByteOffset = I * 5

			local B0 = buffer.readu8(Data, ByteOffset + 0)
			local B1 = buffer.readu8(Data, ByteOffset + 1)
			local B2 = buffer.readu8(Data, ByteOffset + 2)
			local B3 = buffer.readu8(Data, ByteOffset + 3)
			local B4 = buffer.readu8(Data, ByteOffset + 4)

			local T0 = bit32.band(B0, 0x1F)
			local T1 = bit32.bor(bit32.rshift(B0, 5), bit32.lshift(bit32.band(B1, 0x03), 3))
			local T2 = bit32.band(bit32.rshift(B1, 2), 0x1F)
			local T3 = bit32.bor(bit32.rshift(B1, 7), bit32.lshift(bit32.band(B2, 0x0F), 1))
			local T4 = bit32.bor(bit32.rshift(B2, 4), bit32.lshift(bit32.band(B3, 0x01), 4))
			local T5 = bit32.band(bit32.rshift(B3, 1), 0x1F)
			local T6 = bit32.bor(bit32.rshift(B3, 6), bit32.lshift(bit32.band(B4, 0x07), 2))
			local T7 = bit32.rshift(B4, 3)

			buffer.writeu16(Poly, (PolyOffset + 0) * 2, T0)
			buffer.writeu16(Poly, (PolyOffset + 1) * 2, T1)
			buffer.writeu16(Poly, (PolyOffset + 2) * 2, T2)
			buffer.writeu16(Poly, (PolyOffset + 3) * 2, T3)
			buffer.writeu16(Poly, (PolyOffset + 4) * 2, T4)
			buffer.writeu16(Poly, (PolyOffset + 5) * 2, T5)
			buffer.writeu16(Poly, (PolyOffset + 6) * 2, T6)
			buffer.writeu16(Poly, (PolyOffset + 7) * 2, T7)
		end
	elseif L == 10 then
		local IterCount = N / 4

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 4
			local ByteOffset = I * 5

			local B0 = buffer.readu8(Data, ByteOffset + 0)
			local B1 = buffer.readu8(Data, ByteOffset + 1)
			local B2 = buffer.readu8(Data, ByteOffset + 2)
			local B3 = buffer.readu8(Data, ByteOffset + 3)
			local B4 = buffer.readu8(Data, ByteOffset + 4)

			local T0 = bit32.bor(B0, bit32.lshift(bit32.band(B1, 0x03), 8))
			local T1 = bit32.bor(bit32.rshift(B1, 2), bit32.lshift(bit32.band(B2, 0x0F), 6))
			local T2 = bit32.bor(bit32.rshift(B2, 4), bit32.lshift(bit32.band(B3, 0x3F), 4))
			local T3 = bit32.bor(bit32.rshift(B3, 6), bit32.lshift(B4, 2))

			buffer.writeu16(Poly, (PolyOffset + 0) * 2, T0)
			buffer.writeu16(Poly, (PolyOffset + 1) * 2, T1)
			buffer.writeu16(Poly, (PolyOffset + 2) * 2, T2)
			buffer.writeu16(Poly, (PolyOffset + 3) * 2, T3)
		end
	elseif L == 11 then
		local IterCount = N / 8

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 8
			local ByteOffset = I * 11

			local B0 = buffer.readu8(Data, ByteOffset + 0)
			local B1 = buffer.readu8(Data, ByteOffset + 1)
			local B2 = buffer.readu8(Data, ByteOffset + 2)
			local B3 = buffer.readu8(Data, ByteOffset + 3)
			local B4 = buffer.readu8(Data, ByteOffset + 4)
			local B5 = buffer.readu8(Data, ByteOffset + 5)
			local B6 = buffer.readu8(Data, ByteOffset + 6)
			local B7 = buffer.readu8(Data, ByteOffset + 7)
			local B8 = buffer.readu8(Data, ByteOffset + 8)
			local B9 = buffer.readu8(Data, ByteOffset + 9)
			local B10 = buffer.readu8(Data, ByteOffset + 10)

			local T0 = bit32.bor(B0, bit32.lshift(bit32.band(B1, 0x07), 8))
			local T1 = bit32.bor(bit32.rshift(B1, 3), bit32.lshift(bit32.band(B2, 0x3F), 5))
			local T2 = bit32.bor(bit32.rshift(B2, 6), bit32.lshift(B3, 2), bit32.lshift(bit32.band(B4, 0x01), 10))
			local T3 = bit32.bor(bit32.rshift(B4, 1), bit32.lshift(bit32.band(B5, 0x0F), 7))
			local T4 = bit32.bor(bit32.rshift(B5, 4), bit32.lshift(bit32.band(B6, 0x7F), 4))
			local T5 = bit32.bor(bit32.rshift(B6, 7), bit32.lshift(B7, 1), bit32.lshift(bit32.band(B8, 0x03), 9))
			local T6 = bit32.bor(bit32.rshift(B8, 2), bit32.lshift(bit32.band(B9, 0x1F), 6))
			local T7 = bit32.bor(bit32.rshift(B9, 5), bit32.lshift(B10, 3))

			buffer.writeu16(Poly, (PolyOffset + 0) * 2, T0)
			buffer.writeu16(Poly, (PolyOffset + 1) * 2, T1)
			buffer.writeu16(Poly, (PolyOffset + 2) * 2, T2)
			buffer.writeu16(Poly, (PolyOffset + 3) * 2, T3)
			buffer.writeu16(Poly, (PolyOffset + 4) * 2, T4)
			buffer.writeu16(Poly, (PolyOffset + 5) * 2, T5)
			buffer.writeu16(Poly, (PolyOffset + 6) * 2, T6)
			buffer.writeu16(Poly, (PolyOffset + 7) * 2, T7)
		end
	else
		local IterCount = N / 2

		for I = 0, IterCount - 1 do
			local PolyOffset = I * 2
			local ByteOffset = I * 3

			local B0 = buffer.readu8(Data, ByteOffset + 0)
			local B1 = buffer.readu8(Data, ByteOffset + 1)
			local B2 = buffer.readu8(Data, ByteOffset + 2)

			local T0 = bit32.bor(B0, bit32.lshift(bit32.band(B1, 0x0F), 8))
			local T1 = bit32.bor(bit32.rshift(B1, 4), bit32.lshift(B2, 4))

			if T0 >= 3329 or T1 >= 3329 then
				error("Invalid polynomial coefficient encoding")
			end

			buffer.writeu16(Poly, (PolyOffset + 0) * 2, T0)
			buffer.writeu16(Poly, (PolyOffset + 1) * 2, T1)
		end
	end

	return Poly
end

function Serialize.DecodeAt(Data: buffer, DataOffset: number, Poly: buffer, PolyOffset: number, L: number)
	if L == 12 then
		for I = 0, 127 do
			local ByteIdx = DataOffset + I * 3
			local PolyIdx = PolyOffset + I * 4

			local B0 = buffer.readu8(Data, ByteIdx)
			local B1 = buffer.readu8(Data, ByteIdx + 1)
			local B2 = buffer.readu8(Data, ByteIdx + 2)

			local T0 = bit32.bor(B0, bit32.lshift(bit32.band(B1, 0x0F), 8))
			local T1 = bit32.bor(bit32.rshift(B1, 4), bit32.lshift(B2, 4))

			if T0 >= 3329 or T1 >= 3329 then
				error("Invalid polynomial coefficient encoding")
			end

			buffer.writeu16(Poly, PolyIdx, T0)
			buffer.writeu16(Poly, PolyIdx + 2, T1)
		end
	end
end

return Serialize


--[=[
	ML-KEM Utilities Module

	Utility functions for ML-KEM implementation

	Example usage:
		local Utils = require(script.Utils)
		local IsEqual = Utils.CtMemcmp(Buffer1, Buffer2)

		Utils.CtCondMemcpy(Condition, Dest, Src1, Src2)

		local PubKeySize = Utils.GetKemPublicKeyLen(4) -- 1568 bytes
		local SecKeySize = Utils.GetKemSecretKeyLen(4) -- 3168 bytes
--]=]

--!strict
--!optimize 2
--!native

local MlKemParams = require("./Params")

local Utils = {}

function Utils.CtMemcmp(Bytes0: buffer, Bytes1: buffer): number
	local Len0 = buffer.len(Bytes0)
	local Len1 = buffer.len(Bytes1)

	if Len0 ~= Len1 then
		return 0x00000000
	end

	local Diff = 0
	local Len4 = Len0 - (Len0 % 4)

	for Index = 0, Len4 - 4, 4 do
		local Word0 = buffer.readu32(Bytes0, Index)
		local Word1 = buffer.readu32(Bytes1, Index)
		Diff = bit32.bor(Diff, bit32.bxor(Word0, Word1))
	end

	for Index = Len4, Len0 - 1 do
		local Byte0 = buffer.readu8(Bytes0, Index)
		local Byte1 = buffer.readu8(Bytes1, Index)
		Diff = bit32.bor(Diff, bit32.bxor(Byte0, Byte1))
	end

	return if Diff == 0 then 0xFFFFFFFF else 0x00000000
end

function Utils.CtSelect(Cond: number, A: number, B: number): number
	local Mask = bit32.band(Cond, 0xFFFFFFFF)

	return bit32.bor(bit32.band(A, Mask), bit32.band(B, bit32.bnot(Mask)))
end

function Utils.CtCondMemcpy(Cond: number, Sink: buffer, Source0: buffer, Source1: buffer)
	local SinkLen = buffer.len(Sink)
	local Mask = bit32.band(Cond, 0xFFFFFFFF)
	local InvMask = bit32.bnot(Mask)
	local Len4 = SinkLen - (SinkLen % 4)

	for Index = 0, Len4 - 4, 4 do
		local Word0 = buffer.readu32(Source0, Index)
		local Word1 = buffer.readu32(Source1, Index)
		local Selected = bit32.bor(bit32.band(Word0, Mask), bit32.band(Word1, InvMask))
		buffer.writeu32(Sink, Index, Selected)
	end

	for Index = Len4, SinkLen - 1 do
		local Byte0 = buffer.readu8(Source0, Index)
		local Byte1 = buffer.readu8(Source1, Index)
		local Selected = bit32.bor(bit32.band(Byte0, Mask), bit32.band(Byte1, InvMask))
		buffer.writeu8(Sink, Index, Selected)
	end
end

function Utils.BitReverse(Value: number, BitWidth: number): number
	local Reversed = 0
	for I = 0, BitWidth - 1 do
		local Bit = bit32.band(bit32.rshift(Value, I), 1)
		Reversed = bit32.bor(Reversed, bit32.lshift(Bit, BitWidth - 1 - I))
	end

	return Reversed
end

function Utils.GetPkePublicKeyLen(K: number): number
	if not MlKemParams.CheckK(K) then
		error("Invalid K parameter")
	end

	return K * 12 * 32 + 32
end

function Utils.GetPkeSecretKeyLen(K: number): number
	if not MlKemParams.CheckK(K) then
		error("Invalid K parameter")
	end

	return K * 12 * 32
end

function Utils.GetPkeCipherTextLen(K: number, Du: number, Dv: number): number
	if not MlKemParams.CheckK(K) then
		error("Invalid K parameter")
	end
	if not MlKemParams.CheckD(Du) then
		error("Invalid Du parameter")
	end
	if not MlKemParams.CheckD(Dv) then
		error("Invalid Dv parameter")
	end

	return 32 * (K * Du + Dv)
end

function Utils.GetKemPublicKeyLen(K: number): number
	return Utils.GetPkePublicKeyLen(K)
end

function Utils.GetKemSecretKeyLen(K: number): number
	if not MlKemParams.CheckK(K) then
		error("Invalid K parameter")
	end

	return Utils.GetPkeSecretKeyLen(K) + Utils.GetPkePublicKeyLen(K) + 32 + 32
end

function Utils.GetKemCipherTextLen(K: number, Du: number, Dv: number): number
	return Utils.GetPkeCipherTextLen(K, Du, Dv)
end

return Utils