1 files changed, 64 insertions, 16 deletions

diff --git a/packages/base/src/Internal/Numeric.hs b/packages/base/src/Internal/Numeric.hs
index fd0a217..4f7bb82 100644
--- a/packages/base/src/Internal/Numeric.hs
+++ b/packages/base/src/Internal/Numeric.hs
@@ -4,6 +4,7 @@
 {-# LANGUAGE MultiParamTypeClasses #-}
 {-# LANGUAGE FunctionalDependencies #-}
 {-# LANGUAGE UndecidableInstances #-}
+{-# LANGUAGE PatternSynonyms #-}
 
 {-# OPTIONS_GHC -fno-warn-missing-signatures #-}
 
@@ -22,12 +23,18 @@ module Internal.Numeric where
 import Internal.Vector
 import Internal.Matrix
 import Internal.Element
+import Internal.Extract (requires,pattern BAD_SIZE)
 import Internal.ST as ST
 import Internal.Conversion
 import Internal.Vectorized
 import Internal.LAPACK(multiplyR,multiplyC,multiplyF,multiplyQ,multiplyI,multiplyL)
+import Control.Monad
+import Data.Function
+import Data.Int
 import Data.List.Split(chunksOf)
 import qualified Data.Vector.Storable as V
+import Foreign.Ptr
+import Foreign.Storable
 
 --------------------------------------------------------------------------------
 
@@ -44,7 +51,7 @@ type instance ArgOf Matrix a = a -> a -> a
 --------------------------------------------------------------------------------
 
 -- | Basic element-by-element functions for numeric containers
-class Element e => Container c e
+class Storable e => Container c e
   where
     conj'        :: c e -> c e
     size'        :: c e -> IndexOf c
@@ -56,7 +63,7 @@ class Element e => Container c e
     -- | element by element multiplication
     mul         :: c e -> c e -> c e
     equal       :: c e -> c e -> Bool
-    cmap'        :: (Element b) => (e -> b) -> c e -> c b
+    cmap'        :: (Storable b) => (e -> b) -> c e -> c b
     konst'      :: e -> IndexOf c -> c e
     build'       :: IndexOf c -> (ArgOf c e) -> c e
     atIndex'     :: c e -> IndexOf c -> e
@@ -107,7 +114,7 @@ instance Container Vector I
     mul = vectorZipI Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -146,7 +153,7 @@ instance Container Vector Z
     mul = vectorZipL Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -186,7 +193,7 @@ instance Container Vector Float
     mul = vectorZipF Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -223,7 +230,7 @@ instance Container Vector Double
     mul = vectorZipR Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -260,7 +267,7 @@ instance Container Vector (Complex Double)
     mul = vectorZipC Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -296,7 +303,7 @@ instance Container Vector (Complex Float)
     mul = vectorZipQ Mul
     equal = (==)
     scalar' = V.singleton
-    konst' = constantD
+    konst' = constantAux
     build' = buildV
     cmap' = mapVector
     atIndex' = (@>)
@@ -323,7 +330,7 @@ instance Container Vector (Complex Float)
 
 ---------------------------------------------------------------
 
-instance (Num a, Element a, Container Vector a) => Container Matrix a
+instance (Num a, Storable a, Container Vector a) => Container Matrix a
   where
     conj' = liftMatrix conj'
     size' = size
@@ -418,8 +425,8 @@ fromZ = fromZ'
 toZ :: (Container c e) => c e -> c Z
 toZ = toZ'
 
--- | like 'fmap' (cannot implement instance Functor because of Element class constraint)
-cmap :: (Element b, Container c e) => (e -> b) -> c e -> c b
+-- | like 'fmap' (cannot implement instance Functor because of Storable class constraint)
+cmap :: (Storable b, Container c e) => (e -> b) -> c e -> c b
 cmap = cmap'
 
 -- | generic indexing function
@@ -470,7 +477,7 @@ step
 step = step'
 
 
--- | Element by element version of @case compare a b of {LT -> l; EQ -> e; GT -> g}@.
+-- | Storable by element version of @case compare a b of {LT -> l; EQ -> e; GT -> g}@.
 --
 -- Arguments with any dimension = 1 are automatically expanded:
 --
@@ -598,7 +605,7 @@ instance Numeric Z
 --------------------------------------------------------------------------------
 
 -- | Matrix product and related functions
-class (Num e, Element e) => Product e where
+class (Num e, Storable e) => Product e where
     -- | matrix product
     multiply :: Matrix e -> Matrix e -> Matrix e
     -- | sum of absolute value of elements (differs in complex case from @norm1@)
@@ -823,12 +830,12 @@ buildV n f = fromList [f k | k <- ks]
 --------------------------------------------------------
 
 -- | Creates a square matrix with a given diagonal.
-diag :: (Num a, Element a) => Vector a -> Matrix a
+diag :: (Num a, Storable a) => Vector a -> Matrix a
 diag v = diagRect 0 v n n where n = dim v
 
 -- | creates the identity matrix of given dimension
-ident :: (Num a, Element a) => Int -> Matrix a
-ident n = diag (constantD 1 n)
+ident :: (Num a, Storable a) => Int -> Matrix a
+ident n = diag (constantAux 1 n)
 
 --------------------------------------------------------
 
@@ -943,3 +950,44 @@ class Testable t
 
 --------------------------------------------------------------------------------
 
+compareV :: (Storable a, Ord a) => Vector a -> Vector a -> Vector Int32
+compareV = compareG compareStorable
+
+compareStorable :: (Storable a, Ord a) =>
+                    Int32 -> Ptr a
+                    -> Int32 -> Ptr a
+                    -> Int32 -> Ptr Int32
+                    -> IO Int32
+compareStorable xn xp yn yp rn rp = do
+    requires (xn==yn && xn==rn) BAD_SIZE $ do
+    ($ 0) $ fix $ \kloop k -> when (k<xn) $ do
+        xk <- peekElemOff xp (fromIntegral k)
+        yk <- peekElemOff yp (fromIntegral k)
+        pokeElemOff rp (fromIntegral k) $ case compare xk yk of
+            LT -> -1
+            GT -> 1
+            EQ -> 0
+        kloop (succ k)
+    return 0
+
+selectV  :: Storable a => Vector Int32 -> Vector a -> Vector a -> Vector a -> Vector a
+selectV = selectG selectStorable
+
+selectStorable :: Storable a =>
+                    Int32 -> Ptr Int32
+                    -> Int32 -> Ptr a
+                    -> Int32 -> Ptr a
+                    -> Int32 -> Ptr a
+                    -> Int32 -> Ptr a
+                    -> IO Int32
+selectStorable condn condp ltn ltp eqn eqp gtn gtp rn rp = do
+    requires (condn==ltn && ltn==eqn && ltn==gtn && ltn==rn) BAD_SIZE $ do
+    ($ 0) $ fix $ \kloop k -> when (k<condn) $ do
+        condpk <- peekElemOff condp (fromIntegral k)
+        pokeElemOff rp (fromIntegral k) =<< case compare condpk 0 of
+            LT -> peekElemOff ltp (fromIntegral k)
+            GT -> peekElemOff gtp (fromIntegral k)
+            EQ -> peekElemOff eqp (fromIntegral k)
+        kloop (succ k)
+    return 0
+