diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Revision history for mutable-fenwick + +## 0.1.0.0 -- 2025-04-07 + +* First version. Released on an unsuspecting world. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Parsa Alizadeh + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,22 @@ +# mutable-fenwick + +This package provides an implementation of mutable +[Fenwick trees](https://en.wikipedia.org/wiki/Fenwick_tree) in Haskell. + +## Features + +It is maximally generic. Each operation of Fenwick tree is implemented using a subset of constraints +from `Semigroup`, `Monoid`, or `Commutative`, chosen carefully based on the nature of each +operation. This is mostly possible due to how Haskell typeclasses work, and provides different +functionality based on the constraints provided by the underlying element type. + +It is fast and efficient. Every operation is marked as inline, meaning that they will be optimized +for the given element type. With `ArrayC` and `VectorC` from this package, it is possible to use +unboxed arrays and vectors for newtypes that implement a custom algebra (e.g. `Sum`, `Product` or +`Xor`). An implementation using this library can be as fast as a C/C++ implementation. + +It is the only Haskell library (I believe, as of this date) that provides Fenwick trees. +- The [FenwickTree](https://hackage.haskell.org/package/FenwickTree) package is more similar to a + Segment tree, and it does not have a generic interface for the data structure. +- The [binary-indexed-tree](https://hackage.haskell.org/package/binary-indexed-tree) package + has an interface for ST monad, but the implementation is only limited to `Sum` monoid. diff --git a/mutable-fenwick.cabal b/mutable-fenwick.cabal new file mode 100644 --- /dev/null +++ b/mutable-fenwick.cabal @@ -0,0 +1,68 @@ +cabal-version: 2.4 +name: mutable-fenwick +version: 0.1.0.0 + +synopsis: Mutable Fenwick trees +description: + This package provides an implementation of mutable + [Fenwick trees](https://en.wikipedia.org/wiki/Fenwick_tree). + . + It is maximally generic. Each operation of Fenwick tree is implemented using a subset of constraints + from @Semigroup@, @Monoid@, or @Commutative@, chosen based on the nature of each + operation. + . + It is fast and efficient. With @ArrayC@ and @VectorC@ from this package, it is possible to use + unboxed arrays and vectors for newtypes that implement a custom algebra (e.g. @Sum@, @Product@ or + @Xor@). An implementation using this library can be as fast as a C/C++ implementation. + +homepage: https://github.com/ParsaAlizadeh/fenwick-tree +bug-reports: https://github.com/ParsaAlizadeh/fenwick-tree/issues + +license: MIT +license-file: LICENSE +author: Parsa Alizadeh +maintainer: parsa.alizadeh1@gmail.com +copyright: (c) Parsa Alizade, 2025 +category: Data +extra-source-files: + README.md + CHANGELOG.md + +tested-with: GHC == 9.2.8 || == 9.4.8 + +source-repository head + type: git + location: https://github.com/ParsaAlizadeh/fenwick-tree.git + +library + hs-source-dirs: src + default-language: Haskell2010 + exposed-modules: + Data.Array.ArrayC, + Data.Vector.VectorC, + Data.Fenwick.Array, + Data.Fenwick.Vector + + -- Modules included in this library but not exported. + -- other-modules: + + build-depends: + base >= 4.17.2 && < 4.18, + array >= 0.5.4 && < 0.6, + commutative-semigroups >= 0.2.0 && < 0.3, + monoid-subclasses >= 1.2.6 && < 1.3, + vector >= 0.13.1 && < 0.14, + + +test-suite test-fenwick + type: exitcode-stdio-1.0 + hs-source-dirs: tests + default-language: Haskell2010 + build-depends: + base, + hspec, + mutable-fenwick, + array, + vector, + main-is: FenwickTest.hs + diff --git a/src/Data/Array/ArrayC.hs b/src/Data/Array/ArrayC.hs new file mode 100644 --- /dev/null +++ b/src/Data/Array/ArrayC.hs @@ -0,0 +1,65 @@ +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} + +-- | Module : Data.Array.ArrayC +-- +-- This modulo provides 'ArrayC', which given a existing 'MArray' implementation and a +-- representation type, it stores elements by coercing the element type into representation type. +-- Using this, we can have elements with algebraic structures (like 'Data.Monoid.Monoid' used in +-- this package) but stored in efficient unboxed arrays. +module Data.Array.ArrayC ( ArrayC ) where + +import Data.Array.MArray +import Data.Array.Base hiding (array, elems) +import Data.Coerce + +-- | Given an array type (e.g. 'Data.Array.IO.IOUArray') and a representation type (e.g. 'Int'), the +-- type constructor @'ArrayC' 'Data.Array.IO.IOUArray' 'Int'@ is a valid array type, that can store +-- elements that are coercible to 'Int' (the representation type). The same type can be used for +-- immutable and mutable arrays. Use functions from 'Data.Array.IArray' and 'Data.Array.MArray' to +-- create and modify the array. +newtype ArrayC array rep ix elem = ArrayC (array ix rep) + +instance (IArray array rep, Coercible rep elem) => IArray (ArrayC array rep) elem where + bounds (ArrayC array) = bounds array + {-# INLINE bounds #-} + + numElements (ArrayC array) = numElements array + {-# INLINE numElements #-} + + unsafeArray ix elems = ArrayC $ unsafeArray ix (coerce elems) + {-# INLINE unsafeArray #-} + + unsafeAt (ArrayC array) ix = coerce $ unsafeAt array ix + {-# INLINE unsafeAt #-} + + unsafeReplace (ArrayC array) elems = ArrayC $ unsafeReplace array (coerce elems) + {-# INLINE unsafeReplace #-} + + unsafeAccum f (ArrayC array) elems = ArrayC $ unsafeAccum (coerce f) array elems + {-# INLINE unsafeAccum #-} + + unsafeAccumArray f e ix elems = ArrayC $ unsafeAccumArray (coerce f) (coerce e) ix elems + {-# INLINE unsafeAccumArray #-} + +instance (Monad m, MArray arr r m, Coercible r e) => MArray (ArrayC arr r) e m where + getBounds (ArrayC arr) = getBounds arr + {-# INLINE getBounds #-} + + getNumElements (ArrayC arr) = getNumElements arr + {-# INLINE getNumElements #-} + + newArray ix e = ArrayC <$> newArray ix (coerce e) + {-# INLINE newArray #-} + + newArray_ ix = ArrayC <$> newArray_ ix + {-# INLINE newArray_ #-} + + unsafeNewArray_ ix = ArrayC <$> unsafeNewArray_ ix + {-# INLINE unsafeNewArray_ #-} + + unsafeRead (ArrayC arr) i = coerce <$> unsafeRead arr i + {-# INLINE unsafeRead #-} + + unsafeWrite (ArrayC arr) i e = unsafeWrite arr i (coerce e) + {-# INLINE unsafeWrite #-} diff --git a/src/Data/Fenwick/Array.hs b/src/Data/Fenwick/Array.hs new file mode 100644 --- /dev/null +++ b/src/Data/Fenwick/Array.hs @@ -0,0 +1,145 @@ +{-# LANGUAGE BangPatterns #-} + +-- | Module : Data.Fenwick.Array +-- +-- This modulo provides mutable [Fenwick Trees](https://en.wikipedia.org/wiki/Fenwick_tree), using +-- arrays as the underlying data structure. The algebraic structure is given by 'Semigroup' and +-- 'Monoid' instances. Some of the functions require the structure to be +-- 'Data.Semigroup.Commutative.Commutative'. See @monoid-subclasses@ and @commutative-semigroups@. +-- If you want mark @Sum@ or @Product@ as @Commutative@, see @SumCommutative@ and +-- @ProductCommutative@ in the mentioned packages. + +module Data.Fenwick.Array + ( FenMArray + , newFen + , newAccumFen + , newListFen + , getSizeFen + , addFen + , sumPrefixFen + , lowerBoundFen + ) where + +import Data.Array.MArray +import Control.Monad +import Data.Bits +import Data.Semigroup.Cancellative +import Data.Monoid.Cancellative +import Data.Array.Base + +-- | Fenwick tree datatype. @array@ must be a type such that @array 'Int' elem@ refers to a valid +-- mutable array. Fenwick tree is assumed to be data structure over a 1-based array of size @n@. +data FenMArray array elem = FenMArray !Int !(array Int elem) + +-- | least significant bit +lsb :: Int -> Int +lsb node = node .&. (-node) +{-# INLINE lsb #-} + +modifyArray' :: (MArray a e m, Ix i) => a i e -> i -> (e -> e) -> m () +modifyArray' arr i f = do + x <- readArray arr i + let !x' = f x + writeArray arr i x' +{-# INLINABLE modifyArray' #-} + +unsafeModifyArray' :: (MArray a e m, Ix i) => a i e -> Int -> (e -> e) -> m () +unsafeModifyArray' arr i f = do + x <- unsafeRead arr i + let !x' = f x + unsafeWrite arr i x' +{-# INLINABLE unsafeModifyArray' #-} + +-- | Creates a Fenwick tree of size @n@ over 1-based array. All elements of the array are initially +-- @'mempty'@. \( O(n) \) +-- +-- Using type application extension, you can specify the array and element type as the first and +-- second type argument. +-- +-- @ +-- fen <- 'newFen' \@('Data.Array.ArrayC.ArrayC' 'Data.Array.IO.IOUArray' 'Int') \@('Data.Monoid.Sum' 'Int') +-- @ +newFen :: (MArray array elem m, Monoid elem) => Int -> m (FenMArray array elem) +newFen n = do + arr <- newArray (0, n) mempty + pure (FenMArray n arr) +{-# INLINABLE newFen #-} + +-- | @newAccumFen n xs@ creates a Fenwick tree of size @n@ over 1-based array. Every item of @xs@ is +-- pair of an index and an element. The elements of the array are initialized by @'<>'@ing elements +-- for each index. This functions is faster than creating empty array using 'newFen' and adding +-- elements individually using 'addFen'. \( O(n + m) \) where \( m \) is @'Data.Foldable.length' xs@ +newAccumFen :: (MArray array elem m, CommutativeMonoid elem, Foldable t) => Int -> t (Int, elem) -> m (FenMArray array elem) +newAccumFen n xs = do + arr <- newArray (0, n) mempty + forM_ xs $ \(i, e) -> do + modifyArray' arr i (<> e) + forM_ [1..n] $ \i -> do + let j = i + lsb i + when (j <= n) $ do + e <- readArray arr i + unsafeModifyArray' arr j (e <>) + pure (FenMArray n arr) +{-# INLINABLE newAccumFen #-} + +-- | Creates a Fenwick tree and initialize the array based on the elements of the list. \( O(n) \) +newListFen :: (MArray array elem m, CommutativeMonoid elem) => Int -> [elem] -> m (FenMArray array elem) +newListFen n xs = newAccumFen n $ zip [1..n] xs +{-# INLINABLE newListFen #-} + +-- | Get the size of underlying array. \( O(1) \) +getSizeFen :: FenMArray array elem -> Int +getSizeFen (FenMArray n _) = n +{-# INLINE getSizeFen #-} + +-- | Add a value to a cell of the array. The index must be in the range @[1, n]@. \( O(\log n) \) +addFen :: (MArray array elem m, Commutative elem) => FenMArray array elem -> Int -> elem -> m () +addFen (FenMArray n arr) r a = go (check r) where + check i + | i < 1 || i > n = error "index out of range" + | otherwise = i + go i = when (i <= n) $ do + unsafeModifyArray' arr i (a <>) + go (i + lsb i) +{-# INLINABLE addFen #-} + +-- | Given index @r@, get the prefix sum of elements of the array in the range @[1, r]@. It accepts +-- values of @r@ out of the range of indices, assuming that every element out of the range of array +-- is 'mempty'. \( O(\log n) \) +sumPrefixFen :: (MArray array elem m, Monoid elem) => FenMArray array elem -> Int -> m elem +sumPrefixFen (FenMArray n arr) = go mempty . min n where + -- prefix (0, r], for 1 <= r <= n + go !s i + | i <= 0 = pure s + | otherwise = do + x <- unsafeRead arr i + go (x <> s) (i - lsb i) +{-# INLINABLE sumPrefixFen #-} + +-- | Given a prefix sum @q@, find the least index @r@ such that the prefix sum of @[1, r]@ is at +-- least @q@. Only applicable when partial sums are ordered. If the sum of the whole array is +-- less than @query@, it returns @n+1@. This functions is faster than binary search over +-- 'sumPrefixFen'. \( O(\log n) \) +lowerBoundFen :: (MArray array elem m, Monoid elem, Ord elem) => FenMArray array elem -> elem -> m Int +lowerBoundFen (FenMArray n arr) query = go root (n + 1) mempty where + root = bit (finiteBitSize n - countLeadingZeros n - 1) + go node + | odd node = leaf node + | otherwise = nonleaf node + leaf node fallback prepend + | node > n = pure fallback + | otherwise = do + nodeval <- unsafeRead arr node + pure $ if prepend <> nodeval < query + then fallback + else node + nonleaf node fallback prepend + | node > n = go (leftOf node) fallback prepend + | otherwise = do + nodeval <- unsafeRead arr node + if prepend <> nodeval < query + then go (rightOf node) fallback (prepend <> nodeval) + else go (leftOf node) node prepend + leftOf node = node - (lsb node `unsafeShiftR` 1) + rightOf node = node + (lsb node `unsafeShiftR` 1) +{-# INLINABLE lowerBoundFen #-} diff --git a/src/Data/Fenwick/Vector.hs b/src/Data/Fenwick/Vector.hs new file mode 100644 --- /dev/null +++ b/src/Data/Fenwick/Vector.hs @@ -0,0 +1,136 @@ +{-# LANGUAGE BangPatterns #-} + +-- | Module : Data.Fenwick.Vector +-- +-- This modulo provides mutable [Fenwick Trees](https://en.wikipedia.org/wiki/Fenwick_tree), using +-- vectors as the underlying data structure. The algebraic structure is given by 'Semigroup' and +-- 'Monoid' instances. Some of the functions require the structure to be 'Commutative'. See +-- @monoid-subclasses@ and @commutative-semigroups@. If you want mark @Sum@ or @Product@ as +-- @Commutative@, see @SumCommutative@ and @ProductCommutative@ in the mentioned packages. +-- +-- I was hoping to merge this module with "Data.Fenwick.Array" under a common typeclass, although I +-- couldn't find any way through this. I may change this later, if I find the right abstractions. + +module Data.Fenwick.Vector + ( FenMVector + , newFen + , newAccumFen + , newListFen + , getSizeFen + , addFen + , sumPrefixFen + , lowerBoundFen + ) where + +import Control.Monad +import Data.Bits +import Data.Semigroup.Commutative +import Data.Monoid.Cancellative +import qualified Data.Vector.Generic.Mutable as V + +-- | Fenwick tree datatype. @vector@ must be a valid mutable vector type. Fenwick tree is assumed to +-- be data structure over a 1-based array of size @n@. +data FenMVector vector s elem = FenMVector !Int !(vector s elem) + +-- | least significant bit +lsb :: Int -> Int +lsb node = node .&. (-node) +{-# INLINE lsb #-} + +modify' :: (V.PrimMonad m, V.MVector v a) => v (V.PrimState m) a -> (a -> a) -> Int -> m () +modify' vec f = V.modify vec (f $!) +{-# INLINABLE modify' #-} + +unsafeModify' :: (V.PrimMonad m, V.MVector v a) => v (V.PrimState m) a -> (a -> a) -> Int -> m () +unsafeModify' vec f = V.unsafeModify vec (f $!) +{-# INLINABLE unsafeModify' #-} + +-- | Creates a Fenwick tree of size @n@ over 1-based array. All elements of the array are initially +-- @'mempty'@. \( O(n) \) +-- +-- Using type application extension, you can specify the vector and element type as the first and +-- second type argument. +newFen :: (V.MVector vector elem, V.PrimMonad m, Monoid elem) => Int -> m (FenMVector vector (V.PrimState m) elem) +newFen n = do + vec <- V.replicate (n+1) mempty + pure (FenMVector n vec) +{-# INLINABLE newFen #-} + +-- | @newAccumFen n xs@ creates a Fenwick tree of size @n@ over 1-based array. Every item of @xs@ is +-- pair of an index and an element. The elements of the array are initialized by @'<>'@ing elements +-- for each index. This functions is faster than creating empty array using 'newFen' and adding +-- elements individually using 'addFen'. \( O(n + m) \) where \( m \) is @'Data.Foldable.length' xs@ +newAccumFen :: (V.MVector vector elem, V.PrimMonad m, CommutativeMonoid elem, Foldable t) => Int -> t (Int, elem) -> m (FenMVector vector (V.PrimState m) elem) +newAccumFen n xs = do + vec <- V.replicate (n + 1) mempty + forM_ xs $ \(i, e) -> do + modify' vec (<> e) i + forM_ [1..n] $ \i -> do + let j = i + lsb i + when (j <= n) $ do + e <- V.unsafeRead vec i + unsafeModify' vec (e <>) j + pure (FenMVector n vec) +{-# INLINABLE newAccumFen #-} + +-- | Creates a Fenwick tree and initialize the array based on the elements of the list. \( O(n) \) +newListFen :: (V.MVector vector elem, V.PrimMonad m, CommutativeMonoid elem) => Int -> [elem] -> m (FenMVector vector (V.PrimState m) elem) +newListFen n xs = newAccumFen n $ zip [1..n] xs +{-# INLINABLE newListFen #-} + +-- | Get the size of underlying array. \( O(1) \) +getSizeFen :: FenMVector vector s elem -> Int +getSizeFen (FenMVector n _) = n +{-# INLINE getSizeFen #-} + +-- | Add a value to a cell of the array. The index must be in the range @[1, n]@. \( O(\log n) \) +addFen :: (V.PrimMonad m, V.MVector vector elem, Commutative elem) => FenMVector vector (V.PrimState m) elem -> Int -> elem -> m () +addFen (FenMVector n vec) r a = go (check r) where + check i + | i < 1 || i > n = error "index out of range" + | otherwise = i + go i = when (i <= n) $ do + unsafeModify' vec (a <>) i + go (i + lsb i) +{-# INLINABLE addFen #-} + +-- | Given index @r@, get the prefix sum of elements of the array in the range @[1, r]@. It accepts +-- values of @r@ out of the range of indices, assuming that every element out of the range of array +-- is 'mempty'. \( O(\log n) \) +sumPrefixFen :: (Monoid elem, V.PrimMonad m, V.MVector vector elem) => FenMVector vector (V.PrimState m) elem -> Int -> m elem +sumPrefixFen (FenMVector n vec) = go mempty . min n where + -- prefix (0, r], for 1 <= r <= n + go !s i + | i <= 0 = pure s + | otherwise = do + x <- V.unsafeRead vec i + go (x <> s) (i - lsb i) +{-# INLINABLE sumPrefixFen #-} + +-- | Given a prefix sum @q@, find the least index @r@ such that the prefix sum of @[1, r]@ is at +-- least @q@. Only applicable when partial sums are ordered. If the sum of the whole array is +-- less than @query@, it returns @n+1@. This functions is faster than binary search over +-- 'sumPrefixFen'. \( O(\log n) \) +lowerBoundFen :: (Monoid elem, V.PrimMonad m, V.MVector vector elem, Ord elem) => FenMVector vector (V.PrimState m) elem -> elem -> m Int +lowerBoundFen (FenMVector n vec) query = go root (n + 1) mempty where + root = bit (finiteBitSize n - countLeadingZeros n - 1) + go node + | odd node = leaf node + | otherwise = nonleaf node + leaf node fallback prepend + | node > n = pure fallback + | otherwise = do + nodeval <- V.unsafeRead vec node + pure $ if prepend <> nodeval < query + then fallback + else node + nonleaf node fallback prepend + | node > n = go (leftOf node) fallback prepend + | otherwise = do + nodeval <- V.unsafeRead vec node + if prepend <> nodeval < query + then go (rightOf node) fallback (prepend <> nodeval) + else go (leftOf node) node prepend + leftOf node = node - (lsb node `unsafeShiftR` 1) + rightOf node = node + (lsb node `unsafeShiftR` 1) +{-# INLINABLE lowerBoundFen #-} diff --git a/src/Data/Vector/VectorC.hs b/src/Data/Vector/VectorC.hs new file mode 100644 --- /dev/null +++ b/src/Data/Vector/VectorC.hs @@ -0,0 +1,89 @@ +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | Module : Data.Array.VectorC +-- +-- This modulo provides 'VectorC' and 'MVectorC', which given a existing vector implementation and a +-- representation type, it stores elements by coercing the element type into representation type. +-- Using this, we can have elements with algebraic structures (like 'Data.Monoid.Monoid' used in +-- this package) but stored in efficient unboxed vectors. +module Data.Vector.VectorC ( + VectorC, MVectorC +) where + +import qualified Data.Vector.Generic as VG +import qualified Data.Vector.Generic.Mutable as VGM +import Data.Coerce + +-- | Given a mutable vector type @vector@ and representation type @rep@, the type contructor +-- @'MVectorC' vector rep@ is a valid vector type that stores elements that are coercible to @rep@. +-- Use functions in 'Data.Vector.Generic.Mutable.MVector' to create and modify these vectors. +newtype MVectorC vector rep s elem = MVectorC (vector s rep) + +-- | Immutable vectors corresponding to 'MVectorC'. Use functions in 'Data.Vector.Generic.Vector' to create and +-- modify these vectors. +newtype VectorC vector rep elem = VectorC (vector rep) + +type instance VG.Mutable (VectorC vector rep) = MVectorC (VG.Mutable vector) rep + +instance (Coercible elem rep, VGM.MVector vector rep) => VGM.MVector (MVectorC vector rep) elem where + -- basicLength = coerce $ VGM.basicLength @vector @rep + basicLength (MVectorC vec) = VGM.basicLength vec + {-# INLINE basicLength #-} + -- basicUnsafeSlice = coerce $ VGM.basicUnsafeSlice @VU.MVector + basicUnsafeSlice i n (MVectorC vec) = MVectorC (VGM.basicUnsafeSlice i n vec) + {-# INLINE basicUnsafeSlice #-} + -- basicOverlaps = coerce $ VGM.basicOverlaps @VU.MVector + basicOverlaps (MVectorC vec1) (MVectorC vec2) = VGM.basicOverlaps vec1 vec2 + {-# INLINE basicOverlaps #-} + -- basicUnsafeNew = coerce $ VGM.basicUnsafeNew @VU.MVector + basicUnsafeNew n = MVectorC <$> VGM.basicUnsafeNew n + {-# INLINE basicUnsafeNew #-} + -- basicInitialize = coerce $ VGM.basicInitialize @VU.MVector + basicInitialize (MVectorC vec) = VGM.basicInitialize vec + {-# INLINE basicInitialize #-} + -- basicUnsafeReplicate = coerce $ VGM.basicUnsafeReplicate @VU.MVector + basicUnsafeReplicate n a = MVectorC <$> VGM.basicUnsafeReplicate n (coerce a) + {-# INLINE basicUnsafeReplicate #-} + -- basicUnsafeRead = coerce $ VGM.basicUnsafeRead @VU.MVector + basicUnsafeRead (MVectorC vec) i = coerce <$> VGM.basicUnsafeRead vec i + {-# INLINE basicUnsafeRead #-} + -- basicUnsafeWrite = coerce $ VGM.basicUnsafeWrite @VU.MVector + basicUnsafeWrite (MVectorC vec) i a = VGM.basicUnsafeWrite vec i (coerce a) + {-# INLINE basicUnsafeWrite #-} + -- basicClear = coerce $ VGM.basicClear @VU.MVector + basicClear (MVectorC vec) = VGM.basicClear vec + {-# INLINE basicClear #-} + -- basicSet = coerce $ VGM.basicSet @VU.MVector + basicSet (MVectorC vec) a = VGM.basicSet vec (coerce a) + {-# INLINE basicSet #-} + -- basicUnsafeCopy = coerce $ VGM.basicUnsafeCopy @VU.MVector + basicUnsafeCopy (MVectorC vec1) (MVectorC vec2) = VGM.basicUnsafeCopy vec1 vec2 + {-# INLINE basicUnsafeCopy #-} + -- basicUnsafeMove = coerce $ VGM.basicUnsafeMove @VU.MVector + basicUnsafeMove (MVectorC vec1) (MVectorC vec2) = VGM.basicUnsafeMove vec1 vec2 + {-# INLINE basicUnsafeMove #-} + -- basicUnsafeGrow = coerce $ VGM.basicUnsafeGrow @VU.MVector + basicUnsafeGrow (MVectorC vec) n = MVectorC <$> VGM.basicUnsafeGrow vec n + {-# INLINE basicUnsafeGrow #-} + +instance (Coercible elem rep, VG.Vector vector rep) => VG.Vector (VectorC vector rep) elem where + basicUnsafeFreeze = coerce $ VG.basicUnsafeFreeze @vector @rep + {-# INLINE basicUnsafeFreeze #-} + basicUnsafeThaw = coerce $ VG.basicUnsafeThaw @vector @rep + {-# INLINE basicUnsafeThaw #-} + basicLength = coerce $ VG.basicLength @vector @rep + {-# INLINE basicLength #-} + basicUnsafeSlice = coerce $ VG.basicUnsafeSlice @vector @rep + {-# INLINE basicUnsafeSlice #-} + basicUnsafeIndexM = coerce $ VG.basicUnsafeIndexM @vector @rep + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeCopy = coerce $ VG.basicUnsafeCopy @vector @rep + {-# INLINE basicUnsafeCopy #-} + elemseq (VectorC vec) a = VG.elemseq vec (coerce a) + {-# INLINE elemseq #-} diff --git a/tests/FenwickTest.hs b/tests/FenwickTest.hs new file mode 100644 --- /dev/null +++ b/tests/FenwickTest.hs @@ -0,0 +1,87 @@ +{-# LANGUAGE FlexibleInstances, ScopedTypeVariables, TypeApplications, TupleSections #-} + +module Main where + +import Control.Monad +import Control.Monad.ST (RealWorld) +import Data.Array.ArrayC +import Data.Array.IO +import Data.List +import Data.Monoid +import Data.Proxy +import Data.Vector.VectorC +import Test.Hspec +import qualified Data.Fenwick.Array as FA +import qualified Data.Fenwick.Vector as FV +import qualified Data.Vector.Generic.Mutable as VG +import qualified Data.Vector.Unboxed.Mutable as VUM + +class IsFenwick a where + new :: Int -> IO a + newAccum :: Foldable t => Int -> t (Int, Sum Int) -> IO a + fromList :: Int -> [Sum Int] -> IO a + getSize :: a -> Int + increment :: a -> Int -> Sum Int -> IO () + sumPrefix :: a -> Int -> IO (Sum Int) + lowerBound :: a -> Sum Int -> IO Int + +type FenArray = FA.FenMArray (ArrayC IOUArray Int) (Sum Int) +type FenVector = FV.FenMVector (MVectorC VUM.MVector Int) RealWorld (Sum Int) + +instance IsFenwick FenArray where + new = FA.newFen + newAccum = FA.newAccumFen + fromList = FA.newListFen + getSize = FA.getSizeFen + increment = FA.addFen + sumPrefix = FA.sumPrefixFen + lowerBound = FA.lowerBoundFen + +instance IsFenwick FenVector where + new = FV.newFen + newAccum = FV.newAccumFen + fromList = FV.newListFen + getSize = FV.getSizeFen + increment = FV.addFen + sumPrefix = FV.sumPrefixFen + lowerBound = FV.lowerBoundFen + +testFenwick :: forall a. IsFenwick a => Proxy a -> SpecWith () +testFenwick _ = do + it "has a fixed size" $ do + fen <- new @a 100 + getSize fen `shouldBe` 100 + + it "can compute partial sums" $ do + let list = map Sum [3, -1, 4, 0, 9, -10, 8] + n = length list + fen <- fromList @a n list + traverse (sumPrefix fen) [0..n] + `shouldReturn` scanl (<>) mempty list + + it "can be used as a set" $ do + let n = 100 + list = [34, 10, 100, 63, 12, 6, 54, 29, 83] + m = length list + sorted = sort list + fen <- newAccum @a n $ map (, Sum 1) list + forM_ (zip [1..] sorted) $ \(i, x) -> do + lowerBound fen (Sum i) `shouldReturn` x + + it "can be updated" $ do + let n = 5 + allSumPrefix fen = traverse (sumPrefix fen) [1..n] + fen <- new @a n + increment fen 2 (Sum 2) + allSumPrefix fen `shouldReturn` [0, 2, 2, 2, 2] + increment fen 5 (Sum 3) + allSumPrefix fen `shouldReturn` [0, 2, 2, 2, 5] + increment fen 1 (Sum (-1)) + allSumPrefix fen `shouldReturn` [-1, 1, 1, 1, 4] + +main :: IO () +main = hspec $ do + describe "Data.Fenwick.Array" $ + testFenwick (Proxy @FenArray) + describe "Data.Fenwick.Vector" $ + testFenwick (Proxy @FenVector)