willow-data-model 0.7.0

use compact_u64::*;
use ufotofu::{ConsumeAtLeastError, codec_prelude::*};

use super::*;

/// Essentially how to encode a Path, but working with an arbitrary iterator of components. Path encoding consists of calling this directly, relative path encoding consists of first encoding the length of the greatest common suffix and *then* calling this.
async fn encode_from_iterator_of_components<'a, const MCL: usize, C, I>(
    consumer: &mut C,
    path_length: u64,
    component_count: u64,
    components: I,
) -> Result<(), C::Error>
where
    C: BulkConsumer<Item = u8> + ?Sized,
    I: Iterator<Item = &'a Component<MCL>>,
{
    // First byte contains two 4-bit tags of `CompactU64`s:
    // total path length in bytes: 4 bit tag at offset zero
    // total number of components: 4 bit tag at offset four
    let mut header = 0;
    write_tag(&mut header, 4, 0, path_length);
    write_tag(&mut header, 4, 4, component_count);
    consumer.consume_item(header).await?;

    // Next, encode the total path length in compact bytes.
    cu64_encode(path_length, 4, consumer).await?;

    // Next, encode the  total number of components in compact bytes.
    cu64_encode(component_count, 4, consumer).await?;

    // Then, encode the components. Each is prefixed by its length as an 8-bit-tag CompactU64, except for the final component.
    for (i, component) in components.enumerate() {
        // The length of the final component is omitted (because a decoder can infer it from the total length and all prior components' lengths).
        if i as u64 + 1 != component_count {
            cu64_encode_standalone(component.len() as u64, consumer).await?;
        }

        // Each component length (if any) is followed by the raw component data itself.
        consumer
            .bulk_consume_full_slice(component.as_ref())
            .await
            .map_err(ConsumeAtLeastError::into_reason)?;
    }

    Ok(())
}

/// Implements [encode_path](https://willowprotocol.org/specs/encodings/index.html#encode_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize> Encodable for Path<MCL, MCC, MPL> {
    async fn encode<C>(&self, consumer: &mut C) -> Result<(), C::Error>
    where
        C: BulkConsumer<Item = u8> + ?Sized,
    {
        encode_from_iterator_of_components::<MCL, _, _>(
            consumer,
            self.total_length() as u64,
            self.component_count() as u64,
            self.components(),
        )
        .await
    }
}

