Add repeat and replicate to Protocols.Vec

Author: lmbollen

clash-protocols/src/Protocols/Vec.hs

@@ -11,14 +11,26 @@ module Protocols.Vec (
   unzip3,
   concat,
   unconcat,
+  repeat,
+  replicate,
 ) where
 
 -- base
 import Data.Tuple
 import Prelude ()
 
 -- clash-prelude
-import Clash.Prelude hiding (concat, split, unconcat, unzip, unzip3, zip, zip3)
+import Clash.Prelude hiding (
+  concat,
+  repeat,
+  replicate,
+  split,
+  unconcat,
+  unzip,
+  unzip3,
+  zip,
+  zip3,
+ )
 import Clash.Prelude qualified as C
 
 -- clash-protocols-base
@@ -131,3 +143,15 @@ split3Vec ::
 split3Vec v = (v0, v1, v2)
  where
   (v0, splitAtI -> (v1, v2)) = splitAtI v
+
+{- | repeat a circuit for a number of times, the number of times the circuit is repeated
+is determined by the type-level natural number `n`.
+-}
+repeat :: (C.KnownNat n) => Circuit a b -> Circuit (Vec n a) (Vec n b)
+repeat (Circuit function) = Circuit (C.unzip . uncurry (zipWith (curry function)))
+
+{- | replicate a circuit for a given number of times, the number of times the circuit is replicated
+is given by the supplied `SNat n`.
+-}
+replicate :: SNat n -> Circuit a b -> Circuit (Vec n a) (Vec n b)
+replicate SNat = repeat

clash-protocols/tests/Tests/Protocols/Vec.hs

@@ -18,6 +18,7 @@ import Test.Tasty.TH (testGroupGenerator)
 
 -- clash-protocols (me!)
 import Protocols
+import Protocols.Df qualified as Df
 import Protocols.Vec qualified as Vec
 
 import Clash.Hedgehog.Sized.Vector (genVec)
@@ -176,6 +177,34 @@ prop_unconcat =
   dut = Vec.unconcat C.d2
   model = C.unconcat C.d2
 
+prop_repeat :: Property
+prop_repeat =
+  idWithModel
+    @(C.Vec 3 (Df System Int))
+    @(C.Vec 3 (Df System Int))
+    defExpectOptions
+    gen
+    model
+    dut
+ where
+  gen = genVecData genSmallInt
+  dut = Vec.repeat (Df.map succ)
+  model = fmap (fmap succ)
+
+prop_replicate :: Property
+prop_replicate =
+  idWithModel
+    @(C.Vec 3 (Df System Int))
+    @(C.Vec 3 (Df System Int))
+    defExpectOptions
+    gen
+    model
+    dut
+ where
+  gen = genVecData genSmallInt
+  dut = Vec.replicate C.SNat (Df.map succ)
+  model = fmap (fmap succ)
+
 tests :: TestTree
 tests =
   -- TODO: Move timeout option to hedgehog for better error messages.