module Data.Array.Comfort.Boxed ( Array, shape, reshape, mapShape, accessMaybe, (!), Array.toList, Array.fromList, Array.vectorFromList, toAssociations, fromMap, toMap, fromTuple, toTuple, fromRecord, toRecord, fromContainer, toContainer, indices, Array.replicate, Array.map, zipWith, (//), accumulate, fromAssociations, pick, Array.append, Array.take, Array.drop, Array.takeLeft, Array.takeRight, Array.split, Array.takeCenter, ) where import qualified Data.Array.Comfort.Boxed.Unchecked as Array import qualified Data.Array.Comfort.Container as Container import qualified Data.Array.Comfort.Check as Check import qualified Data.Array.Comfort.Shape.Tuple as TupleShape import qualified Data.Array.Comfort.Shape as Shape import Data.Array.Comfort.Boxed.Unchecked (Array(Array)) import qualified Data.Primitive.Array as Prim import qualified Control.Monad.Primitive as PrimM import qualified Control.Monad.Trans.State as MS import Control.Monad.ST (runST) import Control.Applicative ((<$>)) import qualified Data.Foldable as Fold import qualified Data.Map as Map import qualified Data.Set as Set import Data.Map (Map) import Data.Set (Set) import Data.Traversable (Traversable, traverse) import Data.Foldable (forM_) import Data.Either.HT (maybeRight) import Data.Maybe (fromMaybe) import Prelude hiding (zipWith, replicate) shape :: Array.Array sh a -> sh shape :: forall sh a. Array sh a -> sh shape = forall sh a. Array sh a -> sh Array.shape reshape :: (Shape.C sh0, Shape.C sh1) => sh1 -> Array sh0 a -> Array sh1 a reshape :: forall sh0 sh1 a. (C sh0, C sh1) => sh1 -> Array sh0 a -> Array sh1 a reshape = forall sh0 sh1 array0 array1. (C sh0, C sh1) => String -> (array0 -> sh0) -> (sh1 -> array0 -> array1) -> sh1 -> array0 -> array1 Check.reshape String "Boxed" forall sh a. Array sh a -> sh shape forall sh1 sh0 a. sh1 -> Array sh0 a -> Array sh1 a Array.reshape mapShape :: (Shape.C sh0, Shape.C sh1) => (sh0 -> sh1) -> Array sh0 a -> Array sh1 a mapShape :: forall sh0 sh1 a. (C sh0, C sh1) => (sh0 -> sh1) -> Array sh0 a -> Array sh1 a mapShape sh0 -> sh1 f Array sh0 a arr = forall sh0 sh1 a. (C sh0, C sh1) => sh1 -> Array sh0 a -> Array sh1 a reshape (sh0 -> sh1 f forall a b. (a -> b) -> a -> b $ forall sh a. Array sh a -> sh shape Array sh0 a arr) Array sh0 a arr indices :: (Shape.Indexed sh) => sh -> Array.Array sh (Shape.Index sh) indices :: forall sh. Indexed sh => sh -> Array sh (Index sh) indices sh sh = forall sh a. C sh => sh -> [a] -> Array sh a Array.fromList sh sh forall a b. (a -> b) -> a -> b $ forall sh. Indexed sh => sh -> [Index sh] Shape.indices sh sh fromMap :: (Ord k) => Map k a -> Array (Set k) a fromMap :: forall k a. Ord k => Map k a -> Array (Set k) a fromMap Map k a m = forall sh a. C sh => sh -> [a] -> Array sh a Array.fromList (forall k a. Map k a -> Set k Map.keysSet Map k a m) (forall k a. Map k a -> [a] Map.elems Map k a m) toMap :: (Ord k) => Array (Set k) a -> Map k a toMap :: forall k a. Ord k => Array (Set k) a -> Map k a toMap Array (Set k) a arr = forall k a. Eq k => [(k, a)] -> Map k a Map.fromAscList forall a b. (a -> b) -> a -> b $ forall a b. [a] -> [b] -> [(a, b)] zip (forall a. Set a -> [a] Set.toAscList forall a b. (a -> b) -> a -> b $ forall sh a. Array sh a -> sh shape Array (Set k) a arr) (forall sh a. C sh => Array sh a -> [a] Array.toList Array (Set k) a arr) fromTuple :: (TupleShape.NestedTuple tuple) => Shape.DataTuple tuple a -> Array (Shape.NestedTuple ixtype tuple) a fromTuple :: forall tuple a ixtype. NestedTuple tuple => DataTuple tuple a -> Array (NestedTuple ixtype tuple) a fromTuple DataTuple tuple a tuple = case forall s a. State s a -> s -> a MS.evalState (forall shape a. NestedTuple shape => DataTuple shape a -> State Element (shape, [a]) TupleShape.decons DataTuple tuple a tuple) (Int -> Element Shape.Element Int 0) of (tuple sh, [a] xs) -> forall sh a. C sh => sh -> [a] -> Array sh a Array.fromList (forall ixtype tuple. tuple -> NestedTuple ixtype tuple Shape.NestedTuple tuple sh) [a] xs toTuple :: (TupleShape.NestedTuple tuple) => Array (Shape.NestedTuple ixtype tuple) a -> Shape.DataTuple tuple a toTuple :: forall tuple ixtype a. NestedTuple tuple => Array (NestedTuple ixtype tuple) a -> DataTuple tuple a toTuple Array (NestedTuple ixtype tuple) a arr = forall s a. State s a -> s -> a MS.evalState (forall shape a. ElementTuple shape => shape -> State [a] (DataTuple shape a) TupleShape.cons forall a b. (a -> b) -> a -> b $ forall ixtype tuple. NestedTuple ixtype tuple -> tuple Shape.getNestedTuple forall a b. (a -> b) -> a -> b $ forall sh a. Array sh a -> sh shape Array (NestedTuple ixtype tuple) a arr) (forall sh a. C sh => Array sh a -> [a] Array.toList Array (NestedTuple ixtype tuple) a arr) fromRecord :: (Traversable f) => f a -> Array (Shape.Record f) a fromRecord :: forall (f :: * -> *) a. Traversable f => f a -> Array (Record f) a fromRecord f a xs = forall sh a. C sh => sh -> [a] -> Array sh a Array.fromList (forall (f :: * -> *). f Element -> Record f Shape.Record forall a b. (a -> b) -> a -> b $ forall a b c. (a -> b -> c) -> b -> a -> c flip forall s a. State s a -> s -> a MS.evalState (Int -> Element Shape.Element Int 0) forall a b. (a -> b) -> a -> b $ forall (t :: * -> *) (f :: * -> *) a b. (Traversable t, Applicative f) => (a -> f b) -> t a -> f (t b) traverse (forall a b. a -> b -> a const State Element Element TupleShape.next) f a xs) (forall (t :: * -> *) a. Foldable t => t a -> [a] Fold.toList f a xs) toRecord :: (Traversable f) => Array (Shape.Record f) a -> f a toRecord :: forall (f :: * -> *) a. Traversable f => Array (Record f) a -> f a toRecord Array (Record f) a arr = forall s a. State s a -> s -> a MS.evalState (forall (t :: * -> *) (f :: * -> *) a b. (Traversable t, Applicative f) => (a -> f b) -> t a -> f (t b) traverse (forall a b. a -> b -> a const forall a. State [a] a TupleShape.get) forall a b. (a -> b) -> a -> b $ (\(Shape.Record f Element record) -> f Element record) forall a b. (a -> b) -> a -> b $ forall sh a. Array sh a -> sh shape Array (Record f) a arr) (forall sh a. C sh => Array sh a -> [a] Array.toList Array (Record f) a arr) fromContainer :: (Container.C f) => f a -> Array (Container.Shape f) a fromContainer :: forall (f :: * -> *) a. C f => f a -> Array (Shape f) a fromContainer f a xs = forall sh a. C sh => sh -> [a] -> Array sh a Array.fromList (forall (f :: * -> *) a. C f => f a -> Shape f Container.toShape f a xs) (forall (t :: * -> *) a. Foldable t => t a -> [a] Fold.toList f a xs) toContainer :: (Container.C f) => Array (Container.Shape f) a -> f a toContainer :: forall (f :: * -> *) a. C f => Array (Shape f) a -> f a toContainer Array (Shape f) a arr = forall (f :: * -> *) a. C f => Shape f -> [a] -> f a Container.fromList (forall sh a. Array sh a -> sh Array.shape Array (Shape f) a arr) (forall sh a. C sh => Array sh a -> [a] Array.toList Array (Shape f) a arr) infixl 9 ! (!) :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> a ! :: forall sh a. Indexed sh => Array sh a -> Index sh -> a (!) Array sh a arr Index sh ix = forall a. a -> Maybe a -> a fromMaybe (forall a. HasCallStack => String -> a error String "Array.Comfort.Boxed.!: index out of bounds") forall a b. (a -> b) -> a -> b $ forall sh a. Indexed sh => Array sh a -> Index sh -> Maybe a accessMaybe Array sh a arr Index sh ix accessMaybe :: (Shape.Indexed sh) => Array sh a -> Shape.Index sh -> Maybe a accessMaybe :: forall sh a. Indexed sh => Array sh a -> Index sh -> Maybe a accessMaybe (Array sh sh Array a arr) Index sh ix = forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b fmap (forall a. Array a -> Int -> a Prim.indexArray Array a arr) forall a b. (a -> b) -> a -> b $ forall a b. Either a b -> Maybe b maybeRight forall a b. (a -> b) -> a -> b $ forall a. Result Checked a -> Either String a Shape.getChecked forall a b. (a -> b) -> a -> b $ forall sh check. (Indexed sh, Checking check) => sh -> Index sh -> Result check Int Shape.unifiedOffset sh sh Index sh ix zipWith :: (Shape.C sh, Eq sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c zipWith :: forall sh a b c. (C sh, Eq sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c zipWith a -> b -> c f Array sh a a Array sh b b = if forall sh a. Array sh a -> sh shape Array sh a a forall a. Eq a => a -> a -> Bool == forall sh a. Array sh a -> sh shape Array sh b b then forall sh a b c. C sh => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c Array.zipWith a -> b -> c f Array sh a a Array sh b b else forall a. HasCallStack => String -> a error String "zipWith: shapes mismatch" (//) :: (Shape.Indexed sh) => Array sh a -> [(Shape.Index sh, a)] -> Array sh a // :: forall sh a. Indexed sh => Array sh a -> [(Index sh, a)] -> Array sh a (//) (Array sh sh Array a arr) [(Index sh, a)] xs = forall a. (forall s. ST s a) -> a runST (do MutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => Array a -> Int -> Int -> m (MutableArray (PrimState m) a) Prim.thawArray Array a arr Int 0 (forall sh. C sh => sh -> Int Shape.size sh sh) forall (t :: * -> *) (m :: * -> *) a b. (Foldable t, Monad m) => t a -> (a -> m b) -> m () forM_ [(Index sh, a)] xs forall a b. (a -> b) -> a -> b $ \(Index sh ix,a a) -> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> a -> m () Prim.writeArray MutableArray s a marr (forall sh. Indexed sh => sh -> Index sh -> Int Shape.offset sh sh Index sh ix) a a forall sh a. sh -> Array a -> Array sh a Array sh sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> m (Array a) Prim.unsafeFreezeArray MutableArray s a marr) accumulate :: (Shape.Indexed sh) => (a -> b -> a) -> Array sh a -> [(Shape.Index sh, b)] -> Array sh a accumulate :: forall sh a b. Indexed sh => (a -> b -> a) -> Array sh a -> [(Index sh, b)] -> Array sh a accumulate a -> b -> a f (Array sh sh Array a arr) [(Index sh, b)] xs = forall a. (forall s. ST s a) -> a runST (do MutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => Array a -> Int -> Int -> m (MutableArray (PrimState m) a) Prim.thawArray Array a arr Int 0 (forall sh. C sh => sh -> Int Shape.size sh sh) forall (t :: * -> *) (m :: * -> *) a b. (Foldable t, Monad m) => t a -> (a -> m b) -> m () forM_ [(Index sh, b)] xs forall a b. (a -> b) -> a -> b $ \(Index sh ix,b b) -> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> (a -> a) -> m () updateArray MutableArray s a marr (forall sh. Indexed sh => sh -> Index sh -> Int Shape.offset sh sh Index sh ix) forall a b. (a -> b) -> a -> b $ forall a b c. (a -> b -> c) -> b -> a -> c flip a -> b -> a f b b forall sh a. sh -> Array a -> Array sh a Array sh sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> m (Array a) Prim.unsafeFreezeArray MutableArray s a marr) updateArray :: PrimM.PrimMonad m => Prim.MutableArray (PrimM.PrimState m) a -> Int -> (a -> a) -> m () updateArray :: forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> (a -> a) -> m () updateArray MutableArray (PrimState m) a marr Int k a -> a f = forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> a -> m () Prim.writeArray MutableArray (PrimState m) a marr Int k forall b c a. (b -> c) -> (a -> b) -> a -> c . a -> a f forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b =<< forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> m a Prim.readArray MutableArray (PrimState m) a marr Int k toAssociations :: (Shape.Indexed sh) => Array sh a -> [(Shape.Index sh, a)] toAssociations :: forall sh a. Indexed sh => Array sh a -> [(Index sh, a)] toAssociations Array sh a arr = forall a b. [a] -> [b] -> [(a, b)] zip (forall sh. Indexed sh => sh -> [Index sh] Shape.indices forall a b. (a -> b) -> a -> b $ forall sh a. Array sh a -> sh shape Array sh a arr) (forall sh a. C sh => Array sh a -> [a] Array.toList Array sh a arr) fromAssociations :: (Shape.Indexed sh) => a -> sh -> [(Shape.Index sh, a)] -> Array sh a fromAssociations :: forall sh a. Indexed sh => a -> sh -> [(Index sh, a)] -> Array sh a fromAssociations a a sh sh [(Index sh, a)] xs = forall a. (forall s. ST s a) -> a runST (do MutableArray s a marr <- forall (m :: * -> *) a. PrimMonad m => Int -> a -> m (MutableArray (PrimState m) a) Prim.newArray (forall sh. C sh => sh -> Int Shape.size sh sh) a a forall (t :: * -> *) (m :: * -> *) a b. (Foldable t, Monad m) => t a -> (a -> m b) -> m () forM_ [(Index sh, a)] xs forall a b. (a -> b) -> a -> b $ \(Index sh ix,a x) -> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> Int -> a -> m () Prim.writeArray MutableArray s a marr (forall sh. Indexed sh => sh -> Index sh -> Int Shape.offset sh sh Index sh ix) a x forall sh a. sh -> Array a -> Array sh a Array sh sh forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b <$> forall (m :: * -> *) a. PrimMonad m => MutableArray (PrimState m) a -> m (Array a) Prim.unsafeFreezeArray MutableArray s a marr) pick :: (Shape.Indexed sh0, Shape.C sh1) => Array (sh0,sh1) a -> Shape.Index sh0 -> Array sh1 a pick :: forall sh0 sh1 a. (Indexed sh0, C sh1) => Array (sh0, sh1) a -> Index sh0 -> Array sh1 a pick (Array (sh0 sh0,sh1 sh1) Array a x) Index sh0 ix0 = forall sh a. sh -> Array a -> Array sh a Array sh1 sh1 forall a b. (a -> b) -> a -> b $ let k :: Int k = forall sh. C sh => sh -> Int Shape.size sh1 sh1 in forall a. Array a -> Int -> Int -> Array a Prim.cloneArray Array a x (forall sh. Indexed sh => sh -> Index sh -> Int Shape.offset sh0 sh0 Index sh0 ix0 forall a. Num a => a -> a -> a * Int k) Int k