cpu/crates/cpu/src/util/array_vec.rs

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

use fayalite::{
    expr::{
        ops::{ExprCastTo, ExprIndex, ExprPartialEq, ExprPartialOrd},
        ToLiteralBits,
    },
    int::{IntType, SizeType},
    intern::{Intern, Interned},
    prelude::*,
    ty::{MatchVariantWithoutScope, StaticType, TypeProperties},
};
use std::{marker::PhantomData, ops::Index};

#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub struct Length<Max: Size> {
    ty: UInt,
    _phantom: PhantomData<Max>,
}

impl<Max: Size> Length<Max> {
    pub fn new(max: Max::SizeType) -> Self {
        Self {
            ty: UInt::range_inclusive(0..=Max::as_usize(max)),
            _phantom: PhantomData,
        }
    }
    pub fn ty(self) -> UInt {
        self.ty
    }
    pub fn zero(self) -> Expr<Self> {
        Self::from_uint_unchecked(self.ty.zero())
    }
    pub fn from_uint_unchecked(v: impl ToExpr<Type = UInt>) -> Expr<Self> {
        Expr::from_canonical(Expr::canonical(v.to_expr()))
    }
    pub fn cast_from_uint_unchecked<SrcWidth: Size>(
        self,
        v: impl ToExpr<Type = UIntType<SrcWidth>>,
    ) -> Expr<Self> {
        Self::from_uint_unchecked(v.to_expr().cast_to(self.ty))
    }
    pub fn as_uint(this: impl ToExpr<Type = Self>) -> Expr<UInt> {
        let this = this.to_expr();
        this.cast_to(Expr::ty(this).ty)
    }
}

impl<Max: Size, DestWidth: Size> ExprCastTo<UIntType<DestWidth>> for Length<Max> {
    fn cast_to(src: Expr<Self>, to_type: UIntType<DestWidth>) -> Expr<UIntType<DestWidth>> {
        Expr::<UInt>::from_canonical(Expr::canonical(src)).cast_to(to_type)
    }
}

#[allow(non_upper_case_globals)]
pub const Length: __LengthWithoutGenerics = __LengthWithoutGenerics {};

#[non_exhaustive]
pub struct __LengthWithoutGenerics {}

impl<M: SizeType> Index<M> for __LengthWithoutGenerics {
    type Output = Length<M::Size>;

    fn index(&self, max: M) -> &Self::Output {
        Interned::into_inner(Length::new(max).intern_sized())
    }
}

impl<Max: Size> Type for Length<Max> {
    type BaseType = UInt;
    type MaskType = Bool;
    type MatchVariant = Expr<Self>;
    type MatchActiveScope = ();
    type MatchVariantAndInactiveScope = MatchVariantWithoutScope<Self::MatchVariant>;
    type MatchVariantsIter = std::iter::Once<Self::MatchVariantAndInactiveScope>;
    fn match_variants(
        this: Expr<Self>,
        source_location: SourceLocation,
    ) -> Self::MatchVariantsIter {
        let _ = source_location;
        std::iter::once(MatchVariantWithoutScope(this))
    }

    fn mask_type(&self) -> Self::MaskType {
        Bool
    }

    fn canonical(&self) -> CanonicalType {
        self.ty.canonical()
    }

