use crate::Error;
use crate::provider::{Finalized, Finished, Prehash};
use crate::stream::{self, Sink, Source, StreamError};
use core::cmp::Ordering;
use structured_asn1::digest::{Digest, Resolution};
use structured_asn1::memory::{OwnedVec, SortError, Work};
pub type PassError<SE, PE> = StreamError<SE, DigestError<PE>>;
pub type DigestOutput<'a, P> =
Finalized<'a, P, <P as Digest<'a>>::Identity, <P as Digest<'a>>::Output>;
pub type DigestResult<T, E> = Result<T, DigestError<E>>;
pub type PassResult<T, SE, PE> = Result<T, PassError<SE, PE>>;
#[derive(Debug)]
pub enum DigestError<E> {
Processing(Error),
Callback(E),
OutputWidth {
expected: usize,
actual: usize,
},
Unstreamable,
}
fn started<S, E>(start: Result<Resolution<S>, E>) -> Result<S, DigestError<E>> {
match start.map_err(DigestError::Callback)? {
Resolution::Supported(state) => Ok(state),
Resolution::Unsupported | Resolution::Malformed => Err(DigestError::Unstreamable),
}
}
struct Pending<I, S> {
identity: I,
state: S,
}
struct Request<I> {
ordinal: usize,
identity: I,
}
pub struct DigestBatch<'a, P: Digest<'a>> {
provider: &'a P,
pending: OwnedVec<Pending<P::Identity, P::State>>,
mapping: OwnedVec<usize>,
field: &'static str,
at: u64,
}
impl<'a, P: Digest<'a>> DigestBatch<'a, P> {
pub fn prepare<W: Work>(
provider: &'a P,
identities: OwnedVec<P::Identity>,
execution: &W,
field: &'static str,
at: u64,
) -> Result<Self, DigestError<P::Error>> {
crate::provider::require_instance(provider, field, at).map_err(DigestError::Processing)?;
let refusal = |reason| DigestError::Processing(Error::resource(reason, field, at));
execution.charge_work(1).map_err(refusal)?;
let count = identities.len();
let mut ordered = execution.vector();
ordered
.try_reserve_exact(identities.len())
.map_err(refusal)?;
execution
.charge_work(core::mem::size_of_val(&*identities) as u64)
.map_err(refusal)?;
for (ordinal, identity) in identities.into_iter().enumerate() {
ordered
.try_push(Request { ordinal, identity })
.map_err(refusal)?;
}
ordered
.try_sort_unstable_by(|left, right| {
provider.compare(&left.identity, &right.identity, execution)
})
.map_err(|error| match error {
SortError::Resource(reason) => refusal(reason),
SortError::Comparison(error) => DigestError::Callback(error),
})?;
let mut pending: OwnedVec<Pending<P::Identity, P::State>, _> = execution.vector();
let mut mapping = execution.vector();
mapping.try_reserve_exact(count).map_err(refusal)?;
for _ in 0..count {
mapping.try_push(0).map_err(refusal)?;
}
execution
.charge_work(core::mem::size_of_val(&*ordered) as u64)
.map_err(refusal)?;
for Request { ordinal, identity } in ordered.into_owned() {
execution.charge_work(1).map_err(refusal)?;
if let Some(previous) = pending.last()
&& provider
.compare(&previous.identity, &identity, execution)
.map_err(DigestError::Callback)?
== Ordering::Equal
{
mapping[ordinal] = pending.len() - 1;
continue;
}
pending.try_reserve(1).map_err(refusal)?;
let state = started(provider.start(&identity, execution))?;
pending
.try_push(Pending { identity, state })
.map_err(refusal)?;
mapping[ordinal] = pending.len() - 1;
}
Ok(Self {
provider,
pending: pending.into_owned(),
mapping: mapping.into_owned(),
field,
at,
})
}
#[must_use]
pub fn len(&self) -> usize {
self.pending.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.pending.is_empty()
}
pub fn run<S: Source, W: Work>(
mut self,
source: &mut S,
execution: &W,
) -> Result<DigestMemo<'a, P>, PassError<S::Error, P::Error>> {
let field = self.field;
let at = self.at;
let mut outputs = execution.vector();
outputs
.try_reserve_exact(self.pending.len())
.map_err(|reason| StreamError::Processing(Error::resource(reason, field, at)))?;
stream::run(source, &mut BatchSink(&mut self), execution, field, at)?;
self.finish(outputs, execution)
.map_err(StreamError::Callback)
}
pub(crate) fn feed<W: Work>(
&mut self,
bytes: &[u8],
execution: &W,
) -> Result<(), DigestError<P::Error>> {
for pending in &mut self.pending {
execution.charge_work(1).map_err(|reason| {
DigestError::Processing(Error::resource(reason, self.field, self.at))
})?;
self.provider
.update(&mut pending.state, bytes, execution)
.map_err(DigestError::Callback)?;
}
Ok(())
}
pub(crate) fn finish<W: Work>(
self,
mut outputs: OwnedVec<DigestOutput<'a, P>, W>,
execution: &W,
) -> Result<DigestMemo<'a, P>, DigestError<P::Error>> {
let refusal =
|reason| DigestError::Processing(Error::resource(reason, self.field, self.at));
execution
.charge_work(core::mem::size_of_val(&*self.pending) as u64)
.map_err(refusal)?;
for pending in self.pending {
execution.charge_work(1).map_err(refusal)?;
let output = self
.provider
.finish(pending.state, execution)
.map_err(DigestError::Callback)?;
if let Some(expected) = self.provider.width(&pending.identity) {
let actual = output.as_ref().len();
if actual != expected {
return Err(DigestError::OutputWidth { expected, actual });
}
}
outputs
.try_push(Finalized::new(self.provider, pending.identity, output))
.map_err(refusal)?;
}
Ok(DigestMemo {
provider: self.provider,
outputs: outputs.into_owned(),
mapping: self.mapping,
field: self.field,
at: self.at,
})
}
pub(crate) fn provider(&self) -> &'a P {
self.provider
}
pub(crate) fn location(&self) -> (&'static str, u64) {
(self.field, self.at)
}
}
struct BatchSink<'s, 'a, P: Digest<'a>>(&'s mut DigestBatch<'a, P>);
impl<'a, P: Digest<'a>> Sink for BatchSink<'_, 'a, P> {
type Error = DigestError<P::Error>;
fn write<W: Work>(&mut self, bytes: &[u8], execution: &W) -> Result<(), Self::Error> {
self.0.feed(bytes, execution)
}
}
pub struct DigestMemo<'a, P: Digest<'a>> {
provider: &'a P,
outputs: OwnedVec<DigestOutput<'a, P>>,
mapping: OwnedVec<usize>,
field: &'static str,
at: u64,
}
pub struct RelatedDigest<'m, 'a, P: Digest<'a>> {
provider: &'a P,
identity: &'m P::Identity,
output: P::Output,
}
impl<'a, P: Digest<'a>> RelatedDigest<'_, 'a, P> {
#[must_use]
pub fn prehash(&self) -> Prehash<'_, '_, P, P::Identity> {
Prehash::new(self.provider, self.identity, self.output.as_ref())
}
#[must_use]
pub fn finished(&self) -> Finished<'_, P, P::Identity, P::Output> {
Finished::new(self.provider, self.identity, &self.output)
}
}
struct RelatedSink<'s, 'a, P: Digest<'a>> {
provider: &'a P,
state: &'s mut P::State,
}
impl<'a, P: Digest<'a>> Sink for RelatedSink<'_, 'a, P> {
type Error = DigestError<P::Error>;
fn write<W: Work>(&mut self, bytes: &[u8], execution: &W) -> Result<(), Self::Error> {
self.provider
.update(self.state, bytes, execution)
.map_err(DigestError::Callback)
}
}
impl<'a, P: Digest<'a>> DigestMemo<'a, P> {
pub fn digest_related<'m, S: Source, W: Work>(
&'m self,
ordinal: usize,
source: &mut S,
execution: &W,
field: &'static str,
at: u64,
) -> PassResult<Option<RelatedDigest<'m, 'a, P>>, S::Error, P::Error> {
let Some(existing) = self
.for_request(ordinal, execution)
.map_err(StreamError::Callback)?
else {
return Ok(None);
};
let identity = existing.identity();
let mut state =
started(self.provider.start(identity, execution)).map_err(StreamError::Callback)?;
stream::run(
source,
&mut RelatedSink {
provider: self.provider,
state: &mut state,
},
execution,
field,
at,
)?;
let output = self
.provider
.finish(state, execution)
.map_err(|error| StreamError::Callback(DigestError::Callback(error)))?;
if let Some(expected) = self.provider.width(identity) {
let actual = output.as_ref().len();
if actual != expected {
return Err(StreamError::Callback(DigestError::OutputWidth {
expected,
actual,
}));
}
}
Ok(Some(RelatedDigest {
provider: self.provider,
identity,
output,
}))
}
#[must_use]
pub fn len(&self) -> usize {
self.outputs.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.outputs.is_empty()
}
#[must_use]
pub fn request_count(&self) -> usize {
self.mapping.len()
}
pub fn for_request<W: Work>(
&self,
ordinal: usize,
execution: &W,
) -> DigestResult<Option<&DigestOutput<'a, P>>, P::Error> {
execution.charge_work(1).map_err(|reason| {
DigestError::Processing(Error::resource(reason, self.field, self.at))
})?;
let Some(index) = self.mapping.get(ordinal) else {
return Ok(None);
};
debug_assert!(*index < self.outputs.len());
Ok(Some(&self.outputs[*index]))
}
pub fn find<W: Work>(
&self,
identity: &P::Identity,
execution: &W,
) -> DigestResult<Option<&DigestOutput<'a, P>>, P::Error> {
let refusal =
|reason| DigestError::Processing(Error::resource(reason, self.field, self.at));
execution.charge_work(1).map_err(refusal)?;
let mut from = 0;
let mut to = self.outputs.len();
while from < to {
execution.charge_work(1).map_err(refusal)?;
let middle = from + (to - from) / 2;
let held = &self.outputs[middle];
match self
.provider
.compare(held.identity(), identity, execution)
.map_err(DigestError::Callback)?
{
Ordering::Equal => return Ok(Some(held)),
Ordering::Less => from = middle + 1,
Ordering::Greater => to = middle,
}
}
Ok(None)
}
}