move extractor to Element

1 files changed, 79 insertions, 2 deletions

diff --git a/packages/base/src/Internal/Element.hs b/packages/base/src/Internal/Element.hs
index 6fc2981..0f61370 100644
--- a/packages/base/src/Internal/Element.hs
+++ b/packages/base/src/Internal/Element.hs
@@ -24,9 +24,10 @@ module Internal.Element where
 import Internal.Tools
 import Internal.Vector
 import Internal.Matrix
+import Internal.Vectorized
 import qualified Internal.ST as ST
 import Data.Array
-
+import Text.Printf
 import Data.Vector.Storable(fromList)
 import Data.List(transpose,intersperse)
 import Foreign.Storable(Storable)
@@ -78,7 +79,83 @@ instance (Element a, Read a) => Read (Matrix a) where
 breakAt c l = (a++[c],tail b) where
     (a,b) = break (==c) l
 
-------------------------------------------------------------------
+--------------------------------------------------------------------------------
+
+data Extractor
+    = All
+    | Range Int Int Int
+    | Pos (Vector I)
+    | PosCyc (Vector I)
+    | Take Int
+    | TakeLast Int
+    | Drop Int
+    | DropLast Int
+   deriving Show
+
+
+--
+infixl 9 ??
+(??)  :: Element t => Matrix t -> (Extractor,Extractor) -> Matrix t
+
+minEl = toScalarI Min
+maxEl = toScalarI Max
+cmodi = vectorMapValI ModVS
+
+extractError m e = error $ printf "can't extract %s from matrix %dx%d" (show e) (rows m) (cols m)
+
+m ?? (Range a s b,e) | s /= 1 = m ?? (Pos (idxs [a,a+s .. b]), e)
+m ?? (e,Range a s b) | s /= 1 = m ?? (e, Pos (idxs [a,a+s .. b]))
+
+m ?? e@(Range a _ b,_) | a < 0 || b >= rows m = extractError m e
+m ?? e@(_,Range a _ b) | a < 0 || b >= cols m = extractError m e
+
+m ?? e@(Pos vs,_) | dim vs>0 && (minEl vs < 0 || maxEl vs >= fi (rows m)) = extractError m e
+m ?? e@(_,Pos vs) | dim vs>0 && (minEl vs < 0 || maxEl vs >= fi (cols m)) = extractError m e
+
+m ?? (All,All) = m
+
+m ?? (Range a _ b,e) | a > b = m ?? (Take 0,e)
+m ?? (e,Range a _ b) | a > b = m ?? (e,Take 0)
+
+m ?? (Take n,e)
+    | n <= 0      = (0><cols m) [] ?? (All,e)
+    | n >= rows m =              m ?? (All,e)
+
+m ?? (e,Take n)
+    | n <= 0      = (rows m><0) [] ?? (e,All)
+    | n >= cols m =              m ?? (e,All)
+
+m ?? (Drop n,e)
+    | n <= 0      =              m ?? (All,e)
+    | n >= rows m = (0><cols m) [] ?? (All,e)
+
+m ?? (e,Drop n)
+    | n <= 0      =              m ?? (e,All)
+    | n >= cols m = (rows m><0) [] ?? (e,All)
+
+m ?? (TakeLast n, e) = m ?? (Drop (rows m - n), e)
+m ?? (e, TakeLast n) = m ?? (e, Drop (cols m - n))
+
+m ?? (DropLast n, e) = m ?? (Take (rows m - n), e)
+m ?? (e, DropLast n) = m ?? (e, Take (cols m - n))
+
+m ?? (er,ec) = extractR m moder rs modec cs
+  where
+    (moder,rs) = mkExt (rows m) er
+    (modec,cs) = mkExt (cols m) ec
+    ran a b = (0, idxs [a,b])
+    pos ks  = (1, ks)
+    mkExt _ (Pos  ks)     = pos ks
+    mkExt n (PosCyc ks)
+        | n == 0          = mkExt n (Take 0)
+        | otherwise       = pos (cmodi (fi n) ks)
+    mkExt _ (Range mn _ mx) = ran mn mx
+    mkExt _ (Take k)      = ran 0 (k-1)
+    mkExt n (Drop k)      = ran k (n-1)
+    mkExt n _             = ran 0 (n-1) -- All
+
+--------------------------------------------------------------------------------
+
 
 -- | creates a matrix from a vertical list of matrices
 joinVert :: Element t => [Matrix t] -> Matrix t