diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright (c) 2015, cchalmers + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of cchalmers nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,111 @@ +## dense + +[![Build Status](https://travis-ci.org/cchalmers/dense.svg)](https://travis-ci.org/cchalmers/dense) +[![Haddock](https://rawgit.com/cchalmers/dense/gh-pages/haddock.svg)](https://cchalmers.github.io/dense/) +[![Hackage](https://img.shields.io/hackage/v/dense.svg?style=flat)](https://hackage.haskell.org/package/dense) + +[`dense`]: http://hackage.haskell.org/package/dense +[`vector`]: http://hackage.haskell.org/package/vector +[`linear`]: http://hackage.haskell.org/package/linear +[`repa`]: http://hackage.haskell.org/package/repa +[`array`]: http://hackage.haskell.org/package/array +[`yarr`]: http://hackage.haskell.org/package/yarr + +[`dense`] is a multidimensional array library build on top of the +[`vector`] package, using indices from the [`linear`] package. Native +support for mutable arrays, stencils and parallel computation. + +### Array type + +Arrays are just vectors (from [`vector`]) with a shape: + + +```.haskell +data Array v f a = Array !(f Layout) !(v a) +``` + +where `Layout f = f Int` is the shape of the array, given by a vector +from [`linear`] (`V1`, `V2`, `V3` or `V4`). These vectors are also used +to indexing: + +```.haskell +> a ! V3 1 2 3 +``` + +### Delayed arrays + +A delayed array, defined by + +```.haskell +data Delayed f a = Delayed !(Layout f) (Int -> a) +``` + +can be constructing from a normal array via `delay`. It can be useful +for mapping a function over an array and computing the result in +parallel via `manifest`: + +```.haskell +> manifest . fmap (+100) . delay +``` + +or equivalently using the `delayed` isomorphism: + +```.haskell +> delayed +~ 100 +``` + +`Delayed` is an instance of many classes, including `Additive` from +[`linear`](http://hackage.haskell.org/package/linear): + +```.haskell +> manifest $ delay a ^+^ 3 *^ delay b +``` + +### Mutable + +[`dense`] has similar mutable capabilities to [`vector`], supporting +mutable operations over a `PrimMonad` in `Data.Dense.Mutable`. + +### Stencils + +[`dense`] has good stencil support, allowing construction of 1D, 2D or 3D +stencils using template haskell and quasiquoters. + +```.haskell +myStencil = [stencil| + 2/5 8/5 2/5 + 8/5 2 8/5 + 2/5 8/5 2/5 +|] +``` + +Stencils made with template haskell are unrolled at compile time. + +### Comparison to other array libraries + +[`array`] supports multidimensional and mutable arrays but [`dense`] +provides many more high level functions as well as stencils and parallel +computation. + +[`repa`] and [`yarr`] +[`dense`] has a lot of the same features as [`repa`] and [`yarr`]. +Performance should be similar (more benchmarks needed) but [`dense`] also +has support for mutable arrays and multidimensional stencils. + +### Package structure + +Like [`vector`], there is a [`Data.Shaped.Generic`] module for working +over any generic vector as well as [`Data.Shaped.Unboxed`] and +[`Data.Shaped.Storable`] modules. Unlike [`vector`], boxed vectors are +in [`Data.Shaped.Boxed`]. + +The [`Data.Shaped`] module includes a subset of [`Data.Shaped.Generic`] +as well as some extra reexports and is intended to be imported +*unqualified*. + + +[`Data.Shaped`]: https://cchalmers.github.io/dense/Data-Shaped.html +[`Data.Shaped.Boxed`]: https://cchalmers.github.io/dense/Data-Shaped-Boxed.html +[`Data.Shaped.Generic`]: https://cchalmers.github.io/dense/Data-Shaped-Generic.html +[`Data.Shaped.Storable`]: https://cchalmers.github.io/dense/Data-Shaped-Storable.html +[`Data.Shaped.Unboxed`]: https://cchalmers.github.io/dense/Data-Shaped-Unboxed.html diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/dense.cabal b/dense.cabal new file mode 100644 --- /dev/null +++ b/dense.cabal @@ -0,0 +1,83 @@ +name: dense +version: 0.1.0.0 +synopsis: Mutable and immutable dense multidimensional arrays +description: + Multidimentional array library build on top of the vector package, + using indices from the linear package. Native support for mutable + arrays, stencils and parallel computation. +license: BSD3 +license-file: LICENSE +author: cchalmers +maintainer: c.chalmers@me.com +copyright: (c) Christopher Chalmers 2016 +category: Data +build-type: Simple +extra-source-files: README.md +cabal-version: >=1.10 + +source-repository head + type: git + location: https://github.com/cchalmers/dense + +library + exposed-modules: + Data.Dense + Data.Dense.Base + Data.Dense.Boxed + Data.Dense.Generic + Data.Dense.Index + Data.Dense.Stencil + Data.Dense.Mutable + Data.Dense.TH + Data.Dense.Storable + Data.Dense.Unboxed + other-extensions: + BangPatterns CPP ConstraintKinds DeriveDataTypeable DeriveFunctor + DeriveGeneric FlexibleContexts FlexibleInstances + MultiParamTypeClasses MultiWayIf RankNTypes StandaloneDeriving + TypeFamilies + build-depends: + base >=4.6 && <5, + binary, + bytes, + cereal, + comonad, + deepseq, + ghc-prim, + hashable, + lens, + linear >= 1.20 && <1.21, + primitive, + semigroupoids, + template-haskell, + transformers, + transformers-compat, + vector + hs-source-dirs: src + ghc-options: -Wall + default-language: Haskell2010 + +test-suite doctests + default-language: Haskell2010 + type: exitcode-stdio-1.0 + main-is: doctest.hs + hs-source-dirs: tests + build-depends: + base >=4.6 && <5, + binary, + bytes, + cereal, + comonad, + deepseq, + ghc-prim, + hashable, + lens, + linear >= 1.20 && <1.21, + primitive, + semigroupoids, + template-haskell, + transformers, + transformers-compat, + vector, + doctest, + simple-reflect diff --git a/src/Data/Dense.hs b/src/Data/Dense.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense.hs @@ -0,0 +1,191 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- This module provides a large subset of the full functionality of +-- "dense" without exporting names that conflict with names in prelude, +-- so it can often be imported unqualified. It also includes reexported +-- classes and data types from other modules. However it does not +-- contain much functions necessary to construct arrays, for that see +-- "Data.Dense.Generic" or one of the type specific modules intended to +-- be imported qualified. Typical imports for shaped will look like +-- this: +-- +-- @ +-- import "Data.Dense" +-- import qualified "Data.Dense.Unboxed" as U +-- @ +-- +-- For boxed-specific arrays (a la "Data.Vector") see "Data.Dense.Boxed". +----------------------------------------------------------------------------- +module Data.Dense + ( + -- * Array types + Array + , BArray + , UArray + , SArray + , PArray + + -- * Indexing + , Layout + , HasLayout (..) + , Shape + , extent + , size + + -- ** Folds over indexes + , indexes + , indexesBetween + , indexesFrom + + -- ** Lenses + , vector + + -- ** Traversals + , values + , values' + , valuesBetween + + -- * Construction + + -- ** Flat arrays + , flat + + -- ** Shaped from lists + , fromListInto + , fromListInto_ + + -- ** Shaped from vectors + , fromVectorInto + , fromVectorInto_ + + -- ** Generating + -- | See "Data.Shaped.Generic". + + -- * Functions on arrays + + -- ** Empty arrays + -- | See 'Control.Lens.Empty.AsEmpty' class or "Data.Shaped.Generic". + + -- ** Indexing + -- | See 'Control.Lens.At.Ixed' class. + + -- ** Modifying arrays + -- | See "Data.Shaped.Generic". + + -- ** Slices + + -- *** Matrix + , ixRow + , rows + , ixColumn + , columns + + -- *** 3D + , ixPlane + , planes + , flattenPlane + + -- * Mutable + , MArray + , BMArray + , UMArray + , SMArray + , PMArray + + -- * Delayed + + , Delayed + + -- ** Generating delayed + + , delayed + , seqDelayed + , delay + , manifest + , seqManifest + , genDelayed + , indexDelayed + , affirm + , seqAffirm + + -- ** Helpful reexports + , (*^) + , (^*) + , (^/) + , Additive (..) + , Metric (..) + + -- * Focused + + , Focused + + -- ** Generating focused + + , focusOn + , unfocus + , unfocused + , extendFocus + + -- ** Focus location + , locale + , shiftFocus + + -- ** Boundary + , Boundary (..) + , peekB + , peeksB + , peekRelativeB + + -- ** Helpful reexports + , Comonad (..) + , ComonadStore (..) + + -- * Stencils + , Stencil + + -- ** Constructing stencils + , stencil + , mkStencil + , mkStencilTH + + -- ** Using stencils + , stencilSum + + -- * Common shapes + , V1 (..) + , V2 (..) + , V3 (..) + , V4 (..) + , R1 (..) + , R2 (..) + , R3 (..) + , R4 (..) + + -- ** Extra planes + , _xz + , _yz + , _yx + , _zy + , _zx + ) where + +import Data.Dense.Generic +import Control.Comonad.Store +import Linear hiding (vector) +import Data.Dense.TH +import Data.Dense.Stencil + diff --git a/src/Data/Dense/Base.hs b/src/Data/Dense/Base.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Base.hs @@ -0,0 +1,613 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE StandaloneDeriving #-} +{-# LANGUAGE TypeFamilies #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Base +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- Base module for multidimensional arrays. This module exports the +-- constructors for the 'Array' data type. +-- +-- Also, to prevent this module becomming too large, only the data types +-- and the functions nessesary for the instances are defined here. All +-- other functions are defined in "Data.Dense.Generic". +----------------------------------------------------------------------------- +module Data.Dense.Base + ( + -- * Array types + Array (..) + , Boxed + + -- ** Lenses + , vector + , values + + -- ** Conversion to/from mutable arrays + + , unsafeThaw + , unsafeFreeze + + -- * Delayed + + , Delayed (..) + , delay + , manifest + , genDelayed + , indexDelayed + + -- * Focused + + , Focused (..) + + ) where + + +#if __GLASGOW_HASKELL__ <= 708 +import Control.Applicative (pure, (*>)) +import Data.Foldable (Foldable) +import Data.Monoid (Monoid, mappend, mempty) +#endif + +import Control.Applicative (liftA2) +import Control.Comonad +import Control.Comonad.Store +import Control.DeepSeq +import Control.Lens +import Control.Lens.Internal (noEffect) +import Control.Monad (guard, liftM) +import Control.Monad.Primitive +import Data.Binary as Binary +import Data.Bytes.Serial +import Data.Data +import qualified Data.Foldable as F +import Data.Functor.Apply +import Data.Functor.Classes +import Data.Functor.Extend +import Data.Hashable +import Data.Serialize as Cereal +import Data.Traversable (for) +import qualified Data.Vector as B +import Data.Vector.Generic (Vector) +import qualified Data.Vector.Generic as G +import Data.Vector.Generic.Lens (vectorTraverse) +import qualified Data.Vector.Generic.Mutable as GM +import qualified Data.Vector.Generic.New as New +-- import GHC.Generics (Generic, Generic1) +import Linear hiding (vector) +import Text.ParserCombinators.ReadPrec (readS_to_Prec) +import qualified Text.Read as Read + +import Data.Dense.Index +import Data.Dense.Mutable (MArray (..)) + +import Control.Concurrent (forkOn, getNumCapabilities, + newEmptyMVar, putMVar, + takeMVar) +import System.IO.Unsafe (unsafePerformIO) + +import Prelude hiding (null, replicate, + zipWith, zipWith3) + +import GHC.Types (SPEC (..)) + +-- | An 'Array' is a vector with a shape. +data Array v f a = Array !(Layout f) !(v a) + deriving Typeable + +-- Lenses -------------------------------------------------------------- + +-- | Indexed traversal over the elements of an array. The index is the +-- current position in the array. +values :: (Shape f, Vector v a, Vector w b) + => IndexedTraversal (f Int) (Array v f a) (Array w f b) a b +values = \f arr -> reindexed (shapeFromIndex $ extent arr) (vector . vectorTraverse) f arr +{-# INLINE values #-} + +-- | Indexed lens over the underlying vector of an array. The index is +-- the 'extent' of the array. You must _not_ change the length of the +-- vector, otherwise an error will be thrown (even for 'V1' layouts, +-- use 'flat' for 'V1'). +vector :: (Vector v a, Vector w b) => IndexedLens (Layout f) (Array v f a) (Array w f b) (v a) (w b) +vector f (Array l v) = + indexed f l v <&> \w -> + sizeMissmatch (G.length v) (G.length w) + ("vector: trying to replace vector of length " ++ show (G.length v) ++ " with one of length " ++ show (G.length w)) + $ Array l w +{-# INLINE vector #-} + +-- Mutable conversion -------------------------------------------------- + +-- | O(1) Unsafe convert a mutable array to an immutable one without +-- copying. The mutable array may not be used after this operation. +unsafeFreeze :: (PrimMonad m, Vector v a) + => MArray (G.Mutable v) f (PrimState m) a -> m (Array v f a) +unsafeFreeze (MArray l mv) = Array l `liftM` G.unsafeFreeze mv +{-# INLINE unsafeFreeze #-} + +-- | O(1) Unsafely convert an immutable array to a mutable one without +-- copying. The immutable array may not be used after this operation. +unsafeThaw :: (PrimMonad m, Vector v a) + => Array v f a -> m (MArray (G.Mutable v) f (PrimState m) a) +unsafeThaw (Array l v) = MArray l `liftM` G.unsafeThaw v +{-# INLINE unsafeThaw #-} + +------------------------------------------------------------------------ +-- Instances +------------------------------------------------------------------------ + +-- | The 'size' of the 'layout' __must__ remain the same or an error is thrown. +instance Shape f => HasLayout f (Array v f a) where + layout f (Array l v) = f l <&> \l' -> + sizeMissmatch (shapeSize l) (shapeSize l') + ("layout (Array): trying to replace shape " ++ showShape l ++ " with " ++ showShape l') + $ Array l' v + {-# INLINE layout #-} + +-- layout :: (Shape l, Shape t) => Lens (Array v l a) (Array v t a) (Layout l) (Layout t) + +instance (Vector v a, Eq1 f, Eq a) => Eq (Array v f a) where + Array l1 v1 == Array l2 v2 = eq1 l1 l2 && G.eq v1 v2 + {-# INLINE (==) #-} + +instance (Vector v a, Show1 f, Show a) => Show (Array v f a) where + showsPrec p (Array l v2) = showParen (p > 10) $ + showString "Array " . showsPrec1 11 l . showChar ' ' . G.showsPrec 11 v2 + +type instance Index (Array v f a) = f Int +type instance IxValue (Array v f a) = a + +instance (Shape f, Vector v a) => Ixed (Array v f a) where + ix x f (Array l v) + | shapeInRange l x = f (G.unsafeIndex v i) <&> + \a -> Array l (G.modify (\mv -> GM.unsafeWrite mv i a) v) + where i = shapeToIndex l x + ix _ _ arr = pure arr + {-# INLINE ix #-} + +instance (Vector v a, Vector v b) => Each (Array v f a) (Array v f b) a b where + each = vector . vectorTraverse + {-# INLINE each #-} + +instance (Shape f, Vector v a) => AsEmpty (Array v f a) where + _Empty = nearly (Array zero G.empty) (F.all (==0) . extent) + {-# INLINE _Empty #-} + +instance (Vector v a, Read1 f, Read a) => Read (Array v f a) where + readPrec = Read.parens $ Read.prec 10 $ do + Read.Ident "Array" <- Read.lexP + l <- readS_to_Prec readsPrec1 + v <- G.readPrec + return $ Array l v + +instance (NFData (f Int), NFData (v a)) => NFData (Array v f a) where + rnf (Array l v) = rnf l `seq` rnf v + {-# INLINE rnf #-} + +-- Boxed instances ----------------------------------------------------- + +-- | The vector is the boxed vector. +type Boxed v = v ~ B.Vector + +instance Boxed v => Functor (Array v f) where + fmap = over vector . fmap + {-# INLINE fmap #-} + +instance Boxed v => F.Foldable (Array v f) where + foldMap f = F.foldMap f . view vector + {-# INLINE foldMap #-} + +instance Boxed v => Traversable (Array v f) where + traverse = each + {-# INLINE traverse #-} + +#if (MIN_VERSION_transformers(0,5,0)) || !(MIN_VERSION_transformers(0,4,0)) +instance (Boxed v, Eq1 f) => Eq1 (Array v f) where + liftEq f (Array l1 v1) (Array l2 v2) = eq1 l1 l2 && G.and (G.zipWith f v1 v2) + {-# INLINE liftEq #-} + +instance (Boxed v, Read1 f) => Read1 (Array v f) where + liftReadsPrec _ f = readsData $ readsBinaryWith readsPrec1 (const f) "Array" (\c l -> Array c (G.fromList l)) + {-# INLINE liftReadsPrec #-} +#else +instance (Boxed v, Eq1 f) => Eq1 (Array v f) where + eq1 = (==) + {-# INLINE eq1 #-} + +instance (Boxed v, Read1 f) => Read1 (Array v f) where + readsPrec1 = readsPrec + {-# INLINE readsPrec1 #-} +#endif + +instance (Boxed v, Shape f) => FunctorWithIndex (f Int) (Array v f) +instance (Boxed v, Shape f) => FoldableWithIndex (f Int) (Array v f) +instance (Boxed v, Shape f) => TraversableWithIndex (f Int) (Array v f) where + itraverse = itraverseOf values + {-# INLINE itraverse #-} + itraversed = values + {-# INLINE itraversed #-} + +instance (Boxed v, Shape f, Serial1 f) => Serial1 (Array v f) where + serializeWith putF (Array l v) = do + serializeWith serialize l + F.traverse_ putF v + deserializeWith = genGet (deserializeWith deserialize) + +-- deriving instance (Generic1 v, Generic1 f) => Generic1 (Array v f) + +-- instance (v ~ B.Vector, Shape l) => Apply (Array v l) where +-- instance (v ~ B.Vector, Shape l) => Bind (Array v l) where +-- instance (v ~ B.Vector, Shape l) => Additive (Array v l) where +-- instance (v ~ B.Vector, Shape l) => Metric (Array v l) where + +-- V1 instances -------------------------------------------------------- + +-- Array v V1 a is essentially v a with a wrapper. + +type instance G.Mutable (Array v f) = MArray (G.Mutable v) f + +-- | 1D Arrays can be used as a generic 'Vector'. +instance (Vector v a, f ~ V1) => Vector (Array v f) a where + {-# INLINE basicUnsafeFreeze #-} + {-# INLINE basicUnsafeThaw #-} + {-# INLINE basicLength #-} + {-# INLINE basicUnsafeSlice #-} + {-# INLINE basicUnsafeIndexM #-} + basicUnsafeFreeze = unsafeFreeze + basicUnsafeThaw = unsafeThaw + basicLength (Array (V1 n) _) = n + basicUnsafeSlice i n (Array _ v) = Array (V1 n) $ G.basicUnsafeSlice i n v + basicUnsafeIndexM (Array _ v) = G.basicUnsafeIndexM v + +-- Serialise instances ------------------------------------------------- + +instance (Vector v a, Shape f, Serial1 f, Serial a) => Serial (Array v f a) where + serialize (Array l v) = do + serializeWith serialize l + traverseOf_ vectorTraverse serialize v + {-# INLINE serialize #-} + deserialize = genGet (deserializeWith deserialize) deserialize + {-# INLINE deserialize #-} + +instance (Vector v a, Shape f, Binary (f Int), Binary a) => Binary (Array v f a) where + put (Array l v) = do + Binary.put l + traverseOf_ vectorTraverse Binary.put v + {-# INLINE put #-} + get = genGet Binary.get Binary.get + {-# INLINE get #-} + +instance (Vector v a, Shape f, Serialize (f Int), Serialize a) => Serialize (Array v f a) where + put (Array l v) = do + Cereal.put l + traverseOf_ vectorTraverse Cereal.put v + {-# INLINE put #-} + get = genGet Cereal.get Cereal.get + {-# INLINE get #-} + +genGet :: Monad m => (Vector v a, Shape f) => m (f Int) -> m a -> m (Array v f a) +genGet getL getA = do + l <- getL + let n = shapeSize l + nv0 = New.create (GM.new n) + f acc i = (\a -> New.modify (\mv -> GM.write mv i a) acc) `liftM` getA + nv <- F.foldlM f nv0 [0 .. n - 1] + return $! Array l (G.new nv) +{-# INLINE genGet #-} + +instance (Vector v a, Foldable f, Hashable a) => Hashable (Array v f a) where + hashWithSalt s (Array l v) = G.foldl' hashWithSalt s' v + where s' = F.foldl' hashWithSalt s l + {-# INLINE hashWithSalt #-} + +-- deriving instance (Generic (v a), Generic1 f) => Generic (Array v f a) +deriving instance (Typeable f, Typeable v, Typeable a, Data (f Int), Data (v a)) => Data (Array v f a) + + +-- instance (Vector v a, Typeable v, Typeable l, Shape l, Data a) => Data (Array v l a) where +-- gfoldl f z (Array l a) = +-- z (\l' a' -> Array (l & partsOf traverse .~ l') (G.fromList a')) `f` F.toList l `f` G.toList a +-- gunfold k z _ = k (k (z (\l a -> Array (zero & partsOf traverse .~ l) (G.fromList a)))) +-- toConstr _ = con +-- dataTypeOf _ = ty +-- dataCast1 = gcast1 + +-- ty :: DataType +-- ty = mkDataType "Array" [con] + +-- con :: Constr +-- con = mkConstr ty "Array" [] Prefix + +------------------------------------------------------------------------ +-- Delayed +------------------------------------------------------------------------ + +-- | A delayed representation of an array. This useful for mapping over +-- an array in parallel. +data Delayed f a = Delayed !(Layout f) (f Int -> a) + deriving (Typeable, Functor) + +-- | Turn a material array into a delayed one with the same shape. +delay :: (Vector v a, Shape f) => Array v f a -> Delayed f a +delay (Array l v) = Delayed l (G.unsafeIndex v . shapeToIndex l) +{-# INLINE delay #-} + +-- | The 'size' of the 'layout' __must__ remain the same or an error is thrown. +instance Shape f => HasLayout f (Delayed f a) where + layout f (Delayed l ixF) = f l <&> \l' -> + sizeMissmatch (shapeSize l) (shapeSize l') + ("layout (Delayed): trying to replace shape " ++ showShape l ++ " with " ++ showShape l') + $ Delayed l' ixF + {-# INLINE layout #-} + +-- | 'foldMap' in parallel. +instance Shape f => Foldable (Delayed f) where + foldr f b (Delayed l ixF) = foldrOf shapeIndexes (\x -> f (ixF x)) b l + {-# INLINE foldr #-} + + foldMap = foldDelayed . const + +#if __GLASGOW_HASKELL__ >= 710 + length = size + {-# INLINE length #-} +#endif + +instance (Shape f, Show1 f, Show a) => Show (Delayed f a) where + showsPrec p arr@(Delayed l _) = showParen (p > 10) $ + showString "Delayed " . showsPrec1 11 l . showChar ' ' . showsPrec 11 (F.toList arr) + +-- instance (Shape f, Show1 f) => Show1 (Delayed f) where +-- showsPrec1 = showsPrec + +instance Shape f => Traversable (Delayed f) where + traverse f arr = delay <$> traversed f (manifest arr) + +instance Shape f => Apply (Delayed f) where + {-# INLINE (<.>) #-} + {-# INLINE (<. ) #-} + {-# INLINE ( .>) #-} + (<.>) = liftI2 id + (<. ) = liftI2 const + ( .>) = liftI2 (const id) + +instance Shape f => Additive (Delayed f) where + zero = _Empty # () + {-# INLINE zero #-} + + -- This can only be satisfied on if one array is larger than the other + -- in all dimensions, otherwise there will be gaps in the array + liftU2 f (Delayed l ixF) (Delayed k ixG) + | l `eq1` k = Delayed l (liftA2 f ixF ixG) + + -- l > k + | F.all (>= EQ) cmp = Delayed l $ \x -> + if | shapeInRange l x -> liftA2 f ixF ixG x + | otherwise -> ixF x + + -- k > l + | F.all (<= EQ) cmp = Delayed k $ \x -> + if | shapeInRange k x -> liftA2 f ixF ixG x + | otherwise -> ixG x + + -- not possible to union array sizes because there would be gaps, + -- just intersect them instead + | otherwise = Delayed (shapeIntersect l k) $ liftA2 f ixF ixG + where cmp = liftI2 compare l k + + liftI2 f (Delayed l ixF) (Delayed k ixG) = Delayed (shapeIntersect l k) $ liftA2 f ixF ixG + {-# INLINE liftI2 #-} + +instance Shape f => Metric (Delayed f) + +instance FunctorWithIndex (f Int) (Delayed f) where + imap f (Delayed l ixF) = Delayed l $ \x -> f x (ixF x) + {-# INLINE imap #-} + +-- | 'ifoldMap' in parallel. +instance Shape f => FoldableWithIndex (f Int) (Delayed f) where + ifoldr f b (Delayed l ixF) = foldrOf shapeIndexes (\x -> f x (ixF x)) b l + {-# INLINE ifoldr #-} + + ifolded = ifoldring ifoldr + {-# INLINE ifolded #-} + + ifoldMap = foldDelayed + {-# INLINE ifoldMap #-} + +instance Shape f => TraversableWithIndex (f Int) (Delayed f) where + itraverse f arr = delay <$> itraverse f (manifest arr) + {-# INLINE itraverse #-} + +instance Shape f => Each (Delayed f a) (Delayed f b) a b where + each = traversed + {-# INLINE each #-} + +instance Shape f => AsEmpty (Delayed f a) where + _Empty = nearly (Delayed zero (error "empty delayed array")) + (\(Delayed l _) -> F.all (==0) l) + {-# INLINE _Empty #-} + +type instance Index (Delayed f a) = f Int +type instance IxValue (Delayed f a) = a +instance Shape f => Ixed (Delayed f a) where + ix x f arr@(Delayed l ixF) + | shapeInRange l x = f (ixF x) <&> \a -> + let g y | eq1 x y = a + | otherwise = ixF x + in Delayed l g + | otherwise = pure arr + {-# INLINE ix #-} + +-- | Index a delayed array, returning a 'IndexOutOfBounds' exception if +-- the index is out of range. +indexDelayed :: Shape f => Delayed f a -> f Int -> a +indexDelayed (Delayed l ixF) x = + boundsCheck l x $ ixF x +{-# INLINE indexDelayed #-} + +foldDelayed :: (Shape f, Monoid m) => (f Int -> a -> m) -> (Delayed f a) -> m +foldDelayed f (Delayed l ixF) = unsafePerformIO $ do + childs <- for [0 .. threads - 1] $ \c -> do + child <- newEmptyMVar + _ <- forkOn c $ do + let k | c == threads - 1 = q + r + | otherwise = q + x = c * q + m = x + k + go i (Just s) acc + | i >= m = acc + | otherwise = let !acc' = acc `mappend` f s (ixF s) + in go (i+1) (shapeStep l s) acc' + go _ Nothing acc = acc + putMVar child $! go x (Just $ shapeFromIndex l x) mempty + return child + F.fold <$> for childs takeMVar + where + !n = shapeSize l + !(q, r) = n `quotRem` threads + !threads = unsafePerformIO getNumCapabilities +{-# INLINE foldDelayed #-} + +-- | Parallel manifestation of a delayed array into a material one. +manifest :: (Vector v a, Shape f) => Delayed f a -> Array v f a +manifest (Delayed l ixF) = Array l v + where + !v = unsafePerformIO $! do + mv <- GM.new n + childs <- for [0 .. threads - 1] $ \c -> do + child <- newEmptyMVar + _ <- forkOn c $ do + let k | c == threads - 1 = q + r + | otherwise = q + x = c * q + iforOf_ (linearIndexesBetween x (x+k)) l $ \i s -> + GM.unsafeWrite mv i (ixF s) + putMVar child () + return child + F.for_ childs takeMVar + G.unsafeFreeze mv + !n = shapeSize l + !(q, r) = n `quotRem` threads + !threads = unsafePerformIO getNumCapabilities +{-# INLINE manifest #-} + +linearIndexesBetween :: Shape f => Int -> Int -> IndexedFold Int (Layout f) (f Int) +linearIndexesBetween i0 k g l = go SPEC i0 (Just $ shapeFromIndex l i0) + where + go !_ i (Just x) = indexed g i x *> go SPEC (i+1) (guard (i+1 < k) *> shapeStep l x) + go !_ _ _ = noEffect +{-# INLINE linearIndexesBetween #-} + +-- | Generate a 'Delayed' array using the given 'Layout' and +-- construction function. +genDelayed :: Layout f -> (f Int -> a) -> Delayed f a +genDelayed = Delayed +{-# INLINE genDelayed #-} + +------------------------------------------------------------------------ +-- Focused +------------------------------------------------------------------------ + +-- | A delayed representation of an array with a focus on a single +-- element. This element is the target of 'extract'. +data Focused f a = Focused !(f Int) !(Delayed f a) + deriving (Typeable, Functor) + +-- | The 'size' of the 'layout' __must__ remain the same or an error is thrown. +instance Shape f => HasLayout f (Focused f a) where + layout f (Focused x (Delayed l ixF)) = f l <&> \l' -> + sizeMissmatch (shapeSize l) (shapeSize l') + ("layout (Focused): trying to replace shape " ++ showShape l ++ " with " ++ showShape l') + $ Focused x (Delayed l' ixF) + {-# INLINE layout #-} + +instance Shape f => Comonad (Focused f) where + {-# INLINE extract #-} + {-# INLINE extend #-} + extract (Focused x d) = indexDelayed d x + extend f (Focused x d@(Delayed l _)) = + Focused x (genDelayed l $ \i -> f (Focused i d)) + +instance Shape f => Extend (Focused f) where + {-# INLINE extended #-} + extended = extend + +instance Shape f => ComonadStore (f Int) (Focused f) where + {-# INLINE pos #-} + {-# INLINE peek #-} + {-# INLINE peeks #-} + {-# INLINE seek #-} + {-# INLINE seeks #-} + pos (Focused x _) = x + peek x (Focused _ d) = indexDelayed d x + peeks f (Focused x d) = indexDelayed d (f x) + seek x (Focused _ d) = Focused x d + seeks f (Focused x d) = Focused (f x) d + +instance (Shape f, Show1 f, Show a) => Show (Focused f a) where + showsPrec p (Focused l d) = showParen (p > 10) $ + showString "Focused " . showsPrec1 11 l . showChar ' ' . showsPrec 11 d + +-- instance (Shape f, Show1 f) => Show1 (Focused f) where +-- showsPrec1 = showsPrec + +type instance Index (Focused f a) = f Int +type instance IxValue (Focused f a) = a + +instance Shape f => Foldable (Focused f) where + foldr f b (Focused _ d) = F.foldr f b d + {-# INLINE foldr #-} + + foldMap f (Focused _ d) = F.foldMap f d + {-# INLINE foldMap #-} + +#if __GLASGOW_HASKELL__ >= 710 + length = size + {-# INLINE length #-} +#endif + +instance Shape f => Traversable (Focused f) where + traverse f (Focused u d) = Focused u <$> traverse f d + {-# INLINE traverse #-} + +-- | Index relative to focus. +instance Shape f => FunctorWithIndex (f Int) (Focused f) where + imap f (Focused u d) = Focused u (imap (f . (^-^ u)) d) + {-# INLINE imap #-} + +-- | Index relative to focus. +instance Shape f => FoldableWithIndex (f Int) (Focused f) where + ifoldr f b (Focused u d) = ifoldr (f . (^-^ u)) b d + {-# INLINE ifoldr #-} + + ifolded = ifoldring ifoldr + {-# INLINE ifolded #-} + + ifoldMap f (Focused u d) = ifoldMap (f . (^-^) u) d + {-# INLINE ifoldMap #-} + +-- | Index relative to focus. +instance Shape f => TraversableWithIndex (f Int) (Focused f) where + itraverse f (Focused u d) = Focused u <$> itraverse (f . (^-^ u)) d + {-# INLINE itraverse #-} + +-- | Index relative to focus. +instance Shape f => Ixed (Focused f a) where + ix i f (Focused u d) = Focused u <$> ix (i ^-^ u) f d + {-# INLINE ix #-} + diff --git a/src/Data/Dense/Boxed.hs b/src/Data/Dense/Boxed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Boxed.hs @@ -0,0 +1,663 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Boxed +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- Boxed multidimensional arrays. +----------------------------------------------------------------------------- +module Data.Dense.Boxed + ( + -- * BArray types + BArray + , Shape + + -- * Layout of an array + , HasLayout (..) + , Layout + + -- ** Extracting size + , extent + , size + + -- ** Folds over indexes + , indexes + , indexesFrom + , indexesBetween + + -- * Underlying vector + , vector + + -- ** Traversals + , values + , values' + , valuesBetween + + -- * Construction + + -- ** Flat arrays + , flat + , fromList + + -- ** From lists + , fromListInto + , fromListInto_ + + -- ** From vectors + , fromVectorInto + , fromVectorInto_ + + -- ** Initialisation + , replicate + , generate + , linearGenerate + + -- ** Monadic initialisation + , create + , replicateM + , generateM + , linearGenerateM + + -- * Functions on arrays + + -- ** Empty arrays + , empty + , null + + -- ** Indexing + + , (!) + , (!?) + , unsafeIndex + , linearIndex + , unsafeLinearIndex + + -- *** Monadic indexing + , indexM + , unsafeIndexM + , linearIndexM + , unsafeLinearIndexM + + -- ** Modifying arrays + + -- ** Bulk updates + , (//) + + -- ** Accumulations + , accum + + -- ** Mapping + , map + , imap + + -- * Zipping + -- ** Tuples + , zip + , zip3 + + -- ** Zip with function + , zipWith + , zipWith3 + , izipWith + , izipWith3 + + -- ** Slices + + -- *** Matrix + , ixRow + , rows + , ixColumn + , columns + + -- *** 3D + , ixPlane + , planes + , flattenPlane + + -- *** Ordinals + , unsafeOrdinals + + -- * Mutable + , BMArray + + , thaw + , freeze + , unsafeThaw + , unsafeFreeze + + -- * Delayed + + , G.Delayed + + -- ** Generating delayed + + , delayed + , seqDelayed + , delay + , manifest + , seqManifest + , G.genDelayed + , G.indexDelayed + , affirm + , seqAffirm + + -- * Focused + + , G.Focused + + -- ** Generating focused + + , G.focusOn + , G.unfocus + , G.unfocused + , G.extendFocus + + -- ** Focus location + , G.locale + , G.shiftFocus + + ) where + +import Control.Lens hiding (imap) +import Control.Monad.Primitive +import Control.Monad.ST +import qualified Data.Foldable as F +import Data.Vector (Vector) +import Linear hiding (vector) + +import Prelude hiding (map, null, replicate, zip, + zip3, zipWith, zipWith3) + +import Data.Dense.Generic (BArray) +import qualified Data.Dense.Generic as G +import Data.Dense.Index +import Data.Dense.Mutable (BMArray) + +-- Lenses -------------------------------------------------------------- + +-- | Same as 'values' but restrictive in the vector type. +values :: Shape f + => IndexedTraversal (f Int) (BArray f a) (BArray f b) a b +values = G.values' +{-# INLINE values #-} + +-- | Same as 'values' but restrictive in the vector type. +values' :: Shape f + => IndexedTraversal (f Int) (BArray f a) (BArray f b) a b +values' = G.values' +{-# INLINE values' #-} + +-- | Same as 'values' but restrictive in the vector type. +valuesBetween + :: Shape f + => f Int + -> f Int + -> IndexedTraversal' (f Int) (BArray f a) a +valuesBetween = G.valuesBetween +{-# INLINE valuesBetween #-} + +-- | 1D arrays are just vectors. You are free to change the length of +-- the vector when going 'over' this 'Iso' (unlike 'linear'). +-- +-- Note that 'V1' arrays are an instance of 'Vector' so you can use +-- any of the functions in 'Data.Vector.Generic' on them without +-- needing to convert. +flat :: Iso (BArray V1 a) (BArray V1 b) (Vector a) (Vector b) +flat = G.flat +{-# INLINE flat #-} + +-- | Indexed lens over the underlying vector of an array. The index is +-- the 'extent' of the array. You must _not_ change the length of the +-- vector, otherwise an error will be thrown (even for 'V1' layouts, +-- use 'flat' for 'V1'). +vector :: IndexedLens (Layout f) (BArray f a) (BArray f b) (Vector a) (Vector b) +vector = G.vector +{-# INLINE vector #-} + +-- Constructing vectors ------------------------------------------------ + +-- | Contruct a flat array from a list. (This is just 'G.fromList' from +-- 'Data.Vector.Generic'.) +fromList :: [a] -> BArray V1 a +fromList = G.fromList +{-# INLINE fromList #-} + +-- | O(n) Convert the first @n@ elements of a list to an BArrayith the +-- given shape. Returns 'Nothing' if there are not enough elements in +-- the list. +fromListInto :: Shape f => Layout f -> [a] -> Maybe (BArray f a) +fromListInto = G.fromListInto +{-# INLINE fromListInto #-} + +-- | O(n) Convert the first @n@ elements of a list to an BArrayith the +-- given shape. Throw an error if the list is not long enough. +fromListInto_ :: Shape f => Layout f -> [a] -> BArray f a +fromListInto_ = G.fromListInto_ +{-# INLINE fromListInto_ #-} + +-- | Create an array from a 'vector' and a 'layout'. Return 'Nothing' if +-- the vector is not the right shape. +fromVectorInto :: Shape f => Layout f -> Vector a -> Maybe (BArray f a) +fromVectorInto = G.fromVectorInto +{-# INLINE fromVectorInto #-} + +-- | Create an array from a 'vector' and a 'layout'. Throws an error if +-- the vector is not the right shape. +fromVectorInto_ :: Shape f => Layout f -> Vector a -> BArray f a +fromVectorInto_ = G.fromVectorInto_ +{-# INLINE fromVectorInto_ #-} + +-- | The empty 'BArray' with a 'zero' shape. +empty :: (Additive f) => BArray f a +empty = G.empty +{-# INLINE empty #-} + +-- | Test is if the array is 'empty'. +null :: F.Foldable f => BArray f a -> Bool +null = G.null +{-# INLINE null #-} + +-- Indexing ------------------------------------------------------------ + +-- | Index an element of an array. Throws 'IndexOutOfBounds' if the +-- index is out of bounds. +(!) :: Shape f => BArray f a -> f Int -> a +(!) = (G.!) +{-# INLINE (!) #-} + +-- | Safe index of an element. +(!?) :: Shape f => BArray f a -> f Int -> Maybe a +(!?) = (G.!?) +{-# INLINE (!?) #-} + +-- | Index an element of an array without bounds checking. +unsafeIndex :: Shape f => BArray f a -> f Int -> a +unsafeIndex = G.unsafeIndex +{-# INLINE unsafeIndex #-} + +-- | Index an element of an array while ignoring its shape. +linearIndex :: BArray f a -> Int -> a +linearIndex = G.linearIndex +{-# INLINE linearIndex #-} + +-- | Index an element of an array while ignoring its shape, without +-- bounds checking. +unsafeLinearIndex :: BArray f a -> Int -> a +unsafeLinearIndex = G.unsafeLinearIndex +{-# INLINE unsafeLinearIndex #-} + +-- Monadic indexing ---------------------------------------------------- + +-- | /O(1)/ Indexing in a monad. +-- +-- The monad allows operations to be strict in the vector when necessary. +-- Suppose vector copying is implemented like this: +-- +-- > copy mv v = ... write mv i (v ! i) ... +-- +-- For lazy vectors, @v ! i@ would not be evaluated which means that +-- @mv@ would unnecessarily retain a reference to @v@ in each element +-- written. +-- +-- With 'indexM', copying can be implemented like this instead: +-- +-- > copy mv v = ... do +-- > x <- indexM v i +-- > write mv i x +-- +-- Here, no references to @v@ are retained because indexing (but /not/ +-- the elements) is evaluated eagerly. +-- +-- Throws an error if the index is out of range. +indexM :: (Shape f, Monad m) => BArray f a -> f Int -> m a +indexM = G.indexM +{-# INLINE indexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeIndexM :: (Shape f, Monad m) => BArray f a -> f Int -> m a +unsafeIndexM = G.unsafeIndexM +{-# INLINE unsafeIndexM #-} + +-- | /O(1)/ Indexing in a monad. Throws an error if the index is out of +-- range. +linearIndexM :: (Shape f, Monad m) => BArray f a -> Int -> m a +linearIndexM = G.linearIndexM +{-# INLINE linearIndexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeLinearIndexM :: Monad m => BArray f a -> Int -> m a +unsafeLinearIndexM = G.unsafeLinearIndexM +{-# INLINE unsafeLinearIndexM #-} + +-- Initialisation ------------------------------------------------------ + +-- | Execute the monadic action and freeze the resulting array. +create :: (forall s. ST s (BMArray f s a)) -> BArray f a +create m = m `seq` runST (m >>= G.unsafeFreeze) +{-# INLINE create #-} + +-- | O(n) BArray of the given shape with the same value in each position. +replicate :: Shape f => f Int -> a -> BArray f a +replicate = G.replicate +{-# INLINE replicate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +linearGenerate :: Shape f => Layout f -> (Int -> a) -> BArray f a +linearGenerate = G.linearGenerate +{-# INLINE linearGenerate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +generate :: Shape f => Layout f -> (f Int -> a) -> BArray f a +generate = G.generate +{-# INLINE generate #-} + +-- Monadic initialisation ---------------------------------------------- + +-- | O(n) Construct an array of the given shape by filling each position +-- with the monadic value. +replicateM :: (Monad m, Shape f) => Layout f -> m a -> m (BArray f a) +replicateM = G.replicateM +{-# INLINE replicateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +generateM :: (Monad m, Shape f) => Layout f -> (f Int -> m a) -> m (BArray f a) +generateM = G.generateM +{-# INLINE generateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +linearGenerateM :: (Monad m, Shape f) => Layout f -> (Int -> m a) -> m (BArray f a) +linearGenerateM = G.linearGenerateM +{-# INLINE linearGenerateM #-} + +-- Modifying ----------------------------------------------------------- + +-- | /O(n)/ Map a function over an array +map :: (a -> b) -> BArray f a -> BArray f b +map = G.map +{-# INLINE map #-} + +-- | /O(n)/ Apply a function to every element of a vector and its index +imap :: Shape f => (f Int -> a -> b) -> BArray f a -> BArray f b +imap = G.imap +{-# INLINE imap #-} + +-- Bulk updates -------------------------------------------------------- + + +-- | For each pair (i,a) from the list, replace the array element at +-- position i by a. +(//) :: Shape f => BArray f a -> [(f Int, a)] -> BArray f a +(//) = (G.//) +{-# INLINE (//) #-} + +-- Accumilation -------------------------------------------------------- + +-- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the array element +-- @a@ at position @i@ by @f a b@. +-- +accum :: Shape f + => (a -> b -> a) -- ^ accumulating function @f@ + -> BArray f a -- ^ initial array + -> [(f Int, b)] -- ^ list of index/value pairs (of length @n@) + -> BArray f a +accum = G.accum +{-# INLINE accum #-} + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- Tuple zip ----------------------------------------------------------- + +-- | Zip two arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip :: Shape f + => BArray f a + -> BArray f b + -> BArray f (a,b) +zip = G.zip + +-- | Zip three arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip3 :: Shape f + => BArray f a + -> BArray f b + -> BArray f c + -> BArray f (a,b,c) +zip3 = G.zip3 + +-- Zip with function --------------------------------------------------- + +-- | Zip two arrays using the given function. If the array's don't have +-- the same shape, the new array with be the intersection of the two +-- shapes. +zipWith :: Shape f + => (a -> b -> c) + -> BArray f a + -> BArray f b + -> BArray f c +zipWith = G.zipWith +{-# INLINE zipWith #-} + +-- | Zip three arrays using the given function. If the array's don't +-- have the same shape, the new array with be the intersection of the +-- two shapes. +zipWith3 :: Shape f + => (a -> b -> c -> d) + -> BArray f a + -> BArray f b + -> BArray f c + -> BArray f d +zipWith3 = G.zipWith3 +{-# INLINE zipWith3 #-} + +-- Indexed zipping ----------------------------------------------------- + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith :: Shape f + => (f Int -> a -> b -> c) + -> BArray f a + -> BArray f b + -> BArray f c +izipWith = G.izipWith +{-# INLINE izipWith #-} + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith3 :: Shape f + => (f Int -> a -> b -> c -> d) + -> BArray f a + -> BArray f b + -> BArray f c + -> BArray f d +izipWith3 = G.izipWith3 +{-# INLINE izipWith3 #-} + +------------------------------------------------------------------------ +-- Slices +------------------------------------------------------------------------ + +-- $setup +-- >>> import Debug.SimpleReflect +-- >>> import qualified Data.Vector as V +-- >>> let m = fromListInto_ (V2 3 4) [a,b,c,d,e,f,g,h,i,j,k,l] :: BArray V2 Expr + +-- | Indexed traversal over the rows of a matrix. Each row is an +-- efficient 'Data.Vector.Generic.slice' of the original vector. +-- +-- >>> traverseOf_ rows print m +-- [a,b,c,d] +-- [e,f,g,h] +-- [i,j,k,l] +rows :: IndexedTraversal Int (BArray V2 a) (BArray V2 b) (Vector a) (Vector b) +rows = G.rows +{-# INLINE rows #-} + +-- | Affine traversal over a single row in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 . each *~ 2 +-- [a,b,c,d] +-- [e * 2,f * 2,g * 2,h * 2] +-- [i,j,k,l] +-- +-- The row vector should remain the same size to satisfy traversal +-- laws but give reasonable behaviour if the size differs: +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0,1] +-- [a,b,c,d] +-- [0,1,g,h] +-- [i,j,k,l] +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0..100] +-- [a,b,c,d] +-- [0,1,2,3] +-- [i,j,k,l] +ixRow :: Int -> IndexedTraversal' Int (BArray V2 a) (Vector a) +ixRow = G.ixRow +{-# INLINE ixRow #-} + +-- | Indexed traversal over the columns of a matrix. Unlike 'rows', each +-- column is a new separate vector. +-- +-- >>> traverseOf_ columns print m +-- [a,e,i] +-- [b,f,j] +-- [c,g,k] +-- [d,h,l] +-- +-- >>> traverseOf_ rows print $ m & columns . indices odd . each .~ 0 +-- [a,0,c,0] +-- [e,0,g,0] +-- [i,0,k,0] +-- +-- The vectors should be the same size to be a valid traversal. If the +-- vectors are different sizes, the number of rows in the new array +-- will be the length of the smallest vector. +columns :: IndexedTraversal Int (BArray V2 a) (BArray V2 b) (Vector a) (Vector b) +columns = G.columns +{-# INLINE columns #-} + +-- | Affine traversal over a single column in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixColumn 2 . each +~ 1 +-- [a,b,c + 1,d] +-- [e,f,g + 1,h] +-- [i,j,k + 1,l] +ixColumn :: Int -> IndexedTraversal' Int (BArray V2 a) (Vector a) +ixColumn = G.ixColumn +{-# INLINE ixColumn #-} + +-- | Traversal over a single plane of a 3D array given a lens onto that +-- plane (like '_xy', '_yz', '_zx'). +ixPlane :: ALens' (V3 Int) (V2 Int) + -> Int + -> IndexedTraversal' Int (BArray V3 a) (BArray V2 a) +ixPlane = G.ixPlane +{-# INLINE ixPlane #-} + +-- | Traversal over all planes of 3D array given a lens onto that plane +-- (like '_xy', '_yz', '_zx'). +planes :: ALens' (V3 Int) (V2 Int) + -> IndexedTraversal Int (BArray V3 a) (BArray V3 b) (BArray V2 a) (BArray V2 b) +planes = G.planes +{-# INLINE planes #-} + +-- | Flatten a plane by reducing a vector in the third dimension to a +-- single value. +flattenPlane :: ALens' (V3 Int) (V2 Int) + -> (Vector a -> b) + -> BArray V3 a + -> BArray V2 b +flattenPlane = G.flattenPlane +{-# INLINE flattenPlane #-} + +-- Ordinals ------------------------------------------------------------ + +-- | This 'Traversal' should not have any duplicates in the list of +-- indices. +unsafeOrdinals :: Shape f => [f Int] -> IndexedTraversal' (f Int) (BArray f a) a +unsafeOrdinals = G.unsafeOrdinals +{-# INLINE [0] unsafeOrdinals #-} + +-- Mutable ------------------------------------------------------------- + +-- | O(n) Yield a mutable copy of the immutable vector. +freeze :: PrimMonad m => BMArray f (PrimState m) a -> m (BArray f a) +freeze = G.freeze +{-# INLINE freeze #-} + +-- | O(n) Yield an immutable copy of the mutable array. +thaw :: PrimMonad m => BArray f a -> m (BMArray f (PrimState m) a) +thaw = G.thaw +{-# INLINE thaw #-} + +-- | O(1) Unsafe convert a mutable array to an immutable one without +-- copying. The mutable array may not be used after this operation. +unsafeFreeze :: PrimMonad m => BMArray f (PrimState m) a -> m (BArray f a) +unsafeFreeze = G.unsafeFreeze +{-# INLINE unsafeFreeze #-} + +-- | O(1) Unsafely convert an immutable array to a mutable one without +-- copying. The immutable array may not be used after this operation. +unsafeThaw :: PrimMonad m => BArray f a -> m (BMArray f (PrimState m) a) +unsafeThaw = G.unsafeThaw +{-# INLINE unsafeThaw #-} + +------------------------------------------------------------------------ +-- Delayed +------------------------------------------------------------------------ + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in parallel. +delayed :: (Shape f, Shape k) + => Iso (BArray f a) (BArray k b) (G.Delayed f a) (G.Delayed k b) +delayed = G.delayed +{-# INLINE delayed #-} + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in sequence. +seqDelayed :: (Shape f, Shape k) + => Iso (BArray f a) (BArray k b) (G.Delayed f a) (G.Delayed k b) +seqDelayed = G.seqDelayed +{-# INLINE seqDelayed #-} + +-- | Turn a material array into a delayed one with the same shape. +delay :: Shape f => BArray f a -> G.Delayed f a +delay = G.delay +{-# INLINE delay #-} + +-- | Parallel manifestation of a delayed array into a material one. +manifest :: Shape f => G.Delayed f a -> BArray f a +manifest = G.manifest +{-# INLINE manifest #-} + +-- | Sequential manifestation of a delayed array. +seqManifest :: Shape f => G.Delayed f a -> BArray f a +seqManifest = G.seqManifest +{-# INLINE seqManifest #-} + +-- | 'manifest' an array to a 'BArray' and delay again. +affirm :: Shape f => G.Delayed f a -> G.Delayed f a +affirm = delay . manifest +{-# INLINE affirm #-} + +-- | 'seqManifest' an array to a 'BArray' and delay again. +seqAffirm :: Shape f => G.Delayed f a -> G.Delayed f a +seqAffirm = delay . seqManifest +{-# INLINE seqAffirm #-} + diff --git a/src/Data/Dense/Generic.hs b/src/Data/Dense/Generic.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Generic.hs @@ -0,0 +1,1017 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Generic +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- This module provides generic functions over multidimensional arrays. +----------------------------------------------------------------------------- +module Data.Dense.Generic + ( + -- * Array types + Array + , Shape (..) + , BArray + , UArray + , SArray + , PArray + + -- * Layout of an array + , HasLayout (..) + , Layout + + -- ** Extracting size + , extent + , size + + -- ** Folds over indexes + , indexes + , indexesFrom + , indexesBetween + + -- * Underlying vector + , vector + + -- ** Traversals + , values + , values' + , valuesBetween + + -- * Construction + + -- ** Flat arrays + , flat + , fromList + + -- ** From lists + , fromListInto + , fromListInto_ + + -- ** From vectors + , fromVectorInto + , fromVectorInto_ + + -- ** Initialisation + , replicate + , generate + , linearGenerate + + -- ** Monadic initialisation + , create + , replicateM + , generateM + , linearGenerateM + + -- * Functions on arrays + + -- ** Empty arrays + , empty + , null + + -- ** Indexing + + , (!) + , (!?) + , unsafeIndex + , linearIndex + , unsafeLinearIndex + + -- *** Monadic indexing + , indexM + , unsafeIndexM + , linearIndexM + , unsafeLinearIndexM + + -- ** Modifying arrays + + -- ** Bulk updates + , (//) + + -- ** Accumulations + , accum + + -- ** Mapping + , map + , imap + + -- * Zipping + -- ** Tuples + , Data.Dense.Generic.zip + , Data.Dense.Generic.zip3 + + -- ** Zip with function + , zipWith + , zipWith3 + , izipWith + , izipWith3 + + -- ** Slices + + -- *** Matrix + , ixRow + , rows + , ixColumn + , columns + + -- *** 3D + , ixPlane + , planes + , flattenPlane + + -- *** Ordinals + , unsafeOrdinals + + -- * Mutable + , MArray + , M.BMArray + , M.UMArray + , M.SMArray + , M.PMArray + + , thaw + , freeze + , unsafeThaw + , unsafeFreeze + + -- * Delayed + + , Delayed + + -- ** Generating delayed + + , delayed + , seqDelayed + , delay + , manifest + , seqManifest + , genDelayed + , indexDelayed + , affirm + , seqAffirm + + -- * Focused + + , Focused + + -- ** Generating focused + + , focusOn + , unfocus + , unfocused + , extendFocus + + -- ** Focus location + , locale + , shiftFocus + + -- ** Boundary + , Boundary (..) + , peekB + , peeksB + , peekRelativeB + + -- * Fusion + -- ** Streams + , streamGenerate + , streamGenerateM + , streamIndexes + + -- ** Bundles + , bundleGenerate + , bundleGenerateM + , bundleIndexes + + ) where + + +#if __GLASGOW_HASKELL__ <= 708 +import Control.Applicative (Applicative, pure, (<*>)) +import Data.Foldable (Foldable) +#endif + +import Control.Comonad +import Control.Comonad.Store +import Control.Lens hiding (imap) +import Control.Monad (liftM) +import Control.Monad.Primitive +import Control.Monad.ST +import qualified Data.Foldable as F +import Data.Functor.Classes +import qualified Data.List as L +import Data.Maybe (fromMaybe) +import Data.Typeable +import qualified Data.Vector as B +import Data.Vector.Fusion.Bundle (MBundle) +import qualified Data.Vector.Fusion.Bundle as Bundle +import qualified Data.Vector.Fusion.Bundle.Monadic as MBundle +import Data.Vector.Fusion.Bundle.Size +import Data.Vector.Fusion.Stream.Monadic (Step (..), Stream (..)) +import qualified Data.Vector.Fusion.Stream.Monadic as Stream +import Data.Vector.Generic (Vector) +import qualified Data.Vector.Generic as G +import qualified Data.Vector.Generic.Mutable as GM +import qualified Data.Vector.Primitive as P +import qualified Data.Vector.Storable as S +import qualified Data.Vector.Unboxed as U +import Linear hiding (vector) + +import Data.Dense.Base +import Data.Dense.Index +import Data.Dense.Mutable (MArray (..)) +import qualified Data.Dense.Mutable as M + +import Prelude hiding (map, null, replicate, + zipWith, zipWith3) + +-- Aliases ------------------------------------------------------------- + +-- | 'Boxed' array. +type BArray = Array B.Vector + +-- | 'Data.Vector.Unboxed.Unbox'ed array. +type UArray = Array U.Vector + +-- | 'Foreign.Storable.Storeable' array. +type SArray = Array S.Vector + +-- | 'Data.Primitive.Types.Prim' array. +type PArray = Array P.Vector + +-- Lenses -------------------------------------------------------------- + +-- | Same as 'values' but restrictive in the vector type. +values' :: (Shape f, Vector v a, Vector v b) + => IndexedTraversal (f Int) (Array v f a) (Array v f b) a b +values' = values +{-# INLINE values' #-} + +-- | Traverse over the 'values' between two indexes. +valuesBetween :: (Shape f, Vector v a) => f Int -> f Int -> IndexedTraversal' (f Int) (Array v f a) a +valuesBetween a b = unsafeOrdinals (toListOf (shapeIndexesFrom a) b) +{-# INLINE valuesBetween #-} + +-- | 1D arrays are just vectors. You are free to change the length of +-- the vector when going 'over' this 'Iso' (unlike 'linear'). +-- +-- Note that 'V1' arrays are an instance of 'Vector' so you can use +-- any of the functions in "Data.Vector.Generic" on them without +-- needing to convert. +flat :: Vector w b => Iso (Array v V1 a) (Array w V1 b) (v a) (w b) +flat = iso (\(Array _ v) -> v) (\v -> Array (V1 $ G.length v) v) +{-# INLINE flat #-} + +-- Constructing vectors ------------------------------------------------ + +-- | Contruct a flat array from a list. (This is just 'G.fromList' from +-- 'Data.Vector.Generic'.) +fromList :: Vector v a => [a] -> Array v V1 a +fromList = G.fromList +{-# INLINE fromList #-} + +-- | O(n) Convert the first @n@ elements of a list to an Array with the +-- given shape. Returns 'Nothing' if there are not enough elements in +-- the list. +fromListInto :: (Shape f, Vector v a) => Layout f -> [a] -> Maybe (Array v f a) +fromListInto l as + | G.length v == n = Just $ Array l v + | otherwise = Nothing + where v = G.fromListN n as + n = shapeSize l +{-# INLINE fromListInto #-} + +-- | O(n) Convert the first @n@ elements of a list to an Array with the +-- given shape. Throw an error if the list is not long enough. +fromListInto_ :: (Shape f, Vector v a) => Layout f -> [a] -> Array v f a +fromListInto_ l as = fromMaybe err $ fromListInto l as + where + err = error $ "fromListInto_: shape " ++ showShape l ++ " is too large for list" +{-# INLINE fromListInto_ #-} + +-- | Create an array from a 'vector' and a 'layout'. Return 'Nothing' if +-- the vector is not the right shape. +fromVectorInto :: (Shape f, Vector v a) => Layout f -> v a -> Maybe (Array v f a) +fromVectorInto l v + | shapeSize l == G.length v = Just $! Array l v + | otherwise = Nothing +{-# INLINE fromVectorInto #-} + +-- | Create an array from a 'vector' and a 'layout'. Throws an error if +-- the vector is not the right shape. +fromVectorInto_ :: (Shape f, Vector v a) => Layout f -> v a -> Array v f a +fromVectorInto_ l as = fromMaybe err $ fromVectorInto l as + where + err = error $ "fromVectorInto_: shape " ++ showShape l ++ " is too large for the vector" +{-# INLINE fromVectorInto_ #-} + +-- | The empty 'Array' with a 'zero' shape. +empty :: (Vector v a, Additive f) => Array v f a +empty = Array zero G.empty +{-# INLINE empty #-} + +-- | Test is if the array is 'empty'. +null :: Foldable f => Array v f a -> Bool +null (Array l _) = F.all (==0) l +{-# INLINE null #-} + +-- Indexing ------------------------------------------------------------ + +-- | Index an element of an array. Throws 'IndexOutOfBounds' if the +-- index is out of bounds. +(!) :: (Shape f, Vector v a) => Array v f a -> f Int -> a +(!) (Array l v) i = boundsCheck l i $ G.unsafeIndex v (shapeToIndex l i) +{-# INLINE (!) #-} + +-- | Safe index of an element. +(!?) :: (Shape f, Vector v a) => Array v f a -> f Int -> Maybe a +Array l v !? i + | shapeInRange l i = Just $! G.unsafeIndex v (shapeToIndex l i) + | otherwise = Nothing +{-# INLINE (!?) #-} + +-- | Index an element of an array without bounds checking. +unsafeIndex :: (Shape f, Vector v a) => Array v f a -> f Int -> a +unsafeIndex (Array l v) i = G.unsafeIndex v (shapeToIndex l i) +{-# INLINE unsafeIndex #-} + +-- | Index an element of an array while ignoring its shape. +linearIndex :: Vector v a => Array v f a -> Int -> a +linearIndex (Array _ v) i = v G.! i +{-# INLINE linearIndex #-} + +-- | Index an element of an array while ignoring its shape, without +-- bounds checking. +unsafeLinearIndex :: Vector v a => Array v f a -> Int -> a +unsafeLinearIndex (Array _ v) i = G.unsafeIndex v i +{-# INLINE unsafeLinearIndex #-} + +-- Monadic indexing ---------------------------------------------------- + +-- | /O(1)/ Indexing in a monad. +-- +-- The monad allows operations to be strict in the vector when necessary. +-- Suppose vector copying is implemented like this: +-- +-- > copy mv v = ... write mv i (v ! i) ... +-- +-- For lazy vectors, @v ! i@ would not be evaluated which means that +-- @mv@ would unnecessarily retain a reference to @v@ in each element +-- written. +-- +-- With 'indexM', copying can be implemented like this instead: +-- +-- > copy mv v = ... do +-- > x <- indexM v i +-- > write mv i x +-- +-- Here, no references to @v@ are retained because indexing (but /not/ +-- the elements) is evaluated eagerly. +-- +-- Throws an error if the index is out of range. +indexM :: (Shape f, Vector v a, Monad m) => Array v f a -> f Int -> m a +indexM (Array l v) i = boundsCheck l i $ G.unsafeIndexM v (shapeToIndex l i) +{-# INLINE indexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeIndexM :: (Shape f, Vector v a, Monad m) => Array v f a -> f Int -> m a +unsafeIndexM (Array l v) i = G.unsafeIndexM v (shapeToIndex l i) +{-# INLINE unsafeIndexM #-} + +-- | /O(1)/ Indexing in a monad. Throws an error if the index is out of +-- range. +linearIndexM :: (Shape f, Vector v a, Monad m) => Array v f a -> Int -> m a +linearIndexM (Array l v) i = boundsCheck l (shapeFromIndex l i) $ G.unsafeIndexM v i +{-# INLINE linearIndexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeLinearIndexM :: (Vector v a, Monad m) => Array v f a -> Int -> m a +unsafeLinearIndexM (Array _ v) = G.unsafeIndexM v +{-# INLINE unsafeLinearIndexM #-} + +-- Initialisation ------------------------------------------------------ + +-- | Execute the monadic action and freeze the resulting array. +create :: Vector v a => (forall s. ST s (MArray (G.Mutable v) f s a)) -> Array v f a +create m = m `seq` runST (m >>= unsafeFreeze) +{-# INLINE create #-} + +-- | O(n) Array of the given shape with the same value in each position. +replicate :: (Shape f, Vector v a) => f Int -> a -> Array v f a +replicate l a + | n > 0 = Array l $ G.replicate n a + | otherwise = empty + where n = shapeSize l +{-# INLINE replicate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +linearGenerate :: (Shape f, Vector v a) => Layout f -> (Int -> a) -> Array v f a +linearGenerate l f + | n > 0 = Array l $ G.generate n f + | otherwise = empty + where n = shapeSize l +{-# INLINE linearGenerate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +generate :: (Shape f, Vector v a) => Layout f -> (f Int -> a) -> Array v f a +generate l f = Array l $ G.unstream (bundleGenerate l f) +{-# INLINE generate #-} + +-- Monadic initialisation ---------------------------------------------- + +-- | O(n) Construct an array of the given shape by filling each position +-- with the monadic value. +replicateM :: (Monad m, Shape f, Vector v a) => Layout f -> m a -> m (Array v f a) +replicateM l a + | n > 0 = Array l `liftM` G.replicateM n a + | otherwise = return empty + where n = shapeSize l +{-# INLINE replicateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +generateM :: (Monad m, Shape f, Vector v a) => Layout f -> (f Int -> m a) -> m (Array v f a) +generateM l f = Array l `liftM` unstreamM (bundleGenerateM l f) +{-# INLINE generateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +linearGenerateM :: (Monad m, Shape f, Vector v a) => Layout f -> (Int -> m a) -> m (Array v f a) +linearGenerateM l f + | n > 0 = Array l `liftM` G.generateM n f + | otherwise = return empty + where n = shapeSize l +{-# INLINE linearGenerateM #-} + +-- Modifying ----------------------------------------------------------- + +-- | /O(n)/ Map a function over an array +map :: (Vector v a, Vector v b) => (a -> b) -> Array v f a -> Array v f b +map f (Array l a) = Array l (G.map f a) +{-# INLINE map #-} + +-- | /O(n)/ Apply a function to every element of a vector and its index +imap :: (Shape f, Vector v a, Vector v b) => (f Int -> a -> b) -> Array v f a -> Array v f b +imap f (Array l v) = + Array l $ (G.unstream . Bundle.inplace (Stream.zipWith f (streamIndexes l)) id . G.stream) v +{-# INLINE imap #-} + +-- Bulk updates -------------------------------------------------------- + +-- | For each pair (i,a) from the list, replace the array element at +-- position i by a. +(//) :: (G.Vector v a, Shape f) => Array v f a -> [(f Int, a)] -> Array v f a +Array l v // xs = Array l $ v G.// over (each . _1) (shapeToIndex l) xs + +-- Accumilation -------------------------------------------------------- + +-- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the array element +-- @a@ at position @i@ by @f a b@. +-- +accum :: (Shape f, Vector v a) + => (a -> b -> a) -- ^ accumulating function @f@ + -> Array v f a -- ^ initial array + -> [(f Int, b)] -- ^ list of index/value pairs (of length @n@) + -> Array v f a +accum f (Array l v) us = Array l $ G.accum f v (over (mapped . _1) (shapeToIndex l) us) +{-# INLINE accum #-} + +------------------------------------------------------------------------ +-- Streams +------------------------------------------------------------------------ + +-- Copied from Data.Vector.Generic because it isn't exported from there. + +unstreamM :: (Monad m, Vector v a) => Bundle.MBundle m u a -> m (v a) +{-# INLINE [1] unstreamM #-} +unstreamM s = do + xs <- MBundle.toList s + return $ G.unstream $ Bundle.unsafeFromList (MBundle.size s) xs + +unstreamPrimM :: (PrimMonad m, Vector v a) => Bundle.MBundle m u a -> m (v a) +{-# INLINE [1] unstreamPrimM #-} +unstreamPrimM s = GM.munstream s >>= G.unsafeFreeze + +-- FIXME: the next two functions are only necessary for the specialisations +unstreamPrimM_IO :: Vector v a => Bundle.MBundle IO u a -> IO (v a) +{-# INLINE unstreamPrimM_IO #-} +unstreamPrimM_IO = unstreamPrimM + +unstreamPrimM_ST :: Vector v a => Bundle.MBundle (ST s) u a -> ST s (v a) +{-# INLINE unstreamPrimM_ST #-} +unstreamPrimM_ST = unstreamPrimM + +{-# RULES + +"unstreamM[IO]" unstreamM = unstreamPrimM_IO +"unstreamM[ST]" unstreamM = unstreamPrimM_ST #-} + +-- | Generate a stream from a 'Layout''s indices. +streamGenerate :: (Monad m, Shape f) => Layout f -> (f Int -> a) -> Stream m a +streamGenerate l f = streamGenerateM l (return . f) +{-# INLINE streamGenerate #-} + +-- | Generate a stream from a 'Layout''s indices. +streamGenerateM :: (Monad m, Shape f) => Layout f -> (f Int -> m a) -> Stream m a +streamGenerateM l f = l `seq` Stream step (if eq1 l zero then Nothing else Just zero) + where + {-# INLINE [0] step #-} + step (Just i) = do + x <- f i + return $ Yield x (shapeStep l i) + step Nothing = return Done +{-# INLINE [1] streamGenerateM #-} + +-- | Stream a sub-layout of an 'Array'. The layout should be shapeInRange of +-- the array's layout, this is not checked. +unsafeStreamSub :: (Monad m, Shape f, G.Vector v a) => Layout f -> Array v f a -> Stream m a +unsafeStreamSub l2 (Array l1 v) = streamGenerateM l2 $ \x -> G.basicUnsafeIndexM v (shapeToIndex l1 x) +{-# INLINE unsafeStreamSub #-} + +-- | Stream a sub-layout of an 'Array'. +streamSub :: (Monad m, Shape f, G.Vector v a) => Layout f -> Array v f a -> Stream m a +streamSub l2 arr@(Array l1 _) = unsafeStreamSub (shapeIntersect l1 l2) arr +{-# INLINE streamSub #-} + +-- | Make a stream of the indexes of a 'Layout'. +streamIndexes :: (Monad m, Shape f) => Layout f -> Stream m (f Int) +streamIndexes l = Stream step (if eq1 l zero then Nothing else Just zero) + where + {-# INLINE [0] step #-} + step (Just i) = return $ Yield i (shapeStep l i) + step Nothing = return Done +{-# INLINE [1] streamIndexes #-} + +------------------------------------------------------------------------ +-- Bundles +------------------------------------------------------------------------ + +-- | Generate a bundle from 'Layout' indices. +bundleGenerate :: (Monad m, Shape f) => Layout f -> (f Int -> a) -> MBundle m v a +bundleGenerate l f = bundleGenerateM l (return . f) +{-# INLINE bundleGenerate #-} + +-- | Generate a bundle from 'Layout' indices. +bundleGenerateM :: (Monad m, Shape f) => Layout f -> (f Int -> m a) -> MBundle m v a +bundleGenerateM l f = MBundle.fromStream (streamGenerateM l f) (Exact (shapeSize l)) +{-# INLINE [1] bundleGenerateM #-} + +-- | Generate a bundle of indexes for the given 'Layout'. +bundleIndexes :: (Monad m, Shape f) => Layout f -> MBundle m v (f Int) +bundleIndexes l = MBundle.fromStream (streamIndexes l) (Exact (shapeSize l)) +{-# INLINE [1] bundleIndexes #-} + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- Tuple zip ----------------------------------------------------------- + +-- | Zip two arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip :: (Shape f, Vector v a, Vector v b, Vector v (a,b)) + => Array v f a + -> Array v f b + -> Array v f (a,b) +zip = zipWith (,) + +-- | Zip three arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v (a,b,c)) + => Array v f a + -> Array v f b + -> Array v f c + -> Array v f (a,b,c) +zip3 = zipWith3 (,,) + +-- Zip with function --------------------------------------------------- + +-- | Zip two arrays using the given function. If the array's don't have +-- the same shape, the new array with be the intersection of the two +-- shapes. +zipWith :: (Shape f, Vector v a, Vector v b, Vector v c) + => (a -> b -> c) + -> Array v f a + -> Array v f b + -> Array v f c +zipWith f a1@(Array l1 v1) a2@(Array l2 v2) + | eq1 l1 l1 = Array l1 $ G.zipWith f v1 v2 + | otherwise = Array l' $ G.unstream $ + MBundle.fromStream (Stream.zipWith f (streamSub l' a1) (streamSub l' a2)) (Exact (shapeSize l')) + where l' = shapeIntersect l1 l2 +{-# INLINE zipWith #-} + +-- | Zip three arrays using the given function. If the array's don't +-- have the same shape, the new array with be the intersection of the +-- two shapes. +zipWith3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v d) + => (a -> b -> c -> d) + -> Array v f a + -> Array v f b + -> Array v f c + -> Array v f d +zipWith3 f a1@(Array l1 v1) a2@(Array l2 v2) a3@(Array l3 v3) + | eq1 l1 l2 && + eq1 l2 l3 = Array l1 $ G.zipWith3 f v1 v2 v3 + | otherwise = Array l' $ G.unstream $ + MBundle.fromStream (Stream.zipWith3 f (streamSub l' a1) (streamSub l' a2) (streamSub l' a3)) (Exact (shapeSize l')) + where l' = shapeIntersect (shapeIntersect l1 l2) l3 +{-# INLINE zipWith3 #-} + +-- Indexed zipping ----------------------------------------------------- + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith :: (Shape f, Vector v a, Vector v b, Vector v c) + => (f Int -> a -> b -> c) + -> Array v f a + -> Array v f b + -> Array v f c +izipWith f a1@(Array l1 v1) a2@(Array l2 v2) + | eq1 l1 l2 = Array l1 $ G.unstream $ Bundle.zipWith3 f (bundleIndexes l1) (G.stream v1) (G.stream v2) + | otherwise = Array l' $ G.unstream $ + MBundle.fromStream (Stream.zipWith3 f (streamIndexes l') (streamSub l' a1) (streamSub l' a2)) (Exact (shapeSize l')) + where l' = shapeIntersect l1 l2 +{-# INLINE izipWith #-} + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith3 :: (Shape f, Vector v a, Vector v b, Vector v c, Vector v d) + => (f Int -> a -> b -> c -> d) + -> Array v f a + -> Array v f b + -> Array v f c + -> Array v f d +izipWith3 f a1@(Array l1 v1) a2@(Array l2 v2) a3@(Array l3 v3) + | eq1 l1 l2 = Array l1 $ G.unstream $ Bundle.zipWith4 f (bundleIndexes l1) (G.stream v1) (G.stream v2) (G.stream v3) + | otherwise = + Array l' $ G.unstream $ MBundle.fromStream + (Stream.zipWith4 f (streamIndexes l') (streamSub l' a1) (streamSub l' a2) (streamSub l' a3)) (Exact (shapeSize l')) + where l' = shapeIntersect (shapeIntersect l1 l2) l3 +{-# INLINE izipWith3 #-} + +------------------------------------------------------------------------ +-- Slices +------------------------------------------------------------------------ + +-- $setup +-- >>> import Debug.SimpleReflect +-- >>> let m = fromListInto_ (V2 3 4) [a,b,c,d,e,f,g,h,i,j,k,l] :: BArray V2 Expr + +-- | Indexed traversal over the rows of a matrix. Each row is an +-- efficient 'Data.Vector.Generic.slice' of the original vector. +-- +-- >>> traverseOf_ rows print m +-- [a,b,c,d] +-- [e,f,g,h] +-- [i,j,k,l] +rows :: (Vector v a, Vector w b) + => IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b) +rows f (Array l@(V2 x y) v) = Array l . G.concat <$> go 0 0 where + go i a | i >= x = pure [] + | otherwise = (:) <$> indexed f i (G.slice a y v) <*> go (i+1) (a+y) +{-# INLINE rows #-} + +-- | Affine traversal over a single row in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 . each *~ 2 +-- [a,b,c,d] +-- [e * 2,f * 2,g * 2,h * 2] +-- [i,j,k,l] +-- +-- The row vector should remain the same size to satisfy traversal +-- laws but give reasonable behaviour if the size differs: +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ B.fromList [0,1] +-- [a,b,c,d] +-- [0,1,g,h] +-- [i,j,k,l] +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ B.fromList [0..100] +-- [a,b,c,d] +-- [0,1,2,3] +-- [i,j,k,l] +ixRow :: Vector v a => Int -> IndexedTraversal' Int (Array v V2 a) (v a) +ixRow i f m@(Array (l@(V2 x y)) v) + | y >= 0 && i < x = Array l . G.unsafeUpd v . L.zip [a..] . G.toList . G.take y <$> indexed f i (G.slice a y v) + | otherwise = pure m + where a = i * y +{-# INLINE ixRow #-} + +-- | Indexed traversal over the columns of a matrix. Unlike 'rows', each +-- column is a new separate vector. +-- +-- >>> traverseOf_ columns print m +-- [a,e,i] +-- [b,f,j] +-- [c,g,k] +-- [d,h,l] +-- +-- >>> traverseOf_ rows print $ m & columns . indices odd . each .~ 0 +-- [a,0,c,0] +-- [e,0,g,0] +-- [i,0,k,0] +-- +-- The vectors should be the same size to be a valid traversal. If the +-- vectors are different sizes, the number of rows in the new array +-- will be the length of the smallest vector. +columns :: (Vector v a, Vector w b) + => IndexedTraversal Int (Array v V2 a) (Array w V2 b) (v a) (w b) +columns f m@(Array l@(V2 _ y) _) = transposeConcat l <$> go 0 where + go j | j >= y = pure [] + | otherwise = (:) <$> indexed f j (getColumn m j) <*> go (j+1) +{-# INLINE columns #-} + +-- | Affine traversal over a single column in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixColumn 2 . each +~ 1 +-- [a,b,c + 1,d] +-- [e,f,g + 1,h] +-- [i,j,k + 1,l] +ixColumn :: Vector v a => Int -> IndexedTraversal' Int (Array v V2 a) (v a) +ixColumn j f m@(Array (l@(V2 _ y)) v) + | j >= 0 && j < y = Array l . G.unsafeUpd v . L.zip js . G.toList . G.take y <$> indexed f j (getColumn m j) + | otherwise = pure m + where js = [j, j + y .. ] +{-# INLINE ixColumn #-} + +getColumn :: Vector v a => Array v V2 a -> Int -> v a +getColumn (Array (V2 x y) v) j = G.generate x $ \i -> G.unsafeIndex v (i * y + j) +{-# INLINE getColumn #-} + +transposeConcat :: Vector v a => V2 Int -> [v a] -> Array v V2 a +transposeConcat (V2 _ y) vs = Array (V2 x' y) $ G.create $ do + mv <- GM.new (x'*y) + iforM_ vs $ \j v -> + F.for_ [0..x'-1] $ \i -> + GM.write mv (i*y + j) (v G.! i) + return mv + where x' = minimum $ fmap G.length vs +{-# INLINE transposeConcat #-} + +-- | Traversal over a single plane of a 3D array given a lens onto that +-- plane (like '_xy', '_yz', '_zx'). +ixPlane :: Vector v a + => ALens' (V3 Int) (V2 Int) + -> Int + -> IndexedTraversal' Int (Array v V3 a) (Array v V2 a) +ixPlane l32 i f a@(Array l v) + | i < 0 || i >= k = pure a + | otherwise = Array l . (v G.//) . L.zip is . toListOf values + <$> indexed f i (getPlane l32 i a) + where + is = toListOf (cloneLens l32 . shapeIndexes . to (\x -> shapeToIndex l $ pure i & l32 #~ x)) l + k = F.sum $ l & l32 #~ 0 + +-- | Traversal over all planes of 3D array given a lens onto that plane +-- (like '_xy', '_yz', '_zx'). +planes :: (Vector v a, Vector w b) + => ALens' (V3 Int) (V2 Int) + -> IndexedTraversal Int (Array v V3 a) (Array w V3 b) (Array v V2 a) (Array w V2 b) +planes l32 f a@(Array l _) = concatPlanes l l32 <$> go 0 where + go i | i >= k = pure [] + | otherwise = (:) <$> indexed f i (getPlane l32 i a) <*> go (i+1) + k = F.sum $ l & l32 #~ 0 +{-# INLINE planes #-} + +concatPlanes :: Vector v a => V3 Int -> ALens' (V3 Int) (V2 Int) -> [Array v V2 a] -> Array v V3 a +concatPlanes l l32 as = create $ do + arr <- M.new l + iforM_ as $ \i m -> + iforMOf_ values m $ \x a -> do + let w = pure i & l32 #~ x + M.write arr w a + return arr + +getPlane :: Vector v a => ALens' (V3 Int) (V2 Int) -> Int -> Array v V3 a -> Array v V2 a +getPlane l32 i a = generate (a ^# layout . l32) $ \x -> a ! (pure i & l32 #~ x) + +-- | Flatten a plane by reducing a vector in the third dimension to a +-- single value. +flattenPlane :: (Vector v a, Vector w b) + => ALens' (V3 Int) (V2 Int) + -> (v a -> b) + -> Array v V3 a + -> Array w V2 b +flattenPlane l32 f a@(Array l _) = generate l' $ \x -> f (getVector x) + where + getVector x = G.generate n $ \i -> a ! (pure i & l32 #~ x) + n = F.sum $ l & l32 #~ 0 + l' = l ^# l32 +{-# INLINE flattenPlane #-} + +-- Ordinals ------------------------------------------------------------ + +-- | This 'Traversal' should not have any duplicates in the list of +-- indices. +unsafeOrdinals :: (Vector v a, Shape f) => [f Int] -> IndexedTraversal' (f Int) (Array v f a) a +unsafeOrdinals is f (Array l v) = Array l . (v G.//) <$> traverse g is + where g x = let i = shapeToIndex l x in (,) i <$> indexed f x (G.unsafeIndex v i) +{-# INLINE [0] unsafeOrdinals #-} + +setOrdinals :: (Indexable (f Int) p, Vector v a, Shape f) => [f Int] -> p a a -> Array v f a -> Array v f a +setOrdinals is f (Array l v) = Array l $ G.unsafeUpd v (fmap g is) + where g x = let i = shapeToIndex l x in (,) i $ indexed f x (G.unsafeIndex v i) +{-# INLINE setOrdinals #-} + +{-# RULES +"unsafeOrdinals/setOrdinals" forall (is :: [f Int]). + unsafeOrdinals is = sets (setOrdinals is) + :: Vector v a => ASetter' (Array v f a) a; +"unsafeOrdinalts/isetOrdintals" forall (is :: [f Int]). + unsafeOrdinals is = sets (setOrdinals is) + :: Vector v a => AnIndexedSetter' (f Int) (Array v f a) a + #-} + +-- Mutable ------------------------------------------------------------- + +-- | O(n) Yield a mutable copy of the immutable vector. +freeze :: (PrimMonad m, Vector v a) + => MArray (G.Mutable v) f (PrimState m) a -> m (Array v f a) +freeze (MArray l mv) = Array l `liftM` G.freeze mv +{-# INLINE freeze #-} + +-- | O(n) Yield an immutable copy of the mutable array. +thaw :: (PrimMonad m, Vector v a) + => Array v f a -> m (MArray (G.Mutable v) f (PrimState m) a) +thaw (Array l v) = MArray l `liftM` G.thaw v +{-# INLINE thaw #-} + +------------------------------------------------------------------------ +-- Delayed +------------------------------------------------------------------------ + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in parallel. +delayed :: (Vector v a, Vector w b, Shape f, Shape g) + => Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b) +delayed = iso delay manifest +{-# INLINE delayed #-} + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in parallel. +seqDelayed :: (Vector v a, Vector w b, Shape f, Shape g) + => Iso (Array v f a) (Array w g b) (Delayed f a) (Delayed g b) +seqDelayed = iso delay seqManifest +{-# INLINE seqDelayed #-} + +-- | Sequential manifestation of a delayed array. +seqManifest :: (Vector v a, Shape f) => Delayed f a -> Array v f a +seqManifest (Delayed l f) = generate l f +{-# INLINE seqManifest #-} + +-- | 'manifest' an array to a 'UArray' and delay again. See +-- "Data.Dense.Boxed" or "Data.Dense.Storable" to 'affirm' for other +-- types of arrays. +affirm :: (Shape f, U.Unbox a) => Delayed f a -> Delayed f a +affirm = delay . (manifest :: (U.Unbox a, Shape f) => Delayed f a -> UArray f a) +{-# INLINE affirm #-} + +-- | 'seqManifest' an array to a 'UArray' and delay again. See +-- "Data.Dense.Boxed" or "Data.Dense.Storable" to 'affirm' for other +-- types of arrays. +seqAffirm :: (Shape f, U.Unbox a) => Delayed f a -> Delayed f a +seqAffirm = delay . (seqManifest :: (U.Unbox a, Shape f) => Delayed f a -> UArray f a) +{-# INLINE seqAffirm #-} + +------------------------------------------------------------------------ +-- Focused +------------------------------------------------------------------------ + +-- | Focus on a particular element of a delayed array. +focusOn :: f Int -> Delayed f a -> Focused f a +focusOn = Focused -- XXX do range checking +{-# INLINE focusOn #-} + +-- | Discard the focus to retrieve the delayed array. +unfocus :: Focused f a -> Delayed f a +unfocus (Focused _ d) = d +{-# INLINE unfocus #-} + +-- | Indexed lens onto the delayed array, indexed at the focus. +unfocused :: IndexedLens (f Int) (Focused f a) (Focused f b) (Delayed f a) (Delayed f b) +unfocused f (Focused x d) = Focused x <$> indexed f x d +{-# INLINE unfocused #-} + +-- | Modify a 'Delayed' array by extracting a value from a 'Focused' +-- each point. +extendFocus :: Shape f => (Focused f a -> b) -> Delayed f a -> Delayed f b +extendFocus f = unfocus . extend f . focusOn zero +{-# INLINE extendFocus #-} + +-- | Lens onto the position of a 'ComonadStore'. +-- +-- @ +-- 'locale' :: 'Lens'' ('Focused' l a) (l 'Int') +-- @ +locale :: ComonadStore s w => Lens' (w a) s +locale f w = (`seek` w) <$> f (pos w) +{-# INLINE locale #-} + +-- | Focus on a neighbouring element, relative to the current focus. +shiftFocus :: Applicative f => f Int -> Focused f a -> Focused f a +shiftFocus dx (Focused x d@(Delayed l _)) = Focused x' d + where + x' = f <$> l <*> x <*> dx + f k i di + | i' < 0 = k + i' + | i' >= k = i' - k + | otherwise = i' + where i' = i + di +{-# INLINE shiftFocus #-} + +-- Boundary conditions ------------------------------------------------- + +-- | The boundary condition used for indexing relative elements in a +-- 'Focused'. +data Boundary + = Clamp -- ^ clamp coordinates to the extent of the array + | Mirror -- ^ mirror coordinates beyond the array extent + | Wrap -- ^ wrap coordinates around on each dimension + deriving (Show, Read, Typeable) + +-- Peeking ------------------------------------------------------------- + +-- | Index a focused using a 'Boundary' condition. +peekB :: Shape f => Boundary -> f Int -> Focused f a -> a +peekB = \b x -> peeksB b (const x) +{-# INLINE peekB #-} + +-- | Index an element relative to the current focus using a 'Boundary' +-- condition. +peekRelativeB :: Shape f => Boundary -> f Int -> Focused f a -> a +peekRelativeB = \b i -> peeksB b (^+^ i) +{-# INLINE peekRelativeB #-} + +-- | Index an element by applying a function the current position, using +-- a boundary condition. +peeksB :: Shape f => Boundary -> (f Int -> f Int) -> Focused f a -> a +peeksB = \case + Clamp -> clampPeeks + Wrap -> wrapPeeks + Mirror -> mirrorPeeks +{-# INLINE peeksB #-} + +-- After much testing, this seems to be the most reliable method to get +-- stencilSum to inline properly. + +-- Wrap + +wrapPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a +wrapPeeks f (Focused x (Delayed l ixF)) = ixF $! wrapIndex l (f x) +{-# INLINE wrapPeeks #-} + +wrapIndex :: Shape f => Layout f -> f Int -> f Int +wrapIndex !l !x = liftI2 f l x where + f n i + | i < 0 = n + i + | i < n = i + | otherwise = i - n +{-# INLINE wrapIndex #-} + +-- Clamp + +clampPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a +clampPeeks f (Focused x (Delayed l ixF)) = ixF $! clampIndex l (f x) +{-# INLINE clampPeeks #-} + +clampIndex :: Shape f => Layout f -> f Int -> f Int +clampIndex !l !x = liftI2 f l x where + f n i + | i < 0 = 0 + | i >= n = n - 1 + | otherwise = i +{-# INLINE clampIndex #-} + +-- Mirror + +mirrorPeeks :: Shape f => (f Int -> f Int) -> Focused f a -> a +mirrorPeeks f (Focused x (Delayed l ixF)) = ixF $! mirrorIndex l (f x) +{-# INLINE mirrorPeeks #-} + +mirrorIndex :: Shape f => Layout f -> f Int -> f Int +mirrorIndex !l !x = liftI2 f l x where + f n i + | i < 0 = - i + | i < n = i + | otherwise = i - n +{-# INLINE mirrorIndex #-} + diff --git a/src/Data/Dense/Index.hs b/src/Data/Dense/Index.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Index.hs @@ -0,0 +1,370 @@ +{-# LANGUAGE DefaultSignatures #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE Rank2Types #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Mutable +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- This module provides a class for types that can be converted to and +-- from linear indexes. +-- +-- The default instances are defined in row-major order. +----------------------------------------------------------------------------- +module Data.Dense.Index + ( -- * Shape class + Layout + , Shape (..) + , indexIso + , shapeIndexes + , shapeIndexesFrom + , shapeIndexesBetween + + -- * HasLayout + , HasLayout (..) + , extent + , size + , indexes + , indexesBetween + , indexesFrom + + -- * Exceptions + + -- (* Bounds checking + , ArrayException (IndexOutOfBounds) + , _IndexOutOfBounds + , boundsCheck + + -- (* Size missmatch + , SizeMissmatch (..) + , AsSizeMissmatch (..) + , sizeMissmatch + + -- * Utilities + , showShape + ) where + +import Control.Applicative +import Control.Exception +import Control.Exception.Lens +import Control.Lens +import Control.Lens.Internal.Getter +import Data.Foldable as F +import Data.Typeable + +import Data.Functor.Classes +import Data.Traversable +import Linear + +-- | A 'Layout' is the full size of an array. This alias is used to help +-- distinguish between the layout of an array and an index (usually +-- just @l Int@) in a type signature. +type Layout f = f Int + +------------------------------------------------------------------------ +-- Shape class +------------------------------------------------------------------------ + +-- | Class for types that can be converted to and from linear indexes. +class (Eq1 f, Additive f, Traversable f) => Shape f where + -- | Convert a shape to its linear index using the 'Layout'. + shapeToIndex :: Layout f -> f Int -> Int + shapeToIndex l x = F.foldl (\k (e, a) -> k * e + a) 0 (liftI2 (,) l x) + {-# INLINE shapeToIndex #-} + + -- | Convert a linear index to a shape the 'Layout'. + shapeFromIndex :: Layout f -> Int -> f Int + shapeFromIndex l i = snd $ mapAccumR quotRem i l + {-# INLINE shapeFromIndex #-} + + -- | Calculate the intersection of two shapes. + shapeIntersect :: Layout f -> Layout f -> Layout f + shapeIntersect = liftU2 min + {-# INLINE shapeIntersect #-} + + -- | Increment a shape by one. It is assumed that the provided index + -- is 'inRange'. + unsafeShapeStep :: Layout f -> f Int -> f Int + unsafeShapeStep l + = shapeFromIndex l + . (+1) + . shapeToIndex l + {-# INLINE unsafeShapeStep #-} + + -- | Increment a shape by one. It is assumed that the provided index + -- is 'inRange'. + shapeStep :: Layout f -> f Int -> Maybe (f Int) + shapeStep l = fmap (shapeFromIndex l) + . guardPure (< shapeSize l) + . (+1) + . shapeToIndex l + {-# INLINE shapeStep #-} + + -- | Increment a shape by one between the two bounds + shapeStepBetween :: f Int -> Layout f -> f Int -> Maybe (f Int) + shapeStepBetween a l = fmap (^+^ a) . shapeStep l . (^-^ a) + {-# INLINE shapeStepBetween #-} + + -- | @inRange ex i@ checks @i < ex@ for every coordinate of @f@. + shapeInRange :: Layout f -> f Int -> Bool + shapeInRange l i = F.and $ liftI2 (\ii li -> ii >= 0 && ii < li) i l + {-# INLINE shapeInRange #-} + + -- | The number of elements in a shape. + shapeSize :: Layout f -> Int + shapeSize = F.product + {-# INLINE shapeSize #-} + +guardPure :: Alternative f => (a -> Bool) -> a -> f a +guardPure p a = if p a then pure a else empty +{-# INLINE guardPure #-} + +instance Shape V0 + +instance Shape V1 where + {-# INLINE shapeToIndex #-} + {-# INLINE shapeFromIndex #-} + {-# INLINE shapeIntersect #-} + {-# INLINE shapeStep #-} + {-# INLINE shapeInRange #-} + shapeToIndex _ (V1 i) = i + shapeFromIndex _ i = V1 i + shapeIntersect = min + shapeStep l = guardPure (shapeInRange l) . (+1) + shapeStepBetween _a b i = guardPure (> b) i' + where i' = i + 1 + shapeInRange m i = i >= 0 && i < m + +instance Shape V2 where + shapeToIndex (V2 _x y) (V2 i j) = y*i + j + {-# INLINE shapeToIndex #-} + + shapeFromIndex (V2 _x y) n = V2 i j + where (i, j) = n `quotRem` y + {-# INLINE shapeFromIndex #-} + + shapeStep (V2 x y) (V2 i j) + | j + 1 < y = Just (V2 i (j + 1)) + | i + 1 < x = Just (V2 (i + 1) 0 ) + | otherwise = Nothing + {-# INLINE shapeStep #-} + + unsafeShapeStep (V2 _ y) (V2 i j) + | j + 1 < y = V2 i (j + 1) + | otherwise = V2 (i + 1) 0 + {-# INLINE unsafeShapeStep #-} + + shapeStepBetween (V2 _ia ja) (V2 ib jb) (V2 i j) + | j + 1 < jb = Just (V2 i (j + 1)) + | i + 1 < ib = Just (V2 (i + 1) ja ) + | otherwise = Nothing + {-# INLINE shapeStepBetween #-} + +instance Shape V3 where + shapeStep (V3 x y z) (V3 i j k) + | k + 1 < z = Just (V3 i j (k + 1)) + | j + 1 < y = Just (V3 i (j + 1) 0 ) + | i + 1 < x = Just (V3 (i + 1) 0 0 ) + | otherwise = Nothing + {-# INLINE shapeStep #-} + + shapeStepBetween (V3 _ia ja ka) (V3 ib jb kb) (V3 i j k) + | k < kb = Just (V3 i j (k + 1)) + | j < jb = Just (V3 i (j + 1) ka ) + | i < ib = Just (V3 (i + 1) ja ka ) + | otherwise = Nothing + {-# INLINE shapeStepBetween #-} + +instance Shape V4 where + shapeStep (V4 x y z w) (V4 i j k l) + | l + 1 < w = Just (V4 i j k (l + 1)) + | k + 1 < z = Just (V4 i j (k + 1) 0 ) + | j + 1 < y = Just (V4 i (j + 1) 0 0 ) + | i + 1 < x = Just (V4 (i + 1) 0 0 0 ) + | otherwise = Nothing + {-# INLINE shapeStep #-} + + shapeStepBetween (V4 _ia ja ka la) (V4 ib jb kb lb) (V4 i j k l) + | l < lb = Just (V4 i j k (l + 1)) + | k < kb = Just (V4 i j (k + 1) la ) + | j < jb = Just (V4 i (j + 1) ka la ) + | i < ib = Just (V4 (i + 1) ja ka la ) + | otherwise = Nothing + {-# INLINE shapeStepBetween #-} + +-- instance Dim n => Shape (V n) + +-- | @'toIndex' l@ and @'fromIndex' l@ form two halfs of an isomorphism. +indexIso :: Shape f => Layout f -> Iso' (f Int) Int +indexIso l = iso (shapeToIndex l) (shapeFromIndex l) +{-# INLINE indexIso #-} + +------------------------------------------------------------------------ +-- HasLayout +------------------------------------------------------------------------ + +-- | Class of things that have a 'Layout'. This means we can use the +-- same functions for the various different arrays in the library. +class Shape f => HasLayout f a | a -> f where + -- | Lens onto the 'Layout' of something. + layout :: Lens' a (Layout f) + default layout :: (a ~ f Int) => (Layout f -> g (Layout f)) -> a -> g a + layout = id + {-# INLINE layout #-} + +instance i ~ Int => HasLayout V0 (V0 i) +instance i ~ Int => HasLayout V1 (V1 i) +instance i ~ Int => HasLayout V2 (V2 i) +instance i ~ Int => HasLayout V3 (V3 i) +instance i ~ Int => HasLayout V4 (V4 i) + +-- | Get the extent of an array. +-- +-- @ +-- 'extent' :: 'Data.Dense.Base.Array' v f a -> f 'Int' +-- 'extent' :: 'Data.Dense.Mutable.MArray' v f s a -> f 'Int' +-- 'extent' :: 'Data.Dense.Base.Delayed' f a -> f 'Int' +-- 'extent' :: 'Data.Dense.Base.Focused' f a -> f 'Int' +-- @ +extent :: HasLayout f a => a -> f Int +extent = view layout +{-# INLINE extent #-} + +-- | Get the total number of elements in an array. +-- +-- @ +-- 'size' :: 'Data.Dense.Base.Array' v f a -> 'Int' +-- 'size' :: 'Data.Dense.Mutable.MArray' v f s a -> 'Int' +-- 'size' :: 'Data.Dense.Base.Delayed' f a -> 'Int' +-- 'size' :: 'Data.Dense.Base.Focused' f a -> 'Int' +-- @ +size :: HasLayout f a => a -> Int +size = shapeSize . view layout +{-# INLINE size #-} + +-- NB: lens already uses indices so we settle for indexes + +-- | Indexed fold for all the indexes in the layout. +indexes :: HasLayout f a => IndexedFold Int a (f Int) +indexes = layout . shapeIndexes +{-# INLINE indexes #-} + +-- | 'indexes' for a 'Shape'. +shapeIndexes :: Shape f => IndexedFold Int (Layout f) (f Int) +shapeIndexes g l = go (0::Int) (if eq1 l zero then Nothing else Just zero) where + go i (Just x) = indexed g i x *> go (i + 1) (shapeStep l x) + go _ Nothing = noEffect +{-# INLINE shapeIndexes #-} + +-- | Indexed fold starting starting from some point, where the index is +-- the linear index for the original layout. +indexesFrom :: HasLayout f a => f Int -> IndexedFold Int a (f Int) +indexesFrom a = layout . shapeIndexesFrom a +{-# INLINE indexesFrom #-} + +-- | 'indexesFrom' for a 'Shape'. +shapeIndexesFrom :: Shape f => f Int -> IndexedFold Int (Layout f) (f Int) +shapeIndexesFrom a f l = shapeIndexesBetween a l f l +{-# INLINE shapeIndexesFrom #-} + +-- | Indexed fold between the two indexes where the index is the linear +-- index for the original layout. +indexesBetween :: HasLayout f a => f Int -> f Int -> IndexedFold Int a (f Int) +indexesBetween a b = layout . shapeIndexesBetween a b +{-# INLINE indexesBetween #-} + +-- | 'indexesBetween' for a 'Shape'. +shapeIndexesBetween :: Shape f => f Int -> f Int -> IndexedFold Int (Layout f) (f Int) +shapeIndexesBetween a b f l = + go (if eq1 l a || not (shapeInRange l b) then Nothing else Just a) where + go (Just x) = indexed f (shapeToIndex l x) x *> go (shapeStepBetween a b x) + go Nothing = noEffect +{-# INLINE shapeIndexesBetween #-} + +------------------------------------------------------------------------ +-- Exceptions +------------------------------------------------------------------------ + +-- Bounds check -------------------------------------------------------- + +-- | @boundsCheck l i@ performs a bounds check for index @i@ and layout +-- @l@. Throws an 'IndexOutOfBounds' exception when out of range in +-- the form @(i, l)@. This can be caught with the '_IndexOutOfBounds' +-- prism. +-- +-- >>> boundsCheck (V2 3 5) (V2 1 4) "in range" +-- "in range" +-- +-- >>> boundsCheck (V2 10 20) (V2 10 5) "in bounds" +-- "*** Exception: array index out of range: (V2 10 5, V2 10 20) +-- +-- >>> catching _IndexOutOfBounds (boundsCheck (V1 2) (V1 2) (putStrLn "in range")) print +-- "(V1 2, V1 2)" +-- +-- The output format is suitable to be read using the '_Show' prism: +-- +-- >>> trying (_IndexOutOfBounds . _Show) (boundsCheck (V1 2) (V1 20) (putStrLn "in range")) :: IO (Either (V1 Int, V1 Int) ()) +-- Left (V1 20,V1 2) +boundsCheck :: Shape l => Layout l-> l Int -> a -> a +boundsCheck l i + | shapeInRange l i = id + | otherwise = throwing _IndexOutOfBounds $ "(" ++ showShape i ++ ", " ++ showShape l ++ ")" +{-# INLINE boundsCheck #-} + +-- Size missmatch ------------------------------------------------------ + +-- | Thrown when two sizes that should match, don't. +data SizeMissmatch = SizeMissmatch String + deriving Typeable + +instance Exception SizeMissmatch +instance Show SizeMissmatch where + showsPrec _ (SizeMissmatch s) + = showString "size missmatch" + . (if not (null s) then showString ": " . showString s + else id) + +-- | Exception thown from missmatching sizes. +class AsSizeMissmatch t where + -- | Extract information about an 'SizeMissmatch'. + -- + -- @ + -- '_SizeMissmatch' :: 'Prism'' 'SizeMissmatch' 'String' + -- '_SizeMissmatch' :: 'Prism'' 'SomeException' 'String' + -- @ + _SizeMissmatch :: Prism' t String + +instance AsSizeMissmatch SizeMissmatch where + _SizeMissmatch = prism' SizeMissmatch $ (\(SizeMissmatch s) -> Just s) + {-# INLINE _SizeMissmatch #-} + +instance AsSizeMissmatch SomeException where + _SizeMissmatch = exception . (_SizeMissmatch :: Prism' SizeMissmatch String) + {-# INLINE _SizeMissmatch #-} + +-- | Check the sizes are equal. If not, throw 'SizeMissmatch'. +sizeMissmatch :: Int -> Int -> String -> a -> a +sizeMissmatch i j err + | i == j = id + | otherwise = throwing _SizeMissmatch err +{-# INLINE sizeMissmatch #-} + +-- Utilities ----------------------------------------------------------- + +-- | Show a shape in the form @VN i1 i2 .. iN@ where @N@ is the 'length' +-- of the shape. +showShape :: Shape f => f Int -> String +showShape l = "V" ++ show (lengthOf folded l) ++ " " ++ unwords (show <$> F.toList l) + diff --git a/src/Data/Dense/Mutable.hs b/src/Data/Dense/Mutable.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Mutable.hs @@ -0,0 +1,304 @@ +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Mutable +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- This module provides generic functions over mutable multidimensional +-- arrays. +----------------------------------------------------------------------------- +module Data.Dense.Mutable + ( + -- * Mutable array + MArray (..) + , UMArray + , SMArray + , BMArray + , PMArray + + -- * Lenses + , mlayout + , mvector + + -- * Creation + , new + , replicate + , replicateM + , clone + + -- * Standard operations + -- ** Indexing + , read + , linearRead + , unsafeRead + , unsafeLinearRead + + -- ** Writing + , write + , linearWrite + , unsafeWrite + , unsafeLinearWrite + + -- ** Modifying + , modify + , linearModify + , unsafeModify + , unsafeLinearModify + + -- ** Swap + , swap + , linearSwap + , unsafeSwap + , unsafeLinearSwap + + -- ** Exchange + , exchange + , linearExchange + , unsafeExchange + , unsafeLinearExchange + + -- * Misc + , set + , clear + , copy + + ) where + +import Control.Monad (liftM) +import Control.Monad.Primitive +import Control.Lens (IndexedLens, indexed, Lens, (<&>)) +import Data.Foldable as F +import Data.Typeable +import qualified Data.Vector as B +import Data.Vector.Generic.Mutable (MVector) +import qualified Data.Vector.Generic.Mutable as GM +import qualified Data.Vector.Primitive.Mutable as P +import qualified Data.Vector.Storable.Mutable as S +import qualified Data.Vector.Unboxed.Mutable as U +import Linear.V1 + +import Data.Dense.Index + +import Prelude hiding (read, replicate) + +-- | A mutable array with a shape. +data MArray v l s a = MArray !(Layout l) !(v s a) + deriving Typeable + +-- | Boxed mutable array. +type BMArray = MArray B.MVector + +-- | Unboxed mutable array. +type UMArray = MArray U.MVector + +-- | Storable mutable array. +type SMArray = MArray S.MVector + +-- | Primitive mutable array. +type PMArray = MArray P.MVector + +-- | Lens onto the shape of the vector. The total size of the layout +-- _must_ remain the same or an error is thrown. +mlayout :: (Shape l, Shape l') => Lens (MArray v l s a) (MArray v l' s a) (Layout l) (Layout l') +mlayout f (MArray l v) = f l <&> \l' -> + sizeMissmatch (F.product l) (F.product l') + ("mlayout: trying to replace shape " ++ showShape l ++ ", with " ++ showShape l') + $ MArray l' v +{-# INLINE mlayout #-} + +instance Shape f => HasLayout f (MArray v f s a) where + layout = mlayout + {-# INLINE layout #-} + +-- | Indexed lens over the underlying vector of an array. The index is +-- the 'extent' of the array. You must __not__ change the length of +-- the vector, otherwise an error will be thrown. +mvector :: (MVector v a, MVector w b) => IndexedLens (Layout l) (MArray v l s a) (MArray w l t b) (v s a) (w t b) +mvector f (MArray l v) = + indexed f l v <&> \w -> + sizeMissmatch (GM.length v) (GM.length w) + ("mvector: trying to replace vector of length " ++ show (GM.length v) ++ ", with one of length " ++ show (GM.length w)) + $ MArray l w +{-# INLINE mvector #-} + +-- | New mutable array with shape @l@. +new :: (PrimMonad m, Shape l, MVector v a) => Layout l -> m (MArray v l (PrimState m) a) +new l = MArray l `liftM` GM.new (F.product l) +{-# INLINE new #-} + +-- | New mutable array with shape @l@ filled with element @a@. +replicate :: (PrimMonad m, Shape l, MVector v a) => Layout l -> a -> m (MArray v l (PrimState m) a) +replicate l a = MArray l `liftM` GM.replicate (F.product l) a +{-# INLINE replicate #-} + +-- | New mutable array with shape @l@ filled with result of monadic +-- action @a@. +replicateM :: (PrimMonad m, Shape l, MVector v a) => Layout l -> m a -> m (MArray v l (PrimState m) a) +replicateM l a = MArray l `liftM` GM.replicateM (F.product l) a +{-# INLINE replicateM #-} + +-- | Clone a mutable array, making a new, separate mutable array. +clone :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> m (MArray v l (PrimState m) a) +clone (MArray l v) = MArray l `liftM` GM.clone v +{-# INLINE clone #-} + +-- Individual elements ------------------------------------------------- + +-- | Clear the elements of a mutable array. This is usually a no-op for +-- unboxed arrays. +clear :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> m () +clear (MArray _ v) = GM.clear v +{-# INLINE clear #-} + +-- | Read a mutable array at element @l@. +read :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> m a +read (MArray l v) s = boundsCheck l s $ GM.unsafeRead v (shapeToIndex l s) +{-# INLINE read #-} + +-- | Write a mutable array at element @l@. +write :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> a -> m () +write (MArray l v) s a = boundsCheck l s $ GM.unsafeWrite v (shapeToIndex l s) a +{-# INLINE write #-} + +-- | Modify a mutable array at element @l@ by applying a function. +modify :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> (a -> a) -> m () +modify (MArray l v) s f = boundsCheck l s $ GM.unsafeRead v i >>= GM.unsafeWrite v i . f + where i = shapeToIndex l s +{-# INLINE modify #-} + +-- | Swap two elements in a mutable array. +swap :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> l Int -> m () +swap (MArray l v) i j = boundsCheck l i boundsCheck l j $ GM.unsafeSwap v (shapeToIndex l i) (shapeToIndex l j) +{-# INLINE swap #-} + +-- | Replace the element at the give position and return the old +-- element. +exchange :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> a -> m a +exchange (MArray l v) i a = boundsCheck l i $ GM.unsafeExchange v (shapeToIndex l i) a +{-# INLINE exchange #-} + +-- | Read a mutable array at element @i@ by indexing the internal +-- vector. +linearRead :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> m a +linearRead (MArray _ v) = GM.read v +{-# INLINE linearRead #-} + +-- | Write a mutable array at element @i@ by indexing the internal +-- vector. +linearWrite :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> a -> m () +linearWrite (MArray _ v) = GM.write v +{-# INLINE linearWrite #-} + +-- | Swap two elements in a mutable array by indexing the internal +-- vector. +linearSwap :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> Int -> m () +linearSwap (MArray _ v) = GM.swap v +{-# INLINE linearSwap #-} + +-- | Modify a mutable array at element @i@ by applying a function. +linearModify :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> (a -> a) -> m () +linearModify (MArray _ v) i f = GM.read v i >>= GM.unsafeWrite v i . f +{-# INLINE linearModify #-} + +-- | Replace the element at the give position and return the old +-- element. +linearExchange :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> a -> m a +linearExchange (MArray _ v) i a = GM.exchange v i a +{-# INLINE linearExchange #-} + +-- Unsafe varients + +-- | 'read' without bounds checking. +unsafeRead :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> m a +unsafeRead (MArray l v) s = GM.unsafeRead v (shapeToIndex l s) +{-# INLINE unsafeRead #-} + +-- | 'write' without bounds checking. +unsafeWrite :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> a -> m () +unsafeWrite (MArray l v) s = GM.unsafeWrite v (shapeToIndex l s) +{-# INLINE unsafeWrite #-} + +-- | 'swap' without bounds checking. +unsafeSwap :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> l Int -> m () +unsafeSwap (MArray l v) s j = GM.unsafeSwap v (shapeToIndex l s) (shapeToIndex j s) +{-# INLINE unsafeSwap #-} + +-- | 'modify' without bounds checking. +unsafeModify :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> (a -> a) -> m () +unsafeModify (MArray l v) s f = GM.unsafeRead v i >>= GM.unsafeWrite v i . f + where i = shapeToIndex l s +{-# INLINE unsafeModify #-} + +-- | Replace the element at the give position and return the old +-- element. +unsafeExchange :: (PrimMonad m, Shape l, MVector v a) => MArray v l (PrimState m) a -> l Int -> a -> m a +unsafeExchange (MArray l v) i a = GM.unsafeExchange v (shapeToIndex l i) a +{-# INLINE unsafeExchange #-} + +-- | 'linearRead' without bounds checking. +unsafeLinearRead :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> m a +unsafeLinearRead (MArray _ v) = GM.unsafeRead v +{-# INLINE unsafeLinearRead #-} + +-- | 'linearWrite' without bounds checking. +unsafeLinearWrite :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> a -> m () +unsafeLinearWrite (MArray _ v) = GM.unsafeWrite v +{-# INLINE unsafeLinearWrite #-} + +-- | 'linearSwap' without bounds checking. +unsafeLinearSwap :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> Int -> m () +unsafeLinearSwap (MArray _ v) = GM.unsafeSwap v +{-# INLINE unsafeLinearSwap #-} + +-- | 'linearModify' without bounds checking. +unsafeLinearModify :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> (a -> a) -> m () +unsafeLinearModify (MArray _ v) i f = GM.unsafeRead v i >>= GM.unsafeWrite v i . f +{-# INLINE unsafeLinearModify #-} + +-- | Replace the element at the give position and return the old +-- element. +unsafeLinearExchange :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> Int -> a -> m a +unsafeLinearExchange (MArray _ v) i a = GM.unsafeExchange v i a +{-# INLINE unsafeLinearExchange #-} + +-- Filling and copying ------------------------------------------------- + +-- | Set all elements in a mutable array to a constant value. +set :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> a -> m () +set (MArray _ v) = GM.set v +{-# INLINE set #-} + +-- | Copy all elements from one array into another. +copy :: (PrimMonad m, MVector v a) => MArray v l (PrimState m) a -> MArray v l (PrimState m) a -> m () +copy (MArray _ v) (MArray _ u) = GM.copy v u +{-# INLINE copy #-} + +-- V1 instances -------------------------------------------------------- + +-- Array v V1 a is essentially v a with a wrapper. Instance is provided +-- for convience. + +instance (MVector v a, l ~ V1) => MVector (MArray v l) a where + {-# INLINE basicLength #-} + {-# INLINE basicUnsafeSlice #-} + {-# INLINE basicOverlaps #-} + {-# INLINE basicUnsafeNew #-} + {-# INLINE basicUnsafeRead #-} + {-# INLINE basicUnsafeWrite #-} + {-# INLINE basicInitialize #-} + basicLength (MArray (V1 n) _) = n + basicUnsafeSlice i n (MArray _ v) = MArray (V1 n) $ GM.basicUnsafeSlice i n v + basicOverlaps (MArray _ v) (MArray _ w) = GM.basicOverlaps v w + basicUnsafeNew n = MArray (V1 n) `liftM` GM.basicUnsafeNew n + basicUnsafeRead (MArray _ v) = GM.basicUnsafeRead v + basicUnsafeWrite (MArray _ v) = GM.basicUnsafeWrite v + basicInitialize (MArray _ v) = GM.basicInitialize v + diff --git a/src/Data/Dense/Stencil.hs b/src/Data/Dense/Stencil.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Stencil.hs @@ -0,0 +1,110 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TemplateHaskell #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Stencil +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- Stencils can be used to sum (or any fold) over neighbouring sites to +-- the current position on a 'Focused'. +----------------------------------------------------------------------------- +module Data.Dense.Stencil + ( -- * The Stencil type + Stencil (..) + , mkStencil + , mkStencilUnboxed + + -- ** Using stencils + , stencilSum + + ) where + +import Control.Lens +import Data.Dense.Base +import Data.Dense.Generic (Boundary (..), peekRelativeB) +import Data.Dense.Index +import qualified Data.Foldable as F +import Data.Functor.Classes +import qualified Data.Vector.Unboxed as U +import Text.Show + +-- Types --------------------------------------------------------------- + +-- | Stencils are used to fold over neighbouring array sites. To +-- construct a stencil use 'mkStencil', 'mkStencilUnboxed'. For +-- static sized stencils you can use the quasiquoter +-- 'Data.Dense.TH.stencil'. +-- +-- To use a stencil you can use 'stencilSum' or use the 'Foldable' and +-- 'FoldableWithIndex' instances. +newtype Stencil f a = Stencil (forall b. (f Int -> a -> b -> b) -> b -> b) + +instance (Show1 f, Show a) => Show (Stencil f a) where + showsPrec _ s = showListWith g (itoList s) where + g (i,x) = showChar '(' . showsPrec1 0 i . showChar ',' . showsPrec 0 x . showChar ')' + +instance F.Foldable (Stencil f) where + foldr f z (Stencil s) = s (\_ a b -> f a b) z + {-# INLINE foldr #-} + +instance FoldableWithIndex (f Int) (Stencil f) where + ifoldr f b (Stencil s) = s f b + {-# INLINE ifoldr #-} + ifoldMap = ifoldMapOf (ifoldring ifoldr) + {-# INLINE ifoldMap #-} + +instance Functor (Stencil f) where + fmap f (Stencil s) = Stencil $ \g z -> s (\x a b -> g x (f a) b) z + {-# INLINE [0] fmap #-} + +-- | Make a stencil folding over a list. +-- +-- If the list is staticlly known this should expand at compile time +-- via rewrite rules, similar to 'Data.Dense.TH.makeStencilTH' but less reliable. If +-- that does not happen the resulting could be slow. If the list is +-- not know at compile time, 'mkStencilUnboxed' can be signifcantly +-- faster (but isn't subject expending via rewrite rules). +mkStencil :: [(f Int, a)] -> Stencil f a +mkStencil l = Stencil $ \g z -> myfoldr (\(i,a) b -> g i a b) z l +{-# INLINE mkStencil #-} + +-- Version of foldr that recursivly expands the list via rewrite rules. +myfoldr :: (a -> b -> b) -> b -> [a] -> b +myfoldr f b = go where + go [] = b + go (a:as) = f a (go as) +{-# INLINE [0] myfoldr #-} + +{-# RULES +"mkStencil/cons" forall f b a as. + myfoldr f b (a:as) = f a (myfoldr f b as) + #-} + +-- | Make a stencil folding over an unboxed vector from the list. +mkStencilUnboxed :: (U.Unbox (f Int), U.Unbox a) => [(f Int, a)] -> Stencil f a +mkStencilUnboxed l = Stencil $ \g z -> U.foldr (\(i,a) b -> g i a b) z v + where !v = U.fromList l +{-# INLINE mkStencilUnboxed #-} + +-- | Sum the elements around a 'Focused' using a 'Boundary' condition +-- and a 'Stencil'. +-- +-- This is often used in conjunction with 'Data.Dense.extendFocus'. +stencilSum :: (Shape f, Num a) => Boundary -> Stencil f a -> Focused f a -> a +stencilSum bnd s = \w -> + let f i b a = b + a * peekRelativeB bnd i w + {-# INLINE [0] f #-} + in ifoldl' f 0 s +{-# INLINE stencilSum #-} + diff --git a/src/Data/Dense/Storable.hs b/src/Data/Dense/Storable.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Storable.hs @@ -0,0 +1,674 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Storable +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- 'Storeable' multidimentional arrays. +----------------------------------------------------------------------------- +module Data.Dense.Storable + ( + -- * SArray types + SArray + , Storable + , Shape + + -- * Layout of an array + , HasLayout (..) + , Layout + + -- ** Extracting size + , extent + , size + + -- ** Folds over indexes + , indexes + , indexesFrom + , indexesBetween + + -- * Underlying vector + , vector + + -- ** Traversals + , values + , values' + , valuesBetween + + -- * Construction + + -- ** Flat arrays + , flat + , fromList + + -- ** From lists + , fromListInto + , fromListInto_ + + -- ** From vectors + , fromVectorInto + , fromVectorInto_ + + -- ** Initialisation + , replicate + , generate + , linearGenerate + + -- ** Monadic initialisation + , create + , replicateM + , generateM + , linearGenerateM + + -- * Functions on arrays + + -- ** Empty arrays + , empty + , null + + -- ** Indexing + + , (!) + , (!?) + , unsafeIndex + , linearIndex + , unsafeLinearIndex + + -- *** Monadic indexing + , indexM + , unsafeIndexM + , linearIndexM + , unsafeLinearIndexM + + -- ** Modifying arrays + + -- ** Bulk updates + , (//) + + -- ** Accumulations + , accum + + -- ** Mapping + , map + , imap + + -- * Zipping + + -- ** Zip with function + , zipWith + , zipWith3 + , izipWith + , izipWith3 + + -- ** Slices + + -- *** Matrix + , ixRow + , rows + , ixColumn + , columns + + -- *** 3D + , ixPlane + , planes + , flattenPlane + + -- *** Ordinals + , unsafeOrdinals + + -- * Mutable + , SMArray + + , thaw + , freeze + , unsafeThaw + , unsafeFreeze + + -- * Delayed + + , G.Delayed + + -- ** Generating delayed + + , delayed + , seqDelayed + , delay + , manifest + , seqManifest + , G.genDelayed + , G.indexDelayed + , affirm + , seqAffirm + + -- * Focused + + , G.Focused + + -- ** Generating focused + + , G.focusOn + , G.unfocus + , G.unfocused + , G.extendFocus + + -- ** Focus location + , G.locale + , G.shiftFocus + + -- ** Pointers + + , unsafeWithPtr + , unsafeToForeignPtr + , unsafeFromForeignPtr + ) where + +import Control.Lens hiding (imap) +import Control.Monad.Primitive +import Control.Monad.ST +import qualified Data.Foldable as F +import Data.Vector.Storable (Storable, Vector) +import qualified Data.Vector.Storable as S +import Linear hiding (vector) +import Foreign (Ptr, ForeignPtr) + +import Prelude hiding (map, null, replicate, zip, + zip3, zipWith, zipWith3) + +import Data.Dense.Base (Array (..)) +import Data.Dense.Generic (SArray) +import qualified Data.Dense.Generic as G +import Data.Dense.Index +import Data.Dense.Mutable (SMArray) + +-- Lenses -------------------------------------------------------------- + +-- | Same as 'values' but restrictive in the vector type. +values :: (Shape f, Storable a, Storable b) + => IndexedTraversal (f Int) (SArray f a) (SArray f b) a b +values = G.values' +{-# INLINE values #-} + +-- | Same as 'values' but restrictive in the vector type. +values' :: (Shape f, Storable a, Storable b) + => IndexedTraversal (f Int) (SArray f a) (SArray f b) a b +values' = G.values' +{-# INLINE values' #-} + +-- | Same as 'values' but restrictive in the vector type. +valuesBetween + :: (Shape f, Storable a) + => f Int + -> f Int + -> IndexedTraversal' (f Int) (SArray f a) a +valuesBetween = G.valuesBetween +{-# INLINE valuesBetween #-} + +-- | 1D arrays are just vectors. You are free to change the length of +-- the vector when going 'over' this 'Iso' (unlike 'linear'). +-- +-- Note that 'V1' arrays are an instance of 'Vector' so you can use +-- any of the functions in 'Data.Vector.Generic' on them without +-- needing to convert. +flat :: Storable b => Iso (SArray V1 a) (SArray V1 b) (Vector a) (Vector b) +flat = G.flat +{-# INLINE flat #-} + +-- | Indexed lens over the underlying vector of an array. The index is +-- the 'extent' of the array. You must _not_ change the length of the +-- vector, otherwise an error will be thrown (even for 'V1' layouts, +-- use 'flat' for 'V1'). +vector :: (Storable a, Storable b) => IndexedLens (Layout f) (SArray f a) (SArray f b) (Vector a) (Vector b) +vector = G.vector +{-# INLINE vector #-} + +-- Constructing vectors ------------------------------------------------ + +-- | Contruct a flat array from a list. (This is just 'G.fromList' from +-- 'Data.Vector.Generic'.) +fromList :: Storable a => [a] -> SArray V1 a +fromList = G.fromList +{-# INLINE fromList #-} + +-- | O(n) Convert the first @n@ elements of a list to an SArrayith the +-- given shape. Returns 'Nothing' if there are not enough elements in +-- the list. +fromListInto :: (Shape f, Storable a) => Layout f -> [a] -> Maybe (SArray f a) +fromListInto = G.fromListInto +{-# INLINE fromListInto #-} + +-- | O(n) Convert the first @n@ elements of a list to an SArrayith the +-- given shape. Throw an error if the list is not long enough. +fromListInto_ :: (Shape f, Storable a) => Layout f -> [a] -> SArray f a +fromListInto_ = G.fromListInto_ +{-# INLINE fromListInto_ #-} + +-- | Create an array from a 'vector' and a 'layout'. Return 'Nothing' if +-- the vector is not the right shape. +fromVectorInto :: (Shape f, Storable a) => Layout f -> Vector a -> Maybe (SArray f a) +fromVectorInto = G.fromVectorInto +{-# INLINE fromVectorInto #-} + +-- | Create an array from a 'vector' and a 'layout'. Throws an error if +-- the vector is not the right shape. +fromVectorInto_ :: (Shape f, Storable a) => Layout f -> Vector a -> SArray f a +fromVectorInto_ = G.fromVectorInto_ +{-# INLINE fromVectorInto_ #-} + +-- | The empty 'SArray' with a 'zero' shape. +empty :: (Storable a, Additive f) => SArray f a +empty = G.empty +{-# INLINE empty #-} + +-- | Test is if the array is 'empty'. +null :: F.Foldable f => SArray f a -> Bool +null = G.null +{-# INLINE null #-} + +-- Indexing ------------------------------------------------------------ + +-- | Index an element of an array. Throws 'IndexOutOfBounds' if the +-- index is out of bounds. +(!) :: (Shape f, Storable a) => SArray f a -> f Int -> a +(!) = (G.!) +{-# INLINE (!) #-} + +-- | Safe index of an element. +(!?) :: (Shape f, Storable a) => SArray f a -> f Int -> Maybe a +(!?) = (G.!?) +{-# INLINE (!?) #-} + +-- | Index an element of an array without bounds checking. +unsafeIndex :: (Shape f, Storable a) => SArray f a -> f Int -> a +unsafeIndex = G.unsafeIndex +{-# INLINE unsafeIndex #-} + +-- | Index an element of an array while ignoring its shape. +linearIndex :: Storable a => SArray f a -> Int -> a +linearIndex = G.linearIndex +{-# INLINE linearIndex #-} + +-- | Index an element of an array while ignoring its shape, without +-- bounds checking. +unsafeLinearIndex :: Storable a => SArray f a -> Int -> a +unsafeLinearIndex = G.unsafeLinearIndex +{-# INLINE unsafeLinearIndex #-} + +-- Monadic indexing ---------------------------------------------------- + +-- | /O(1)/ Indexing in a monad. +-- +-- The monad allows operations to be strict in the vector when necessary. +-- Suppose vector copying is implemented like this: +-- +-- > copy mv v = ... write mv i (v ! i) ... +-- +-- For lazy vectors, @v ! i@ would not be evaluated which means that +-- @mv@ would unnecessarily retain a reference to @v@ in each element +-- written. +-- +-- With 'indexM', copying can be implemented like this instead: +-- +-- > copy mv v = ... do +-- > x <- indexM v i +-- > write mv i x +-- +-- Here, no references to @v@ are retained because indexing (but /not/ +-- the elements) is evaluated eagerly. +-- +-- Throws an error if the index is out of range. +indexM :: (Shape f, Storable a, Monad m) => SArray f a -> f Int -> m a +indexM = G.indexM +{-# INLINE indexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeIndexM :: (Shape f, Storable a, Monad m) => SArray f a -> f Int -> m a +unsafeIndexM = G.unsafeIndexM +{-# INLINE unsafeIndexM #-} + +-- | /O(1)/ Indexing in a monad. Throws an error if the index is out of +-- range. +linearIndexM :: (Shape f, Storable a, Monad m) => SArray f a -> Int -> m a +linearIndexM = G.linearIndexM +{-# INLINE linearIndexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeLinearIndexM :: (Storable a, Monad m) => SArray f a -> Int -> m a +unsafeLinearIndexM = G.unsafeLinearIndexM +{-# INLINE unsafeLinearIndexM #-} + +-- Initialisation ------------------------------------------------------ + +-- | Execute the monadic action and freeze the resulting array. +create :: Storable a => (forall s. ST s (SMArray f s a)) -> SArray f a +create m = m `seq` runST (m >>= G.unsafeFreeze) +{-# INLINE create #-} + +-- | O(n) SArray of the given shape with the same value in each position. +replicate :: (Shape f, Storable a) => f Int -> a -> SArray f a +replicate = G.replicate +{-# INLINE replicate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +linearGenerate :: (Shape f, Storable a) => Layout f -> (Int -> a) -> SArray f a +linearGenerate = G.linearGenerate +{-# INLINE linearGenerate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +generate :: (Shape f, Storable a) => Layout f -> (f Int -> a) -> SArray f a +generate = G.generate +{-# INLINE generate #-} + +-- Monadic initialisation ---------------------------------------------- + +-- | O(n) Construct an array of the given shape by filling each position +-- with the monadic value. +replicateM :: (Monad m, Shape f, Storable a) => Layout f -> m a -> m (SArray f a) +replicateM = G.replicateM +{-# INLINE replicateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +generateM :: (Monad m, Shape f, Storable a) => Layout f -> (f Int -> m a) -> m (SArray f a) +generateM = G.generateM +{-# INLINE generateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +linearGenerateM :: (Monad m, Shape f, Storable a) => Layout f -> (Int -> m a) -> m (SArray f a) +linearGenerateM = G.linearGenerateM +{-# INLINE linearGenerateM #-} + +-- Modifying ----------------------------------------------------------- + +-- | /O(n)/ Map a function over an array +map :: (Storable a, Storable b) => (a -> b) -> SArray f a -> SArray f b +map = G.map +{-# INLINE map #-} + +-- | /O(n)/ Apply a function to every element of a vector and its index +imap :: (Shape f, Storable a, Storable b) => (f Int -> a -> b) -> SArray f a -> SArray f b +imap = G.imap +{-# INLINE imap #-} + +-- Bulk updates -------------------------------------------------------- + + +-- | For each pair (i,a) from the list, replace the array element at +-- position i by a. +(//) :: (Storable a, Shape f) => SArray f a -> [(f Int, a)] -> SArray f a +(//) = (G.//) +{-# INLINE (//) #-} + +-- Accumilation -------------------------------------------------------- + +-- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the array element +-- @a@ at position @i@ by @f a b@. +-- +accum :: (Shape f, Storable a) + => (a -> b -> a) -- ^ accumulating function @f@ + -> SArray f a -- ^ initial array + -> [(f Int, b)] -- ^ list of index/value pairs (of length @n@) + -> SArray f a +accum = G.accum +{-# INLINE accum #-} + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- Zip with function --------------------------------------------------- + +-- | Zip two arrays using the given function. If the array's don't have +-- the same shape, the new array with be the intersection of the two +-- shapes. +zipWith :: (Shape f, Storable a, Storable b, Storable c) + => (a -> b -> c) + -> SArray f a + -> SArray f b + -> SArray f c +zipWith = G.zipWith +{-# INLINE zipWith #-} + +-- | Zip three arrays using the given function. If the array's don't +-- have the same shape, the new array with be the intersection of the +-- two shapes. +zipWith3 :: (Shape f, Storable a, Storable b, Storable c, Storable d) + => (a -> b -> c -> d) + -> SArray f a + -> SArray f b + -> SArray f c + -> SArray f d +zipWith3 = G.zipWith3 +{-# INLINE zipWith3 #-} + +-- Indexed zipping ----------------------------------------------------- + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith :: (Shape f, Storable a, Storable b, Storable c) + => (f Int -> a -> b -> c) + -> SArray f a + -> SArray f b + -> SArray f c +izipWith = G.izipWith +{-# INLINE izipWith #-} + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith3 :: (Shape f, Storable a, Storable b, Storable c, Storable d) + => (f Int -> a -> b -> c -> d) + -> SArray f a + -> SArray f b + -> SArray f c + -> SArray f d +izipWith3 = G.izipWith3 +{-# INLINE izipWith3 #-} + +------------------------------------------------------------------------ +-- Slices +------------------------------------------------------------------------ + +-- $setup +-- >>> import qualified Data.Vector.Storable as V +-- >>> let m = fromListInto_ (V2 2 3) [1..] :: SArray V2 Int + +-- | Indexed traversal over the rows of a matrix. Each row is an +-- efficient 'Data.Vector.Generic.slice' of the original vector. +-- +-- >>> traverseOf_ rows print m +-- [1,2,3] +-- [4,5,6] +rows :: (Storable a, Storable b) + => IndexedTraversal Int (SArray V2 a) (SArray V2 b) (Vector a) (Vector b) +rows = G.rows +{-# INLINE rows #-} + +-- | Affine traversal over a single row in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 . each +~ 2 +-- [1,2,3] +-- [6,7,8] +-- +-- The row vector should remain the same size to satisfy traversal +-- laws but give reasonable behaviour if the size differs: +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0,1] +-- [1,2,3] +-- [0,1,6] +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0..100] +-- [1,2,3] +-- [0,1,2] +ixRow :: Storable a => Int -> IndexedTraversal' Int (SArray V2 a) (Vector a) +ixRow = G.ixRow +{-# INLINE ixRow #-} + +-- | Indexed traversal over the columns of a matrix. Unlike 'rows', each +-- column is a new separate vector. +-- +-- >>> traverseOf_ columns print m +-- [1,4] +-- [2,5] +-- [3,6] +-- +-- >>> traverseOf_ rows print $ m & columns . indices odd . each .~ 0 +-- [1,0,3] +-- [4,0,6] +-- +-- The vectors should be the same size to be a valid traversal. If the +-- vectors are different sizes, the number of rows in the new array +-- will be the length of the smallest vector. +columns :: (Storable a, Storable b) + => IndexedTraversal Int (SArray V2 a) (SArray V2 b) (Vector a) (Vector b) +columns = G.columns +{-# INLINE columns #-} + +-- | Affine traversal over a single column in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixColumn 2 . each *~ 10 +-- [1,2,30] +-- [4,5,60] +ixColumn :: Storable a => Int -> IndexedTraversal' Int (SArray V2 a) (Vector a) +ixColumn = G.ixColumn +{-# INLINE ixColumn #-} + +-- | Traversal over a single plane of a 3D array given a lens onto that +-- plane (like '_xy', '_yz', '_zx'). +ixPlane :: Storable a + => ALens' (V3 Int) (V2 Int) + -> Int + -> IndexedTraversal' Int (SArray V3 a) (SArray V2 a) +ixPlane = G.ixPlane +{-# INLINE ixPlane #-} + +-- | Traversal over all planes of 3D array given a lens onto that plane +-- (like '_xy', '_yz', '_zx'). +planes :: (Storable a, Storable b) + => ALens' (V3 Int) (V2 Int) + -> IndexedTraversal Int (SArray V3 a) (SArray V3 b) (SArray V2 a) (SArray V2 b) +planes = G.planes +{-# INLINE planes #-} + +-- | Flatten a plane by reducing a vector in the third dimension to a +-- single value. +flattenPlane :: (Storable a, Storable b) + => ALens' (V3 Int) (V2 Int) + -> (Vector a -> b) + -> SArray V3 a + -> SArray V2 b +flattenPlane = G.flattenPlane +{-# INLINE flattenPlane #-} + +-- Ordinals ------------------------------------------------------------ + +-- | This 'Traversal' should not have any duplicates in the list of +-- indices. +unsafeOrdinals :: (Storable a, Shape f) => [f Int] -> IndexedTraversal' (f Int) (SArray f a) a +unsafeOrdinals = G.unsafeOrdinals +{-# INLINE [0] unsafeOrdinals #-} + +-- Mutable ------------------------------------------------------------- + +-- | O(n) Yield a mutable copy of the immutable vector. +freeze :: (PrimMonad m, Storable a) + => SMArray f (PrimState m) a -> m (SArray f a) +freeze = G.freeze +{-# INLINE freeze #-} + +-- | O(n) Yield an immutable copy of the mutable array. +thaw :: (PrimMonad m, Storable a) + => SArray f a -> m (SMArray f (PrimState m) a) +thaw = G.thaw +{-# INLINE thaw #-} + +-- | O(1) Unsafe convert a mutable array to an immutable one without +-- copying. The mutable array may not be used after this operation. +unsafeFreeze :: (PrimMonad m, Storable a) + => SMArray f (PrimState m) a -> m (SArray f a) +unsafeFreeze = G.unsafeFreeze +{-# INLINE unsafeFreeze #-} + +-- | O(1) Unsafely convert an immutable array to a mutable one without +-- copying. The immutable array may not be used after this operation. +unsafeThaw :: (PrimMonad m, Storable a) + => SArray f a -> m (SMArray f (PrimState m) a) +unsafeThaw = G.unsafeThaw +{-# INLINE unsafeThaw #-} + +------------------------------------------------------------------------ +-- Delayed +------------------------------------------------------------------------ + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in parallel. +delayed :: (Storable a, Storable b, Shape f, Shape k) + => Iso (SArray f a) (SArray k b) (G.Delayed f a) (G.Delayed k b) +delayed = G.delayed +{-# INLINE delayed #-} + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in sequence. +seqDelayed :: (Storable a, Storable b, Shape f, Shape k) + => Iso (SArray f a) (SArray k b) (G.Delayed f a) (G.Delayed k b) +seqDelayed = G.seqDelayed +{-# INLINE seqDelayed #-} + +-- | Turn a material array into a delayed one with the same shape. +delay :: (Storable a, Shape f) => SArray f a -> G.Delayed f a +delay = G.delay +{-# INLINE delay #-} + +-- | Parallel manifestation of a delayed array into a material one. +manifest :: (Storable a, Shape f) => G.Delayed f a -> SArray f a +manifest = G.manifest +{-# INLINE manifest #-} + +-- | Sequential manifestation of a delayed array. +seqManifest :: (Storable a, Shape f) => G.Delayed f a -> SArray f a +seqManifest = G.seqManifest +{-# INLINE seqManifest #-} + +-- | 'manifest' an array to a 'SArray' and delay again. +affirm :: (Shape f, Storable a) => G.Delayed f a -> G.Delayed f a +affirm = delay . manifest +{-# INLINE affirm #-} + +-- | 'seqManifest' an array to a 'SArray' and delay again. +seqAffirm :: (Shape f, Storable a) => G.Delayed f a -> G.Delayed f a +seqAffirm = delay . seqManifest +{-# INLINE seqAffirm #-} + +-- Pointer operations -------------------------------------------------- + +-- | Pass a pointer to the array's data to the IO action. Modifying +-- data through the 'Ptr' is unsafe. +unsafeWithPtr :: Storable a => SArray f a -> (Ptr a -> IO b) -> IO b +unsafeWithPtr (Array _ v) = S.unsafeWith v +{-# INLINE unsafeWithPtr #-} + +-- | Yield the underlying ForeignPtr. Modifying the data through the +-- 'ForeignPtr' is unsafe. +unsafeToForeignPtr :: Storable a => SArray f a -> ForeignPtr a +unsafeToForeignPtr (Array _ v) = fp + where (fp, _, _) = S.unsafeToForeignPtr v +{-# INLINE unsafeToForeignPtr #-} + +-- | O(1) Create an array from a layout and 'ForeignPtr'. It is +-- assumed the pointer points directly to the data (no offset). +-- Modifying data through the 'ForeignPtr' afterwards is unsafe. +unsafeFromForeignPtr + :: (Shape f, Storable a) => Layout f -> ForeignPtr a -> SArray f a +unsafeFromForeignPtr l fp = Array l (S.unsafeFromForeignPtr0 fp (shapeSize l)) +{-# INLINE unsafeFromForeignPtr #-} + diff --git a/src/Data/Dense/TH.hs b/src/Data/Dense/TH.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/TH.hs @@ -0,0 +1,729 @@ +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE QuasiQuotes #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TemplateHaskell #-} + +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.TH +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : experimental +-- Portability : non-portable +-- +-- Contains QuasiQuotes and TemplateHaskell utilities for creating dense +-- arrays, stencils and fixed length vectors. +-- +-- The parser for the QuasiQuotes is still a work in progress. +----------------------------------------------------------------------------- +module Data.Dense.TH + ( -- * Creating dense arrays + dense + + -- * Fixed length vector + , v + + -- * Stencils + , stencil + + -- ** Stencils from lists + , ShapeLift (..) + , mkStencilTH + , mkStencilTHBy + + ) where + +import Control.Applicative hiding (many, empty) +import Control.Lens +import Control.Monad +import Data.Char +import Data.Foldable as F +import Data.Function (on) +import qualified Data.List as List +import Data.Maybe +import Data.Monoid (Endo) +import qualified Data.Vector as Vector +import Language.Haskell.TH +import Language.Haskell.TH.Quote +import Language.Haskell.TH.Syntax +import Linear +import qualified Linear.V as V +import Text.ParserCombinators.ReadP +import qualified Text.Read.Lex as Lex + +import Data.Dense.Generic (empty, fromListInto_) +import Data.Dense.Index +import Data.Dense.Stencil + +-- | QuasiQuoter for producing a dense arrays using a custom parser. +-- Values are space separated, while also allowing infix expressions +-- (like @5/7@). If you want to apply a function, it should be done in +-- brackets. Supports 1D, 2D and 3D arrays. +-- +-- The number of rows/columns must be consistent thought out the +-- array. +-- +-- === __Examples__ +-- +-- - 1D arrays are of the following form form. Note these can be +-- used as 'V1', 'V2' or 'V3' arrays. +-- +-- @ +-- ['dense'| 5 -3 1 -3 5 |] :: ('R1' f, 'Vector.Vector' v a, 'Num' a) => 'Data.Dense.Array' v f a +-- @ +-- +-- +-- - 2D arrays are of the following form. Note these can be used as +-- 'V2' or 'V3' arrays. +-- +-- @ +-- chars :: 'Data.Dense.UArray' 'V2' 'Char' +-- chars :: ['dense'| +-- \'a\' \'b\' \'c\' +-- \'d\' \'e\' \'f\' +-- \'g\' \'h\' \'i\' +-- |] +-- @ +-- +-- - 3D arrays are of the following form. Note the order in which +-- 'dense' formats the array. The array @a@ is such that @a ! 'V3' +-- x y z = "xyz"@ +-- +-- @ +-- a :: 'Data.Dense.BArray' 'V3' 'String' +-- a = ['dense'| +-- "000" "100" "200" +-- "010" "110" "210" +-- "020" "120" "220" +-- +-- "001" "101" "201" +-- "011" "111" "211" +-- "021" "121" "221" +-- +-- "002" "102" "202" +-- "012" "112" "212" +-- "022" "122" "222" +-- |] +-- @ +-- +dense :: QuasiQuoter +dense = QuasiQuoter + { quoteExp = parseDense + , quotePat = error "dense can't be used in pattern" + , quoteType = error "dense can't be used in type" + , quoteDec = error "dense can't be used in dec" + } + +-- | List of expressions forming a stencil. To be either turned into +-- either a real list or an unfolded stencil. +parseDense :: String -> Q Exp +parseDense str = + case mapM (mapM parseLine) (map lines ps) of + Left err -> fail err + Right [] -> [| empty |] + Right [[as]] -> uncurry mkArray $ parse1D as + Right [ass] -> uncurry mkArray $ parse2D ass + Right asss -> uncurry mkArray $ parse3D asss + where ps = paragraphs str + +-- | Split a string up into paragraphs separated by a new line. Extra +-- newlines inbetween paragraphs are stripped. +paragraphs :: String -> [String] +paragraphs = go [] . strip where + go ps ('\n':'\n':xs) = [reverse ps] ++ go [] (strip xs) + go ps (x:xs) = go (x:ps) xs + go [] [] = [] + go ps [] = [reverse $ strip ps] + + strip = dropWhile (\x -> x == '\n' || x == ' ') + +-- Creating arrays ----------------------------------------------------- + +mkArray :: ShapeLift f => Layout f -> [Exp] -> Q Exp +mkArray l as = do + lE <- liftShape' l + let fromListE = AppE (VarE 'fromListInto_) lE + pure $ AppE fromListE (ListE as) + +------------------------------------------------------------------------ +-- V n a +------------------------------------------------------------------------ + +-- | Type safe 'QuasiQuoter' for fixed length vectors 'V.V'. Values are +-- space separated. Can be used as expressions or patterns. +-- +-- @ +-- [v| x y z w q r |] :: 'V.V' 6 a +-- @ +-- +-- Note this requires @DataKinds@. Also requires @ViewPatterns@ if 'v' +-- is used as a pattern. +-- +-- === __Examples__ +-- +-- @ +-- >>> let a = [v| 1 2 3 4 5 |] +-- >>> :t a +-- a :: Num a => V 5 a +-- >>> a +-- V {toVector = [1,2,3,4,5]} +-- >>> let f [v| a b c d e |] = (a,b,c,d,e) +-- >>> :t f +-- f :: V 5 t -> (t, t, t, t, t) +-- >>> f a +-- (1,2,3,4,5) +-- @ +-- +-- Variables and infix expressions are also allowed. Negative values can +-- be expressed by a leading @-@ with a space before but no space +-- after. +-- +-- @ +-- >>> let b x = [v| 1\/x 2 \/ x (succ x)**2 x-2 x - 3 -x |] +-- >>> b Debug.SimpleReflect.a +-- V {toVector = [1 \/ a,2 \/ a,succ a**2,a - 2,a - 3,negate a]} +-- @ +v :: QuasiQuoter +v = QuasiQuoter + { quoteExp = parseV + , quotePat = patternV + , quoteType = error "v can't be used as type" + , quoteDec = error "v can't be used as dec" + } + +parseV :: String -> Q Exp +parseV s = case parseLine s of + Right as -> + let e = pure $ ListE as + n = pure . LitT $ NumTyLit (toInteger $ length as) + in [| (V.V :: Vector.Vector a -> V.V $n a) (Vector.fromList $e) |] + Left err -> fail $ "v: " ++ err + +------------------------------------------------------------------------ +-- Stencils +------------------------------------------------------------------------ + +parseStencilLine :: String -> Either String [Maybe Exp] +parseStencilLine s = + case List.sortBy (compare `on` (length . snd)) rs of + (xs,"") : _ -> Right xs + (_ , x) : _ -> Left $ "parse error on input " ++ head (words x) + _ -> Left "no parse" + where + rs = readP_to_S (many $ mExp <* skipSpaces) s + mExp = fmap Just noAppExpression <|> skip + skip = do + Lex.Ident "_" <- Lex.lex + pure Nothing + +-- | List of expressions forming a stencil. To be either turned into +-- either a real list or an unfolded stencil. +parseStencil :: String -> Q Exp +parseStencil str = + case mapM (mapM parseStencilLine) (map lines ps) of + Left err -> fail err + Right [] -> [| mkStencil [] |] + Right [[as]] -> uncurry mkStencilE $ parse1D as + Right [ass] -> uncurry mkStencilE $ parse2D ass + Right asss -> uncurry mkStencilE $ parse3D asss + where ps = paragraphs str + +mkStencilE :: ShapeLift f => Layout f -> [Maybe Exp] -> Q Exp +mkStencilE l as = do + when (F.any even l) $ reportWarning + "stencil has an even size in some dimension, the centre element may be incorrect" + + let ixes = map (^-^ fmap (`div` 2) l) (toListOf shapeIndexes l) + -- indexes zipped with expressions, discarding 'Nothing's + xs = mapMaybe (sequenceOf _2) (zip ixes as) + + mkStencilTHBy pure xs + +-- | QuasiQuoter for producing a static stencil definition. This is a +-- versatile parser for 1D, 2D and 3D stencils. The parsing is similar +-- to 'dense' but 'stencil' also supports @_@, which means ignore this +-- element. Also, stencils should have an odd length in all dimensions +-- so there is always a center element (which is used as 'zero'). +-- +-- === __Examples__ +-- +-- - 1D stencils are of the form +-- +-- @ +-- ['stencil'| 5 -3 1 -3 5 |] :: 'Num' a => 'Stencil' 'V1' a +-- @ +-- +-- - 2D stencils are of the form +-- +-- @ +-- myStencil2 :: 'Num' a => 'Stencil' 'V2' a +-- myStencil2 = ['stencil'| +-- 0 1 0 +-- 1 0 1 +-- 0 1 0 +-- |] +-- @ +-- +-- - 3D stencils have gaps between planes. +-- +-- @ +-- myStencil3 :: 'Fractional' a => 'Stencil' 'V3' a +-- myStencil3 :: ['stencil'| +-- 1\/20 3\/10 1\/20 +-- 3\/10 1 3\/10 +-- 1\/20 3\/10 1\/20 +-- +-- 3\/10 1 3\/10 +-- 1 _ 1 +-- 3\/10 1 3\/10 +-- +-- 1\/20 3\/10 1\/20 +-- 3\/10 1 3\/10 +-- 1\/20 3\/10 1\/20 +-- |] +-- @ +-- +-- Variables can also be used +-- +-- @ +-- myStencil2' :: a -> a -> a -> 'Stencil' 'V2' a +-- myStencil2' a b c = ['stencil'| +-- c b c +-- b a b +-- c b c +-- |] +-- @ +-- +-- +stencil :: QuasiQuoter +stencil = QuasiQuoter + { quoteExp = parseStencil + , quotePat = error "stencil can't be used in pattern" + , quoteType = error "stencil can't be used in type" + , quoteDec = error "stencil can't be used in dec" + } + +-- | Construct a 'Stencil' by unrolling the list at compile time. For +-- example +-- +-- @ +-- 'ifoldr' f b $('mkStencilTH' [('V1' (-1), 5), ('V1' 0, 3), ('V1' 1, 5)]) +-- @ +-- +-- will be get turned into +-- +-- @ +-- f ('V1' (-1)) 5 (f ('V1' 0) 3 (f ('V1' 1) 5 b)) +-- @ +-- +-- at compile time. Since there are no loops and all target indexes +-- are known at compile time, this can lead to more optimisations and +-- faster execution times. This can lead to around a 2x speed up +-- compared to folding over unboxed vectors. +-- +-- @ +-- myStencil = $('mkStencilTH' (as :: [(f 'Int', a)])) :: 'Stencil' f a +-- @ +mkStencilTH :: (ShapeLift f, Lift a) => [(f Int, a)] -> Q Exp +mkStencilTH = mkStencilTHBy lift + +-- | 'mkStencilTH' with a custom 'lift' function for @a@. +mkStencilTHBy :: ShapeLift f => (a -> Q Exp) -> [(f Int, a)] -> Q Exp +mkStencilTHBy aLift as = do + -- See Note [mkName-capturing] + f <- newName "mkStencilTHBy_f" + b <- newName "mkStencilTHBy_b" + let appF (i,a) e = do + iE <- liftShape' i + aE <- aLift a + pure $ AppE (AppE (AppE (VarE f) iE) aE) e + + e <- foldrM appF (VarE b) as + pure $ AppE (ConE 'Stencil) (LamE [VarP f,VarP b] e) + +{- +~~~~ Note [mkName-capturing] + +Since 'newName' will capture any other names below it with the same +name. So if we simply used @newName "b"@, [stencil| a b c |] where +a=1; b=2; c=3 would convert @b@ to @b_a5y0@ (or w/e the top level b +is) and fail. To prevent this I've used a name that's unlikely +conflict. + +Another solution would be to use lookupValueName on all variables. +But this would either require traversing over all 'Name's in every +'Exp' (shown below) or parse in the Q monad. + +-- | Lookup and replace all names made with 'mkName' using +-- 'lookupValueName'; failing if not in scope. +replaceMkName :: Exp -> Q Exp +replaceMkName = template f where + f (Name (OccName s) NameS) = + lookupValueName s >>= \case + Just nm -> pure nm + -- Sometimes a variable may not be in scope yet because it's + -- generated in a TH splice that hasn't been run yet. + Nothing -> fail $ "Not in scope: ‘" ++ s ++ "’" + f nm = pure nm + +-} + +------------------------------------------------------------------------ +-- Parsing expressions +------------------------------------------------------------------------ + +parseLine :: String -> Either String [Exp] +parseLine s = + case List.sortBy (compare `on` (length . snd)) rs of + (xs,"") : _ -> Right xs + (_ , x) : _ -> Left $ "parse error on input " ++ head (words x) + _ -> Left "no parse" + where + rs = readP_to_S (many noAppExpression <* skipSpaces) s + +-- | Fail the parser if the next non-space is a @-@ directly followed by +-- a non-space. +closeNegateFail :: ReadP () +closeNegateFail = do + s <- look + case s of + ' ' : s' -> case dropWhile isSpace s' of + '-' : c : _ -> if isSpace c then pure () else pfail + _ -> pure () + _ -> pure () + +-- | If there is a space before but not after a @-@, it is treated as a +-- separate expression. +-- +-- @ +-- "1 2 -3 4" -> [1, 2, -3, 4] +-- "1 2 - 3 4" -> [1, -1, 4] +-- "1 2-3 4" -> [1, -1, 4] +-- "11 -3/2 -3/2 4" -> [1, -1, 4] +-- "1 -3/2 4" -> [1.0,-1.5,4.0] +-- @ +noAppExpression :: ReadP Exp +noAppExpression = do + aE <- anExpr True + + option aE $ do + closeNegateFail + i <- infixExp + bE <- noAppExpression + pure $ UInfixE aE i bE + +-- | Parse an express without any top level application. Infix functions +-- are still permitted. +-- +-- This is only a small subset of the full haskell syntax. The +-- following syntax is supported: +-- +-- - Variables/constructors: @a@, @'Just'@ etc. +-- - Numbers: @3@, @-6@, @7.8@, @1e-6@, @0x583fa@ +-- - Parenthesis/tuples: @()@ @(f a)@, @(a,b)@ +-- - Lists +-- - Strings +-- - Function application +-- - Infix operators: symbols (@+@, @/@ etc.) and blackticked (like @`mod`@) +-- +-- More advanced haskell syntax are not yet supported: +-- +-- - let bindings +-- - lambdas +-- - partial infix application (+) (1+) (+2) +-- - type signatures +-- - comments +-- +-- This could be replaced by haskell-src-meta but since I want a +-- custom parser for 'noAppExpression' it doesn't seem worth the extra +-- dependencies. +expression :: ReadP Exp +expression = do + f <- anExpr True + args <- many (anExpr False) + let aE = F.foldl AppE f args + + option aE $ do + -- if the next lex isn't a symbol, we move on to the next statement + i <- infixExp + bE <- expression + pure $ UInfixE aE i bE + +-- | Parse an infix expression. Either a symbol or a name wrapped in @`@. +infixExp :: ReadP Exp +infixExp = do + a <- Lex.lex + case a of + Lex.Symbol s -> pure $ symbol s + Lex.Punc "`" -> do + Lex.Ident x <- Lex.lex + Lex.Punc "`" <- Lex.lex + ident x + _ -> pfail + +-- Lexing -------------------------------------------------------------- + +-- | Parse a single expression. +anExpr + :: Bool -- ^ Allow a leading @-@ to mean 'negate' + -> ReadP Exp +anExpr new = do + a <- Lex.lex + case a of + Lex.Char c -> pure $ LitE (CharL c) + Lex.String s -> pure $ LitE (StringL s) + Lex.Punc s -> punc s + Lex.Ident s -> ident s + Lex.Symbol s -> if new then prefix s else pfail + Lex.Number n -> pure $ LitE (number n) + Lex.EOF -> pfail + +-- | Convert a name to an expression. +ident :: String -> ReadP Exp +ident "_" = pfail +ident s@(x:_) | isUpper x = pure $ ConE (mkName s) +ident s = pure $ VarE (mkName s) + +-- | Convert a symbol to an expression. +symbol :: String -> Exp +symbol s@(':':_) = ConE (mkName s) +symbol s = VarE (mkName s) + +-- | Parse from some punctuation. +punc :: String -> ReadP Exp +punc = \case + -- parenthesis / tuples + "(" -> do as <- expression `sepBy` comma + Lex.Punc ")" <- Lex.lex + pure $ TupE as + -- lists + "[" -> do as <- expression `sepBy` comma + Lex.Punc "]" <- Lex.lex + pure $ ListE as + _ -> pfail + +prefix :: String -> ReadP Exp +prefix "-" = do + e <- anExpr False + pure $ AppE (VarE 'negate) e +prefix _ = pfail + +comma :: ReadP () +comma = do + Lex.Punc "," <- Lex.lex + pure () + +-- | Turn a 'Number' into a literal 'Integer' if possible, otherwise +-- make a literal `Rational`. +number :: Lex.Number -> Lit +number n = + maybe (RationalL $ Lex.numberToRational n) + IntegerL + (Lex.numberToInteger n) + +------------------------------------------------------------------------ +-- Parsing patterns +------------------------------------------------------------------------ + +patternV :: String -> Q Pat +patternV s = do + case parsePattern s of + Left err -> fail err + Right pats -> do + fE <- vTuple (length pats) + pure $ ViewP fE (TupP pats) + +parsePattern :: String -> Either String [Pat] +parsePattern s = + case List.sortBy (compare `on` (length . snd)) rs of + (xs,"") : _ -> Right xs + (_ , x) : _ -> Left $ "parse error on input " ++ head (words x) + _ -> Left "no parse" + where rs = readP_to_S (many pattern <* skipSpaces) s + +pattern :: ReadP Pat +pattern = do + a <- Lex.lex + case a of + Lex.Char c -> pure $ LitP (CharL c) + Lex.String s -> pure $ LitP (StringL s) + Lex.Punc s -> puncP s + Lex.Ident n -> pure $ identP n + Lex.Symbol s -> prefixP s + Lex.Number n -> pure $ LitP (number n) + Lex.EOF -> pfail + +-- | Convert a name to an expression. +identP :: String -> Pat +identP "_" = WildP +identP s@(x:_) | isUpper x = ConP (mkName s) [] +identP s = VarP (mkName s) + +-- | Parse from some punctuation. +puncP :: String -> ReadP Pat +puncP = \case + "~" -> TildeP <$> pattern + + "(" -> do as <- pattern `sepBy` comma + Lex.Punc ")" <- Lex.lex + pure $ TupP as + + "[" -> do as <- pattern `sepBy` comma + Lex.Punc "]" <- Lex.lex + pure $ ListP as + _ -> pfail + +prefixP :: String -> ReadP Pat +prefixP "!" = do + c:_ <- look + when (isSpace c) pfail + BangP <$> pattern +prefixP "~" = TildeP <$> pattern +prefixP _ = pfail + +-- | Create an expression for converting a (V n a) to an n-tuple. +vTuple :: Int -> Q Exp +vTuple n + | n > 62 = error "max supported length is 62 for v pattern" + | otherwise = do + vN <- newName "v" + let idx i = AppE (AppE (VarE 'Vector.unsafeIndex) (VarE vN)) (intE i) + let xs = TupE $ map idx [0..n-1] + a <- newName "a" + let tup = iterate (\x -> AppT x (VarT a)) (TupleT n) !! n + typ = ForallT [PlainTV a] [] + (AppT (AppT ArrowT (AppT (AppT (ConT ''V.V) (intT n)) (VarT a))) tup) + + [| (\(V.V $(pure $ VarP vN)) -> $(pure xs)) :: $(pure typ) |] + where + intE = LitE . IntegerL . toInteger + intT = LitT . NumTyLit . toInteger + +-- Parsing specific dimensions ----------------------------------------- + +-- | Parse a 1D list. If the system is not valid, return a string +-- with error message. +parse1D :: [a] -> (V1 Int, [a]) +parse1D as = (V1 x, as) where + x = length as + +-- | Parse a 2D list of lists. If the system is not valid, returns an +-- error +parse2D :: [[a]] -> (V2 Int, [a]) +parse2D as + | Just e <- badX = error ("parse2D: " ++ errMsg e) + | otherwise = (V2 x y, F.concat $ List.transpose as) + where + x = head xs + y = length as + xs = map length as + + badX = ifind (const (/= x)) xs + errMsg (i,j) = + "row " ++ show i ++ " has " ++ show j ++ " columns but the first" + ++ " row has " ++ show x ++ " columns" + +-- | Parse a 3D list of list of lists. If the system is not valid, +-- return a string with error message. +-- +-- The element are reordered in the appropriate way for the array: +-- +-- @ +-- >>> parse3D [["abc","def","ghi"],["jkl","mno","pqr"],["stu","vwx","yz!"]] +-- ((V3 3 3 3), "ajsdmvgpybktenwhqzclufoxir!") +-- @ +-- +parse3D :: [[[a]]] -> (V3 Int, [a]) +parse3D as + | nullOf (each.each.each) as = (zero, []) + | Just e <- badX = error $ errorCol e + | Just e <- badY = error $ errorRow e + | otherwise = (V3 x y z, as') + where + z = length as + y = length (head as) + x = length (head (head as)) + + -- reorder and concatenate so it's the correct order for the array + as' = F.concatMap F.concat (List.transpose $ map List.transpose $ List.transpose as) + + -- check for inconsistencies + badY = ifind (const (/= y)) (map length as) + badX = ifindOf' (traversed <.> traversed <. to length) (const (/= x)) as + + -- error messages for inconsistent rows/columns + errorCol ((k,j),i) = + "plane " ++ show k ++ ", row " ++ show j ++ " has " ++ show i ++ + " columns" ++ ", but the first row has " ++ show x ++ " columns" + errorRow (k,j) = + "plane " ++ show k ++ " has " ++ show j ++ " rows but the first" + ++ " plane has " ++ show x ++ " rows" + +-- | Version of ifindOf which is consistent with ifind (in that it also returns the index). +ifindOf' :: IndexedGetting i (Endo (Maybe (i, a))) s a -> (i -> a -> Bool) -> s -> Maybe (i, a) +ifindOf' l p = ifoldrOf l (\i a y -> if p i a then Just (i, a) else y) Nothing +{-# INLINE ifindOf' #-} + +-- Shape lift class ---------------------------------------------------- + +-- | Class of shapes that can be 'lift'ed. +-- +-- This is to prevent orphans for the 'Lift' class. +class Shape f => ShapeLift f where + -- | 'lift' for 'Shape's. + liftShape :: Lift a => f a -> Q Exp + + -- | Polymorphic 'lift' for a 'Shape's. + liftShape' :: Lift a => f a -> Q Exp + +instance ShapeLift V1 where + liftShape (V1 x) = [| V1 x |] + liftShape' (V1 x) = [| v1 x |] + +instance ShapeLift V2 where + liftShape (V2 x y) = [| V2 x y |] + liftShape' (V2 x y) = [| v2 x y |] + +instance ShapeLift V3 where + liftShape (V3 x y z) = [| V3 x y z |] + liftShape' (V3 x y z) = [| v3 x y z |] + +instance ShapeLift V4 where + liftShape (V4 x y z w) = [| V4 x y z w |] + liftShape' (V4 x y z w) = [| v4 x y z w |] + +v1 :: (R1 f, Shape f, Num a) => a -> f a +v1 x = set _x x one +{-# INLINE [0] v1 #-} + +v2 :: (R2 f, Shape f, Num a) => a -> a -> f a +v2 x y = set _xy (V2 x y) one +{-# INLINE [0] v2 #-} + +v3 :: (R3 f, Shape f, Num a) => a -> a -> a -> f a +v3 x y z = set _xyz (V3 x y z) one +{-# INLINE [0] v3 #-} + +v4 :: (R4 f, Shape f, Num a) => a -> a -> a -> a -> f a +v4 x y z w = set _xyzw (V4 x y z w) one +{-# INLINE [0] v4 #-} + +one :: (Shape f, Num a) => f a +one = 1 <$ (zero :: Additive f => f Int) + +-- are these nessesary? +{-# RULES + "v1/V1" v1 = V1; + "v1/V2" forall a. v1 a = V2 a 1; + "v1/V3" forall a. v1 a = V3 a 1 1; + "v2/V2" v2 = V2; + "v2/V3" forall a b. v2 a b = V3 a b 1; + "v3/V3" v3 = V3; + "v4/V4" v4 = V4 + #-} + diff --git a/src/Data/Dense/Unboxed.hs b/src/Data/Dense/Unboxed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Dense/Unboxed.hs @@ -0,0 +1,666 @@ +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE RankNTypes #-} +----------------------------------------------------------------------------- +-- | +-- Module : Data.Dense.Unboxed +-- Copyright : (c) Christopher Chalmers +-- License : BSD3 +-- +-- Maintainer : Christopher Chalmers +-- Stability : provisional +-- Portability : non-portable +-- +-- Unboxed multidimensional arrays. +----------------------------------------------------------------------------- +module Data.Dense.Unboxed + ( + -- * UArray types + UArray + , Unbox + , Shape + + -- * Layout of an array + , HasLayout (..) + , Layout + + -- ** Extracting size + , extent + , size + + -- ** Folds over indexes + , indexes + , indexesFrom + , indexesBetween + + -- * Underlying vector + , vector + + -- ** Traversals + , values + , values' + , valuesBetween + + -- * Construction + + -- ** Flat arrays + , flat + , fromList + + -- ** From lists + , fromListInto + , fromListInto_ + + -- ** From vectors + , fromVectorInto + , fromVectorInto_ + + -- ** Initialisation + , replicate + , generate + , linearGenerate + + -- ** Monadic initialisation + , create + , replicateM + , generateM + , linearGenerateM + + -- * Functions on arrays + + -- ** Empty arrays + , empty + , null + + -- ** Indexing + + , (!) + , (!?) + , unsafeIndex + , linearIndex + , unsafeLinearIndex + + -- *** Monadic indexing + , indexM + , unsafeIndexM + , linearIndexM + , unsafeLinearIndexM + + -- ** Modifying arrays + + -- ** Bulk updates + , (//) + + -- ** Accumulations + , accum + + -- ** Mapping + , map + , imap + + -- * Zipping + -- ** Tuples + , zip + , zip3 + + -- ** Zip with function + , zipWith + , zipWith3 + , izipWith + , izipWith3 + + -- ** Slices + + -- *** Matrix + , ixRow + , rows + , ixColumn + , columns + + -- *** 3D + , ixPlane + , planes + , flattenPlane + + -- *** Ordinals + , unsafeOrdinals + + -- * Mutable + , UMArray + + , thaw + , freeze + , unsafeThaw + , unsafeFreeze + + -- * Delayed + + , G.Delayed + + -- ** Generating delayed + + , delayed + , seqDelayed + , delay + , manifest + , seqManifest + , G.genDelayed + , G.indexDelayed + , affirm + , seqAffirm + + -- * Focused + + , G.Focused + + -- ** Generating focused + + , G.focusOn + , G.unfocus + , G.unfocused + , G.extendFocus + + -- ** Focus location + , G.locale + , G.shiftFocus + + ) where + +import Control.Lens hiding (imap) +import Control.Monad.Primitive +import Control.Monad.ST +import qualified Data.Foldable as F +import Data.Vector.Unboxed (Unbox, Vector) +import Linear hiding (vector) + +import Prelude hiding (map, null, replicate, zip, + zip3, zipWith, zipWith3) + +import Data.Dense.Generic (UArray) +import qualified Data.Dense.Generic as G +import Data.Dense.Index +import Data.Dense.Mutable (UMArray) + +-- Lenses -------------------------------------------------------------- + +-- | Same as 'values' but restrictive in the vector type. +values :: (Shape f, Unbox a, Unbox b) + => IndexedTraversal (f Int) (UArray f a) (UArray f b) a b +values = G.values' +{-# INLINE values #-} + +-- | Same as 'values' but restrictive in the vector type. +values' :: (Shape f, Unbox a, Unbox b) + => IndexedTraversal (f Int) (UArray f a) (UArray f b) a b +values' = G.values' +{-# INLINE values' #-} + +-- | Same as 'values' but restrictive in the vector type. +valuesBetween + :: (Shape f, Unbox a) + => f Int + -> f Int + -> IndexedTraversal' (f Int) (UArray f a) a +valuesBetween = G.valuesBetween +{-# INLINE valuesBetween #-} + +-- | 1D arrays are just vectors. You are free to change the length of +-- the vector when going 'over' this 'Iso' (unlike 'linear'). +-- +-- Note that 'V1' arrays are an instance of 'Vector' so you can use +-- any of the functions in 'Data.Vector.Generic' on them without +-- needing to convert. +flat :: Unbox b => Iso (UArray V1 a) (UArray V1 b) (Vector a) (Vector b) +flat = G.flat +{-# INLINE flat #-} + +-- | Indexed lens over the underlying vector of an array. The index is +-- the 'extent' of the array. You must _not_ change the length of the +-- vector, otherwise an error will be thrown (even for 'V1' layouts, +-- use 'flat' for 'V1'). +vector :: (Unbox a, Unbox b) => IndexedLens (Layout f) (UArray f a) (UArray f b) (Vector a) (Vector b) +vector = G.vector +{-# INLINE vector #-} + +-- Constructing vectors ------------------------------------------------ + +-- | Contruct a flat array from a list. (This is just 'G.fromList' from +-- 'Data.Vector.Generic'.) +fromList :: Unbox a => [a] -> UArray V1 a +fromList = G.fromList +{-# INLINE fromList #-} + +-- | O(n) Convert the first @n@ elements of a list to an UArrayith the +-- given shape. Returns 'Nothing' if there are not enough elements in +-- the list. +fromListInto :: (Shape f, Unbox a) => Layout f -> [a] -> Maybe (UArray f a) +fromListInto = G.fromListInto +{-# INLINE fromListInto #-} + +-- | O(n) Convert the first @n@ elements of a list to an UArrayith the +-- given shape. Throw an error if the list is not long enough. +fromListInto_ :: (Shape f, Unbox a) => Layout f -> [a] -> UArray f a +fromListInto_ = G.fromListInto_ +{-# INLINE fromListInto_ #-} + +-- | Create an array from a 'vector' and a 'layout'. Return 'Nothing' if +-- the vector is not the right shape. +fromVectorInto :: (Shape f, Unbox a) => Layout f -> Vector a -> Maybe (UArray f a) +fromVectorInto = G.fromVectorInto +{-# INLINE fromVectorInto #-} + +-- | Create an array from a 'vector' and a 'layout'. Throws an error if +-- the vector is not the right shape. +fromVectorInto_ :: (Shape f, Unbox a) => Layout f -> Vector a -> UArray f a +fromVectorInto_ = G.fromVectorInto_ +{-# INLINE fromVectorInto_ #-} + +-- | The empty 'UArray' with a 'zero' shape. +empty :: (Unbox a, Additive f) => UArray f a +empty = G.empty +{-# INLINE empty #-} + +-- | Test is if the array is 'empty'. +null :: F.Foldable f => UArray f a -> Bool +null = G.null +{-# INLINE null #-} + +-- Indexing ------------------------------------------------------------ + +-- | Index an element of an array. Throws 'IndexOutOfBounds' if the +-- index is out of bounds. +(!) :: (Shape f, Unbox a) => UArray f a -> f Int -> a +(!) = (G.!) +{-# INLINE (!) #-} + +-- | Safe index of an element. +(!?) :: (Shape f, Unbox a) => UArray f a -> f Int -> Maybe a +(!?) = (G.!?) +{-# INLINE (!?) #-} + +-- | Index an element of an array without bounds checking. +unsafeIndex :: (Shape f, Unbox a) => UArray f a -> f Int -> a +unsafeIndex = G.unsafeIndex +{-# INLINE unsafeIndex #-} + +-- | Index an element of an array while ignoring its shape. +linearIndex :: Unbox a => UArray f a -> Int -> a +linearIndex = G.linearIndex +{-# INLINE linearIndex #-} + +-- | Index an element of an array while ignoring its shape, without +-- bounds checking. +unsafeLinearIndex :: Unbox a => UArray f a -> Int -> a +unsafeLinearIndex = G.unsafeLinearIndex +{-# INLINE unsafeLinearIndex #-} + +-- Monadic indexing ---------------------------------------------------- + +-- | /O(1)/ Indexing in a monad. +-- +-- The monad allows operations to be strict in the vector when necessary. +-- Suppose vector copying is implemented like this: +-- +-- > copy mv v = ... write mv i (v ! i) ... +-- +-- For lazy vectors, @v ! i@ would not be evaluated which means that +-- @mv@ would unnecessarily retain a reference to @v@ in each element +-- written. +-- +-- With 'indexM', copying can be implemented like this instead: +-- +-- > copy mv v = ... do +-- > x <- indexM v i +-- > write mv i x +-- +-- Here, no references to @v@ are retained because indexing (but /not/ +-- the elements) is evaluated eagerly. +-- +-- Throws an error if the index is out of range. +indexM :: (Shape f, Unbox a, Monad m) => UArray f a -> f Int -> m a +indexM = G.indexM +{-# INLINE indexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeIndexM :: (Shape f, Unbox a, Monad m) => UArray f a -> f Int -> m a +unsafeIndexM = G.unsafeIndexM +{-# INLINE unsafeIndexM #-} + +-- | /O(1)/ Indexing in a monad. Throws an error if the index is out of +-- range. +linearIndexM :: (Shape f, Unbox a, Monad m) => UArray f a -> Int -> m a +linearIndexM = G.linearIndexM +{-# INLINE linearIndexM #-} + +-- | /O(1)/ Indexing in a monad without bounds checks. See 'indexM' for an +-- explanation of why this is useful. +unsafeLinearIndexM :: (Unbox a, Monad m) => UArray f a -> Int -> m a +unsafeLinearIndexM = G.unsafeLinearIndexM +{-# INLINE unsafeLinearIndexM #-} + +-- Initialisation ------------------------------------------------------ + +-- | Execute the monadic action and freeze the resulting array. +create :: Unbox a + => (forall s. ST s (UMArray f s a)) -> UArray f a +create m = m `seq` runST (m >>= G.unsafeFreeze) +{-# INLINE create #-} + +-- | O(n) UArray of the given shape with the same value in each position. +replicate :: (Shape f, Unbox a) => f Int -> a -> UArray f a +replicate = G.replicate +{-# INLINE replicate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +linearGenerate :: (Shape f, Unbox a) => Layout f -> (Int -> a) -> UArray f a +linearGenerate = G.linearGenerate +{-# INLINE linearGenerate #-} + +-- | O(n) Construct an array of the given shape by applying the +-- function to each index. +generate :: (Shape f, Unbox a) => Layout f -> (f Int -> a) -> UArray f a +generate = G.generate +{-# INLINE generate #-} + +-- Monadic initialisation ---------------------------------------------- + +-- | O(n) Construct an array of the given shape by filling each position +-- with the monadic value. +replicateM :: (Monad m, Shape f, Unbox a) => Layout f -> m a -> m (UArray f a) +replicateM = G.replicateM +{-# INLINE replicateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +generateM :: (Monad m, Shape f, Unbox a) => Layout f -> (f Int -> m a) -> m (UArray f a) +generateM = G.generateM +{-# INLINE generateM #-} + +-- | O(n) Construct an array of the given shape by applying the monadic +-- function to each index. +linearGenerateM :: (Monad m, Shape f, Unbox a) => Layout f -> (Int -> m a) -> m (UArray f a) +linearGenerateM = G.linearGenerateM +{-# INLINE linearGenerateM #-} + +-- Modifying ----------------------------------------------------------- + +-- | /O(n)/ Map a function over an array +map :: (Unbox a, Unbox b) => (a -> b) -> UArray f a -> UArray f b +map = G.map +{-# INLINE map #-} + +-- | /O(n)/ Apply a function to every element of a vector and its index +imap :: (Shape f, Unbox a, Unbox b) => (f Int -> a -> b) -> UArray f a -> UArray f b +imap = G.imap +{-# INLINE imap #-} + +-- Bulk updates -------------------------------------------------------- + + +-- | For each pair (i,a) from the list, replace the array element at +-- position i by a. +(//) :: (Unbox a, Shape f) => UArray f a -> [(f Int, a)] -> UArray f a +(//) = (G.//) +{-# INLINE (//) #-} + +-- Accumilation -------------------------------------------------------- + +-- | /O(m+n)/ For each pair @(i,b)@ from the list, replace the array element +-- @a@ at position @i@ by @f a b@. +-- +accum :: (Shape f, Unbox a) + => (a -> b -> a) -- ^ accumulating function @f@ + -> UArray f a -- ^ initial array + -> [(f Int, b)] -- ^ list of index/value pairs (of length @n@) + -> UArray f a +accum = G.accum +{-# INLINE accum #-} + +------------------------------------------------------------------------ +-- Zipping +------------------------------------------------------------------------ + +-- Tuple zip ----------------------------------------------------------- + +-- | Zip two arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip :: (Shape f, Unbox a, Unbox b) + => UArray f a + -> UArray f b + -> UArray f (a,b) +zip = G.zip + +-- | Zip three arrays element wise. If the array's don't have the same +-- shape, the new array with be the intersection of the two shapes. +zip3 :: (Shape f, Unbox a, Unbox b, Unbox c) + => UArray f a + -> UArray f b + -> UArray f c + -> UArray f (a,b,c) +zip3 = G.zip3 + +-- Zip with function --------------------------------------------------- + +-- | Zip two arrays using the given function. If the array's don't have +-- the same shape, the new array with be the intersection of the two +-- shapes. +zipWith :: (Shape f, Unbox a, Unbox b, Unbox c) + => (a -> b -> c) + -> UArray f a + -> UArray f b + -> UArray f c +zipWith = G.zipWith +{-# INLINE zipWith #-} + +-- | Zip three arrays using the given function. If the array's don't +-- have the same shape, the new array with be the intersection of the +-- two shapes. +zipWith3 :: (Shape f, Unbox a, Unbox b, Unbox c, Unbox d) + => (a -> b -> c -> d) + -> UArray f a + -> UArray f b + -> UArray f c + -> UArray f d +zipWith3 = G.zipWith3 +{-# INLINE zipWith3 #-} + +-- Indexed zipping ----------------------------------------------------- + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith :: (Shape f, Unbox a, Unbox b, Unbox c) + => (f Int -> a -> b -> c) + -> UArray f a + -> UArray f b + -> UArray f c +izipWith = G.izipWith +{-# INLINE izipWith #-} + +-- | Zip two arrays using the given function with access to the index. +-- If the array's don't have the same shape, the new array with be the +-- intersection of the two shapes. +izipWith3 :: (Shape f, Unbox a, Unbox b, Unbox c, Unbox d) + => (f Int -> a -> b -> c -> d) + -> UArray f a + -> UArray f b + -> UArray f c + -> UArray f d +izipWith3 = G.izipWith3 +{-# INLINE izipWith3 #-} + +------------------------------------------------------------------------ +-- Slices +------------------------------------------------------------------------ + +-- $setup +-- >>> import qualified Data.Vector.Unboxed as V +-- >>> let m = fromListInto_ (V2 2 3) [1..] :: UArray V2 Int + +-- | Indexed traversal over the rows of a matrix. Each row is an +-- efficient 'Data.Vector.Generic.slice' of the original vector. +-- +-- >>> traverseOf_ rows print m +-- [1,2,3] +-- [4,5,6] +rows :: (Unbox a, Unbox b) + => IndexedTraversal Int (UArray V2 a) (UArray V2 b) (Vector a) (Vector b) +rows = G.rows +{-# INLINE rows #-} + +-- | Affine traversal over a single row in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 . each +~ 2 +-- [1,2,3] +-- [6,7,8] +-- +-- The row vector should remain the same size to satisfy traversal +-- laws but give reasonable behaviour if the size differs: +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0,1] +-- [1,2,3] +-- [0,1,6] +-- +-- >>> traverseOf_ rows print $ m & ixRow 1 .~ V.fromList [0..100] +-- [1,2,3] +-- [0,1,2] +ixRow :: Unbox a => Int -> IndexedTraversal' Int (UArray V2 a) (Vector a) +ixRow = G.ixRow +{-# INLINE ixRow #-} + +-- | Indexed traversal over the columns of a matrix. Unlike 'rows', each +-- column is a new separate vector. +-- +-- >>> traverseOf_ columns print m +-- [1,4] +-- [2,5] +-- [3,6] +-- +-- >>> traverseOf_ rows print $ m & columns . indices odd . each .~ 0 +-- [1,0,3] +-- [4,0,6] +-- +-- The vectors should be the same size to be a valid traversal. If the +-- vectors are different sizes, the number of rows in the new array +-- will be the length of the smallest vector. +columns :: (Unbox a, Unbox b) + => IndexedTraversal Int (UArray V2 a) (UArray V2 b) (Vector a) (Vector b) +columns = G.columns +{-# INLINE columns #-} + +-- | Affine traversal over a single column in a matrix. +-- +-- >>> traverseOf_ rows print $ m & ixColumn 2 . each *~ 10 +-- [1,2,30] +-- [4,5,60] +ixColumn :: Unbox a => Int -> IndexedTraversal' Int (UArray V2 a) (Vector a) +ixColumn = G.ixColumn +{-# INLINE ixColumn #-} + +-- | Traversal over a single plane of a 3D array given a lens onto that +-- plane (like '_xy', '_yz', '_zx'). +ixPlane :: Unbox a + => ALens' (V3 Int) (V2 Int) + -> Int + -> IndexedTraversal' Int (UArray V3 a) (UArray V2 a) +ixPlane = G.ixPlane +{-# INLINE ixPlane #-} + +-- | Traversal over all planes of 3D array given a lens onto that plane +-- (like '_xy', '_yz', '_zx'). +planes :: (Unbox a, Unbox b) + => ALens' (V3 Int) (V2 Int) + -> IndexedTraversal Int (UArray V3 a) (UArray V3 b) (UArray V2 a) (UArray V2 b) +planes = G.planes +{-# INLINE planes #-} + +-- | Flatten a plane by reducing a vector in the third dimension to a +-- single value. +flattenPlane :: (Unbox a, Unbox b) + => ALens' (V3 Int) (V2 Int) + -> (Vector a -> b) + -> UArray V3 a + -> UArray V2 b +flattenPlane = G.flattenPlane +{-# INLINE flattenPlane #-} + +-- Ordinals ------------------------------------------------------------ + +-- | This 'Traversal' should not have any duplicates in the list of +-- indices. +unsafeOrdinals :: (Unbox a, Shape f) => [f Int] -> IndexedTraversal' (f Int) (UArray f a) a +unsafeOrdinals = G.unsafeOrdinals +{-# INLINE [0] unsafeOrdinals #-} + +-- Mutable ------------------------------------------------------------- + +-- | O(n) Yield a mutable copy of the immutable vector. +freeze :: (PrimMonad m, Unbox a) + => UMArray f (PrimState m) a -> m (UArray f a) +freeze = G.freeze +{-# INLINE freeze #-} + +-- | O(n) Yield an immutable copy of the mutable array. +thaw :: (PrimMonad m, Unbox a) + => UArray f a -> m (UMArray f (PrimState m) a) +thaw = G.thaw +{-# INLINE thaw #-} + +-- | O(1) Unsafe convert a mutable array to an immutable one without +-- copying. The mutable array may not be used after this operation. +unsafeFreeze :: (PrimMonad m, Unbox a) + => UMArray f (PrimState m) a -> m (UArray f a) +unsafeFreeze = G.unsafeFreeze +{-# INLINE unsafeFreeze #-} + +-- | O(1) Unsafely convert an immutable array to a mutable one without +-- copying. The immutable array may not be used after this operation. +unsafeThaw :: (PrimMonad m, Unbox a) + => UArray f a -> m (UMArray f (PrimState m) a) +unsafeThaw = G.unsafeThaw +{-# INLINE unsafeThaw #-} + +------------------------------------------------------------------------ +-- Delayed +------------------------------------------------------------------------ + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in parallel. +delayed :: (Unbox a, Unbox b, Shape f, Shape k) + => Iso (UArray f a) (UArray k b) (G.Delayed f a) (G.Delayed k b) +delayed = G.delayed +{-# INLINE delayed #-} + +-- | Isomorphism between an array and its delayed representation. +-- Conversion to the array is done in sequence. +seqDelayed :: (Unbox a, Unbox b, Shape f, Shape k) + => Iso (UArray f a) (UArray k b) (G.Delayed f a) (G.Delayed k b) +seqDelayed = G.seqDelayed +{-# INLINE seqDelayed #-} + +-- | Turn a material array into a delayed one with the same shape. +delay :: (Unbox a, Shape f) => UArray f a -> G.Delayed f a +delay = G.delay +{-# INLINE delay #-} + +-- | Parallel manifestation of a delayed array into a material one. +manifest :: (Unbox a, Shape f) => G.Delayed f a -> UArray f a +manifest = G.manifest +{-# INLINE manifest #-} + +-- | Sequential manifestation of a delayed array. +seqManifest :: (Unbox a, Shape f) => G.Delayed f a -> UArray f a +seqManifest = G.seqManifest +{-# INLINE seqManifest #-} + +-- | 'manifest' an array to a 'UArray' and delay again. +affirm :: (Shape f, Unbox a) => G.Delayed f a -> G.Delayed f a +affirm = delay . manifest +{-# INLINE affirm #-} + +-- | 'seqManifest' an array to a 'UArray' and delay again. +seqAffirm :: (Shape f, Unbox a) => G.Delayed f a -> G.Delayed f a +seqAffirm = delay . seqManifest +{-# INLINE seqAffirm #-} + diff --git a/tests/doctest.hs b/tests/doctest.hs new file mode 100644 --- /dev/null +++ b/tests/doctest.hs @@ -0,0 +1,11 @@ +import Test.DocTest + +main = doctest + [ "src/Data/Dense/Index.hs" + , "src/Data/Dense/Mutable.hs" + , "src/Data/Dense/Unboxed.hs" + , "src/Data/Dense/Storable.hs" + , "src/Data/Dense/Boxed.hs" + , "src/Data/Dense/Generic.hs" + , "src/Data/Dense/Base.hs" + ]