Move scanlPar, scanrPar to RTree

export properly
clash-lang · Apr 22, 2022 · a0e0565 · a0e0565
1 parent f0a5e75
commit a0e0565
Show file tree

Hide file tree

Showing 7 changed files with 70 additions and 101 deletions.
diff --git a/clash-prelude/clash-prelude.cabal b/clash-prelude/clash-prelude.cabal
@@ -274,7 +274,6 @@ Library
                       Clash.Sized.Signed
                       Clash.Sized.Unsigned
                       Clash.Sized.Vector
-                      Clash.Sized.Vector.Par
 
                       Clash.Sized.Internal.BitVector
                       Clash.Sized.Internal.Index

diff --git a/clash-prelude/src/Clash/Explicit/Prelude.hs b/clash-prelude/src/Clash/Explicit/Prelude.hs
@@ -103,7 +103,6 @@ module Clash.Explicit.Prelude
   , module Clash.Sized.Fixed
     -- *** Fixed size vectors
   , module Clash.Sized.Vector
-  , module Clash.Sized.Vector.Par
     -- *** Perfect depth trees
   , module Clash.Sized.RTree
     -- ** Annotations
@@ -187,7 +186,6 @@ import Clash.Sized.RTree
 import Clash.Sized.Signed
 import Clash.Sized.Unsigned
 import Clash.Sized.Vector hiding (fromList, unsafeFromList)
