use super::MlDsaParameters;
use crate::{
DType, Error, Matrix, OpAttribute, Result,
runtime::{BufferBinding, ComputeDispatch, KernelId, PushConstant, SemanticDispatch},
};
pub fn verify_batch(
messages: &Matrix,
msg_offsets: &Matrix,
msg_lengths: &Matrix,
signatures: &Matrix,
public_keys: &Matrix,
) -> Result<Matrix> {
verify_batch_with_parameters(
MlDsaParameters::MlDsa65,
messages,
msg_offsets,
msg_lengths,
signatures,
public_keys,
)
}
pub fn verify_batch_with_parameters(
parameters: MlDsaParameters,
messages: &Matrix,
msg_offsets: &Matrix,
msg_lengths: &Matrix,
signatures: &Matrix,
public_keys: &Matrix,
) -> Result<Matrix> {
verify_batch_with_context(
parameters,
messages,
msg_offsets,
msg_lengths,
signatures,
public_keys,
0..0,
)
}
pub fn verify_batch_with_context(
parameters: MlDsaParameters,
messages: &Matrix,
msg_offsets: &Matrix,
msg_lengths: &Matrix,
signatures: &Matrix,
public_keys: &Matrix,
context: std::ops::Range<usize>,
) -> Result<Matrix> {
let (contract, kernel) = match parameters {
MlDsaParameters::MlDsa65 => (
crate::core::operation::cryptography::pqc::ML_DSA_65_VERIFY_BATCH,
KernelId::CryptographyMlDsa65VerifyU8,
),
MlDsaParameters::MlDsa44 => (
crate::core::operation::cryptography::pqc::ML_DSA_44_VERIFY_BATCH,
KernelId::CryptographyMlDsa44VerifyU8,
),
MlDsaParameters::MlDsa87 => (
crate::core::operation::cryptography::pqc::ML_DSA_87_VERIFY_BATCH,
KernelId::CryptographyMlDsa87VerifyU8,
),
};
let operation = contract.name();
let sig_bytes = parameters.signature_size();
let pk_bytes = parameters.public_key_size();
if messages.dtype() != DType::U8 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be a U8 matrix; found {}",
messages.dtype().token()
)));
}
if messages.shape().len() != 1 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be rank-1 [total_bytes]; found {:?}",
messages.shape()
)));
}
let context_length = context
.end
.checked_sub(context.start)
.ok_or_else(|| Error::invalid_argument(format!("{operation} context range is reversed")))?;
if context.end > messages.num_elements() || context_length > 255 {
return Err(Error::invalid_argument(format!(
"{operation} context must be within messages and at most 255 bytes"
)));
}
let context_offset = u32::try_from(context.start)
.map_err(|_| Error::invalid_argument(format!("{operation} context offset exceeds u32")))?;
let context_length = u32::try_from(context_length)
.map_err(|_| Error::invalid_argument(format!("{operation} context length exceeds u32")))?;
let batch_size = validate_u32_vector(msg_offsets, "msg_offsets", None, operation)?;
validate_u32_vector(msg_lengths, "msg_lengths", Some(batch_size), operation)?;
validate_packed_u8(signatures, "signatures", batch_size, sig_bytes, operation)?;
validate_packed_u8(public_keys, "public_keys", batch_size, pk_bytes, operation)?;
let engine = messages.engine_handle();
for (label, matrix) in [
("msg_offsets", msg_offsets),
("msg_lengths", msg_lengths),
("signatures", signatures),
("public_keys", public_keys),
] {
if !engine.same_as(matrix.engine_handle()) {
return Err(Error::invalid_argument(format!(
"{operation} {label} must share an engine with messages"
)));
}
}
let batch_u32 = u32::try_from(batch_size)
.map_err(|_| Error::invalid_argument(format!("{operation} batch_size exceeds u32")))?;
let message_bytes = u32::try_from(messages.num_elements())
.map_err(|_| Error::invalid_argument(format!("{operation} message byte size exceeds u32")))?;
for matrix in [msg_offsets, msg_lengths, signatures, public_keys] {
let bytes = matrix
.num_elements()
.checked_mul(matrix.dtype().size_bytes())
.ok_or_else(|| Error::invalid_argument(format!("{operation} buffer byte size overflows")))?;
u32::try_from(bytes)
.map_err(|_| Error::invalid_argument(format!("{operation} buffer byte size exceeds u32")))?;
}
let results = Matrix::allocate_output(engine, vec![batch_size], batch_size, DType::U32)?;
if batch_size == 0 {
return Ok(results);
}
let buffers = [
BufferBinding::read(messages.storage()),
BufferBinding::read(msg_offsets.storage()),
BufferBinding::read(msg_lengths.storage()),
BufferBinding::read(signatures.storage()),
BufferBinding::read(public_keys.storage()),
BufferBinding::write(results.storage()),
];
let push_constants = [
PushConstant::U32(batch_u32),
PushConstant::U32(message_bytes),
PushConstant::U32(context_offset),
PushConstant::U32(context_length),
];
let inputs = [messages, msg_offsets, msg_lengths, signatures, public_keys];
let outputs = [&results];
let attributes = [
OpAttribute::UnsignedInteger {
name: "batch_size".into(),
value: u64::from(batch_u32),
},
OpAttribute::UnsignedInteger {
name: "context_offset".into(),
value: u64::from(context_offset),
},
OpAttribute::UnsignedInteger {
name: "context_length".into(),
value: u64::from(context_length),
},
];
engine.record_semantic(
ComputeDispatch {
kernel,
buffers: &buffers,
push_constants: &push_constants,
workgroups: kernel.linear_workgroups(batch_u32),
},
SemanticDispatch {
contract,
inputs: &inputs,
outputs: &outputs,
attributes: &attributes,
},
)?;
Ok(results)
}
#[cfg_attr(
not(test),
expect(dead_code, reason = "private verifier awaits runtime admission")
)]
#[expect(
clippy::too_many_arguments,
reason = "private adapter preserves the five-buffer verifier contract"
)]
pub(crate) fn verify_prehashed_batch(
parameters: MlDsaParameters,
algorithm: u32,
messages: &Matrix,
msg_offsets: &Matrix,
msg_lengths: &Matrix,
signatures: &Matrix,
public_keys: &Matrix,
context: std::ops::Range<usize>,
) -> Result<Matrix> {
if KernelId::mldsa_prehash_bytes(algorithm).is_none() {
return Err(Error::invalid_argument("unknown ML-DSA prehash identifier"));
}
let (contract, kernel) = match parameters {
MlDsaParameters::MlDsa65 => (
crate::core::operation::cryptography::pqc::ML_DSA_65_VERIFY_PREHASHED_BATCH,
KernelId::CryptographyMlDsa65VerifyPrehashedU8,
),
MlDsaParameters::MlDsa44 => (
crate::core::operation::cryptography::pqc::ML_DSA_44_VERIFY_PREHASHED_BATCH,
KernelId::CryptographyMlDsa44VerifyPrehashedU8,
),
MlDsaParameters::MlDsa87 => (
crate::core::operation::cryptography::pqc::ML_DSA_87_VERIFY_PREHASHED_BATCH,
KernelId::CryptographyMlDsa87VerifyPrehashedU8,
),
};
let operation = contract.name();
let sig_bytes = parameters.signature_size();
let pk_bytes = parameters.public_key_size();
if messages.dtype() != DType::U8 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be a U8 matrix; found {}",
messages.dtype().token()
)));
}
if messages.shape().len() != 1 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be rank-1 [total_bytes]; found {:?}",
messages.shape()
)));
}
let context_length = context
.end
.checked_sub(context.start)
.ok_or_else(|| Error::invalid_argument(format!("{operation} context range is reversed")))?;
if context.end > messages.num_elements() || context_length > 255 {
return Err(Error::invalid_argument(format!(
"{operation} context must be within messages and at most 255 bytes"
)));
}
let context_offset = u32::try_from(context.start)
.map_err(|_| Error::invalid_argument(format!("{operation} context offset exceeds u32")))?;
let context_length = u32::try_from(context_length)
.map_err(|_| Error::invalid_argument(format!("{operation} context length exceeds u32")))?;
let batch_size = validate_u32_vector(msg_offsets, "msg_offsets", None, operation)?;
validate_u32_vector(msg_lengths, "msg_lengths", Some(batch_size), operation)?;
validate_packed_u8(signatures, "signatures", batch_size, sig_bytes, operation)?;
validate_packed_u8(public_keys, "public_keys", batch_size, pk_bytes, operation)?;
let engine = messages.engine_handle();
for (label, matrix) in [
("msg_offsets", msg_offsets),
("msg_lengths", msg_lengths),
("signatures", signatures),
("public_keys", public_keys),
] {
if !engine.same_as(matrix.engine_handle()) {
return Err(Error::invalid_argument(format!(
"{operation} {label} must share an engine with messages"
)));
}
}
let batch_u32 = u32::try_from(batch_size)
.map_err(|_| Error::invalid_argument(format!("{operation} batch_size exceeds u32")))?;
let message_bytes = u32::try_from(messages.num_elements())
.map_err(|_| Error::invalid_argument(format!("{operation} message byte size exceeds u32")))?;
for matrix in [msg_offsets, msg_lengths, signatures, public_keys] {
let bytes = matrix
.num_elements()
.checked_mul(matrix.dtype().size_bytes())
.ok_or_else(|| Error::invalid_argument(format!("{operation} buffer byte size overflows")))?;
u32::try_from(bytes)
.map_err(|_| Error::invalid_argument(format!("{operation} buffer byte size exceeds u32")))?;
}
let results = Matrix::allocate_output(engine, vec![batch_size], batch_size, DType::U32)?;
if batch_size == 0 {
return Ok(results);
}
let buffers = [
BufferBinding::read(messages.storage()),
BufferBinding::read(msg_offsets.storage()),
BufferBinding::read(msg_lengths.storage()),
BufferBinding::read(signatures.storage()),
BufferBinding::read(public_keys.storage()),
BufferBinding::write(results.storage()),
];
let push_constants = [
PushConstant::U32(batch_u32),
PushConstant::U32(message_bytes),
PushConstant::U32(context_offset),
PushConstant::U32(context_length),
PushConstant::U32(algorithm),
];
let inputs = [messages, msg_offsets, msg_lengths, signatures, public_keys];
let outputs = [&results];
let attributes = [
OpAttribute::UnsignedInteger {
name: "batch_size".into(),
value: u64::from(batch_u32),
},
OpAttribute::UnsignedInteger {
name: "context_offset".into(),
value: u64::from(context_offset),
},
OpAttribute::UnsignedInteger {
name: "context_length".into(),
value: u64::from(context_length),
},
OpAttribute::UnsignedInteger {
name: "algorithm".into(),
value: u64::from(algorithm),
},
];
engine.record_semantic(
ComputeDispatch {
kernel,
buffers: &buffers,
push_constants: &push_constants,
workgroups: kernel.linear_workgroups(batch_u32),
},
SemanticDispatch {
contract,
inputs: &inputs,
outputs: &outputs,
attributes: &attributes,
},
)?;
Ok(results)
}
#[cfg_attr(
not(test),
expect(dead_code, reason = "private verifier awaits runtime admission")
)]
#[expect(
clippy::too_many_arguments,
reason = "private adapter preserves the five-buffer verifier contract"
)]
pub(crate) fn verify_hash_message_batch(
parameters: MlDsaParameters,
algorithm: u32,
messages: &Matrix,
msg_offsets: &Matrix,
msg_lengths: &Matrix,
signatures: &Matrix,
public_keys: &Matrix,
context: std::ops::Range<usize>,
) -> Result<Matrix> {
if KernelId::mldsa_prehash_bytes(algorithm).is_none() {
return Err(Error::invalid_argument("unknown ML-DSA prehash identifier"));
}
let (contract, kernel) = match parameters {
MlDsaParameters::MlDsa65 => (
crate::core::operation::cryptography::pqc::ML_DSA_65_VERIFY_HASH_MESSAGE_BATCH,
KernelId::CryptographyMlDsa65VerifyHashMessageU8,
),
MlDsaParameters::MlDsa44 => (
crate::core::operation::cryptography::pqc::ML_DSA_44_VERIFY_HASH_MESSAGE_BATCH,
KernelId::CryptographyMlDsa44VerifyHashMessageU8,
),
MlDsaParameters::MlDsa87 => (
crate::core::operation::cryptography::pqc::ML_DSA_87_VERIFY_HASH_MESSAGE_BATCH,
KernelId::CryptographyMlDsa87VerifyHashMessageU8,
),
};
let operation = contract.name();
let sig_bytes = parameters.signature_size();
let pk_bytes = parameters.public_key_size();
if messages.dtype() != DType::U8 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be a U8 matrix; found {}",
messages.dtype().token()
)));
}
if messages.shape().len() != 1 {
return Err(Error::invalid_argument(format!(
"{operation} messages must be rank-1 [total_bytes]; found {:?}",
messages.shape()
)));
}
let context_length = context
.end
.checked_sub(context.start)
.ok_or_else(|| Error::invalid_argument(format!("{operation} context range is reversed")))?;
if context.end > messages.num_elements() || context_length > 255 {
return Err(Error::invalid_argument(format!(
"{operation} context must be within messages and at most 255 bytes"
)));
}
let context_offset = u32::try_from(context.start)
.map_err(|_| Error::invalid_argument(format!("{operation} context offset exceeds u32")))?;
let context_length = u32::try_from(context_length)
.map_err(|_| Error::invalid_argument(format!("{operation} context length exceeds u32")))?;
let batch_size = validate_u32_vector(msg_offsets, "msg_offsets", None, operation)?;
validate_u32_vector(msg_lengths, "msg_lengths", Some(batch_size), operation)?;
validate_packed_u8(signatures, "signatures", batch_size, sig_bytes, operation)?;
validate_packed_u8(public_keys, "public_keys", batch_size, pk_bytes, operation)?;
let engine = messages.engine_handle();
for (label, matrix) in [
("msg_offsets", msg_offsets),
("msg_lengths", msg_lengths),
("signatures", signatures),
("public_keys", public_keys),
] {
if !engine.same_as(matrix.engine_handle()) {
return Err(Error::invalid_argument(format!(
"{operation} {label} must share an engine with messages"
)));
}
}
let batch_u32 = u32::try_from(batch_size)
.map_err(|_| Error::invalid_argument(format!("{operation} batch_size exceeds u32")))?;
let message_bytes = u32::try_from(messages.num_elements())
.map_err(|_| Error::invalid_argument(format!("{operation} message byte size exceeds u32")))?;
for matrix in [msg_offsets, msg_lengths, signatures, public_keys] {
let bytes = matrix
.num_elements()
.checked_mul(matrix.dtype().size_bytes())
.ok_or_else(|| Error::invalid_argument(format!("{operation} buffer byte size overflows")))?;
u32::try_from(bytes)
.map_err(|_| Error::invalid_argument(format!("{operation} buffer byte size exceeds u32")))?;
}
let results = Matrix::allocate_output(engine, vec![batch_size], batch_size, DType::U32)?;
if batch_size == 0 {
return Ok(results);
}
let buffers = [
BufferBinding::read(messages.storage()),
BufferBinding::read(msg_offsets.storage()),
BufferBinding::read(msg_lengths.storage()),
BufferBinding::read(signatures.storage()),
BufferBinding::read(public_keys.storage()),
BufferBinding::write(results.storage()),
];
let push_constants = [
PushConstant::U32(batch_u32),
PushConstant::U32(message_bytes),
PushConstant::U32(context_offset),
PushConstant::U32(context_length),
PushConstant::U32(algorithm),
];
let inputs = [messages, msg_offsets, msg_lengths, signatures, public_keys];
let outputs = [&results];
let attributes = [
OpAttribute::UnsignedInteger {
name: "batch_size".into(),
value: u64::from(batch_u32),
},
OpAttribute::UnsignedInteger {
name: "context_offset".into(),
value: u64::from(context_offset),
},
OpAttribute::UnsignedInteger {
name: "context_length".into(),
value: u64::from(context_length),
},
OpAttribute::UnsignedInteger {
name: "algorithm".into(),
value: u64::from(algorithm),
},
];
engine.record_semantic(
ComputeDispatch {
kernel,
buffers: &buffers,
push_constants: &push_constants,
workgroups: kernel.linear_workgroups(batch_u32),
},
SemanticDispatch {
contract,
inputs: &inputs,
outputs: &outputs,
attributes: &attributes,
},
)?;
Ok(results)
}
fn validate_u32_vector(
matrix: &Matrix,
label: &str,
expected_len: Option<usize>,
operation: &str,
) -> Result<usize> {
if matrix.dtype() != DType::U32 {
return Err(Error::invalid_argument(format!(
"{operation} {label} must be a U32 matrix; found {}",
matrix.dtype().token()
)));
}
let [len] = matrix.shape() else {
return Err(Error::invalid_argument(format!(
"{operation} {label} must be rank-1; found {:?}",
matrix.shape()
)));
};
if let Some(expected) = expected_len
&& *len != expected
{
return Err(Error::invalid_argument(format!(
"{operation} {label} length {len} does not match batch_size {expected}"
)));
}
Ok(*len)
}
fn validate_packed_u8(
matrix: &Matrix,
label: &str,
batch_size: usize,
stride: usize,
operation: &str,
) -> Result<()> {
if matrix.dtype() != DType::U8 {
return Err(Error::invalid_argument(format!(
"{operation} {label} must be a U8 matrix; found {}",
matrix.dtype().token()
)));
}
let expected = batch_size
.checked_mul(stride)
.ok_or_else(|| Error::invalid_argument(format!("{operation} {label} size overflows")))?;
let [actual] = matrix.shape() else {
return Err(Error::invalid_argument(format!(
"{operation} {label} must be rank-1 [batch * stride]; found {:?}",
matrix.shape()
)));
};
if *actual != expected {
return Err(Error::invalid_argument(format!(
"{operation} {label} must contain {expected} bytes for batch {batch_size}; found {actual}"
)));
}
Ok(())
}