Add MonadFix instance

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 #-}

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)