-import Clash.Sized.Vector.Par
 import Clash.XException
 
 {- $setup

diff --git a/clash-prelude/src/Clash/Explicit/Prelude/Safe.hs b/clash-prelude/src/Clash/Explicit/Prelude/Safe.hs
@@ -78,7 +78,6 @@ module Clash.Explicit.Prelude.Safe
   , module Clash.Sized.Fixed
     -- *** Fixed size vectors
   , module Clash.Sized.Vector
-  , module Clash.Sized.Vector.Par
     -- *** Perfect depth trees
   , module Clash.Sized.RTree
     -- ** Annotations
@@ -150,7 +149,6 @@ import Clash.Sized.RTree
 import Clash.Sized.Signed
 import Clash.Sized.Unsigned
 import Clash.Sized.Vector hiding (fromList, unsafeFromList)
-import Clash.Sized.Vector.Par
 import Clash.XException
 
 {- $setup

diff --git a/clash-prelude/src/Clash/Prelude.hs b/clash-prelude/src/Clash/Prelude.hs
@@ -126,7 +126,6 @@ module Clash.Prelude
   , module Clash.Sized.Fixed
     -- *** Fixed size vectors
   , module Clash.Sized.Vector
-  , module Clash.Sized.Vector.Par
     -- *** Perfect depth trees
   , module Clash.Sized.RTree
     -- ** Annotations
@@ -212,7 +211,6 @@ import           Clash.Sized.RTree
 import           Clash.Sized.Signed
 import           Clash.Sized.Unsigned
 import           Clash.Sized.Vector hiding (fromList, unsafeFromList)
-import           Clash.Sized.Vector.Par
 import           Clash.Signal hiding
   (HiddenClockName, HiddenResetName, HiddenEnableName)
 import           Clash.Signal.Delayed

diff --git a/clash-prelude/src/Clash/Prelude/Safe.hs b/clash-prelude/src/Clash/Prelude/Safe.hs
@@ -94,7 +94,6 @@ module Clash.Prelude.Safe
   , module Clash.Sized.Fixed
     -- *** Fixed size vectors
   , module Clash.Sized.Vector
-  , module Clash.Sized.Vector.Par
     -- *** Perfect depth trees
   , module Clash.Sized.RTree
     -- ** Annotations
@@ -163,7 +162,6 @@ import           Clash.Sized.RTree
 import           Clash.Sized.Signed
 import           Clash.Sized.Unsigned
 import           Clash.Sized.Vector hiding (fromList, unsafeFromList)
-import           Clash.Sized.Vector.Par
 import           Clash.Signal
 import           Clash.Signal.Delayed
 import           Clash.XException

diff --git a/clash-prelude/src/Clash/Sized/RTree.hs b/clash-prelude/src/Clash/Sized/RTree.hs
@@ -17,6 +17,8 @@ Maintainer :  QBayLogic B.V. <devops@qbaylogic.com>
 
 {-# OPTIONS_GHC -fplugin GHC.TypeLits.Normalise -fplugin GHC.TypeLits.KnownNat.Solver #-}
 
+{-# OPTIONS_GHC -fno-warn-incomplete-patterns #-}
+
 module Clash.Sized.RTree
   ( -- * 'RTree' data type
     RTree (LR, BR)
@@ -41,6 +43,9 @@ module Clash.Sized.RTree
   , tfold
     -- ** Specialised folds
   , tdfold
+    -- ** Prefix sums (scans)
+  , scanlPar
+  , scanrPar
     -- * Conversions
   , v2t
   , t2v
@@ -58,21 +63,23 @@ import Data.Foldable               (toList)
 import Data.Kind                   (Type)
 import Data.Singletons             (Apply, TyFun, type (@@))
 import Data.Proxy                  (Proxy (..))
-import GHC.TypeLits                (KnownNat, Nat, type (+), type (^), type (*))
+import Data.Type.Equality          ((:~:) (..))
+import GHC.TypeLits                (KnownNat, Nat, type (+), type (^), type (*),
+                                    sameNat)
 import Language.Haskell.TH.Syntax  (Lift(..))
 #if MIN_VERSION_template_haskell(2,16,0)
 import Language.Haskell.TH.Compat
 #endif
-import Prelude                     hiding ((++), (!!))
+import Prelude                     hiding ((++), (!!), map)
 import Test.QuickCheck             (Arbitrary (..), CoArbitrary (..))
 
 import Clash.Annotations.Primitive (hasBlackBox)
 import Clash.Class.BitPack         (BitPack (..), packXWith)
-import Clash.Promoted.Nat          (SNat (..), UNat (..), pow2SNat, snatToNum,
-                                    subSNat, toUNat)
+import Clash.Promoted.Nat          (SNat (..), SNatLE (..), UNat (..), compareSNat,
+                                    pow2SNat, snatToNum, subSNat, toUNat)
 import Clash.Promoted.Nat.Literals (d1)
 import Clash.Sized.Index           (Index)
-import Clash.Sized.Vector          (Vec (..), (!!), (++), dtfold, replace)
+import Clash.Sized.Vector          (Vec (..), (!!), (++), dtfold, map, replace)
 import Clash.XException
   (ShowX (..), NFDataX (..), isX, showsX, showsPrecXWith)
 
@@ -537,3 +544,61 @@ lazyT :: KnownNat d
       => RTree d a
       -> RTree d a
 lazyT = tzipWith (flip const) (trepeat ())
+
+-- | Low-depth (left) scan, see 'Clash.Sized.Vector.scanl'.
+--
+-- >>> scanlPar (+) (1 :> 2 :> 3 :> 4 :> Nil)
+-- 1 :> 3 :> 6 :> 10 :> Nil
+--
+-- The operator must be associative.
+scanlPar :: KnownNat n
+         => (a -> a -> a) -- ^ Must be associative
+         -> Vec (2^n) a -> Vec (2^n) a
+scanlPar op v = scanlInductiveRTree op (v2t v)
+{-# INLINE scanlPar #-}
+
+-- | Low-depth (right) scan, see 'Clash.Sized.Vector.scanr'.
+--
+-- >>> scanrPar (+) (1 :> 2 :> 3 :> 4 :> Nil)
+-- 10 :> 9 :> 7 :> 4 :> Nil
+--
+-- The operator must be associative.
+scanrPar :: KnownNat n
+         => (a -> a -> a) -- ^ Must be associative
+         -> Vec (2^n) a -> Vec (2^n) a
+scanrPar op v = scanrInductiveRTree op (v2t v)
+{-# INLINE scanrPar #-}
+
+scanlInductiveRTree
+  :: forall a n. KnownNat n
+  => (a -> a -> a)
+  -> RTree n a
+  -> Vec (2^n) a
+scanlInductiveRTree op tr =
+  -- I have to use sameNat and compareSNat both; the type checker cannot infer
+  -- that n <= 0 means n ~ 0.
+  case (sameNat (Proxy @n) (Proxy @0), compareSNat (SNat @n) (SNat @0), tr) of
+    (Just Refl, _, LR x) -> x :> Nil
+    (_, SNatGT, BR x y) ->
+        let
+          x' = scanlInductiveRTree op x
+          y' = scanlInductiveRTree op y
+          l = x' !! (length x'-1) -- 'last' doesn't work here
+        in x' ++ map (l `op`) y'
+
+scanrInductiveRTree
+  :: forall a n. KnownNat n
+  => (a -> a -> a)
+  -> RTree n a
+  -> Vec (2^n) a
+scanrInductiveRTree op tr =
+  -- I have to use sameNat and compareSNat both; the type checker cannot infer
+  -- that n <= 0 means n ~ 0.
+  case (sameNat (Proxy @n) (Proxy @0), compareSNat (SNat @n) (SNat @0), tr) of
+    (Just Refl, _, LR x) -> x :> Nil
+    (_, SNatGT, BR x y) ->
+        let
+          x' = scanrInductiveRTree op x
+          y' = scanrInductiveRTree op y
+          l = y' !! (0::Int)  -- `head` doesn't work here
+        in map (l `op`) x' ++ y'
diff --git a/clash-prelude/src/Clash/Sized/Vector/Par.hs b/clash-prelude/src/Clash/Sized/Vector/Par.hs