[Haskell-cafe] Optimizing array operation?

Xiao-Yong Jin xj2106 at columbia.edu
Wed Mar 4 20:00:43 EST 2009


Dear list,

I decided to extract some common code I use in my data
analysis code and end up with some interesting pattern.

> -- | Confined array operation.  The type Arr here is actually used as
> -- a C-like array with index runs from 0 to n-1.
> type Arr = Array Int
> 
> -- | Element fetching function, no bounds checking.  Works with array
> -- with index [0..n-1]
> (!>) :: Arr e -> Int -> e
> (!>) = unsafeAt
> {-# INLINE (!>) #-}
> 
> -- | Length of the array.
> arrLength :: Arr e -> Int
> arrLength = numElements
> {-# INLINE arrLength #-}
> 
> -- | Sequentially loop over all elements in an array.
> loopArrM :: (Monad m) =>
>             (e -> m a)          -- ^ function to apply to elements
>          -> Arr e               -- ^ the array
>          -> m ()
> loopArrM = loopArrMp (\ _ -> True)
> {-# INLINE loopArrM #-}
> 
> -- | Same as loopArrM but only to those with index satisfies predicate.
> loopArrMp :: (Monad m) =>
>              (Int -> Bool)      -- ^ predicate
>           -> (e -> m a)         -- ^ function to apply to elements
>           -> Arr e              -- ^ the array
>           -> m ()
> loopArrMp p f arr = loopArrMG 0 end p f arr
>     where
>       end = arrLength arr
> {-# INLINE loopArrMp #-}
> 
> -- | Same as loopArrM but only to indices within a range.
> loopArrMr :: (Monad m) =>
>              Int                -- ^ start index
>           -> Int                -- ^ end index
>           -> (e -> m a)         -- ^ function to apply to elements
>           -> Arr e              -- ^ the array
>           -> m ()
> loopArrMr start end = loopArrMG start end (\ _ -> True)
> {-# INLINE loopArrMr #-}
> 
> -- | Generic loopArrM for indices with in a range [start, end-1] and
> -- satisfy predicate.
> loopArrMG :: (Monad m) =>
>              Int                -- ^ start index
>           -> Int                -- ^ end index
>           -> (Int -> Bool)      -- ^ predicate
>           -> (e -> m a)         -- ^ function to apply to elements
>           -> Arr e              -- ^ the array
>           -> m ()
> loopArrMG start end p f arr = go start
>     where
>       go !j | j == end  = return ()
>             | p j       = f (arr !> j) >> go (j + 1)
>             | otherwise = go (j + 1)
> {-# INLINE loopArrMG #-}

It is quite frequently to loop over Arr with various
conditions, so I wrote loopArrMG to do it.  I know a list
comprehension would just do the trick, but the extra time
and heap allocation is not what I want.

My question here is that whether there is a better data
structure available to accomplish such task?  If not, is
there a better way to write the above code?

Thanks,
Xiao-Yong
-- 
    c/*    __o/*
    <\     * (__
    */\      <


More information about the Haskell-Cafe mailing list