fayalite/crates/fayalite/src/array.rs

// SPDX-License-Identifier: LGPL-3.0-or-later
// See Notices.txt for copyright information

use crate::{
    expr::{
        CastToBits, Expr, HdlPartialEq, ReduceBits, ToExpr,
        ops::{ArrayLiteral, ExprFromIterator, ExprIntoIterator, ExprPartialEq},
    },
    int::{Bool, DYN_SIZE, DynSize, KnownSize, Size, SizeType},
    intern::{Intern, Interned, LazyInterned},
    module::transform::visit::{Fold, Folder, Visit, Visitor},
    sim::value::{SimValue, SimValuePartialEq},
    source_location::SourceLocation,
    ty::{
        CanonicalType, MatchVariantWithoutScope, OpaqueSimValueSlice, OpaqueSimValueWriter,
        OpaqueSimValueWritten, StaticType, Type, TypeProperties, TypeWithDeref,
        serde_impls::SerdeCanonicalType,
    },
    util::ConstUsize,
};
use serde::{Deserialize, Deserializer, Serialize, Serializer, de::Error};
use std::{iter::FusedIterator, ops::Index};

#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub struct ArrayType<T: Type = CanonicalType, Len: Size = DynSize> {
    element: LazyInterned<T>,
    len: Len::SizeType,
    type_properties: TypeProperties,
}

impl<T: Type, Len: Size> std::fmt::Debug for ArrayType<T, Len> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Array<{:?}, {}>", self.element, self.len())
    }
}

pub type Array<T = CanonicalType, const LEN: usize = DYN_SIZE> = ArrayType<T, ConstUsize<LEN>>;

#[allow(non_upper_case_globals)]
pub const Array: ArrayWithoutGenerics = ArrayWithoutGenerics;
#[allow(non_upper_case_globals)]
pub const ArrayType: ArrayWithoutGenerics = ArrayWithoutGenerics;

impl<T: Type, Len: Size> ArrayType<T, Len> {
    const fn make_type_properties(element: TypeProperties, len: usize) -> TypeProperties {
        let TypeProperties {
            is_passive,
            is_storable,
            is_castable_from_bits,
            bit_width,
            sim_only_values_len,
        } = element;
        let Some(bit_width) = bit_width.checked_mul(len) else {
            panic!("array too big");
        };
        let Some(sim_only_values_len) = sim_only_values_len.checked_mul(len) else {
            panic!("array too big");
        };
        TypeProperties {
            is_passive,
            is_storable,
            is_castable_from_bits,
            bit_width,
            sim_only_values_len,
        }
    }
    pub fn new(element: T, len: Len::SizeType) -> Self {
        let type_properties =
            Self::make_type_properties(element.canonical().type_properties(), Len::as_usize(len));
        Self {
            element: LazyInterned::Interned(element.intern_sized()),
            len,
            type_properties,
        }
    }
    pub fn element(&self) -> T {
        *self.element
    }
    pub fn len(self) -> usize {
        Len::as_usize(self.len)
    }
    pub fn is_empty(self) -> bool {
        self.len() == 0
    }
    pub fn type_properties(self) -> TypeProperties {
        self.type_properties
    }
    pub fn as_dyn_array(self) -> Array {
        Array::new_dyn(self.element().canonical(), self.len())
    }
    pub fn can_connect<T2: Type, Len2: Size>(self, rhs: ArrayType<T2, Len2>) -> bool {
        self.len() == rhs.len()
            && self
                .element()
                .canonical()
                .can_connect(rhs.element().canonical())
    }
}

impl<T: Type, Len: KnownSize + Size<SizeType = Len>> ArrayType<T, Len> {
    pub fn new_static(element: T) -> Self {
        Self::new(element, Len::SizeType::default())
    }
}

impl<T: StaticType, Len: KnownSize> Default for ArrayType<T, Len> {
    fn default() -> Self {
        Self::TYPE
    }
}

