willow25 0.7.3

use core::cmp::min;
use core::fmt::Debug;
use core::hash::Hash;
use core::ops::Bound;
use core::{cmp::Ordering, ops::RangeBounds};

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

use order_theory::{GreatestElement, UpperSemilattice};

use compact_u64::*;
use ufotofu::codec_prelude::*;

use crate::is_bitflagged;
use crate::paths::path_extends_path::*;
use crate::prelude::*;

/// An [`Area`](https://willowprotocol.org/specs/grouping-entries/#Area) is a box in three-dimensional willow space, consisting of all entries matching either a single [`SubspaceId`] or entries of arbitrary [`SubspaceIds`](SubspaceId), prefixed by some [`Path`], and a [`TimeRange`](super::TimeRange).
///
/// Areas are the default way by which application developers should aggregate entries. See [the specification](https://willowprotocol.org/specs/grouping-entries/#areas) for more details.
///
/// ```
/// use willow25::prelude::*;
///
/// let a1 = Area::new(None, Path::new(), TimeRange::new_closed(0.into(), 17.into()));
///
/// assert!(a1.includes(&([5; 32].into(), Path::new(), Timestamp::from(9))));
/// assert_eq!(a1.subspace(), None);
///
/// let a2 = Area::new(Some([42; 32].into()), Path::new(), TimeRange::new_open(15.into()));
/// assert_eq!(
///     a1.intersection(&a2),
///     Ok(Area::new(
///         Some([42; 32].into()),
///         Path::new(),
///         TimeRange::new_closed(15.into(), 17.into()),
///     )),
/// );
/// ```
///
/// [Specification](https://willowprotocol.org/specs/grouping-entries/#Area)
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
#[cfg_attr(feature = "dev", derive(Arbitrary))]
pub struct Area {
    subspace: Option<SubspaceId>,
    path: Path,
    times: TimeRange,
}

impl Grouping for Area {
    fn includes<Coord>(&self, coord: &Coord) -> bool
    where
        Coord: Coordinatelike + ?Sized,
    {
        self.times().contains(&coord.timestamp())
            && self
                .subspace()
                .map(|subspace_id| subspace_id == coord.subspace_id())
                .unwrap_or(true)
            && coord.path().is_prefixed_by(self.path())
    }

    fn intersection(&self, other: &Self) -> Result<Self, EmptyGrouping> {
        if self.path.is_related_to(other.path()) {
            let times = self.times().intersection_willow_range(other.times())?;
            let path = self.path().least_upper_bound(other.path());

            match (self.subspace(), other.subspace()) {
                (None, None) => Ok(Self::new(None, path, times)),
                (Some(subspace), None) | (None, Some(subspace)) => {
                    Ok(Area::new(Some(subspace.clone()), path, times))
                }
                (Some(self_subspace), Some(other_subspace)) => {
                    if self_subspace == other_subspace {
                        Ok(Area::new(Some(self_subspace.clone()), path, times))
                    } else {
                        Err(EmptyGrouping)
                    }
                }
            }
        } else {
            Err(EmptyGrouping)
        }
    }
}

impl RangeBounds<SubspaceId> for Area {
    fn start_bound(&self) -> Bound<&SubspaceId> {
        match self.subspace() {
            None => Bound::Unbounded,
            Some(s) => Bound::Included(s),
        }
    }

    fn end_bound(&self) -> Bound<&SubspaceId> {
        match self.subspace() {
            None => Bound::Unbounded,
            Some(s) => Bound::Included(s),
        }
    }
}

impl RangeBounds<Timestamp> for Area {
    fn start_bound(&self) -> Bound<&Timestamp> {
        self.times().start_bound()
    }

    fn end_bound(&self) -> Bound<&Timestamp> {
        self.times().end_bound()
    }
}

