Search code examples
arrayshaskellfold

Is there a function in haskell working like a mixture of accumArray and foldr?


let me call the function accumrArray.

accumrArray :: 
              (e' -> e -> e) An accumulating function
           -> e              A default element 
           -> (i, i)         The bounds of the array 
           -> [(i, e')]      List of associations 
           -> a i e          The array

accumrArray  (:) [] (1,2) [(1,1),(2,2),(2,3)]  === array [(1,[1]), (2,[2,3])]
head $ (accumrArray (:) [] (1,1) [(1,x)|x<-[4..]]) ! 1 === 4

Solution

  • How strange... I wrote this function a few days ago for someone else. The function first appeared in LML (I believe), but never made it into the Haskell array library.

    Here you go:

    {-# LANGUAGE ScopedTypeVariables #-}
    import Data.Array
    import System.IO.Unsafe
    import Data.IORef
    import Data.Array.MArray
    import Data.Array.Base
    import Control.Monad
    import Data.Array.IO
    
    accumArrayR :: forall a e i. Ix i => (a -> e -> e) -> e -> (i,i) -> [(i,a)] -> Array i e
    accumArrayR f e bounds@(l,u) assocs = unsafePerformIO $ do
      ref <- newIORef assocs
      arr <- newArray_ bounds
      let _ = arr :: IOArray i e
      let n = safeRangeSize (l,u)
      let elem x = unsafePerformIO $ do
                      ass <- readIORef ref
                      let loop [] = writeIORef ref [] >> return e
                          loop ((y,a):rest) = do
                             let ix = safeIndex bounds n y
                             let r = f a (elem x)
                             unsafeWrite arr ix r
                             if (ix == x)
                                then writeIORef ref rest >> return r
                                else loop rest
                      loop ass
      forM_ [0..n] $ \ix -> unsafeWrite arr ix (elem ix)
      unsafeFreeze arr
    

    A challenge for the reader: use accumArrayR to implement linear-time depth-first-search of a graph.

    Edit I should mention that the function isn't thread-safe as written. Turning the IORef into an MVar would fix it, but there might be better ways.