use alloc::vec::Vec;
use crate::{
try_compact_note_decryption_inner, try_note_decryption_inner, BatchDomain, EphemeralKeyBytes,
ShieldedOutput, COMPACT_NOTE_SIZE, ENC_CIPHERTEXT_SIZE,
};
#[allow(clippy::type_complexity)]
pub fn try_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, ENC_CIPHERTEXT_SIZE>>(
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
) -> Vec<Option<((D::Note, D::Recipient, D::Memo), usize)>> {
batch_note_decryption(ivks, outputs, try_note_decryption_inner)
}
#[allow(clippy::type_complexity)]
pub fn try_compact_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, COMPACT_NOTE_SIZE>>(
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
) -> Vec<Option<((D::Note, D::Recipient), usize)>> {
batch_note_decryption(ivks, outputs, try_compact_note_decryption_inner)
}
fn batch_note_decryption<D: BatchDomain, Output: ShieldedOutput<D, CS>, F, FR, const CS: usize>(
ivks: &[D::IncomingViewingKey],
outputs: &[(D, Output)],
decrypt_inner: F,
) -> Vec<Option<(FR, usize)>>
where
F: Fn(&D, &D::IncomingViewingKey, &EphemeralKeyBytes, &Output, &D::SymmetricKey) -> Option<FR>,
{
if ivks.is_empty() {
return (0..outputs.len()).map(|_| None).collect();
};
let ephemeral_keys = D::batch_epk(outputs.iter().map(|(_, output)| output.ephemeral_key()));
let items = ephemeral_keys.iter().flat_map(|(epk, ephemeral_key)| {
ivks.iter().map(move |ivk| {
(
epk.as_ref().map(|epk| D::ka_agree_dec(ivk, epk)),
ephemeral_key,
)
})
});
let keys = D::batch_kdf(items);
keys.chunks(ivks.len())
.zip(ephemeral_keys.iter().zip(outputs.iter()))
.map(|(key_chunk, ((_, ephemeral_key), (domain, output)))| {
key_chunk
.iter()
.zip(ivks.iter().enumerate())
.filter_map(|(key, (i, ivk))| {
key.as_ref()
.and_then(|key| decrypt_inner(domain, ivk, ephemeral_key, output, key))
.map(|out| (out, i))
})
.next()
})
.collect::<Vec<Option<_>>>()
}