use crate::{
expr::{ops::VariantAccess, Expr, ToExpr},
hdl,
int::Bool,
intern::{Intern, Interned},
module::{
connect, enum_match_variants_helper, incomplete_wire, wire,
EnumMatchVariantAndInactiveScopeImpl, EnumMatchVariantsIterImpl, Scope,
},
source_location::SourceLocation,
ty::{CanonicalType, MatchVariantAndInactiveScope, StaticType, Type, TypeProperties},
};
use hashbrown::HashMap;
use std::{convert::Infallible, fmt, iter::FusedIterator};
#[derive(Copy, Clone, PartialEq, Eq, Hash, Debug)]
pub struct EnumVariant {
pub name: Interned<str>,
pub ty: Option<CanonicalType>,
}
impl EnumVariant {
pub fn fmt_debug_in_enum(self) -> FmtDebugInEnum {
FmtDebugInEnum(self)
}
}
#[derive(Copy, Clone)]
pub struct FmtDebugInEnum(EnumVariant);
impl fmt::Debug for FmtDebugInEnum {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let EnumVariant { name, ty } = self.0;
if let Some(ty) = ty {
write!(f, "{name}({ty:?})")
} else {
write!(f, "{name}")
}
}
}
impl fmt::Display for FmtDebugInEnum {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
fmt::Debug::fmt(self, f)
}
}
#[derive(Clone, Eq)]
struct EnumImpl {
variants: Interned<[EnumVariant]>,
name_indexes: HashMap<Interned<str>, usize>,
type_properties: TypeProperties,
}
impl std::fmt::Debug for EnumImpl {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("Enum ")?;
f.debug_set()
.entries(
self.variants
.iter()
.map(|variant| variant.fmt_debug_in_enum()),
)
.finish()
}
}
impl std::hash::Hash for EnumImpl {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.variants.hash(state);
}
}
impl PartialEq for EnumImpl {
fn eq(&self, other: &Self) -> bool {
self.variants == other.variants
}
}
#[derive(Copy, Clone, PartialEq, Eq, Hash)]
pub struct Enum(Interned<EnumImpl>);
impl fmt::Debug for Enum {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
const fn discriminant_bit_width_impl(variant_count: usize) -> usize {
match variant_count.next_power_of_two().checked_ilog2() {
Some(x) => x as usize,
None => 0,
}
}
#[derive(Clone)]
pub struct EnumTypePropertiesBuilder {
type_properties: TypeProperties,
variant_count: usize,
}
impl EnumTypePropertiesBuilder {
#[must_use]
pub const fn new() -> Self {
Self {
type_properties: TypeProperties {
is_passive: true,
is_storable: true,
is_castable_from_bits: true,
bit_width: 0,
},
variant_count: 0,
}
}
pub const fn clone(&self) -> Self {
Self { ..*self }
}
#[must_use]
pub const fn variant(self, field_props: Option<TypeProperties>) -> Self {
let Self {
mut type_properties,
variant_count,
} = self;
if let Some(TypeProperties {
is_passive,
is_storable,
is_castable_from_bits,
bit_width,
}) = field_props
{
assert!(is_passive, "variant type must be a passive type");
type_properties = TypeProperties {
is_passive: true,
is_storable: type_properties.is_storable & is_storable,
is_castable_from_bits: type_properties.is_castable_from_bits
& is_castable_from_bits,
bit_width: if type_properties.bit_width < bit_width {
bit_width
} else {
type_properties.bit_width
},
};
}
Self {
type_properties,
variant_count: variant_count + 1,
}
}
pub const fn finish(self) -> TypeProperties {
assert!(
self.variant_count != 0,
"zero-variant enums aren't yet supported: \
https://github.com/chipsalliance/firrtl-spec/issues/208",
);
let Some(bit_width) = self
.type_properties
.bit_width
.checked_add(discriminant_bit_width_impl(self.variant_count))
else {
panic!("enum is too big: bit-width overflowed");
};
TypeProperties {
bit_width,
..self.type_properties
}
}
}
impl Default for EnumTypePropertiesBuilder {
fn default() -> Self {
Self::new()
}
}
impl Enum {
#[track_caller]
pub fn new(variants: Interned<[EnumVariant]>) -> Self {
let mut name_indexes = HashMap::with_capacity(variants.len());
let mut type_props_builder = EnumTypePropertiesBuilder::new();
for (index, EnumVariant { name, ty }) in variants.iter().enumerate() {
if let Some(old_index) = name_indexes.insert(*name, index) {
panic!("duplicate variant name {name:?}: at both index {old_index} and {index}");
}
type_props_builder = type_props_builder.variant(ty.map(CanonicalType::type_properties));
}
Self(
EnumImpl {
variants,
name_indexes,
type_properties: type_props_builder.finish(),
}
.intern_sized(),
)
}
pub fn discriminant_bit_width(self) -> usize {
discriminant_bit_width_impl(self.variants().len())
}
pub fn type_properties(self) -> TypeProperties {
self.0.type_properties
}
pub fn name_indexes(&self) -> &HashMap<Interned<str>, usize> {
&self.0.name_indexes
}
pub fn can_connect(self, rhs: Self) -> bool {
if self.0.variants.len() != rhs.0.variants.len() {
return false;
}
for (
&EnumVariant {
name: lhs_name,
ty: lhs_ty,
},
&EnumVariant {
name: rhs_name,
ty: rhs_ty,
},
) in self.0.variants.iter().zip(rhs.0.variants.iter())
{
if lhs_name != rhs_name {
return false;
}
match (lhs_ty, rhs_ty) {
(None, None) => {}
(None, Some(_)) | (Some(_), None) => return false,
(Some(lhs_ty), Some(rhs_ty)) => {
if !lhs_ty.can_connect(rhs_ty) {
return false;
}
}
}
}
true
}
}
pub trait EnumType:
Type<
BaseType = Enum,
MaskType = Bool,
MatchActiveScope = Scope,
MatchVariantAndInactiveScope = EnumMatchVariantAndInactiveScope<Self>,
MatchVariantsIter = EnumMatchVariantsIter<Self>,
>
{
fn variants(&self) -> Interned<[EnumVariant]>;
fn match_activate_scope(
v: Self::MatchVariantAndInactiveScope,
) -> (Self::MatchVariant, Self::MatchActiveScope);
}
pub struct EnumMatchVariantAndInactiveScope<T: EnumType>(EnumMatchVariantAndInactiveScopeImpl<T>);
impl<T: EnumType> MatchVariantAndInactiveScope for EnumMatchVariantAndInactiveScope<T> {
type MatchVariant = T::MatchVariant;
type MatchActiveScope = Scope;
fn match_activate_scope(self) -> (Self::MatchVariant, Self::MatchActiveScope) {
T::match_activate_scope(self)
}
}
impl<T: EnumType> EnumMatchVariantAndInactiveScope<T> {
pub fn variant_access(&self) -> VariantAccess {
self.0.variant_access()
}
pub fn activate(self) -> (VariantAccess, Scope) {
self.0.activate()
}
}
#[derive(Clone)]
pub struct EnumMatchVariantsIter<T: EnumType> {
pub(crate) inner: EnumMatchVariantsIterImpl<T>,
pub(crate) variant_index: std::ops::Range<usize>,
}
impl<T: EnumType> Iterator for EnumMatchVariantsIter<T> {
type Item = EnumMatchVariantAndInactiveScope<T>;
fn next(&mut self) -> Option<Self::Item> {
self.variant_index.next().map(|variant_index| {
EnumMatchVariantAndInactiveScope(self.inner.for_variant_index(variant_index))
})
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.variant_index.size_hint()
}
}
impl<T: EnumType> ExactSizeIterator for EnumMatchVariantsIter<T> {
fn len(&self) -> usize {
self.variant_index.len()
}
}
impl<T: EnumType> FusedIterator for EnumMatchVariantsIter<T> {}
impl<T: EnumType> DoubleEndedIterator for EnumMatchVariantsIter<T> {
fn next_back(&mut self) -> Option<Self::Item> {
self.variant_index.next_back().map(|variant_index| {
EnumMatchVariantAndInactiveScope(self.inner.for_variant_index(variant_index))
})
}
}
impl EnumType for Enum {
fn match_activate_scope(
v: Self::MatchVariantAndInactiveScope,
) -> (Self::MatchVariant, Self::MatchActiveScope) {
let (expr, scope) = v.0.activate();
(expr.variant_type().map(|_| expr.to_expr()), scope)
}
fn variants(&self) -> Interned<[EnumVariant]> {
self.0.variants
}
}
impl Type for Enum {
type BaseType = Enum;
type MaskType = Bool;
type MatchVariant = Option<Expr<CanonicalType>>;
type MatchActiveScope = Scope;
type MatchVariantAndInactiveScope = EnumMatchVariantAndInactiveScope<Self>;
type MatchVariantsIter = EnumMatchVariantsIter<Self>;
fn match_variants(
this: Expr<Self>,
source_location: SourceLocation,
) -> Self::MatchVariantsIter {
enum_match_variants_helper(this, source_location)
}
fn mask_type(&self) -> Self::MaskType {
Bool
}
fn canonical(&self) -> CanonicalType {
CanonicalType::Enum(*self)
}
#[track_caller]
fn from_canonical(canonical_type: CanonicalType) -> Self {
let CanonicalType::Enum(retval) = canonical_type else {
panic!("expected enum");
};
retval
}
fn source_location() -> SourceLocation {
SourceLocation::builtin()
}
}
#[hdl]
pub enum HdlOption<T: Type> {
HdlNone,
HdlSome(T),
}
#[allow(non_snake_case)]
pub fn HdlNone<T: StaticType>() -> Expr<HdlOption<T>> {
HdlOption[T::TYPE].HdlNone()
}
#[allow(non_snake_case)]
pub fn HdlSome<T: Type>(value: impl ToExpr<Type = T>) -> Expr<HdlOption<T>> {
let value = value.to_expr();
HdlOption[Expr::ty(value)].HdlSome(value)
}
impl<T: Type> HdlOption<T> {
#[track_caller]
pub fn try_map<R: Type, E>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Result<Expr<R>, E>,
) -> Result<Expr<HdlOption<R>>, E> {
Self::try_and_then(expr, |v| Ok(HdlSome(f(v)?)))
}
#[track_caller]
pub fn map<R: Type>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Expr<R>,
) -> Expr<HdlOption<R>> {
Self::and_then(expr, |v| HdlSome(f(v)))
}
#[hdl]
#[track_caller]
pub fn try_and_then<R: Type, E>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Result<Expr<HdlOption<R>>, E>,
) -> Result<Expr<HdlOption<R>>, E> {
type Wrap<T> = T;
#[hdl]
let mut and_then_out = incomplete_wire();
let mut iter = Self::match_variants(expr, SourceLocation::caller());
let none = iter.next().unwrap();
let some = iter.next().unwrap();
assert!(iter.next().is_none());
let (Wrap::<<Self as Type>::MatchVariant>::HdlSome(value), some_scope) =
Self::match_activate_scope(some)
else {
unreachable!();
};
let value = f(value).inspect_err(|_| {
and_then_out.complete(()); })?;
let and_then_out = and_then_out.complete(Expr::ty(value));
connect(and_then_out, value);
drop(some_scope);
let (Wrap::<<Self as Type>::MatchVariant>::HdlNone, none_scope) =
Self::match_activate_scope(none)
else {
unreachable!();
};
connect(and_then_out, Expr::ty(and_then_out).HdlNone());
drop(none_scope);
Ok(and_then_out)
}
#[track_caller]
pub fn and_then<R: Type>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Expr<HdlOption<R>>,
) -> Expr<HdlOption<R>> {
match Self::try_and_then(expr, |v| Ok::<_, Infallible>(f(v))) {
Ok(v) => v,
Err(e) => match e {},
}
}
#[hdl]
#[track_caller]
pub fn and<U: Type>(expr: Expr<Self>, opt_b: Expr<HdlOption<U>>) -> Expr<HdlOption<U>> {
#[hdl]
let and_out = wire(Expr::ty(opt_b));
connect(and_out, Expr::ty(opt_b).HdlNone());
#[hdl]
if let HdlSome(_) = expr {
connect(and_out, opt_b);
}
and_out
}
#[hdl]
#[track_caller]
pub fn try_filter<E>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Result<Expr<Bool>, E>,
) -> Result<Expr<Self>, E> {
#[hdl]
let filtered = wire(Expr::ty(expr));
connect(filtered, Expr::ty(expr).HdlNone());
let mut f = Some(f);
#[hdl]
if let HdlSome(v) = expr {
#[hdl]
if f.take().unwrap()(v)? {
connect(filtered, HdlSome(v));
}
}
Ok(filtered)
}
#[hdl]
#[track_caller]
pub fn filter(expr: Expr<Self>, f: impl FnOnce(Expr<T>) -> Expr<Bool>) -> Expr<Self> {
match Self::try_filter(expr, |v| Ok::<_, Infallible>(f(v))) {
Ok(v) => v,
Err(e) => match e {},
}
}
#[hdl]
#[track_caller]
pub fn try_inspect<E>(
expr: Expr<Self>,
f: impl FnOnce(Expr<T>) -> Result<(), E>,
) -> Result<Expr<Self>, E> {
let mut f = Some(f);
#[hdl]
if let HdlSome(v) = expr {
f.take().unwrap()(v)?;
}
Ok(expr)
}
#[hdl]
#[track_caller]
pub fn inspect(expr: Expr<Self>, f: impl FnOnce(Expr<T>)) -> Expr<Self> {
let mut f = Some(f);
#[hdl]
if let HdlSome(v) = expr {
f.take().unwrap()(v);
}
expr
}
#[hdl]
#[track_caller]
pub fn is_none(expr: Expr<Self>) -> Expr<Bool> {
#[hdl]
let is_none_out: Bool = wire();
connect(is_none_out, false);
#[hdl]
if let HdlNone = expr {
connect(is_none_out, true);
}
is_none_out
}
#[hdl]
#[track_caller]
pub fn is_some(expr: Expr<Self>) -> Expr<Bool> {
#[hdl]
let is_some_out: Bool = wire();
connect(is_some_out, false);
#[hdl]
if let HdlSome(_) = expr {
connect(is_some_out, true);
}
is_some_out
}
#[hdl]
#[track_caller]
pub fn map_or<R: Type>(
expr: Expr<Self>,
default: Expr<R>,
f: impl FnOnce(Expr<T>) -> Expr<R>,
) -> Expr<R> {
#[hdl]
let mapped = wire(Expr::ty(default));
let mut f = Some(f);
#[hdl]
match expr {
HdlSome(v) => connect(mapped, f.take().unwrap()(v)),
HdlNone => connect(mapped, default),
}
mapped
}
#[hdl]
#[track_caller]
pub fn map_or_else<R: Type>(
expr: Expr<Self>,
default: impl FnOnce() -> Expr<R>,
f: impl FnOnce(Expr<T>) -> Expr<R>,
) -> Expr<R> {
#[hdl]
let mut mapped = incomplete_wire();
let mut default = Some(default);
let mut f = Some(f);
let mut retval = None;
#[hdl]
match expr {
HdlSome(v) => {
let v = f.take().unwrap()(v);
let mapped = *retval.get_or_insert_with(|| mapped.complete(Expr::ty(v)));
connect(mapped, v);
}
HdlNone => {
let v = default.take().unwrap()();
let mapped = *retval.get_or_insert_with(|| mapped.complete(Expr::ty(v)));
connect(mapped, v);
}
}
retval.unwrap()
}
#[hdl]
#[track_caller]
pub fn or(expr: Expr<Self>, opt_b: Expr<Self>) -> Expr<Self> {
#[hdl]
let or_out = wire(Expr::ty(expr));
connect(or_out, opt_b);
#[hdl]
if let HdlSome(_) = expr {
connect(or_out, expr);
}
or_out
}
#[hdl]
#[track_caller]
pub fn or_else(expr: Expr<Self>, f: impl FnOnce() -> Expr<Self>) -> Expr<Self> {
#[hdl]
let or_else_out = wire(Expr::ty(expr));
connect(or_else_out, f());
#[hdl]
if let HdlSome(_) = expr {
connect(or_else_out, expr);
}
or_else_out
}
#[hdl]
#[track_caller]
pub fn unwrap_or(expr: Expr<Self>, default: Expr<T>) -> Expr<T> {
#[hdl]
let unwrap_or_else_out = wire(Expr::ty(default));
connect(unwrap_or_else_out, default);
#[hdl]
if let HdlSome(v) = expr {
connect(unwrap_or_else_out, v);
}
unwrap_or_else_out
}
#[hdl]
#[track_caller]
pub fn unwrap_or_else(expr: Expr<Self>, f: impl FnOnce() -> Expr<T>) -> Expr<T> {
#[hdl]
let unwrap_or_else_out = wire(Expr::ty(expr).HdlSome);
connect(unwrap_or_else_out, f());
#[hdl]
if let HdlSome(v) = expr {
connect(unwrap_or_else_out, v);
}
unwrap_or_else_out
}
#[hdl]
#[track_caller]
pub fn xor(expr: Expr<Self>, opt_b: Expr<Self>) -> Expr<Self> {
#[hdl]
let xor_out = wire(Expr::ty(expr));
#[hdl]
if let HdlSome(_) = expr {
#[hdl]
if let HdlNone = opt_b {
connect(xor_out, expr);
} else {
connect(xor_out, Expr::ty(expr).HdlNone());
}
} else {
connect(xor_out, opt_b);
}
xor_out
}
#[hdl]
#[track_caller]
pub fn zip<U: Type>(expr: Expr<Self>, other: Expr<HdlOption<U>>) -> Expr<HdlOption<(T, U)>> {
#[hdl]
let zip_out = wire(HdlOption[(Expr::ty(expr).HdlSome, Expr::ty(other).HdlSome)]);
connect(zip_out, Expr::ty(zip_out).HdlNone());
#[hdl]
if let HdlSome(l) = expr {
#[hdl]
if let HdlSome(r) = other {
connect(zip_out, HdlSome((l, r)));
}
}
zip_out
}
}
impl<T: Type> HdlOption<HdlOption<T>> {
#[hdl]
#[track_caller]
pub fn flatten(expr: Expr<Self>) -> Expr<HdlOption<T>> {
#[hdl]
let flattened = wire(Expr::ty(expr).HdlSome);
#[hdl]
match expr {
HdlSome(v) => connect(flattened, v),
HdlNone => connect(flattened, Expr::ty(expr).HdlSome.HdlNone()),
}
flattened
}
}
impl<T: Type, U: Type> HdlOption<(T, U)> {
#[hdl]
#[track_caller]
pub fn unzip(expr: Expr<Self>) -> Expr<(HdlOption<T>, HdlOption<U>)> {
let (t, u) = Expr::ty(expr).HdlSome;
#[hdl]
let unzipped = wire((HdlOption[t], HdlOption[u]));
connect(unzipped, (HdlOption[t].HdlNone(), HdlOption[u].HdlNone()));
#[hdl]
if let HdlSome(v) = expr {
connect(unzipped.0, HdlSome(v.0));
connect(unzipped.1, HdlSome(v.1));
}
unzipped
}
}