mandulia-0.8.0.1: src/Sort.hsc

{-# LANGUAGE ForeignFunctionInterface #-}
{-
Mandulia -- Mandelbrot/Julia explorer
Copyright (C) 2010  Claude Heiland-Allen <claude@mathr.co.uk>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
-}

module Sort where

import Data.Array (array, (!))
import Foreign.Storable
import Foreign.Ptr (Ptr, castPtr)
import Foreign.C.Types (CInt(..), CDouble(..))
import Foreign.Marshal.Array (withArray, peekArray)
import Unsafe.Coerce (unsafeCoerce)

#include "qsort.h"

foreign import ccall safe "qsort.h sort" c_sort :: Ptr () -> CInt -> IO ()

data Sortee = Sortee{ sKey :: !Double, sValue :: !Int }

instance Storable Sortee where
  sizeOf _ = (#size struct sortee)
  alignment _ = alignment (error "alignment" :: CInt)
  peek ptr = do
    key   <- (#peek struct sortee, key  ) ptr :: IO CDouble
    value <- (#peek struct sortee, value) ptr :: IO CInt
    return Sortee{ sKey = unsafeCoerce key, sValue = fromIntegral value }
  poke ptr (Sortee{ sKey = key, sValue = value }) = do
    (#poke struct sortee, key  ) ptr (unsafeCoerce key :: CDouble)
    (#poke struct sortee, value) ptr (fromIntegral value :: CInt)

sortIO :: (a -> Double) -> [a] -> IO [a]
sortIO p xs = do -- return (sortOn p xs)
  let n  = length xs
      is = [0 .. n - 1]
      ps = map p xs
      a  = array (0, n - 1) (zip is xs)
      ss = zipWith Sortee ps is
  ss' <- withArray ss $ \ptr -> do
    c_sort (castPtr ptr) (fromIntegral n)
    peekArray n ptr
  let is' = map sValue ss'
  return $ map (a !) is'