From 25ddc837c1547fae2197c1c8a1e14dffb59de937 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Sun, 2 Mar 2025 16:11:05 -0800
Subject: [PATCH] add Expr<ArrayType<T, Len>>: IntoIterator and Expr<Array<T>>:
 FromIterator<T>

---
 crates/fayalite/src/array.rs    | 139 ++++++++++++++++++++++++++------
 crates/fayalite/src/expr/ops.rs |  44 ++++++++++
 2 files changed, 159 insertions(+), 24 deletions(-)

diff --git a/crates/fayalite/src/array.rs b/crates/fayalite/src/array.rs
index 647b2e2..0d9b63f 100644
--- a/crates/fayalite/src/array.rs
+++ b/crates/fayalite/src/array.rs
@@ -3,7 +3,7 @@
 
 use crate::{
     expr::{
-        ops::{ArrayIndex, ArrayLiteral, ExprPartialEq},
+        ops::{ArrayLiteral, ExprFromIterator, ExprIntoIterator, ExprPartialEq},
         CastToBits, Expr, HdlPartialEq, ReduceBits, ToExpr,
     },
     int::{Bool, DynSize, KnownSize, Size, SizeType, DYN_SIZE},
@@ -15,7 +15,7 @@ use crate::{
     },
     util::ConstUsize,
 };
-use std::ops::Index;
+use std::{iter::FusedIterator, ops::Index};
 
 #[derive(Copy, Clone, PartialEq, Eq, Hash)]
 pub struct ArrayType<T: Type = CanonicalType, Len: Size = DynSize> {
@@ -151,10 +151,8 @@ impl<T: Type, Len: Size> Type for ArrayType<T, Len> {
         this: Expr<Self>,
         source_location: SourceLocation,
     ) -> Self::MatchVariantsIter {
-        let base = Expr::as_dyn_array(this);
-        let base_ty = Expr::ty(base);
         let _ = source_location;
-        let retval = Vec::from_iter((0..base_ty.len()).map(|i| ArrayIndex::new(base, i).to_expr()));
+        let retval = Vec::from_iter(this);
         std::iter::once(MatchVariantWithoutScope(
             Len::ArrayMatch::<T>::try_from(retval)
                 .ok()
@@ -187,9 +185,7 @@ impl<T: Type, Len: Size> Type for ArrayType<T, Len> {
 
 impl<T: Type, Len: Size> TypeWithDeref for ArrayType<T, Len> {
     fn expr_deref(this: &Expr<Self>) -> &Self::MatchVariant {
-        let base = Expr::as_dyn_array(*this);
-        let base_ty = Expr::ty(base);
-        let retval = Vec::from_iter((0..base_ty.len()).map(|i| ArrayIndex::new(base, i).to_expr()));
+        let retval = Vec::from_iter(*this);
         Interned::into_inner(Intern::intern_sized(
             Len::ArrayMatch::<T>::try_from(retval)
                 .ok()
@@ -230,27 +226,122 @@ where
         let lhs_ty = Expr::ty(lhs);
         let rhs_ty = Expr::ty(rhs);
         assert_eq!(lhs_ty.len(), rhs_ty.len());
-        ArrayLiteral::<Bool, DynSize>::new(
-            Bool,
-            (0..lhs_ty.len())
-                .map(|i| Expr::canonical(lhs[i].cmp_eq(rhs[i])))
-                .collect(),
-        )
-        .cast_to_bits()
-        .all_one_bits()
+        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());
-        ArrayLiteral::<Bool, DynSize>::new(
-            Bool,
-            (0..lhs_ty.len())
-                .map(|i| Expr::canonical(lhs[i].cmp_ne(rhs[i])))
-                .collect(),
-        )
-        .cast_to_bits()
-        .any_one_bits()
+        lhs.into_iter()
+            .zip(rhs)
+            .map(|(l, r)| l.cmp_ne(r))
+            .collect::<Expr<Array<Bool>>>()
+            .cast_to_bits()
+            .any_one_bits()
+    }
+}
+
+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()
     }
 }
diff --git a/crates/fayalite/src/expr/ops.rs b/crates/fayalite/src/expr/ops.rs
index 15c195e..c502fd5 100644
--- a/crates/fayalite/src/expr/ops.rs
+++ b/crates/fayalite/src/expr/ops.rs
@@ -2708,3 +2708,47 @@ impl<T: Type> ToExpr for Uninit<T> {
         }
     }
 }
+
+pub trait ExprIntoIterator: Type {
+    type Item: Type;
+    type ExprIntoIter: Iterator<Item = Expr<Self::Item>>;
+
+    fn expr_into_iter(e: Expr<Self>) -> Self::ExprIntoIter;
+}
+
+impl<T: ExprIntoIterator> IntoIterator for Expr<T> {
+    type Item = Expr<T::Item>;
+    type IntoIter = T::ExprIntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        T::expr_into_iter(self)
+    }
+}
+
+impl<T: ExprIntoIterator> IntoIterator for &'_ Expr<T> {
+    type Item = Expr<T::Item>;
+    type IntoIter = T::ExprIntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        T::expr_into_iter(*self)
+    }
+}
+
+impl<T: ExprIntoIterator> IntoIterator for &'_ mut Expr<T> {
+    type Item = Expr<T::Item>;
+    type IntoIter = T::ExprIntoIter;
+
+    fn into_iter(self) -> Self::IntoIter {
+        T::expr_into_iter(*self)
+    }
+}
+
+pub trait ExprFromIterator<A>: Type {
+    fn expr_from_iter<T: IntoIterator<Item = A>>(iter: T) -> Expr<Self>;
+}
+
+impl<This: ExprFromIterator<A>, A> FromIterator<A> for Expr<This> {
+    fn from_iter<T: IntoIterator<Item = A>>(iter: T) -> Self {
+        This::expr_from_iter(iter)
+    }
+}