    fn from_canonical(canonical_type: CanonicalType) -> Self {
        let ty = <UInt>::from_canonical(canonical_type);
        if let Some(known_max) = Max::KNOWN_VALUE {
            assert_eq!(ty, UInt::range_inclusive(0..=known_max));
        }
        Self {
            ty,
            _phantom: PhantomData,
        }
    }

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

impl<Max: KnownSize> StaticType for Length<Max> {
    const TYPE: Self = Self {
        ty: UInt {
            width: Max::VALUE.next_power_of_two().ilog2() as usize,
        },
        _phantom: PhantomData,
    };
    const MASK_TYPE: Self::MaskType = Bool;
    const TYPE_PROPERTIES: TypeProperties = {
        let mut p = <UInt<1>>::TYPE_PROPERTIES;
        p.bit_width = Self::TYPE.ty.width;
        p
    };
    const MASK_TYPE_PROPERTIES: TypeProperties = Bool::TYPE_PROPERTIES;
}

impl<Max: Size> ExprPartialEq<Self> for Length<Max> {
    fn cmp_eq(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_eq(Self::as_uint(rhs))
    }
    fn cmp_ne(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_ne(Self::as_uint(rhs))
    }
}

impl<Max: Size> ExprPartialOrd<Self> for Length<Max> {
    fn cmp_lt(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_lt(Self::as_uint(rhs))
    }
    fn cmp_le(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_le(Self::as_uint(rhs))
    }
    fn cmp_gt(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_gt(Self::as_uint(rhs))
    }
    fn cmp_ge(lhs: Expr<Self>, rhs: Expr<Self>) -> Expr<Bool> {
        Self::as_uint(lhs).cmp_ge(Self::as_uint(rhs))
    }
}

/// like [`std::vec::Vec`], except with a [`Expr`] for [`len()`][`Self::len()`] and a fixed capacity
#[hdl]
pub struct ArrayVec<T: Type, N: Size> {
    elements: ArrayType<T, N>,
    len: Length<N>,
}

impl<T: Type, N: Size> ArrayVec<T, N> {
    #[hdl]
    pub fn new(self) -> Expr<Self> {
        #[hdl]
        ArrayVec {
            elements: self.elements.uninit(),
            len: self.len.zero(),
        }
    }
    pub fn element(self) -> T {
        self.elements.element()
    }
    pub fn elements_ty(self) -> ArrayType<T, N> {
        self.elements
    }
    pub fn elements_unchecked(this: impl ToExpr<Type = Self>) -> Expr<ArrayType<T, N>> {
        this.to_expr().elements
    }
    #[hdl]
    pub fn from_parts_unchecked(
        elements: impl ToExpr<Type = ArrayType<T, N>>,
        len: impl ToExpr<Type = Length<N>>,
    ) -> Expr<Self> {
        let elements = elements.to_expr();
        let len = len.to_expr();
        assert_eq!(
            Length::new(N::from_usize(Expr::ty(elements).len())),
            Expr::ty(len),
            "len type mismatch",
        );
        #[hdl]
        Self { elements, len }
    }
    pub fn len_ty(self) -> Length<N> {
        self.len
    }
    pub fn len(this: impl ToExpr<Type = Self>) -> Expr<Length<N>> {
        this.to_expr().len
    }
    pub fn is_empty(this: impl ToExpr<Type = Self>) -> Expr<Bool> {
        let len = Self::len(this);
        len.cmp_eq(Expr::ty(len).zero())
    }
    pub fn capacity(self) -> usize {
        self.elements.len()
    }
    pub fn get_unchecked<Idx>(this: impl ToExpr<Type = Self>, index: Idx) -> Expr<T>
    where
        ArrayType<T, N>: ExprIndex<Idx, Output = T>,
    {
        this.to_expr().elements[index]
    }
    #[hdl]
    pub fn for_each(this: impl ToExpr<Type = Self>, mut f: impl FnMut(usize, Expr<T>)) {
        let this = this.to_expr();
        for (index, element) in this.elements.into_iter().enumerate() {
            #[hdl]
            if index.cmp_lt(Length::as_uint(this.len)) {
                f(index, element);
            }
        }
    }
    pub fn mapped_ty<U: Type>(self, new_element_ty: U) -> ArrayVec<U, N> {
        ArrayVec {
            elements: ArrayType[new_element_ty][N::from_usize(self.elements.len())],
            len: self.len,
        }
    }
    #[hdl]
    pub fn map<U: Type>(
        this: impl ToExpr<Type = Self>,
        new_element_ty: U,
        mut f: impl FnMut(usize, Expr<T>) -> Expr<U>,
    ) -> Expr<ArrayVec<U, N>> {
        let this = this.to_expr();
        #[hdl]
        let mapped_array_vec = wire(Expr::ty(this).mapped_ty(new_element_ty));
        connect(mapped_array_vec.len, this.len);
        Self::for_each(this, |index, element| {
            connect(mapped_array_vec[index], f(index, element));
        });
        mapped_array_vec
    }
    #[hdl]
    pub fn as_array_of_options(this: impl ToExpr<Type = Self>) -> Expr<ArrayType<HdlOption<T>, N>> {
        let this = this.to_expr();
        #[hdl]
        let array_vec_as_array_of_options = wire(
            ArrayType[HdlOption[Expr::ty(this).element()]]
                [N::from_usize(Expr::ty(this).capacity())],
        );
        for element in array_vec_as_array_of_options {
            connect(element, Expr::ty(element).HdlNone());
        }
        Self::for_each(this, |index, element| {
            connect(array_vec_as_array_of_options[index], HdlSome(element))
        });
        array_vec_as_array_of_options
    }
    #[hdl]
    pub fn get<Idx: IntType<Dyn = UInt>>(
        this: impl ToExpr<Type = Self>,
        index: impl ToExpr<Type = Idx>,
    ) -> Expr<HdlOption<T>> {
        let this = this.to_expr();
        let index = Expr::as_dyn_int(index.to_expr());
        let never_in_bounds = index.cmp_ge(Expr::ty(this).capacity());
        if let Ok(never_in_bounds) = never_in_bounds.to_literal_bits() {
            if never_in_bounds[0] {
                // avoid error from out-of-bounds constant index
                return HdlOption[Expr::ty(this).element()].HdlNone();
            }
        }
        #[hdl]
        let array_vec_get = wire(HdlOption[Expr::ty(this).element()]);
        connect(array_vec_get, Expr::ty(array_vec_get).HdlNone());
        #[hdl]
        if index.cmp_lt(Length::as_uint(Self::len(this))) {
            connect(array_vec_get, HdlSome(this.elements[index]));
        }
        array_vec_get
    }
}

impl<T: Type, N: Size, Idx, IdxWidth: Size> ExprIndex<Idx> for ArrayVec<T, N>
where
    ArrayType<T, N>: ExprIndex<Idx, Output = T>,
    Idx: ToExpr<Type = UIntType<IdxWidth>>,
{
    type Output = T;

    fn expr_index(this: &Expr<Self>, index: Idx) -> &Expr<Self::Output> {
        // TODO: add assert that index is in-bounds
        <ArrayType<T, N> as ExprIndex<Idx>>::expr_index(&this.elements, index)
    }
}

#[hdl]
pub struct ReadyValidArray<T: Type, N: Size> {
    pub data: ArrayVec<T, N>,
    #[hdl(flip)]
    pub ready: Length<N>,
}

impl<T: Type, N: Size> ReadyValidArray<T, N> {
    #[hdl]
    pub fn firing_len(this: impl ToExpr<Type = Self>) -> Expr<Length<N>> {
        let this = this.to_expr();
        assert_eq!(Expr::ty(this).data.len_ty(), Expr::ty(this).ready);
        #[hdl]
        let firing_len = wire(Expr::ty(this).data.len);
        connect(firing_len, this.data.len);
        #[hdl]
        if this.data.len.cmp_gt(this.ready) {
            connect(firing_len, this.ready);
        }
        firing_len
    }
    #[hdl]
    pub fn firing_data(this: impl ToExpr<Type = Self>) -> Expr<ArrayVec<T, N>> {
        let this = this.to_expr();
        #[hdl]
        let firing_data = wire(Expr::ty(this).data);
        connect(firing_data, this.data);
        connect(firing_data.len, Self::firing_len(this));
        firing_data
    }
}