impl<T: StaticType, Len: KnownSize> StaticType for ArrayType<T, Len> {
    const TYPE: Self = Self {
        element: LazyInterned::new_lazy(&|| T::TYPE.intern_sized()),
        len: Len::SIZE,
        type_properties: Self::TYPE_PROPERTIES,
    };
    const MASK_TYPE: Self::MaskType = ArrayType::<T::MaskType, Len> {
        element: LazyInterned::new_lazy(&|| T::MASK_TYPE.intern_sized()),
        len: Len::SIZE,
        type_properties: Self::MASK_TYPE_PROPERTIES,
    };
    const TYPE_PROPERTIES: TypeProperties =
        Self::make_type_properties(T::TYPE_PROPERTIES, Len::VALUE);
    const MASK_TYPE_PROPERTIES: TypeProperties =
        Self::make_type_properties(T::MASK_TYPE_PROPERTIES, Len::VALUE);
}

impl<T: Type> Array<T> {
    pub fn new_dyn(element: T, len: usize) -> Self {
        Self::new(element, len)
    }
}

impl<T: Type + Fold<State>, Len: Size, State: Folder + ?Sized> Fold<State> for ArrayType<T, Len> {
    fn fold(self, state: &mut State) -> Result<Self, State::Error> {
        state.fold_array_type(self)
    }

    fn default_fold(self, state: &mut State) -> Result<Self, <State as Folder>::Error> {
        Ok(ArrayType::new(self.element().fold(state)?, self.len))
    }
}

impl<T: Type + Visit<State>, Len: Size, State: Visitor + ?Sized> Visit<State>
    for ArrayType<T, Len>
{
    fn visit(&self, state: &mut State) -> Result<(), State::Error> {
        state.visit_array_type(self)
    }

    fn default_visit(&self, state: &mut State) -> Result<(), State::Error> {
        self.element().visit(state)
    }
}

impl<T: Type, Len: Size> Type for ArrayType<T, Len> {
    type BaseType = Array;
    type MaskType = ArrayType<T::MaskType, Len>;
    type SimValue = Len::ArraySimValue<T>;
    type MatchVariant = Len::ArrayMatch<T>;
    type MatchActiveScope = ();
    type MatchVariantAndInactiveScope = MatchVariantWithoutScope<Len::ArrayMatch<T>>;
    type MatchVariantsIter = std::iter::Once<Self::MatchVariantAndInactiveScope>;

    fn match_variants(
        this: Expr<Self>,
        source_location: SourceLocation,
    ) -> Self::MatchVariantsIter {
        let _ = source_location;
        let retval = Vec::from_iter(this);
        std::iter::once(MatchVariantWithoutScope(
            Len::ArrayMatch::<T>::try_from(retval)
                .ok()
                .expect("unreachable"),
        ))
    }

    fn mask_type(&self) -> Self::MaskType {
        ArrayType::new(self.element().mask_type(), self.len)
    }

    fn canonical(&self) -> CanonicalType {
        CanonicalType::Array(Array::new_dyn(self.element().canonical(), self.len()))
    }

    #[track_caller]
    fn from_canonical(canonical_type: CanonicalType) -> Self {
        let CanonicalType::Array(array) = canonical_type else {
            panic!("expected array");
        };
        Self::new(
            T::from_canonical(array.element()),
            Len::from_usize(array.len()),
        )
    }

    fn source_location() -> SourceLocation {
        SourceLocation::builtin()
    }

    fn sim_value_from_opaque(&self, mut opaque: OpaqueSimValueSlice<'_>) -> Self::SimValue {
        let element_ty = self.element();
        let element_size = element_ty.canonical().size();
        let mut value = Vec::with_capacity(self.len());
        for _ in 0..self.len() {
            let (element_opaque, rest) = opaque.split_at(element_size);
            value.push(SimValue::from_opaque(element_ty, element_opaque.to_owned()));
            opaque = rest;
        }
        value.try_into().ok().expect("used correct length")
    }

    fn sim_value_clone_from_opaque(
        &self,
        value: &mut Self::SimValue,
        mut opaque: OpaqueSimValueSlice<'_>,
    ) {
        let element_ty = self.element();
        let element_size = element_ty.canonical().size();
        let value = AsMut::<[SimValue<T>]>::as_mut(value);
        assert_eq!(self.len(), value.len());
        for element_value in value {
            assert_eq!(SimValue::ty(element_value), element_ty);
            let (element_opaque, rest) = opaque.split_at(element_size);
            SimValue::opaque_mut(element_value).clone_from_slice(element_opaque);
            opaque = rest;
        }
    }

