Add MonadFix instance for boxed vectors (#312)

It's #179 with merged into latest master and documentation tweaks Originally PR authored by David Feurer I *believe* this is equivalent to the instance for `[]`. Writing QuickCheck properties for `mfix` seems pretty tricky, so I just added a small unit test. Co-authored-by: David Feuer <[email protected]>
haskell · Jan 16, 2021 · 1da7566 · 1da7566
1 parent 3182ffd
commit 1da7566
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 2 deletions.
diff --git a/Data/Vector.hs b/Data/Vector.hs
@@ -177,11 +177,12 @@ 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 qualified Control.Monad.Fail as Fail
-
+import Control.Monad.Fix ( MonadFix (mfix) )
 import Control.Monad.Zip
+import Data.Function ( fix )
 
 import Prelude hiding ( length, null,
                         replicate, (++), concat,
@@ -382,6 +383,29 @@ instance MonadZip Vector where
   {-# INLINE munzip #-}
   munzip = unzip
 
+-- | Instance has same semantics as one for lists
+--
+--  @since 0.13.0.0
+instance MonadFix Vector where
+  -- We take care to dispose of v0 as soon as possible (see headM docs).
+  --
+  -- It's perfectly safe to use non-monadic indexing within generate
+  -- call since intermediate vector won't be created until result's
+  -- value is demanded.
+  {-# INLINE mfix #-}
+  mfix f
+    | null v0 = empty
+    -- We take first element of resulting vector from v0 and create
+    -- rest using generate. Note that cons should fuse with generate
+    | otherwise = runST $ do
+        h <- headM v0
+        return $ cons h $
+          generate (lv0 - 1) $
+            \i -> fix (\a -> f a ! (i + 1))
+    where
+      -- Used to calculate size of resulting vector
+      v0 = fix (f . head)
+      !lv0 = length v0
 
 instance Applicative.Applicative Vector where
   {-# INLINE pure #-}

diff --git a/changelog.md b/changelog.md
@@ -4,6 +4,7 @@
  * Export `unstreamM` from`from Data.Vector.Generic`: [#70](https://github.com/haskell/vector/issues/70)
  * New functions: `unfoldrExactN` and `unfoldrExactNM`: [#140](https://github.com/haskell/vector/issues/140)
  * Added `iforM` and `iforM_`: [#262](https://github.com/haskell/vector/issues/262)
+ * Added `MonadFix` instance for boxed vectors: [#178](https://github.com/haskell/vector/issues/178)
 
 # Changes in version 0.12.1.2
 

diff --git a/tests/Tests/Vector/UnitTests.hs b/tests/Tests/Vector/UnitTests.hs
@@ -6,6 +6,8 @@ module Tests.Vector.UnitTests (tests) where
 import Control.Applicative as Applicative
 import Control.Exception
 import Control.Monad.Primitive
+import Control.Monad.Fix (mfix)
+import qualified Data.Vector as Vector
 import Data.Int
 import Data.Word
 import Data.Typeable
@@ -80,6 +82,9 @@ tests =
       , testCase "Unboxed" $ testTakeOutOfMemory Unboxed.take
       ]
     ]
+  , testGroup "Data.Vector"
+    [ testCase "MonadFix" checkMonadFix
+    ]
   ]
 
 testsSliceOutOfBounds ::
@@ -157,3 +162,15 @@ _f :: (Generic.Vector v a, Generic.Vector w a, PrimMonad f)
    => Generic.Mutable v (PrimState f) a -> f (w a)
 _f v = Generic.convert `fmap` Generic.unsafeFreeze v
 #endif
+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)