impl Area {
    /// Creates a new `Area` from its constituent optional [`SubspaceId`], [`Path`], and [`TimeRange`](super::TimeRange).
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::new(None, Path::new(), TimeRange::new_closed(0.into(), 17.into()));
    ///
    /// assert!(a.includes(&([5; 32].into(), Path::new(), Timestamp::from(9))));
    /// assert_eq!(a.subspace(), None);
    /// ```
    pub fn new(subspace: Option<SubspaceId>, path: Path, times: TimeRange) -> Self {
        Self {
            subspace,
            path,
            times,
        }
    }

    /// Returns a reference to the inner [`SubspaceId`], if any.
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a1 = Area::new(
    ///     Some([17; 32].into()),
    ///     Path::new(),
    ///     TimeRange::new_closed(0.into(), 17.into()),
    /// );
    /// assert_eq!(a1.subspace(), Some(&[17; 32].into()));
    ///
    /// let a2 = Area::new(None, Path::new(), TimeRange::new_closed(0.into(), 17.into()));
    /// assert_eq!(a2.subspace(), None);
    /// ```
    ///
    /// [Definition](https://willowprotocol.org/specs/grouping-entries/#AreaSubspace).
    pub fn subspace(&self) -> Option<&SubspaceId> {
        self.subspace.as_ref()
    }

    /// Returns a reference to the inner [`Path`].
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::new(
    ///     Some([17; 32].into()),
    ///     Path::new(),
    ///     TimeRange::new_closed(0.into(), 17.into()),
    /// );
    /// assert_eq!(a.path(), &Path::new());
    /// ```
    ///
    /// [Definition](https://willowprotocol.org/specs/grouping-entries/#AreaPath).
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Returns a reference to the inner [`TimeRange`](super::TimeRange).
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::new(
    ///     Some([17; 32].into()),
    ///     Path::new(),
    ///     TimeRange::new_closed(0.into(), 17.into()),
    /// );
    /// assert_eq!(a.times(), &WillowRange::from(TimeRange::new_closed(0.into(), 17.into())));
    /// ```
    ///
    /// [Definition](https://willowprotocol.org/specs/grouping-entries/#AreaTime).
    pub fn times(&self) -> &TimeRange {
        &self.times
    }

    /// Sets the inner subspace id.
    pub fn set_subspace(&mut self, new_subspace: Option<SubspaceId>) {
        self.subspace = new_subspace;
    }

    /// Sets the inner [`Path`].
    pub fn set_path(&mut self, new_path: Path) {
        self.path = new_path;
    }

    /// Sets the inner [`TimeRange`].
    pub fn set_times<TR>(&mut self, new_range: TR)
    where
        TR: Into<TimeRange>,
    {
        self.times = new_range.into();
    }

    /// Returns the [subspace area](https://willowprotocol.org/specs/grouping-entries/#subspace_area) for the given subspace id, i.e., the area which includes exactly the entries of the given subspace id.
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::new_subspace_area([17; 32].into());
    ///
    /// assert!(a.includes(&([17; 32].into(), Path::new(), Timestamp::from(9))));
    /// assert!(!a.includes(&([18; 32].into(), Path::new(), Timestamp::from(9))));
    /// assert_eq!(a.subspace(), Some(&[17; 32].into()));
    /// ```
    pub fn new_subspace_area(subspace_id: SubspaceId) -> Self {
        Self::new(Some(subspace_id), Path::new(), TimeRange::full())
    }

    /// Returns whether an [`Entry`] of the given [coordinate](Coordinatelike) could possibly cause [prefix pruning](https://willowprotocol.org/specs/data-model/index.html#prefix_pruning) in this area.
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::new(
    ///     Some([17; 32].into()),
    ///     Path::new(),
    ///     TimeRange::new_closed(0.into(), 17.into()),
    /// );
    ///
    /// assert!(a.admits_pruning_by(&([17; 32].into(), Path::new(), Timestamp::from(9))));
    /// assert!(!a.admits_pruning_by(&([18; 32].into(), Path::new(), Timestamp::from(9))));
    /// ```
    pub fn admits_pruning_by<Coord>(&self, coord: &Coord) -> bool
    where
        Coord: Coordinatelike,
    {
        if let Some(s) = self.subspace()
            && s != coord.subspace_id()
        {
            return false;
        }

        if coord.timestamp() < *self.times().start() {
            return false;
        }

        coord.path().is_related_to(self.path())
    }

    /// Returns the [`Area`] which [includes](Area::includes) every [coordinate](Coordinatelike).
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// let a = Area::full();
    ///
    /// assert!(a.includes(&([5; 32].into(), Path::new(), Timestamp::from(9))));
    /// assert!(a.includes(&(SubspaceId::from([16; 32]), Path::new(), Timestamp::from(9))));
    /// ```
    pub fn full() -> Self {
        Self {
            subspace: None,
            path: Path::new(),
            times: WillowRange::full(),
        }
    }

    /// Returns whether `self` is the full area, i.e., the area which [includes](Area::includes) every [coordinate](Coordinatelike).
    ///
    /// ```
    /// use willow25::prelude::*;
    ///
    /// assert!(Area::full().is_full());
    /// assert!(!Area::new(
    ///     Some([17; 32].into()),
    ///     Path::new(),
    ///     TimeRange::new_closed(0.into(), 17.into()),
    /// ).is_full());
    /// ```
    pub fn is_full(&self) -> bool {
        self.subspace().is_none() && self.path.is_empty() && self.times.is_full()
    }
}

/// An area is less than another iff all values included in the first are also included in the other.
///
/// This implementation assumes that `S` is inhabited by more than one value.
impl PartialOrd<Self> for Area {
    /// An area is less than another iff all values included in the first are also included in the other.
    ///
    /// This implementation assumes that `S` is inhabited by more than one value.
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        match (
            cmp_subspace(self.subspace(), other.subspace())?,
            other.path().prefix_cmp(self.path())?,
            self.times().partial_cmp(other.times())?,
        ) {
            (Ordering::Equal, Ordering::Equal, Ordering::Equal) => Some(Ordering::Equal),
            (subspaces_cmp, paths_cmp, times_cmp) => {
                if subspaces_cmp.is_le() && paths_cmp.is_le() && times_cmp.is_le() {
                    Some(Ordering::Less)
                } else if subspaces_cmp.is_ge() && paths_cmp.is_ge() && times_cmp.is_ge() {
                    Some(Ordering::Greater)
                } else {
                    None
                }
            }
        }
    }
}

fn cmp_subspace(s1: Option<&SubspaceId>, s2: Option<&SubspaceId>) -> Option<Ordering> {
    match (s1, s2) {
        (None, None) => Some(Ordering::Equal),
        (Some(_), None) => Some(Ordering::Less),
        (None, Some(_)) => Some(Ordering::Greater),
        (Some(s1), Some(s2)) => {
            if s1 == s2 {
                Some(Ordering::Equal)
            } else {
                None
            }
        }
    }
}

impl GreatestElement for Area {
    fn greatest() -> Self {
        Self::new(None, Path::new(), TimeRange::full())
    }

    fn is_greatest(&self) -> bool {
        self.subspace().is_none() && self.times().is_full() && self.path().is_empty()
    }
}

///////////////////////
// Codec Stuff Below //
///////////////////////

/// Implements [encode_area_in_area](https://willowprotocol.org/specs/encodings/index.html#encode_area_in_area).
impl RelativeEncodable<Area> for Area {
    async fn relative_encode<C>(&self, rel: &Area, consumer: &mut C) -> Result<(), C::Error>
    where
        C: BulkConsumer<Item = u8> + ?Sized,
    {
        debug_assert!(self.can_be_encoded_relative_to(rel));

        let self_times_start = u64::from(*self.times().start());
        let self_times_end = self.times().end().map(|t| u64::from(*t));
        let rel_times_start = u64::from(*rel.times().start());
        let rel_times_end = rel.times().end().map(|t| u64::from(*t));

        let (start_diff, start_from_start) = match rel_times_end {
            None => (self_times_start - rel_times_start, true),
            Some(rel_times_end) => {
                if self_times_start - rel_times_start < rel_times_end - self_times_start {
                    (self_times_start - rel_times_start, true)
                } else {
                    (rel_times_end - self_times_start, false)
                }
            }
        };

        let (end_diff, end_from_start) = match rel_times_end {
            None => match self.times.end() {
                None => (None, false),
                Some(self_times_end) => (Some(u64::from(*self_times_end) - rel_times_start), true),
            },

            Some(rel_times_end) => {
                // self is contained in rel, so self_time_end is `Some`.
                let self_times_end: u64 = self_times_end.unwrap();

                if self_times_end - rel_times_start < rel_times_end - self_times_end {
                    (Some(self_times_end - rel_times_start), true)
                } else {
                    (Some(rel_times_end - self_times_end), false)
                }
            }
        };

        let mut header = 0;

        if self.subspace() != rel.subspace() {
            header |= 0b1000_0000;
        }

        if self.times().is_open() {
            header |= 0b0100_0000;
        }

        if start_from_start {
            header |= 0b0010_0000;
        }

        if self.times().is_closed() && end_from_start {
            header |= 0b0001_0000;
        }

        write_tag(&mut header, 2, 4, start_diff);
        write_tag(&mut header, 2, 6, end_diff.unwrap_or(0));

        consumer.consume_item(header).await?;

        if let (Some(self_subspace_id), None) = (self.subspace(), rel.subspace()) {
            consumer.consume_encoded(self_subspace_id).await?;
        }

        cu64_encode(start_diff, 2, consumer).await?;

        if let Some(end_diff) = end_diff {
            cu64_encode(end_diff, 2, consumer).await?;
        }

        encode_path_extends_path(self.path(), rel.path(), consumer).await?;

        Ok(())
    }

    /// Returns `true` iff `rel` includes `self`.
    fn can_be_encoded_relative_to(&self, rel: &Area) -> bool {
        rel.includes_grouping(self)
    }
}

/// Implements [EncodeAreaInArea](https://willowprotocol.org/specs/encodings/index.html#EncodeAreaInArea).
impl RelativeDecodable<Area> for Area {
    type ErrorReason = Blame;

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

/// Implements [encode_area_in_area](https://willowprotocol.org/specs/encodings/index.html#encode_area_in_area).
impl RelativeDecodableCanonic<Area> for Area {
    type ErrorCanonic = Blame;

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

/// Implements [encode_area_in_area](https://willowprotocol.org/specs/encodings/index.html#encode_area_in_area).
impl RelativeEncodableKnownLength<Area> for Area {
    fn len_of_relative_encoding(&self, rel: &Area) -> usize {
        let mut encoding_len = 1; // The header byte.

        let self_times_start = u64::from(*self.times().start());
        let self_times_end = self.times().end().map(|t| u64::from(*t));
        let rel_times_start = u64::from(*rel.times().start());
        let rel_times_end = rel.times().end().map(|t| u64::from(*t));

        let start_diff = match rel_times_end {
            None => self_times_start - rel_times_start,
            Some(rel_times_end) => min(
                self_times_start - rel_times_start,
                rel_times_end - self_times_start,
            ),
        };

        let end_diff = match rel_times_end {
            None => self_times_end.map(|self_times_end| self_times_end - rel_times_start),
            Some(rel_times_end) => {
                // self is contained in rel, so self_time_end is `Some`.
                let self_times_end: u64 = self_times_end.unwrap();
                Some(min(
                    self_times_end - rel_times_start,
                    rel_times_end - self_times_end,
                ))
            }
        };

        if let (Some(self_subspace_id), None) = (self.subspace(), rel.subspace()) {
            encoding_len += self_subspace_id.len_of_encoding();
        }

        encoding_len += cu64_len_of_encoding(2, start_diff);

        if let Some(end_diff) = end_diff {
            encoding_len += cu64_len_of_encoding(2, end_diff);
        }

        encoding_len += path_extends_path_encoding_len(self.path(), rel.path());

        encoding_len
    }
}

async fn relative_decode_maybe_canonic<const CANONIC: bool, P>(
    rel: &Area,
    producer: &mut P,
) -> Result<Area, DecodeError<P::Final, P::Error, Blame>>
where
    P: BulkProducer<Item = u8> + ?Sized,
{
    let header = producer.produce_item().await?;

    // Decode subspace?
    let is_subspace_encoded = is_bitflagged(header, 0);

    // Decode end value of times?
    let is_times_end_open = is_bitflagged(header, 1);

    // Add start_diff to rel.get_times().start, or subtract from rel.get_times().end?
    let start_from_start = is_bitflagged(header, 2);

    // Add end_diff to rel.get_times().start, or subtract from rel.get_times().end?
    let end_from_start = is_bitflagged(header, 3);

    // === Necessary to produce canonic encodings. ===
    // Verify the last two header bits are zero if is_times_end_open
    if CANONIC && is_times_end_open && (is_bitflagged(header, 6) || is_bitflagged(header, 7)) {
        return Err(DecodeError::Other(Blame::TheirFault));
    }
    // ===============================================

    let subspace = if is_subspace_encoded {
        let subspace_id = producer
            .produce_decoded_canonic()
            .await
            .map_err(|err| err.map_other(Into::into))?;

        if let Some(rel_subspace_id) = rel.subspace() {
            if &subspace_id != rel_subspace_id {
                // They encoded an area not contained in rel.
                return Err(DecodeError::Other(Blame::TheirFault));
            } else if CANONIC {
                // They should not have encoded the subspace id explicitly.
                return Err(DecodeError::Other(Blame::TheirFault));
            }
        }

        Some(subspace_id)
    } else {
        rel.subspace.clone()
    };

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

    let rel_times_start = u64::from(*rel.times().start());
    let rel_times_end = rel.times().end().map(|t| u64::from(*t));

    let start = if start_from_start {
        rel_times_start
            .checked_add(start_diff)
            .ok_or(DecodeError::Other(Blame::TheirFault))?
    } else {
        match rel_times_end {
            None => {
                // They should have set `start_from_start` to true.
                return Err(DecodeError::Other(Blame::TheirFault));
            }
            Some(rel_times_end) => rel_times_end
                .checked_sub(start_diff)
                .ok_or(DecodeError::Other(Blame::TheirFault))?,
        }
    };

    if CANONIC {
        let should_have_set_start_from_start = match rel_times_end {
            None => true,
            Some(rel_times_end) => {
                let start_diff = start
                    .checked_sub(rel_times_start)
                    .ok_or(DecodeError::Other(Blame::TheirFault))?;
                let end_diff = rel_times_end
                    .checked_sub(start)
                    .ok_or(DecodeError::Other(Blame::TheirFault))?;
                start_diff < end_diff
            }
        };

        if start_from_start != should_have_set_start_from_start {
            return Err(DecodeError::Other(Blame::TheirFault));
        }
    }

    let end = if is_times_end_open {
        if end_from_start {
            return Err(DecodeError::Other(Blame::TheirFault));
        }

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

        let end = if end_from_start {
            rel_times_start
                .checked_add(end_diff)
                .ok_or(DecodeError::Other(Blame::TheirFault))?
        } else {
            match rel_times_end {
                None => {
                    // They should have set `end_from_start` to true.
                    return Err(DecodeError::Other(Blame::TheirFault));
                }
                Some(rel_times_end) => rel_times_end
                    .checked_sub(end_diff)
                    .ok_or(DecodeError::Other(Blame::TheirFault))?,
            }
        };

        if CANONIC {
            let should_have_set_end_from_start = match rel_times_end {
                None => true,
                Some(rel_times_end) => {
                    let start_diff = end
                        .checked_sub(rel_times_start)
                        .ok_or(DecodeError::Other(Blame::TheirFault))?;
                    let end_diff = rel_times_end
                        .checked_sub(end)
                        .ok_or(DecodeError::Other(Blame::TheirFault))?;
                    start_diff < end_diff
                }
            };

            if end_from_start != should_have_set_end_from_start {
                return Err(DecodeError::Other(Blame::TheirFault));
            }
        }

        Some(end)
    };

    let times = WillowRange::try_new(Timestamp::from(start), end.map(Timestamp::from))
        .map_err(|_| DecodeError::Other(Blame::TheirFault))?;

    let path = if CANONIC {
        decode_path_extends_path_canonic(rel.path(), producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    } else {
        decode_path_extends_path(rel.path(), producer)
            .await
            .map_err(|err| err.map_other(|_| Blame::TheirFault))?
    };

    if !rel.times().includes_willow_range(&times) {
        return Err(DecodeError::Other(Blame::TheirFault));
    }

    Ok(Area::new(subspace, path, times))
}

/// Returns an [`Arbitrary`] area which is guaranteed to be included in the reference area.
#[cfg(feature = "dev")]
pub fn arbitrary_area_in_area<'a>(
    reference: &Area,
    u: &mut arbitrary::Unstructured<'a>,
) -> arbitrary::Result<Area> {
    let subspace = match reference.subspace() {
        None => Option::<SubspaceId>::arbitrary(u)?,
        Some(s) => Some(s.clone()),
    };

    let path_suffix = Path::arbitrary(u)?;
    let path = reference
        .path()
        .append_path(&path_suffix)
        .unwrap_or_else(|_| reference.path().clone());

    let times_candidate = TimeRange::arbitrary(u)?;
    let times = if reference.times().includes_willow_range(&times_candidate) {
        times_candidate
    } else {
        *reference.times()
    };

    Ok(Area {
        subspace,
        path,
        times,
    })
}

/// This is a simple, "inofficial" encoding not on the encodings spec page.
///
/// - First byte is a header of bitflags:
///     - most significant bit: `1` iff the subspace is `any`.
///     - second-most significant bit: `1` iff the timestamp range is open.
///     - remaining six bits: arbitrary.
/// - encoding of the subspace id, or empty string if the subspace is `any`
/// - encoding of the path
/// - encoding of the start of the timestamp range: first an 8-bit [compact u64 tag](https://willowprotocol.org/specs/encodings/index.html#c64_tag), then the [corresponding compact u64 encoding](https://willowprotocol.org/specs/encodings/index.html#c64_corresponding)
/// - encoding of the end of the timestamp range (first an 8-bit [compact u64 tag](https://willowprotocol.org/specs/encodings/index.html#c64_tag), then the [corresponding compact u64 encoding](https://willowprotocol.org/specs/encodings/index.html#c64_corresponding)), or empty string if the timestamp range is open
impl Encodable for Area {
    async fn encode<C>(&self, consumer: &mut C) -> Result<(), C::Error>
    where
        C: BulkConsumer<Item = u8> + ?Sized,
    {
        let mut header = 0u8;

        if self.subspace().is_none() {
            header |= 0b1000_0000;
        }
        if self.times().is_open() {
            header |= 0b0100_0000;
        }

        consumer.consume_item(header).await?;

        if let Some(id) = self.subspace() {
            consumer.consume_encoded(id).await?;
        }

        consumer.consume_encoded(self.path()).await?;

        cu64_encode_standalone(u64::from(*self.times().start()), consumer).await?;

        if let Some(end) = self.times().end() {
            cu64_encode_standalone(u64::from(*end), consumer).await?;
        }

        Ok(())
    }
}