    fn sim_value_to_opaque<'w>(
        &self,
        value: &Self::SimValue,
        mut writer: OpaqueSimValueWriter<'w>,
    ) -> OpaqueSimValueWritten<'w> {
        let element_ty = self.element();
        let element_size = element_ty.canonical().size();
        let value = AsRef::<[SimValue<T>]>::as_ref(value);
        assert_eq!(self.len(), value.len());
        for element_value in value {
            assert_eq!(SimValue::ty(element_value), element_ty);
            writer.fill_prefix_with(element_size, |writer| {
                writer.fill_cloned_from_slice(SimValue::opaque(element_value).as_slice())
            });
        }
        writer.fill_cloned_from_slice(OpaqueSimValueSlice::empty())
    }
}

impl<T: Type + Serialize, Len: Size> Serialize for ArrayType<T, Len> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        SerdeCanonicalType::<T>::Array {
            element: self.element(),
            len: self.len(),
        }
        .serialize(serializer)
    }
}

impl<'de, T: Type + Deserialize<'de>, Len: Size> Deserialize<'de> for ArrayType<T, Len> {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        let name = |len| -> String {
            if let Some(len) = len {
                format!("an Array<_, {len}>")
            } else {
                "an Array<_>".to_string()
            }
        };
        match SerdeCanonicalType::<T>::deserialize(deserializer)? {
            SerdeCanonicalType::Array { element, len } => {
                if let Some(len) = Len::try_from_usize(len) {
                    Ok(Self::new(element, len))
                } else {
                    Err(Error::invalid_value(
                        serde::de::Unexpected::Other(&name(Some(len))),
                        &&*name(Len::KNOWN_VALUE),
                    ))
                }
            }
            ty => Err(Error::invalid_value(
                serde::de::Unexpected::Other(ty.as_serde_unexpected_str()),
                &&*name(Len::KNOWN_VALUE),
            )),
        }
    }
}

impl<T: Type, Len: Size> TypeWithDeref for ArrayType<T, Len> {
    fn expr_deref(this: &Expr<Self>) -> &Self::MatchVariant {
        let retval = Vec::from_iter(*this);
        Interned::into_inner(Intern::intern_sized(
            Len::ArrayMatch::<T>::try_from(retval)
                .ok()
                .expect("unreachable"),
        ))
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug, Default)]
pub struct ArrayWithoutGenerics;

impl<T: Type> Index<T> for ArrayWithoutGenerics {
    type Output = ArrayWithoutLen<T>;

