From 96b012ae19efe199720f1602bb4ef655d81955f4 Mon Sep 17 00:00:00 2001
From: Ben Gamari <ben@smart-cactus.org>
Date: Mon, 9 Sep 2024 10:55:01 -0400
Subject: [PATCH] BiDf: Introduce fanin

This is a useful operation for servicing requests from a number of
sources with a single sink.
---
 clash-protocols/src/Protocols/BiDf.hs         | 37 +++++++++++++
 clash-protocols/tests/Tests/Protocols/BiDf.hs | 53 ++++++++++++++++++-
 2 files changed, 89 insertions(+), 1 deletion(-)

diff --git a/clash-protocols/src/Protocols/BiDf.hs b/clash-protocols/src/Protocols/BiDf.hs
index 175cc8f8..7079e8ac 100644
--- a/clash-protocols/src/Protocols/BiDf.hs
+++ b/clash-protocols/src/Protocols/BiDf.hs
@@ -13,6 +13,8 @@ module Protocols.BiDf (
   loopback,
     -- * Mapping
   dimap,
+  -- * Fan-in
+  fanin
 ) where
 
 import Prelude ()
@@ -100,3 +102,38 @@ dimap f g = circuit $ \biDf -> do
   resp' <- Df.map g -< resp
   (biDf', resp) <- fromDfs -< req'
   idC -< biDf'
+
+-- | Merge a number of 'BiDf's, preferring requests from the last channel.
+fanin
+  :: forall n dom req resp.
+     ( KnownNat n
+     , 1 <= n
+     , NFDataX req
+     , NFDataX resp
+     , HiddenClockResetEnable dom
+     )
+  => Circuit (Vec n (BiDf dom req resp)) (BiDf dom req resp)
+fanin = fromSignals $ \(upFwds, (reqAck, respData)) ->
+  let reqDatas :: Vec n (Signal dom (Df.Data req))
+      reqDatas = map fst upFwds
+      respAcks :: Vec n (Signal dom Ack)
+      respAcks = map snd upFwds
+
+      ((reqAcks, respAck), (respDatas, reqData)) =
+           toSignals fanin' ((reqDatas, respData), (respAcks, reqAck))
+  in (zip reqAcks respDatas, (reqData, respAck))
+ where
+  fanin'
+      :: Circuit (Vec n (Df dom req), Df dom resp)
+                 (Vec n (Df dom resp), Df dom req)
+  fanin' = circuit $ \(reqs, resp) -> do
+    [fwd0, fwd1]
+        <- Df.fanout
+        <| Df.roundrobinCollect @n Df.Parallel
+        <| repeatWithIndexC (\i -> Df.map (\x -> (i,x)))
+        -< reqs
+
+    activeN <- Df.map fst -< fwd1
+    resps <- Df.route <| Df.zip -< (activeN, resp)
+    req <- Df.map snd -< fwd0
+    idC -< (resps, req)
diff --git a/clash-protocols/tests/Tests/Protocols/BiDf.hs b/clash-protocols/tests/Tests/Protocols/BiDf.hs
index b05d6fdf..9095c20d 100644
--- a/clash-protocols/tests/Tests/Protocols/BiDf.hs
+++ b/clash-protocols/tests/Tests/Protocols/BiDf.hs
@@ -3,8 +3,11 @@
 
 module Tests.Protocols.BiDf (tests) where
 
+import Prelude as P
+import qualified Data.List as L
+
 -- clash-prelude
-import Clash.Prelude
+import Clash.Prelude as C
 import qualified Clash.Sized.Vector as Vector
 import Clash.Hedgehog.Sized.Vector
 
@@ -51,6 +54,54 @@ prop_loopback_id =
     BiDf.loopback id -< biDf
     idC -< resp
 
+-- | Test that 'BiDf.fanin' on a single 'BiDf' channel behaves as an identity.
+prop_fanin_id :: Property
+prop_fanin_id =
+  idWithModelSingleDomain @System defExpectOptions gen (\_ _ _ -> id) (exposeClockResetEnable impl)
+ where
+  gen :: Gen [Int]
+  gen = Gen.list (Range.linear 0 10) (Gen.integral (Range.linear 0 100))
+
+  impl
+   :: forall dom a. (HiddenClockResetEnable dom, NFDataX a)
+   => Circuit (Df dom a) (Df dom a)
+  impl = circuit $ \req -> do
+    (biDf, resp) <- BiDf.fromDfs -< req
+    BiDf.loopback id <| BiDf.fanin @1 -< [biDf]
+    idC -< resp
+
+-- | Test that 'BiDf.fanin' on a number of 'BiDf' channels behaves as an
+-- identity on each channel.
+prop_fanin :: Property
+prop_fanin =
+  idWithModelSingleDomain @System expectOpts
+    (gen @3)
+    (\_ _ _ -> id)
+    (exposeClockResetEnable impl)
+ where
+  expectOpts = defExpectOptions
+
+  gen :: forall n. KnownNat n => Gen (Vec n [(Index n, Int)])
+  gen = do
+    xs <- genVec @Gen @n $ Gen.list (Range.linear 0 10) (Gen.integral (Range.linear 0 100))
+    return $ C.zipWith (\i -> fmap (\x -> (i,x))) indicesI xs
+
+  impl
+    :: forall n dom a.
+      (HiddenClockResetEnable dom, KnownNat n, 1 <= n, NFDataX a)
+    => Circuit (Vec n (Df dom a)) (Vec n (Df dom a))
+  impl = circuit $ \reqs -> do
+    (biDfs, resps) <- unbundleC <| repeatC BiDf.fromDfs -< reqs
+    BiDf.loopback id <| BiDf.fanin @n -< biDfs
+    idC -< resps
+
+unbundleC :: forall n a b. Circuit (Vec n (a, b)) (Vec n a, Vec n b)
+unbundleC = fromSignals $ \(fwd, (bwdA, bwdB)) ->
+  let fwdA :: Vec n (Fwd a)
+      fwdB :: Vec n (Fwd b)
+      (fwdA, fwdB) = Vector.unzip fwd
+  in (Vector.zip bwdA bwdB, (fwdA, fwdB))
+
 tests :: TestTree
 tests =
   $(testGroupGenerator)
\ No newline at end of file