pipecircuit 0.0.18

use core::{
    cmp::Ordering,
    fmt,
    hash::{Hash, Hasher},
    marker::PhantomData,
    num::{NonZeroU8, NonZeroU16, NonZeroU32, NonZeroU64},
};

use super::IndexKey;

const SEQUENTIAL_INDEX_MULTIPLIER: u64 = 0x9E37_79B9_7F4A_7C15;

/// A compact, sequential identifier scoped by its row type.
#[repr(transparent)]
pub struct WarpId<Row> {
    value: NonZeroU64,
    row: PhantomData<fn() -> Row>,
}

impl<Row> WarpId<Row> {
    pub const MIN: Self = Self::new(NonZeroU64::MIN);

    pub const fn new(value: NonZeroU64) -> Self {
        Self {
            value,
            row: PhantomData,
        }
    }

    pub const fn value(self) -> NonZeroU64 {
        self.value
    }
}

#[cfg(feature = "serde")]
impl<Row> serde::Serialize for WarpId<Row> {
    fn serialize<Serializer: serde::Serializer>(
        &self,
        serializer: Serializer,
    ) -> Result<Serializer::Ok, Serializer::Error> {
        serde::Serialize::serialize(&self.value, serializer)
    }
}

#[cfg(feature = "serde")]
impl<'message, Row> serde::Deserialize<'message> for WarpId<Row> {
    fn deserialize<Deserializer: serde::Deserializer<'message>>(
        deserializer: Deserializer,
    ) -> Result<Self, Deserializer::Error> {
        <NonZeroU64 as serde::Deserialize<'message>>::deserialize(deserializer).map(Self::new)
    }
}

#[cfg(all(feature = "postcard-bindings", not(target_arch = "wasm32")))]
impl<Row> postcard_bindgen::__private::GenBinding for WarpId<Row> {
    fn get_type() -> postcard_bindgen::__private::ValueType {
        <NonZeroU64 as postcard_bindgen::__private::GenBinding>::get_type()
    }
}

impl<Row> Clone for WarpId<Row> {
    fn clone(&self) -> Self {
        *self
    }
}

impl<Row> Copy for WarpId<Row> {}

impl<Row> fmt::Debug for WarpId<Row> {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.debug_tuple("WarpId").field(&self.value).finish()
    }
}

impl<Row> PartialEq for WarpId<Row> {
    fn eq(&self, other: &Self) -> bool {
        self.value == other.value
    }
}

impl<Row> Eq for WarpId<Row> {}

impl<Row> PartialOrd for WarpId<Row> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<Row> Ord for WarpId<Row> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.value.cmp(&other.value)
    }
}

impl<Row> Hash for WarpId<Row> {
    fn hash<State: Hasher>(&self, state: &mut State) {
        self.value.hash(state);
    }
}

impl<Row> IndexKey for WarpId<Row> {
    fn index_word(self) -> u64 {
        sequential_index_word(self.value)
    }
}

/// Distributes a sequential nonzero integer across the complete `u64` range.
pub const fn sequential_index_word(sequence: NonZeroU64) -> u64 {
    integer_index_word(sequence.get())
}

/// Distributes an unsigned integer across the complete `u64` range.
pub const fn integer_index_word(value: u64) -> u64 {
    value.wrapping_mul(SEQUENTIAL_INDEX_MULTIPLIER)
}

impl IndexKey for u8 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u16 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u32 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self))
    }
}

impl IndexKey for u64 {
    fn index_word(self) -> u64 {
        integer_index_word(self)
    }
}

impl IndexKey for NonZeroU8 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU16 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU32 {
    fn index_word(self) -> u64 {
        integer_index_word(u64::from(self.get()))
    }
}

impl IndexKey for NonZeroU64 {
    fn index_word(self) -> u64 {
        sequential_index_word(self)
    }
}

#[cfg(test)]
mod tests {
    use core::{mem::size_of, num::NonZeroU64};

    use super::{WarpId, integer_index_word, sequential_index_word};
    use crate::storage::warp_table::IndexKey;

    struct Row;

    #[test]
    fn warp_id_is_one_word_and_preserves_its_sequence() {
        let id = WarpId::<Row>::new(NonZeroU64::MIN);

        assert_eq!(size_of::<WarpId<Row>>(), size_of::<u64>());
        assert_eq!(size_of::<Option<WarpId<Row>>>(), size_of::<u64>());
        assert_eq!(id.value(), NonZeroU64::MIN);
        assert_eq!(id.index_word(), sequential_index_word(NonZeroU64::MIN));
        assert_eq!(7_u64.index_word(), integer_index_word(7));
    }

    #[cfg(feature = "serde")]
    #[test]
    fn warp_id_encoding_preserves_the_sequence_without_serializing_the_row()
    -> Result<(), postcard::Error> {
        for value in [NonZeroU64::MIN, NonZeroU64::MAX] {
            let id = WarpId::<Row>::new(value);
            let bytes = postcard::to_allocvec(&id)?;
            assert_eq!(bytes, postcard::to_allocvec(&value)?);
            assert_eq!(postcard::from_bytes::<WarpId<Row>>(&bytes)?, id);
        }
        assert!(postcard::from_bytes::<WarpId<Row>>(&[0]).is_err());
        Ok(())
    }

    #[cfg(all(feature = "postcard-bindings", not(target_arch = "wasm32")))]
    #[test]
    fn warp_id_bindings_preserve_the_nonzero_wire_bound() {
        use postcard_bindgen::__private::GenBinding;

        assert_eq!(WarpId::<Row>::get_type(), NonZeroU64::get_type());
    }

    #[cfg(feature = "serde")]
    proptest::proptest! {
        #[test]
        fn warp_id_round_trips_every_nonzero_sequence(raw in 1_u64..=u64::MAX) {
            let value = NonZeroU64::new(raw)
                .ok_or_else(|| proptest::test_runner::TestCaseError::fail("generated zero"))?;
            let id = WarpId::<Row>::new(value);
            let bytes = postcard::to_allocvec(&id)?;
            let decoded = postcard::from_bytes::<WarpId<Row>>(&bytes)?;
            proptest::prop_assert_eq!(decoded, id);
            proptest::prop_assert_eq!(decoded.index_word(), id.index_word());
        }
    }
}