maxIndex of an MVector - haskell

Question

Data.Vector includes a function maxIndex with type maxIndex :: (Ord a) => Vector a -> Int that returns the index of the maximum value in that Vector. I'm working with mutable Vectors, however, and MVector doesn't have maxIndex defined for it.
What's the best way of getting the data I want out of the MVector I have? My code currently is:
import qualified Data.Vector.Unboxed.Mutable as MV
import Control.Monad.ST
import Control.Monad (mapM_)
type MaxIndex = Int
step :: forall s. MV.MVector s Int -> MaxIndex -> ST s ()
step vec i = do
n <- MV.unsafeRead vec i
MV.write vec i 0
let l = MV.length vec
(k, x) = n `divMod` l
mapM_ (\j -> MV.modify vec (+k) j) [0..l-1] -- side note, this is just
-- fmap (+k) vec, but MVector is not
-- a functor. Is there a better way?
mapM_ (\j -> MV.modify vec (+1) (j `mod` l)) [i+1..i+x]
where i is the index I'm looking to derive inside step. I'm doing this because the actions here need to eventually be wrapped inside an until and repeated until a predicate is satisfied, and freezing and thawing every cycle sounds ludicrously expensive.

Answer 1

I see lots of talk about unsafe freezing which seems suspect since you plan to mutate this memory later, thus violating the assurance you are implicitly giving when calling unsafeFreeze.
My suggestion is to just write an imperative-style maxIndex function. The below is typed but not tested:
import qualified Data.Vector.Unboxed.Mutable as MV
import Control.Monad.ST
import Control.Monad (mapM_)
maxIndex :: (Ord a, MV.Unbox a) => MV.MVector s a -> ST s (Maybe Int)
maxIndex mv | len == 0 = pure Nothing
| otherwise = Just <$> go 0 0
where
len = MV.length mv
go n i | i >=len = pure n
| otherwise = do
nVal <- MV.unsafeRead mv n
iVal <- MV.unsafeRead mv i
if nVal < iVal then go i (i+1)
else go n (i+1)

Answer 2

Have you considered freezing the vector with unsafeFreeze which is supposed to be fast (i.e. Θ(1))? For example you can define maxIndex for mutable vectors like this:
maxIndex = fmap V.maxIndex . V.unsafeFreeze
This assumes that you have imported the following:
import qualified Data.Vector.Unboxed as V
unsafeFreeze doesn't actually copy any data and should be fast, but it would be interesting to run a criterion benchmark to see if this approach is actually faster compared to an explicit loop.