    fn index(&self, element: T) -> &Self::Output {
        Interned::into_inner(Intern::intern_sized(ArrayWithoutLen { element }))
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub struct ArrayWithoutLen<T: Type> {
    element: T,
}

impl<T: Type, L: SizeType> Index<L> for ArrayWithoutLen<T> {
    type Output = ArrayType<T, L::Size>;

    fn index(&self, len: L) -> &Self::Output {
        Interned::into_inner(Intern::intern_sized(ArrayType::new(self.element, len)))
    }
}

impl<Lhs: Type, Rhs: Type, Len: Size> ExprPartialEq<ArrayType<Rhs, Len>> for ArrayType<Lhs, Len>
where
    Lhs: ExprPartialEq<Rhs>,
{
    fn cmp_eq(lhs: Expr<Self>, rhs: Expr<ArrayType<Rhs, Len>>) -> Expr<Bool> {
        let lhs_ty = Expr::ty(lhs);
        let rhs_ty = Expr::ty(rhs);
        assert_eq!(lhs_ty.len(), rhs_ty.len());
        lhs.into_iter()
            .zip(rhs)
            .map(|(l, r)| l.cmp_eq(r))
            .collect::<Expr<Array<Bool>>>()
            .cast_to_bits()
            .all_one_bits()
    }

    fn cmp_ne(lhs: Expr<Self>, rhs: Expr<ArrayType<Rhs, Len>>) -> Expr<Bool> {
        let lhs_ty = Expr::ty(lhs);
        let rhs_ty = Expr::ty(rhs);
        assert_eq!(lhs_ty.len(), rhs_ty.len());
        lhs.into_iter()
            .zip(rhs)
            .map(|(l, r)| l.cmp_ne(r))
            .collect::<Expr<Array<Bool>>>()
            .cast_to_bits()
            .any_one_bits()
    }
}

impl<Lhs: Type, Rhs: Type, Len: Size> SimValuePartialEq<ArrayType<Rhs, Len>> for ArrayType<Lhs, Len>
where
    Lhs: SimValuePartialEq<Rhs>,
{
    fn sim_value_eq(this: &SimValue<Self>, other: &SimValue<ArrayType<Rhs, Len>>) -> bool {
        AsRef::<[_]>::as_ref(&**this)
            .iter()
            .zip(AsRef::<[_]>::as_ref(&**other))
            .all(|(l, r)| SimValuePartialEq::sim_value_eq(l, r))
    }
}

impl<T: Type, Len: Size> ExprIntoIterator for ArrayType<T, Len> {
    type Item = T;
    type ExprIntoIter = ExprArrayIter<T, Len>;

    fn expr_into_iter(e: Expr<Self>) -> Self::ExprIntoIter {
        ExprArrayIter {
            base: e,
            indexes: 0..Expr::ty(e).len(),
        }
    }
}

#[derive(Clone, Debug)]
pub struct ExprArrayIter<T: Type, Len: Size> {
    base: Expr<ArrayType<T, Len>>,
    indexes: std::ops::Range<usize>,
}

impl<T: Type, Len: Size> ExprArrayIter<T, Len> {
    pub fn base(&self) -> Expr<ArrayType<T, Len>> {
        self.base
    }
    pub fn indexes(&self) -> std::ops::Range<usize> {
        self.indexes.clone()
    }
}

impl<T: Type, Len: Size> Iterator for ExprArrayIter<T, Len> {
    type Item = Expr<T>;

    fn next(&mut self) -> Option<Self::Item> {
        self.indexes.next().map(|i| self.base[i])
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.indexes.size_hint()
    }

    fn count(self) -> usize {
        self.indexes.count()
    }

    fn last(mut self) -> Option<Self::Item> {
        self.next_back()
    }

    fn nth(&mut self, n: usize) -> Option<Self::Item> {
        self.indexes.nth(n).map(|i| self.base[i])
    }

    fn fold<B, F>(self, init: B, mut f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.indexes.fold(init, |b, i| f(b, self.base[i]))
    }
}

impl<T: Type, Len: Size> DoubleEndedIterator for ExprArrayIter<T, Len> {
    fn next_back(&mut self) -> Option<Self::Item> {
        self.indexes.next_back().map(|i| self.base[i])
    }

    fn nth_back(&mut self, n: usize) -> Option<Self::Item> {
        self.indexes.nth_back(n).map(|i| self.base[i])
    }

    fn rfold<B, F>(self, init: B, mut f: F) -> B
    where
        F: FnMut(B, Self::Item) -> B,
    {
        self.indexes.rfold(init, |b, i| f(b, self.base[i]))
    }
}

impl<T: Type, Len: Size> ExactSizeIterator for ExprArrayIter<T, Len> {
    fn len(&self) -> usize {
        self.indexes.len()
    }
}

impl<T: Type, Len: Size> FusedIterator for ExprArrayIter<T, Len> {}

impl<A: StaticType> ExprFromIterator<Expr<A>> for Array<A> {
    fn expr_from_iter<T: IntoIterator<Item = Expr<A>>>(iter: T) -> Expr<Self> {
        ArrayLiteral::new(
            A::TYPE,
            iter.into_iter().map(|v| Expr::canonical(v)).collect(),
        )
        .to_expr()
    }
}

impl<'a, A: StaticType> ExprFromIterator<&'a Expr<A>> for Array<A> {
    fn expr_from_iter<T: IntoIterator<Item = &'a Expr<A>>>(iter: T) -> Expr<Self> {
        iter.into_iter().copied().collect()
    }
}