// Decodes the path length and component count as expected at the start of a path encoding, generic over whether the encoding must be canonic or not.
// Implemented as a dedicated function so that it can be used in both absolute and relative decoding.
async fn decode_total_length_and_component_count_maybe_canonic<const CANONIC: bool, P>(
    producer: &mut P,
) -> Result<(usize, usize), DecodeError<P::Final, P::Error, Blame>>
where
    P: BulkProducer<Item = u8> + ?Sized,
{
    // Decode the first byte - the two compact width tags for the path length and component count.
    let header = producer.produce_item().await?;

    // Next, decode the total path length and the component count.
    let total_length = if CANONIC {
        cu64_decode_canonic(header, 4, 0, producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    } else {
        cu64_decode(header, 4, 0, producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    };

    let component_count = if CANONIC {
        cu64_decode_canonic(header, 4, 4, producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    } else {
        cu64_decode(header, 4, 4, producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    };

    // Convert them from u64 to usize, error if usize cannot represent the number.
    let total_length = Blame::u64_to_usize(total_length)?;
    let component_count = Blame::u64_to_usize(component_count)?;

    Ok((total_length, component_count))
}

// Decodes the components of a path encoding, generic over whether the encoding must be canonic or not. Appends them into a PathBuilder. Needs to know the total length of components that had already been appended to that PathBuilder before.
// Implemented as a dedicated function so that it can be used in both absolute and relative decoding.
async fn decode_components_maybe_canonic<
    const CANONIC: bool,
    const MCL: usize,
    const MCC: usize,
    const MPL: usize,
    P,
>(
    producer: &mut P,
    mut builder: PathBuilder<MCL, MCC, MPL>,
    initial_accumulated_component_length: usize,
    remaining_component_count: usize,
    expected_total_length: usize,
) -> Result<Path<MCL, MCC, MPL>, DecodeError<P::Final, P::Error, Blame>>
where
    P: BulkProducer<Item = u8> + ?Sized,
{
    // Decode the actual components.
    // We track the sum of the lengths of all decoded components so far, because we need it to determine the length of the final component.
    let mut accumulated_component_length = initial_accumulated_component_length;

    // Handle decoding of the empty path with dedicated logic to prevent underflows in loop counters =S
    if remaining_component_count == 0 {
        if expected_total_length > accumulated_component_length {
            // Claimed length is incorrect
            Err(DecodeError::Other(Blame::TheirFault))
        } else {
            // Nothing more to do, decoding an empty path turns out to be simple!
            Ok(builder.build())
        }
    } else {
        // We have at least one component.

        // Decode all but the final one (because the final one is encoded without its lenght and hence requires dedicated logic to decode).
        for _ in 1..remaining_component_count {
            let component_len = Blame::u64_to_usize(if CANONIC {
                cu64_decode_canonic_standalone(producer)
                    .await
                    .map_err(|err| err.map_other(|_| Blame::TheirFault))?
            } else {
                cu64_decode_standalone(producer)
                    .await
                    .map_err(|err| err.map_other(|_| Blame::TheirFault))?
            })?;

            if component_len > MCL {
                // Decoded path must respect the MCL.
                return Err(DecodeError::Other(Blame::TheirFault));
            } else {
                // Increase the accumulated length, accounting for errors.
                accumulated_component_length = accumulated_component_length
                    .checked_add(component_len)
                    .ok_or(DecodeError::Other(Blame::TheirFault))?;

                // Copy the component bytes into the Path.
                builder
                    .append_component_from_bulk_producer(component_len, producer)
                    .await?;
            }
        }

        // For the final component, compute its length. If the computation result would be negative, then the encoding was invalid.
        let final_component_length = expected_total_length
            .checked_sub(accumulated_component_length)
            .ok_or(DecodeError::Other(Blame::TheirFault))?;

        if final_component_length > MCL {
            // Decoded path must respect the MCL.
            Err(DecodeError::Other(Blame::TheirFault))
        } else {
            // Copy the final component bytes into the Path.
            builder
                .append_component_from_bulk_producer(final_component_length, producer)
                .await?;

            // What a journey. We are done!
            Ok(builder.build())
        }
    }
}

// Decodes a path, generic over whether the encoding must be canonic or not.
async fn decode_maybe_canonic<
    const CANONIC: bool,
    const MCL: usize,
    const MCC: usize,
    const MPL: usize,
    P,
>(
    producer: &mut P,
) -> Result<Path<MCL, MCC, MPL>, DecodeError<P::Final, P::Error, Blame>>
where
    P: BulkProducer<Item = u8> + ?Sized,
{
    let (total_length, component_count) =
        decode_total_length_and_component_count_maybe_canonic::<CANONIC, _>(producer).await?;

    // Preallocate all storage for the path. Error if the total length or component_count are greater than MPL and MCC respectively allow.
    let builder = PathBuilder::new(total_length, component_count)
        .map_err(|_| DecodeError::Other(Blame::TheirFault))?;

    decode_components_maybe_canonic::<CANONIC, MCL, MCC, MPL, _>(
        producer,
        builder,
        0,
        component_count,
        total_length,
    )
    .await
}

/// Implements [EncodePath](https://willowprotocol.org/specs/encodings/index.html#EncodePath).
impl<const MCL: usize, const MCC: usize, const MPL: usize> Decodable for Path<MCL, MCC, MPL> {
    type ErrorReason = Blame;

    async fn decode<P>(
        producer: &mut P,
    ) -> Result<Self, DecodeError<P::Final, P::Error, Self::ErrorReason>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        decode_maybe_canonic::<false, MCL, MCC, MPL, _>(producer).await
    }
}

/// Implements [encode_path](https://willowprotocol.org/specs/encodings/index.html#encode_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize> DecodableCanonic
    for Path<MCL, MCC, MPL>
{
    type ErrorCanonic = Blame;

    async fn decode_canonic<P>(
        producer: &mut P,
    ) -> Result<Self, DecodeError<P::Final, P::Error, Self::ErrorCanonic>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        decode_maybe_canonic::<true, MCL, MCC, MPL, _>(producer).await
    }
}

// Separate function to allow for reuse in relative encoding.
fn encoding_len_from_iterator_of_components<'a, const MCL: usize, I>(
    path_length: u64,
    component_count: usize,
    components: I,
) -> usize
where
    I: Iterator<Item = &'a Component<MCL>>,
{
    let mut total_enc_len = 1; // First byte for the two four-bit tags at the start of the encoding.

    total_enc_len += cu64_len_of_encoding(4, path_length);
    total_enc_len += cu64_len_of_encoding(4, component_count as u64);

    for (i, comp) in components.enumerate() {
        if i + 1 < component_count {
            total_enc_len += cu64_len_of_encoding(8, comp.len() as u64) + 1;
        }

        total_enc_len += comp.len();
    }

    total_enc_len
}

/// Implements [encode_path](https://willowprotocol.org/specs/encodings/index.html#encode_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize> EncodableKnownLength
    for Path<MCL, MCC, MPL>
{
    fn len_of_encoding(&self) -> usize {
        encoding_len_from_iterator_of_components::<MCL, _>(
            self.total_length() as u64,
            self.component_count(),
            self.components(),
        )
    }
}

////////////////////////////////////
// Relative encoding path <> path //
////////////////////////////////////

/// Implements [path_rel_path](https://willowprotocol.org/specs/encodings/index.html#path_rel_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize> RelativeEncodable<Path<MCL, MCC, MPL>>
    for Path<MCL, MCC, MPL>
{
    async fn relative_encode<Consumer>(
        &self,
        rel: &Path<MCL, MCC, MPL>,
        consumer: &mut Consumer,
    ) -> Result<(), Consumer::Error>
    where
        Consumer: BulkConsumer<Item = u8> + ?Sized,
    {
        let lcp = self.longest_common_prefix(rel);

        cu64_encode_standalone(lcp.component_count() as u64, consumer).await?;

        let suffix_length = self.total_length() - lcp.total_length();
        let suffix_component_count = self.component_count() - lcp.component_count();

        encode_from_iterator_of_components::<MCL, _, _>(
            consumer,
            suffix_length as u64,
            suffix_component_count as u64,
            self.suffix_components(lcp.component_count()),
        )
        .await
    }

    /// Any path can be encoded relative to every path.
    fn can_be_encoded_relative_to(&self, _rel: &Path<MCL, MCC, MPL>) -> bool {
        true
    }
}

// Decodes a path relative to another path, generic over whether the encoding must be canonic or not.
async fn relative_decode_maybe_canonic<
    const CANONIC: bool,
    const MCL: usize,
    const MCC: usize,
    const MPL: usize,
    P,
>(
    producer: &mut P,
    rel: &Path<MCL, MCC, MPL>,
) -> Result<Path<MCL, MCC, MPL>, DecodeError<P::Final, P::Error, Blame>>
where
    P: BulkProducer<Item = u8> + ?Sized,
{
    let prefix_component_count = Blame::u64_to_usize(if CANONIC {
        cu64_decode_canonic_standalone(producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    } else {
        cu64_decode_standalone(producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    })?;

    let (suffix_length, suffix_component_count) =
        decode_total_length_and_component_count_maybe_canonic::<CANONIC, _>(producer).await?;

    if prefix_component_count > rel.component_count() {
        return Err(DecodeError::Other(Blame::TheirFault));
    }

    let prefix_path_length = rel.total_length_of_prefix(prefix_component_count);

    let total_length = prefix_path_length
        .checked_add(suffix_length)
        .ok_or(DecodeError::Other(Blame::TheirFault))?;
    let total_component_count = prefix_component_count
        .checked_add(suffix_component_count)
        .ok_or(DecodeError::Other(Blame::TheirFault))?;

    // Preallocate all storage for the path. Error if the total length or component_count are greater than MPL and MCC respectively allow.
    let builder = PathBuilder::new_from_prefix(
        total_length,
        total_component_count,
        rel,
        prefix_component_count,
    )
    .map_err(|_| DecodeError::Other(Blame::TheirFault))?;

    // Decode the remaining components, add them to the builder, then build.
    let decoded = decode_components_maybe_canonic::<CANONIC, MCL, MCC, MPL, _>(
        producer,
        builder,
        prefix_path_length,
        suffix_component_count,
        total_length,
    )
    .await?;

    if CANONIC {
        // Did the encoding use the *longest* common prefix?
        if prefix_component_count == rel.component_count() {
            // Could not have taken a longer prefix of `r`, i.e., the prefix was maximal.
            Ok(decoded)
        } else if prefix_component_count == decoded.component_count() {
            // The prefix was the full path to decode, so it clearly was chosen maximally.
            Ok(decoded)
        } else {
            // We check whether the next-longer prefix of `r` could have also been used for encoding. If so, error.
            // To efficiently check, we check whether the next component of `r` is equal to its counterpart in what we decoded.
            // Both next components exist, otherwise we would have been in an earlier branch of the `if` expression.
            if rel.component(prefix_component_count).unwrap()
                == decoded.component(prefix_component_count).unwrap()
            {
                // Could have used a longer prefix for decoding. Not canonic!
                Err(DecodeError::Other(Blame::TheirFault))
            } else {
                // Encoding was minimal, yay =)
                Ok(decoded)
            }
        }
    } else {
        // No additional canonicity checks needed.
        Ok(decoded)
    }
}

/// Implements [EncodePathRelativePath](https://willowprotocol.org/specs/encodings/index.html#EncodePathRelativePath).
impl<const MCL: usize, const MCC: usize, const MPL: usize> RelativeDecodable<Path<MCL, MCC, MPL>>
    for Path<MCL, MCC, MPL>
{
    type ErrorReason = Blame;

    async fn relative_decode<P>(
        rel: &Path<MCL, MCC, MPL>,
        producer: &mut P,
    ) -> Result<Self, DecodeError<P::Final, P::Error, Blame>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        relative_decode_maybe_canonic::<false, MCL, MCC, MPL, _>(producer, rel).await
    }
}

/// Implements [path_relative_path](https://willowprotocol.org/specs/encodings/index.html#path_rel_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize>
    RelativeDecodableCanonic<Path<MCL, MCC, MPL>> for Path<MCL, MCC, MPL>
{
    type ErrorCanonic = Blame;

    async fn relative_decode_canonic<P>(
        rel: &Path<MCL, MCC, MPL>,
        producer: &mut P,
    ) -> Result<Self, DecodeError<P::Final, P::Error, Blame>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        relative_decode_maybe_canonic::<true, MCL, MCC, MPL, _>(producer, rel).await
    }
}

/// /// Implements [path_rel_path](https://willowprotocol.org/specs/encodings/index.html#path_rel_path).
impl<const MCL: usize, const MCC: usize, const MPL: usize>
    RelativeEncodableKnownLength<Path<MCL, MCC, MPL>> for Path<MCL, MCC, MPL>
{
    fn len_of_relative_encoding(&self, rel: &Path<MCL, MCC, MPL>) -> usize {
        let lcp = self.longest_common_prefix(rel);
        let path_len_of_suffix = self.total_length() - lcp.total_length();
        let component_count_of_suffix = self.component_count() - lcp.component_count();

        let mut total_enc_len = 0;

        // Number of components in the longest common prefix, encoded as a CompactU64 with an 8-bit tag.
        total_enc_len += cu64_len_of_encoding(8, lcp.component_count() as u64) + 1;

        total_enc_len += encoding_len_from_iterator_of_components::<MCL, _>(
            path_len_of_suffix as u64,
            component_count_of_suffix,
            self.suffix_components(lcp.component_count()),
        );

        total_enc_len
    }
}

////////////////////////////////
// path extends path encoding //
////////////////////////////////

/// Implementations of the [EncodePathExtendsPath](https://willowprotocol.org/specs/encodings/index.html#EncodePathExtendsPath) encoding relation.
pub mod path_extends_path {
    use super::*;

    #[cfg(feature = "dev")]
    use arbitrary::Arbitrary;

    /// Implements encoding for the [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path) encoding function.
    pub async fn encode_path_extends_path<const MCL: usize, const MCC: usize, const MPL: usize, C>(
        path: &Path<MCL, MCC, MPL>,
        prefix: &Path<MCL, MCC, MPL>,
        consumer: &mut C,
    ) -> Result<(), C::Error>
    where
        C: BulkConsumer<Item = u8> + ?Sized,
    {
        if !path.is_prefixed_by(prefix) {
            panic!("Tried to encode relative to a non-prefix with PathExtendsPath");
        }

        let extends_count = prefix.component_count();

        let path_len = path.total_length() - prefix.total_length();
        let diff = path.component_count() - extends_count;

        encode_from_iterator_of_components(
            consumer,
            path_len as u64,
            diff as u64,
            path.suffix_components(extends_count),
        )
        .await?;

        Ok(())
    }

    /// Returns the length of the encoding of `path` relative to `prefix` according to the [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path) encoding function.
    pub fn path_extends_path_encoding_len<const MCL: usize, const MCC: usize, const MPL: usize>(
        path: &Path<MCL, MCC, MPL>,
        prefix: &Path<MCL, MCC, MPL>,
    ) -> usize {
        let prefix_count = prefix.component_count();

        let path_len = path.total_length() - prefix.total_length();
        let diff = path.component_count() - prefix_count;

        encoding_len_from_iterator_of_components(
            path_len as u64,
            diff,
            path.suffix_components(prefix_count),
        )
    }

    /// Implements decoding for the [EncodePathExtendsPath](https://willowprotocol.org/specs/encodings/index.html#EncodePathExtendsPath) encoding relation.
    pub async fn decode_path_extends_path<const MCL: usize, const MCC: usize, const MPL: usize, P>(
        prefix: &Path<MCL, MCC, MPL>,
        producer: &mut P,
    ) -> Result<Path<MCL, MCC, MPL>, DecodeError<P::Final, P::Error, Blame>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        let suffix = Path::<MCL, MCC, MPL>::decode(producer)
            .await
            .map_err(|err| err.map_other(From::from))?;

        let prefix_count = prefix.component_count();

        let total_length = prefix.total_length() + suffix.total_length();
        let total_count = prefix_count + suffix.component_count();

        let mut path_builder =
            PathBuilder::new_from_prefix(total_length, total_count, prefix, prefix_count)
                .map_err(|_err| DecodeError::Other(Blame::TheirFault))?;

        for component in suffix.components() {
            path_builder.append_component(component);
        }

        Ok(path_builder.build())
    }

    /// Implements canonic decoding for the [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path) encoding function.
    pub async fn decode_path_extends_path_canonic<
        const MCL: usize,
        const MCC: usize,
        const MPL: usize,
        P,
    >(
        prefix: &Path<MCL, MCC, MPL>,
        producer: &mut P,
    ) -> Result<Path<MCL, MCC, MPL>, DecodeError<P::Final, P::Error, Blame>>
    where
        P: BulkProducer<Item = u8> + ?Sized,
    {
        let suffix = Path::<MCL, MCC, MPL>::decode_canonic(producer)
            .await
            .map_err(|err| err.map_other(From::from))?;

        let prefix_count = prefix.component_count();

        let total_length = prefix.total_length() + suffix.total_length();
        let total_count = prefix_count + suffix.component_count();

        let mut path_builder =
            PathBuilder::new_from_prefix(total_length, total_count, prefix, prefix_count)
                .map_err(|_err| DecodeError::Other(Blame::TheirFault))?;

        for component in suffix.components() {
            path_builder.append_component(component);
        }

        Ok(path_builder.build())
    }

    /// A wrapper type around [`Path`] whose implementations of the [codec_relative] traits implement the [EncodePathExtendsPath](https://willowprotocol.org/specs/encodings/index.html#encsec_EncodePathExtendsPath) encoding relation.
    #[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Default, Debug)]
    #[cfg_attr(feature = "dev", derive(Arbitrary))]
    pub struct CodecPathExtendsPath<const MCL: usize, const MCC: usize, const MPL: usize>(
        pub Path<MCL, MCC, MPL>,
    );

    /// Implements [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path).
    impl<const MCL: usize, const MCC: usize, const MPL: usize>
        RelativeEncodable<Path<MCL, MCC, MPL>> for CodecPathExtendsPath<MCL, MCC, MPL>
    {
        async fn relative_encode<C>(
            &self,
            rel: &Path<MCL, MCC, MPL>,
            consumer: &mut C,
        ) -> Result<(), C::Error>
        where
            C: BulkConsumer<Item = u8> + ?Sized,
        {
            encode_path_extends_path(&self.0, rel, consumer).await
        }

        /// Returns `true` iff `rel` is a prefix of `self`.
        fn can_be_encoded_relative_to(&self, rel: &Path<MCL, MCC, MPL>) -> bool {
            rel.is_prefix_of(&self.0)
        }
    }

    /// Implements [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path).
    impl<const MCL: usize, const MCC: usize, const MPL: usize>
        RelativeEncodableKnownLength<Path<MCL, MCC, MPL>> for CodecPathExtendsPath<MCL, MCC, MPL>
    {
        fn len_of_relative_encoding(&self, rel: &Path<MCL, MCC, MPL>) -> usize {
            path_extends_path_encoding_len(&self.0, rel)
        }
    }

    /// Implements [EncodePathExtendsPath](https://willowprotocol.org/specs/encodings/index.html#EncodePathExtendsPath).
    impl<const MCL: usize, const MCC: usize, const MPL: usize>
        RelativeDecodable<Path<MCL, MCC, MPL>> for CodecPathExtendsPath<MCL, MCC, MPL>
    {
        type ErrorReason = Blame;

        async fn relative_decode<P>(
            rel: &Path<MCL, MCC, MPL>,
            producer: &mut P,
        ) -> Result<Self, DecodeError<P::Final, P::Error, Self::ErrorReason>>
        where
            P: BulkProducer<Item = u8> + ?Sized,
            Self: Sized,
        {
            Ok(Self(decode_path_extends_path(rel, producer).await?))
        }
    }

    /// Implements [path_extends_path](https://willowprotocol.org/specs/encodings/index.html#path_extends_path).
    impl<const MCL: usize, const MCC: usize, const MPL: usize>
        RelativeDecodableCanonic<Path<MCL, MCC, MPL>> for CodecPathExtendsPath<MCL, MCC, MPL>
    {
        type ErrorCanonic = Blame;

        async fn relative_decode_canonic<P>(
            rel: &Path<MCL, MCC, MPL>,
            producer: &mut P,
        ) -> Result<Self, DecodeError<P::Final, P::Error, Self::ErrorReason>>
        where
            P: BulkProducer<Item = u8> + ?Sized,
            Self: Sized,
        {
            Ok(Self(decode_path_extends_path_canonic(rel, producer).await?))
        }
    }
}