diff --git a/Data/Vector.hs b/Data/Vector.hs index 5c04b412..4deb4de9 100644 --- a/Data/Vector.hs +++ b/Data/Vector.hs @@ -171,11 +171,12 @@ import qualified Data.Vector.Fusion.Bundle as Bundle import Control.DeepSeq ( NFData, rnf ) import Control.Monad ( MonadPlus(..), liftM, ap ) -import Control.Monad.ST ( ST ) +import Control.Monad.ST ( ST, runST ) import Control.Monad.Primitive - +import Control.Monad.Fix ( MonadFix (mfix) ) import Control.Monad.Zip +import Data.Function ( fix ) import Prelude hiding ( length, null, replicate, (++), concat, @@ -361,6 +362,24 @@ instance MonadZip Vector where {-# INLINE munzip #-} munzip = unzip +instance MonadFix Vector where + -- We take care to dispose of v0 as soon as possible. + -- We also avoid setting up the result vector to refer to + -- itself. These measures should prevent memory leaks. + -- It's perfectly safe to use non-monadic indexing within + -- each element, as the result of indexing will be demanded + -- as soon as the vector is produced. + {-# INLINE mfix #-} + mfix f + | null v0 = empty + | otherwise = runST $ do + h <- headM v0 + return $ cons h $ + generate (lv0 - 1) $ + \i -> fix (\a -> f a ! (i + 1)) + where + v0 = fix (f . head) + !lv0 = length v0 instance Applicative.Applicative Vector where {-# INLINE pure #-} diff --git a/tests/Tests/Vector/UnitTests.hs b/tests/Tests/Vector/UnitTests.hs index 5827640d..34ec3e95 100644 --- a/tests/Tests/Vector/UnitTests.hs +++ b/tests/Tests/Vector/UnitTests.hs @@ -3,6 +3,8 @@ module Tests.Vector.UnitTests (tests) where import Control.Applicative as Applicative +import Control.Monad.Fix (mfix) +import qualified Data.Vector as Vector import qualified Data.Vector.Storable as Storable import Foreign.Ptr import Foreign.Storable @@ -20,6 +22,19 @@ instance (Storable a) => Storable (Aligned a) where peek ptr = Aligned Applicative.<$> peek (castPtr ptr) poke ptr = poke (castPtr ptr) . getAligned +tests :: [Test] +tests = + [ testGroup "Data.Vector.Storable.Vector" + [ testCase "Aligned Double" $ + checkAddressAlignment alignedDoubleVec + , testCase "Aligned Int" $ + checkAddressAlignment alignedIntVec + ] + , testGroup "Data.Vector" + [ testCase "MonadFix" checkMonadFix + ] + ] + checkAddressAlignment :: forall a. (Storable a) => Storable.Vector a -> Assertion checkAddressAlignment xs = Storable.unsafeWith xs $ \ptr -> do let ptr' = ptrToWordPtr ptr @@ -31,18 +46,22 @@ checkAddressAlignment xs = Storable.unsafeWith xs $ \ptr -> do dummy :: a dummy = undefined -tests :: [Test] -tests = - [ testGroup "Data.Vector.Storable.Vector Alignment" - [ testCase "Aligned Double" $ - checkAddressAlignment alignedDoubleVec - , testCase "Aligned Int" $ - checkAddressAlignment alignedIntVec - ] - ] - alignedDoubleVec :: Storable.Vector (Aligned Double) alignedDoubleVec = Storable.fromList $ map Aligned [1, 2, 3, 4, 5] alignedIntVec :: Storable.Vector (Aligned Int) alignedIntVec = Storable.fromList $ map Aligned [1, 2, 3, 4, 5] + +checkMonadFix :: Assertion +checkMonadFix = assertBool "checkMonadFix" $ + Vector.toList fewV == fewL && + Vector.toList none == [] + where + facty _ 0 = 1; facty f n = n * f (n - 1) + fewV :: Vector.Vector Int + fewV = fmap ($ 12) $ mfix (\i -> Vector.fromList [facty i, facty (+1), facty (+2)]) + fewL :: [Int] + fewL = fmap ($ 12) $ mfix (\i -> [facty i, facty (+1), facty (+2)]) + + none :: Vector.Vector Int + none = mfix (const Vector.empty)