1 files changed, 93 insertions, 3 deletions

diff --git a/packages/base/src/Internal/Container.hs b/packages/base/src/Internal/Container.hs
index 41b8214..0f2e7d5 100644
--- a/packages/base/src/Internal/Container.hs
+++ b/packages/base/src/Internal/Container.hs
@@ -4,6 +4,8 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
 {-# OPTIONS_GHC -fno-warn-simplifiable-class-constraints #-}
 
@@ -30,8 +32,15 @@ module Internal.Container where
 import Internal.Vector
 import Internal.Matrix
 import Internal.Element
+import Internal.Extract(requires,pattern BAD_SIZE)
 import Internal.Numeric
 import Internal.Algorithms(Field,linearSolveSVD,Herm,mTm)
+import Control.Monad(when)
+import Data.Function
+import Data.Int
+import Foreign.Ptr
+import Foreign.Storable
+import Foreign.Marshal.Array
 #if MIN_VERSION_base(4,11,0)
 import Prelude hiding ((<>))
 #endif
@@ -227,7 +236,7 @@ meanCov x = (med,cov) where
 
 --------------------------------------------------------------------------------
 
-sortVector :: (Ord t, Element t) => Vector t -> Vector t
+sortVector :: (Ord t, Storable t) => Vector t -> Vector t
 sortVector = sortV
 
 {- |
@@ -248,7 +257,7 @@ sortVector = sortV
 -2.20   0.11  -1.58  -0.01   0.19  -0.29   1.04   1.06  -2.09  -0.75
 
 -}
-sortIndex :: (Ord t, Element t) => Vector t -> Vector I
+sortIndex :: (Ord t, Storable t) => Vector t -> Vector I
 sortIndex = sortI
 
 ccompare :: (Ord t, Container c t) => c t -> c t -> c I
@@ -296,10 +305,91 @@ The indexes are autoconformable.
  , 10, 16, 22 ]
 
 -}
-remap :: Element t => Matrix I -> Matrix I -> Matrix t -> Matrix t
+remap :: Storable t => Matrix I -> Matrix I -> Matrix t -> Matrix t
 remap i j m
     | minElement i >= 0 && maxElement i < fromIntegral (rows m) &&
       minElement j >= 0 && maxElement j < fromIntegral (cols m) = remapM i' j' m
     | otherwise = error $ "out of range index in remap"
   where
     [i',j'] = conformMs [i,j]
+
+sortI :: (Storable a, Ord a) => Vector a -> Vector Int32
+sortI = sortG sort_index
+
+type C_Compare a = Ptr a -> Ptr a -> IO Int32
+
+foreign import ccall "wrapper" wrapCompare :: C_Compare a -> IO (FunPtr (C_Compare a))
+
+foreign import ccall "qsort"
+    c_qsort :: Ptr a                   -- ^ base
+               -> Word                 -- ^ nmemb
+               -> Word                 -- ^ size
+               -> FunPtr (C_Compare a) -- ^ compar
+               -> IO ()
+
+sizeOfElem :: forall a. Storable a => Ptr a -> Int
+sizeOfElem _ = sizeOf (undefined :: a)
+
+sort_index :: (Storable a, Ord a) =>
+                Int32 -> Ptr a
+                -> Int32 -> Ptr Int32
+                -> IO Int32
+sort_index vn vp rn rp = do
+    requires (vn == rn) BAD_SIZE $ do
+    comp <- wrapCompare $ \ap bp -> do
+        a <- peekElemOff vp . fromIntegral =<< peek (ap :: Ptr Int32)
+        b <- peekElemOff vp . fromIntegral =<< peek bp
+        return $ case compare a b of
+            LT -> -1
+            GT -> 1
+            EQ -> 0
+    sequence_ [ pokeElemOff rp (fromIntegral i) i | i <- [0 .. rn-1] ]
+    c_qsort rp (fromIntegral rn) 4 comp
+    freeHaskellFunPtr comp
+    return 0
+
+sortV :: (Storable a, Ord a) => Vector a -> Vector a
+sortV = sortG sortStorable
+
+sortStorable :: (Storable a, Ord a) =>
+                Int32 -> Ptr a
+                -> Int32 -> Ptr a
+                -> IO Int32
+sortStorable vn vp rn rp = do
+    requires (vn == rn) BAD_SIZE $ do
+    copyArray rp vp (fromIntegral vn * sizeOfElem vp)
+    comp <- wrapCompare $ \ap bp -> do
+        a <- peek ap
+        b <- peek bp
+        return $ case compare a b of
+            LT -> -1
+            GT -> 1
+            EQ -> 0
+    c_qsort rp (fromIntegral rn) (fromIntegral $ sizeOfElem rp) comp
+    freeHaskellFunPtr comp
+    return 0
+
+remapM :: Storable a => Matrix Int32 -> Matrix Int32 -> Matrix a -> Matrix a
+remapM = remapG remapStorable
+
+remapStorable :: Storable a =>
+    Int32 -> Int32 -> Int32 -> Int32 -> Ptr Int32 -- i
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr Int32 -- j
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr a -- m
+    -> Int32 -> Int32 -> Int32 -> Int32 -> Ptr a -- r
+    -> IO Int32
+remapStorable ir ic iXr iXc ip
+              jr jc jXr jXc jp
+              mr mc mXr mXc mp
+              rr rc rXr rXc rp = do
+    requires (ir==jr && ic==jc && ir==rr && ic==rc) BAD_SIZE $ do
+    ($ 0) $ fix $ \aloop a -> when (a<rr) $ do
+        ($ 0) $ fix $ \bloop b -> when (b<rc) $ do
+            iab <- peekElemOff ip (fromIntegral $ iXr*a + iXc*b)
+            jab <- peekElemOff jp (fromIntegral $ jXr*a + jXc*b)
+            when (0 <= iab && iab < mr && 0 <= jab && jab < mc) $
+                pokeElemOff rp (fromIntegral $ rXr*a + rXc*b)
+                        =<< peekElemOff mp (fromIntegral $ mXr*iab + mXc*jab)
+            bloop (succ b)
+        aloop (succ a)
+    return 0