diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +## Changes + +#### 0.1.0.0 + +- Initial version diff --git a/Data/Array/Convert.hs b/Data/Array/Convert.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Convert.hs @@ -0,0 +1,191 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +module Data.Array.Convert(Convert(..)) where +import Data.Proxy +import qualified Data.Vector as V +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Storable as VS +import GHC.TypeLits(KnownNat) + +import qualified Data.Array.Internal as I +import qualified Data.Array.Internal.Dynamic as D +import qualified Data.Array.Internal.DynamicG as DG +import qualified Data.Array.Internal.DynamicS as DS +import qualified Data.Array.Internal.DynamicU as DU +import qualified Data.Array.Internal.Ranked as R +import qualified Data.Array.Internal.RankedG as RG +import qualified Data.Array.Internal.RankedS as RS +import qualified Data.Array.Internal.RankedU as RU +import qualified Data.Array.Internal.Shaped as S +import qualified Data.Array.Internal.ShapedG as SG +import qualified Data.Array.Internal.ShapedS as SS +import qualified Data.Array.Internal.ShapedU as SU +import Data.Array.Internal.Shape(Shape(..)) + +class Convert a b where + -- | Convert between two array types. + -- If both arrays have the same flavor of boxing this is a O(1) + -- operation. Conversion between different boxing is a O(n) operation. + convert :: a -> b + -- | Convert between two array types if possible, or return error message. + convertE :: a -> Either String b + {-# INLINE convert #-} + convert = either error id . convertE + {-# INLINE convertE #-} + convertE = Right . convert + {-# MINIMAL convert | convertE #-} + +-- We write these instances like "instance (a ~ b) => Convert (X a) (Y b)" so +-- that instance resolution will commit to them as soon as it knows the array +-- types, and then introduce an equality constraint on the element types. This +-- way, if you call 'convert' on an array, type inference can propagate the +-- element type across the conversion. If we wrote them like "instance +-- Convert (X a) (Y a)" instead, they'd not be selected until the element types +-- of both arrays are determined to be the same by other means, which would +-- lead to unnecessary ambiguity errors. + +----- + +instance (a ~ b, DU.Unbox a) => Convert (D.Array a) (DU.Array b) where + convert (D.A (DG.A sh (I.T s o v))) = DU.A (DG.A sh (I.T s o (V.convert v))) + +instance (a ~ b, DU.Unbox a) => Convert (DU.Array a) (D.Array b) where + convert (DU.A (DG.A sh (I.T s o v))) = D.A (DG.A sh (I.T s o (V.convert v))) + +----- + +instance (a ~ b, DS.Unbox a) => Convert (D.Array a) (DS.Array b) where + convert (D.A (DG.A sh (I.T s o v))) = DS.A (DG.A sh (I.T s o (V.convert v))) + +instance (a ~ b, DS.Unbox a) => Convert (DS.Array a) (D.Array b) where + convert (DS.A (DG.A sh (I.T s o v))) = D.A (DG.A sh (I.T s o (V.convert v))) + +----- + +-- As above with the element types, we arrange to select the Convert instance +-- before the ranks are known to be equal, then constrain them to be equal. + +instance (a ~ b, n ~ m, RU.Unbox a) => Convert (R.Array n a) (RU.Array m b) where + convert (R.A (RG.A sh (I.T s o v))) = RU.A (RG.A sh (I.T s o (V.convert v))) + +instance (a ~ b, n ~ m, RU.Unbox a) => Convert (RU.Array n a) (R.Array m b) where + convert (RU.A (RG.A sh (I.T s o v))) = R.A (RG.A sh (I.T s o (V.convert v))) + +----- + +instance (a ~ b, n ~ m, RS.Unbox a) => Convert (R.Array n a) (RS.Array m b) where + convert (R.A (RG.A sh (I.T s o v))) = RS.A (RG.A sh (I.T s o (V.convert v))) + +instance (a ~ b, n ~ m, RS.Unbox a) => Convert (RS.Array n a) (R.Array m b) where + convert (RS.A (RG.A sh (I.T s o v))) = R.A (RG.A sh (I.T s o (V.convert v))) + +----- + +instance (a ~ b, n ~ m, SU.Unbox a) => Convert (S.Array n a) (SU.Array m b) where + convert (S.A (SG.A (I.T s o v))) = SU.A (SG.A (I.T s o (V.convert v))) + +instance (a ~ b, n ~ m, SU.Unbox a) => Convert (SU.Array n a) (S.Array m b) where + convert (SU.A (SG.A (I.T s o v))) = S.A (SG.A (I.T s o (V.convert v))) + +----- + +instance (a ~ b, n ~ m, SS.Unbox a) => Convert (S.Array n a) (SS.Array m b) where + convert (S.A (SG.A (I.T s o v))) = SS.A (SG.A (I.T s o (V.convert v))) + +instance (a ~ b, n ~ m, SS.Unbox a) => Convert (SS.Array n a) (S.Array m b) where + convert (SS.A (SG.A (I.T s o v))) = S.A (SG.A (I.T s o (V.convert v))) + +----- + +instance (a ~ b) => Convert (R.Array n a) (D.Array b) where + convert (R.A (RG.A sh t)) = D.A (DG.A sh t) + +instance (a ~ b, KnownNat n) => Convert (D.Array a) (R.Array n b) where + convertE (D.A (DG.A sh t)) | length sh /= I.valueOf @n = Left "rank mismatch" + | otherwise = Right $ R.A (RG.A sh t) + +instance (a ~ b, S.Shape sh) => Convert (S.Array sh a) (D.Array b) where + convert (S.A a@(SG.A t)) = D.A (DG.A (SG.shapeL a) t) + +instance (a ~ b, S.Rank sh ~ n, S.Shape sh) => Convert (S.Array sh a) (R.Array n b) where + convert (S.A a@(SG.A t)) = R.A (RG.A (SG.shapeL a) t) + +instance (a ~ b, S.Shape sh) => Convert (D.Array a) (S.Array sh b) where + convertE (D.A (DG.A sh t)) | sh == shapeP (Proxy :: Proxy sh) = Right $ S.A (SG.A t) + convertE _ = Left "shape mismatch" + +instance (a ~ b, S.Rank sh ~ n, S.Shape sh) => Convert (R.Array n a) (S.Array sh b) where + convertE (R.A (RG.A sh t)) | sh == shapeP (Proxy :: Proxy sh) = Right $ S.A (SG.A t) + convertE _ = Left "shape mismatch" + +------ + +instance (a ~ b, KnownNat n) => Convert (DS.Array a) (RS.Array n b) where + convertE (DS.A (DG.A sh t)) | length sh /= I.valueOf @n = Left "rank mismatch" + | otherwise = Right $ RS.A (RG.A sh t) + +instance (a ~ b, SS.Rank sh ~ n, SS.Shape sh) => Convert (RS.Array n a) (SS.Array sh b) where + convertE (RS.A (RG.A sh t)) | sh == shapeP (Proxy :: Proxy sh) = Right $ SS.A (SG.A t) + convertE _ = Left "shape mismatch" + +instance (a ~ b, SS.Shape sh) => Convert (DS.Array a) (SS.Array sh b) where + convertE (DS.A (DG.A sh t)) | sh == sh' = Right $ SS.A (SG.A t) + | otherwise = Left $ "shape mismatch: " ++ show (sh, sh') + where sh' = shapeP (Proxy :: Proxy sh) + +instance (a ~ b, SS.Shape sh) => Convert (SS.Array sh a) (DS.Array b) where + convert (SS.A a@(SG.A t)) = DS.A (DG.A (SG.shapeL a) t) + +------ + +instance (a ~ b) => Convert (D.Array a) (DG.Array V.Vector b) where + convert (D.A a) = a + +instance (a ~ b) => Convert (DU.Array a) (DG.Array VU.Vector b) where + convert (DU.A a) = a + +instance (a ~ b) => Convert (DS.Array a) (DG.Array VS.Vector b) where + convert (DS.A a) = a + +------ + +instance (a ~ b, n ~ m) => Convert (R.Array n a) (RG.Array m V.Vector b) where + convert (R.A a) = a + +instance (a ~ b, n ~ m) => Convert (RU.Array n a) (RG.Array m VU.Vector b) where + convert (RU.A a) = a + +instance (a ~ b, n ~ m) => Convert (RS.Array n a) (RG.Array m VS.Vector b) where + convert (RS.A a) = a + +------ + +instance (a ~ b, s ~ t) => Convert (S.Array s a) (SG.Array t V.Vector b) where + convert (S.A a) = a + +instance (a ~ b, s ~ t) => Convert (SU.Array s a) (SG.Array t VU.Vector b) where + convert (SU.A a) = a + +instance (a ~ b, s ~ t) => Convert (SS.Array s a) (SG.Array t VS.Vector b) where + convert (SS.A a) = a + +------ + +-- XXX need more conversions. diff --git a/Data/Array/Dynamic.hs b/Data/Array/Dynamic.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Dynamic.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.Dynamic(module A) where +import Data.Array.Internal.Dynamic as A hiding(Array(..)) +import Data.Array.Internal.Dynamic as A(Array) diff --git a/Data/Array/DynamicG.hs b/Data/Array/DynamicG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/DynamicG.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.DynamicG(module A) where +import Data.Array.Internal.DynamicG as A hiding(Array(..)) +import Data.Array.Internal.DynamicG as A(Array) diff --git a/Data/Array/DynamicS.hs b/Data/Array/DynamicS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/DynamicS.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.DynamicS(module A) where +import Data.Array.Internal.DynamicS as A hiding(Array(..)) +import Data.Array.Internal.DynamicS as A(Array) diff --git a/Data/Array/DynamicU.hs b/Data/Array/DynamicU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/DynamicU.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.DynamicU(module A) where +import Data.Array.Internal.DynamicU as A hiding(Array(..)) +import Data.Array.Internal.DynamicU as A(Array) diff --git a/Data/Array/Internal.hs b/Data/Array/Internal.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal.hs @@ -0,0 +1,529 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE QuantifiedConstraints #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +{-# LANGUAGE UndecidableSuperClasses #-} +module Data.Array.Internal(module Data.Array.Internal) where +import Control.DeepSeq +import Data.Data(Data) +import qualified Data.DList as DL +import Data.Kind (Type) +import Data.List(foldl', zipWith4, zipWith5, sortBy, sortOn) +import Data.Proxy +import GHC.Exts(Constraint, build) +import GHC.Generics(Generic) +import GHC.TypeLits(KnownNat, natVal) +import Text.PrettyPrint +import Text.PrettyPrint.HughesPJClass + +{- HLINT ignore "Reduce duplication" -} + +-- The underlying storage of values must be an instance of Vector. +-- For some types, like unboxed vectors, we require an extra +-- constraint on the elements, which VecElem allows you to express. +-- For vector types that don't need the constraint it can be set +-- to some dummy class. +class Vector v where + type VecElem v :: Type -> Constraint + vIndex :: (VecElem v a) => v a -> Int -> a + vLength :: (VecElem v a) => v a -> Int + vToList :: (VecElem v a) => v a -> [a] + vFromList :: (VecElem v a) => [a] -> v a + vSingleton:: (VecElem v a) => a -> v a + vReplicate:: (VecElem v a) => Int -> a -> v a + vMap :: (VecElem v a, VecElem v b) => (a -> b) -> v a -> v b + vZipWith :: (VecElem v a, VecElem v b, VecElem v c) => (a -> b -> c) -> v a -> v b -> v c + vZipWith3 :: (VecElem v a, VecElem v b, VecElem v c, VecElem v d) => (a -> b -> c -> d) -> v a -> v b -> v c -> v d + vZipWith4 :: (VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e) => (a -> b -> c -> d -> e) -> v a -> v b -> v c -> v d -> v e + vZipWith5 :: (VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e, VecElem v f) => (a -> b -> c -> d -> e -> f) -> v a -> v b -> v c -> v d -> v e -> v f + vAppend :: (VecElem v a) => v a -> v a -> v a + vConcat :: (VecElem v a) => [v a] -> v a + vFold :: (VecElem v a) => (a -> a -> a) -> a -> v a -> a + vSlice :: (VecElem v a) => Int -> Int -> v a -> v a + vSum :: (VecElem v a, Num a) => v a -> a + vProduct :: (VecElem v a, Num a) => v a -> a + vMaximum :: (VecElem v a, Ord a) => v a -> a + vMinimum :: (VecElem v a, Ord a) => v a -> a + vUpdate :: (VecElem v a) => v a -> [(Int, a)] -> v a + vGenerate :: (VecElem v a) => Int -> (Int -> a) -> v a + vAll :: (VecElem v a) => (a -> Bool) -> v a -> Bool + vAny :: (VecElem v a) => (a -> Bool) -> v a -> Bool + +class None a +instance None a + +-- This instance is not used anywheer. It serves more as a reference semantics. +instance Vector [] where + type VecElem [] = None + vIndex = (!!) + vLength = length + vToList = id + vFromList = id + vSingleton = pure + vReplicate = replicate + vMap = map + vZipWith = zipWith + vZipWith3 = zipWith3 + vZipWith4 = zipWith4 + vZipWith5 = zipWith5 + vAppend = (++) + vConcat = concat + vFold = foldl' + vSlice o n = take n . drop o + vSum = sum + vProduct = product + vMaximum = maximum + vMinimum = minimum + vUpdate xs us = loop xs (sortOn fst us) 0 + where + loop [] [] _ = [] + loop [] (_:_) _ = error "vUpdate: out of bounds" + loop as [] _ = as + loop (a:as) ias@((i,a'):ias') n = + case compare i n of + LT -> error "vUpdate: bad index" + EQ -> a' : loop as ias' (n+1) + GT -> a : loop as ias (n+1) + vGenerate n f = map f [0 .. n-1] + vAll = all + vAny = any + +prettyShowL :: (Pretty a) => PrettyLevel -> a -> String +prettyShowL l = render . pPrintPrec l 0 + +-- We expect all N to be non-negative, but we use Int for convenience. +type N = Int + +-- | The type /T/ is the internal type of arrays. In general, +-- operations on /T/ do no sanity checking as that should be done +-- at the point of call. +-- +-- To avoid manipulating the data the indexing into the vector containing +-- the data is somewhat complex. To find where item /i/ of the outermost +-- dimension starts you calculate vector index @offset + i*strides[0]@. +-- To find where item /i,j/ of the two outermost dimensions is you +-- calculate vector index @offset + i*strides[0] + j*strides[1]@, etc. +data T v a = T + { strides :: [N] -- length is tensor rank + , offset :: !N -- offset into vector of values + , values :: !(v a) -- actual values + } + deriving (Show, Generic, Data) + +instance NFData (v a) => NFData (T v a) + +type ShapeL = [Int] + +badShape :: ShapeL -> Bool +badShape = any (< 0) + +-- When shapes match, we can be efficient and use loop-fused comparisons instead +-- of materializing a list. +equalT :: (Vector v, VecElem v a, Eq a, Eq (v a)) + => ShapeL -> T v a -> T v a -> Bool +equalT s x y | strides x == strides y + && offset x == offset y + && values x == values y = True + | otherwise = toVectorT s x == toVectorT s y + +-- Note this assumes the shape is the same for both Vectors. +compareT :: (Vector v, VecElem v a, Ord a, Ord (v a)) + => ShapeL -> T v a -> T v a -> Ordering +compareT s x y = compare (toVectorT s x) (toVectorT s y) + +-- Given the dimensions, return the stride in the underlying vector +-- for each dimension. The first element of the list is the total length. +{-# INLINE getStridesT #-} +getStridesT :: ShapeL -> [N] +getStridesT = scanr (*) 1 + +-- Convert an array to a list by indexing through all the elements. +-- The first argument is the array shape. +-- XXX Copy special cases from Tensor. +{-# INLINE toListT #-} +toListT :: (Vector v, VecElem v a) => ShapeL -> T v a -> [a] +toListT sh a@(T ss0 o0 v) + | isCanonicalT (getStridesT sh) a = vToList v + | otherwise = build $ \cons nil -> + -- TODO: because unScalarT uses vIndex, this has unnecessary bounds + -- checks. We should expose an unchecked indexing function in the Vector + -- class, add top-level bounds checks to cover the full range we'll + -- access, and then do all accesses with the unchecked version. + let go [] ss o rest = cons (unScalarT (T ss o v)) rest + go (n:ns) ss o rest = foldr + (\i -> case indexT (T ss o v) i of T ss' o' _ -> go ns ss' o') + rest + [0..n-1] + in go sh ss0 o0 nil + +-- | Check if the strides are canonical, i.e., if the vector have the natural layout. +-- XXX Copy special cases from Tensor. +{-# INLINE isCanonicalT #-} +isCanonicalT :: (Vector v, VecElem v a) => [N] -> T v a -> Bool +isCanonicalT (n:ss') (T ss o v) = + o == 0 && -- Vector offset is 0 + ss == ss' && -- All strides are normal + vLength v == n -- The vector is the right size +isCanonicalT _ _ = error "impossible" + +-- Convert a value to a scalar array. +{-# INLINE scalarT #-} +scalarT :: (Vector v, VecElem v a) => a -> T v a +scalarT = T [] 0 . vSingleton + +-- Convert a scalar array to the actual value. +{-# INLINE unScalarT #-} +unScalarT :: (Vector v, VecElem v a) => T v a -> a +unScalarT (T _ o v) = vIndex v o + +-- Make a constant array. +{-# INLINE constantT #-} +constantT :: (Vector v, VecElem v a) => ShapeL -> a -> T v a +constantT sh x = T (map (const 0) sh) 0 (vSingleton x) + +-- TODO: change to return a list of vectors. +-- Convert an array to a vector in the natural order. +{-# INLINE toVectorT #-} +toVectorT :: (Vector v, VecElem v a) => ShapeL -> T v a -> v a +toVectorT sh a@(T ats ao v) = + let l : ts' = getStridesT sh + -- Are strides ok from this point? + oks = scanr (&&) True (zipWith (==) ats ts') + loop _ [] _ o = + DL.singleton (vSlice o 1 v) + loop (b:bs) (s:ss) (t:ts) o = + if b then + -- All strides normal from this point, + -- so just take a slice of the underlying vector. + DL.singleton (vSlice o (s*t) v) + else + -- Strides are not normal, collect slices. + DL.concat [ loop bs ss ts (i*t + o) | i <- [0 .. s-1] ] + loop _ _ _ _ = error "impossible" + in if head oks && vLength v == l then + -- All strides are normal, return entire vector + v + else if oks !! length sh then -- Special case for speed. + -- Innermost dimension is normal, so slices are non-trivial. + vConcat $ DL.toList $ loop oks sh ats ao + else + -- All slices would have length 1, going via a list is faster. + vFromList $ toListT sh a + +-- Convert to a vector containing the right elements, +-- but not necessarily in the right order. +{-# INLINE toUnorderedVectorT #-} +toUnorderedVectorT :: (Vector v, VecElem v a) => ShapeL -> T v a -> v a +toUnorderedVectorT sh a@(T ats ao v) = + -- Figure out if the array maps onto some contiguous slice of the vector. + -- Do this by checking if a transposition of the array corresponds to + -- normal strides. + -- First sort the strides in descending order, amnd rearrange the shape the same way. + -- Then compute the strides from this rearranged shape; these will be the normal + -- strides for this shape. If these strides agree with the sorted actual strides + -- it is a transposition, and we can just slice out the relevant piece of the vector. + let + (ats', sh') = unzip $ sortBy (flip compare) $ zip ats sh + l : ts' = getStridesT sh' + in + if ats' == ts' then + vSlice ao l v + else + toVectorT sh a + +-- Convert from a vector. +{-# INLINE fromVectorT #-} +fromVectorT :: ShapeL -> v a -> T v a +fromVectorT sh = T (tail $ getStridesT sh) 0 + +-- Convert from a list +{-# INLINE fromListT #-} +fromListT :: (Vector v, VecElem v a) => [N] -> [a] -> T v a +fromListT sh = fromVectorT sh . vFromList + +-- Index into the outermost dimension of an array. +{-# INLINE indexT #-} +indexT :: T v a -> N -> T v a +indexT (T (s : ss) o v) i = T ss (o + i * s) v +indexT _ _ = error "impossible" + +-- Stretch the given dimensions to have arbitrary size. +-- The stretched dimensions must have size 1, and stretching is +-- done by setting the stride to 0. +{-# INLINE stretchT #-} +stretchT :: [Bool] -> T v a -> T v a +stretchT bs (T ss o v) = T (zipWith (\ b s -> if b then 0 else s) bs ss) o v + +-- Map over the array elements. +{-# INLINE mapT #-} +mapT :: (Vector v, VecElem v a, VecElem v b) => ShapeL -> (a -> b) -> T v a -> T v b +mapT sh f (T ss o v) | product sh >= vLength v = T ss o (vMap f v) +mapT sh f t = fromVectorT sh $ vMap f $ toVectorT sh t + +-- Zip two arrays with a function. +{-# INLINE zipWithT #-} +zipWithT :: (Vector v, VecElem v a, VecElem v b, VecElem v c) => + ShapeL -> (a -> b -> c) -> T v a -> T v b -> T v c +zipWithT sh f t@(T ss _ v) t'@(T _ _ v') = + case (vLength v, vLength v') of + (1, 1) -> + -- If both vectors have length 1, then it's a degenerate case and it's better + -- to operate on the single element directly. + T ss 0 $ vSingleton $ f (vIndex v 0) (vIndex v' 0) + (1, _) -> + -- First vector has length 1, so use a map instead. + mapT sh (vIndex v 0 `f` ) t' + (_, 1) -> + -- Second vector has length 1, so use a map instead. + mapT sh (`f` vIndex v' 0) t + (_, _) -> + let cv = toVectorT sh t + cv' = toVectorT sh t' + in fromVectorT sh $ vZipWith f cv cv' + +-- Zip three arrays with a function. +{-# INLINE zipWith3T #-} +zipWith3T :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d) => + ShapeL -> (a -> b -> c -> d) -> T v a -> T v b -> T v c -> T v d +zipWith3T _ f (T ss _ v) (T _ _ v') (T _ _ v'') | + -- If all vectors have length 1, then it's a degenerate case and it's better + -- to operate on the single element directly. + vLength v == 1, vLength v' == 1, vLength v'' == 1 = + T ss 0 $ vSingleton $ f (vIndex v 0) (vIndex v' 0) (vIndex v'' 0) +zipWith3T sh f t t' t'' = fromVectorT sh $ vZipWith3 f v v' v'' + where v = toVectorT sh t + v' = toVectorT sh t' + v'' = toVectorT sh t'' + +-- Zip four arrays with a function. +{-# INLINE zipWith4T #-} +zipWith4T :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e) => ShapeL -> (a -> b -> c -> d -> e) -> T v a -> T v b -> T v c -> T v d -> T v e +zipWith4T sh f t t' t'' t''' = fromVectorT sh $ vZipWith4 f v v' v'' v''' + where v = toVectorT sh t + v' = toVectorT sh t' + v'' = toVectorT sh t'' + v'''= toVectorT sh t''' + +-- Zip five arrays with a function. +{-# INLINE zipWith5T #-} +zipWith5T :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e, VecElem v f) => ShapeL -> (a -> b -> c -> d -> e -> f) -> T v a -> T v b -> T v c -> T v d -> T v e -> T v f +zipWith5T sh f t t' t'' t''' t'''' = fromVectorT sh $ vZipWith5 f v v' v'' v''' v'''' + where v = toVectorT sh t + v' = toVectorT sh t' + v'' = toVectorT sh t'' + v'''= toVectorT sh t''' + v''''= toVectorT sh t'''' + +-- Do an arbitrary transposition. The first argument should be +-- a permutation of the dimension, i.e., the numbers [0..r-1] in some order +-- (where r is the rank of the array). +{-# INLINE transposeT #-} +transposeT :: [Int] -> T v a -> T v a +transposeT is (T ss o v) = T (permute is ss) o v + +-- Return all subarrays n dimensions down. +-- The shape argument should be a prefix of the array shape. +{-# INLINE subArraysT #-} +subArraysT :: ShapeL -> T v a -> [T v a] +subArraysT sh ten = sub sh ten [] + where sub [] t = (t :) + sub (n:ns) t = foldr (.) id [sub ns (indexT t i) | i <- [0..n-1]] + +-- Reverse the given dimensions. +{-# INLINE reverseT #-} +reverseT :: [N] -> ShapeL -> T v a -> T v a +reverseT rs sh (T ats ao v) = T rts ro v + where (ro, rts) = rev 0 sh ats + rev !_ [] [] = (ao, []) + rev r (m:ms) (t:ts) | r `elem` rs = (o + (m-1)*t, -t : ts') + | otherwise = (o, t : ts') + where (o, ts') = rev (r+1) ms ts + rev _ _ _ = error "reverseT: impossible" + +-- Reduction of all array elements. +{-# INLINE reduceT #-} +reduceT :: (Vector v, VecElem v a) => + ShapeL -> (a -> a -> a) -> a -> T v a -> T v a +reduceT sh f z = scalarT . vFold f z . toVectorT sh + +-- Right fold via toListT. +{-# INLINE foldrT #-} +foldrT + :: (Vector v, VecElem v a) => ShapeL -> (a -> b -> b) -> b -> T v a -> b +foldrT sh f z a = foldr f z (toListT sh a) + +-- Traversal via toListT/fromListT. +{-# INLINE traverseT #-} +traverseT + :: (Vector v, VecElem v a, VecElem v b, Applicative f) + => ShapeL -> (a -> f b) -> T v a -> f (T v b) +traverseT sh f a = fmap (fromListT sh) (traverse f (toListT sh a)) + +-- Fast check if all elements are equal. +allSameT :: (Vector v, VecElem v a, Eq a) => ShapeL -> T v a -> Bool +allSameT sh t@(T _ _ v) + | vLength v <= 1 = True + | otherwise = + let !v' = toVectorT sh t + !x = vIndex v' 0 + in vAll (x ==) v' + +ppT + :: (Vector v, VecElem v a, Pretty a) + => PrettyLevel -> Rational -> ShapeL -> T v a -> Doc +ppT l p sh = maybeParens (p > 10) . vcat . map text . box prettyBoxMode . ppT_ (prettyShowL l) sh + +ppT_ + :: (Vector v, VecElem v a) + => (a -> String) -> ShapeL -> T v a -> String +ppT_ show_ sh t = revDropWhile (== '\n') $ showsT sh t' "" + where ss = map show_ $ toListT sh t + n = maximum $ map length ss + ss' = map padSP ss + padSP s = replicate (n - length s) ' ' ++ s + t' :: T [] String + t' = T (tail (getStridesT sh)) 0 ss' + +showsT :: [N] -> T [] String -> ShowS +showsT (0:_) _ = showString "EMPTY" +showsT [] t = showString $ unScalarT t +showsT s@[_] t = showString $ unwords $ toListT s t +showsT (n:ns) t = + foldr (.) id [ showsT ns (indexT t i) . showString "\n" | i <- [0..n-1] ] + +data BoxMode = BoxMode { _bmBars, _bmUnicode, _bmHeader :: Bool } + +prettyBoxMode :: BoxMode +prettyBoxMode = BoxMode False False False + +box :: BoxMode -> String -> [String] +box BoxMode{..} s = + let bar | _bmUnicode = '\x2502' + | otherwise = '|' + ls = lines s + ls' | _bmBars = map (\ l -> if null l then l else [bar] ++ l ++ [bar]) ls + | otherwise = ls + h = "+" ++ replicate (length (head ls)) '-' ++ "+" + ls'' | _bmHeader = [h] ++ ls' ++ [h] + | otherwise = ls' + in ls'' + +zipWithLong2 :: (a -> b -> b) -> [a] -> [b] -> [b] +zipWithLong2 f (a:as) (b:bs) = f a b : zipWithLong2 f as bs +zipWithLong2 _ _ bs = bs + +padT :: forall v a . (Vector v, VecElem v a) => a -> [(Int, Int)] -> ShapeL -> T v a -> ([Int], T v a) +padT v aps ash at = (ss, fromVectorT ss $ vConcat $ pad' aps ash st at) + where pad' :: [(Int, Int)] -> ShapeL -> [Int] -> T v a -> [v a] + pad' [] sh _ t = [toVectorT sh t] + pad' ((l,h):ps) (s:sh) (n:ns) t = + [vReplicate (n*l) v] ++ concatMap (pad' ps sh ns . indexT t) [0..s-1] ++ [vReplicate (n*h) v] + pad' _ _ _ _ = error $ "pad: rank mismatch: " ++ show (length aps, length ash) + _ : st = getStridesT ss + ss = zipWithLong2 (\ (l,h) s -> l+s+h) aps ash + +-- Check if a reshape is just adding/removing some dimensions of +-- size 1, in which case it can be done by just manipulating +-- the strides. Given the old strides, the old shapes, and the +-- new shape it will return the possible new strides. +simpleReshape :: [N] -> ShapeL -> ShapeL -> Maybe [N] +simpleReshape osts os ns + | filter (1 /=) os == filter (1 /=) ns = Just $ loop ns sts' + -- Old and new dimensions agree where they are not 1. + where + -- Get old strides for non-1 dimensions + sts' = [ st | (st, s) <- zip osts os, s /= 1 ] + -- Insert stride 0 for all 1 dimensions in new shape. + loop [] [] = [] + loop (1:ss) sts = 0 : loop ss sts + loop (_:ss) (st:sts) = st : loop ss sts + loop _ _ = error $ "simpleReshape: shouldn't happen: " ++ show (osts, os, ns) +simpleReshape _ _ _ = Nothing + +{-# INLINE sumT #-} +sumT :: (Vector v, VecElem v a, Num a) => ShapeL -> T v a -> a +sumT sh = vSum . toUnorderedVectorT sh + +{-# INLINE productT #-} +productT :: (Vector v, VecElem v a, Num a) => ShapeL -> T v a -> a +productT sh = vProduct . toUnorderedVectorT sh + +{-# INLINE maximumT #-} +maximumT :: (Vector v, VecElem v a, Ord a) => ShapeL -> T v a -> a +maximumT sh = vMaximum . toUnorderedVectorT sh + +{-# INLINE minimumT #-} +minimumT :: (Vector v, VecElem v a, Ord a) => ShapeL -> T v a -> a +minimumT sh = vMinimum . toUnorderedVectorT sh + +{-# INLINE anyT #-} +anyT :: (Vector v, VecElem v a) => ShapeL -> (a -> Bool) -> T v a -> Bool +anyT sh p = vAny p . toUnorderedVectorT sh + +{-# INLINE allT #-} +allT :: (Vector v, VecElem v a) => ShapeL -> (a -> Bool) -> T v a -> Bool +allT sh p = vAll p . toUnorderedVectorT sh + +{-# INLINE updateT #-} +updateT :: (Vector v, VecElem v a) => ShapeL -> T v a -> [([Int], a)] -> T v a +updateT sh t us = T ss 0 $ vUpdate (toVectorT sh t) $ map ix us + where _ : ss = getStridesT sh + ix (is, a) = (sum $ zipWith (*) is ss, a) + +{-# INLINE generateT #-} +generateT :: (Vector v, VecElem v a) => ShapeL -> ([Int] -> a) -> T v a +generateT sh f = T ss 0 $ vGenerate s g + where s : ss = getStridesT sh + g i = f (toIx ss i) + toIx [] _ = [] + toIx (n:ns) i = q : toIx ns r where (q, r) = quotRem i n + +{-# INLINE iterateNT #-} +iterateNT :: (Vector v, VecElem v a) => Int -> (a -> a) -> a -> T v a +iterateNT n f x = fromListT [n] $ take n $ iterate f x + +{-# INLINE iotaT #-} +iotaT :: (Vector v, VecElem v a, Enum a, Num a) => Int -> T v a +iotaT n = fromListT [n] [0 .. fromIntegral n - 1] -- TODO: should use V.enumFromTo instead + +------- + +-- | Permute the elements of a list, the first argument is indices into the original list. +permute :: [Int] -> [a] -> [a] +permute is xs = map (xs!!) is + +-- | Like 'dropWhile' but at the end of the list. +revDropWhile :: (a -> Bool) -> [a] -> [a] +revDropWhile p = reverse . dropWhile p . reverse + +allSame :: (Eq a) => [a] -> Bool +allSame [] = True +allSame (x : xs) = all (x ==) xs + +-- | Get the value of a type level Nat. +-- Use with explicit type application, i.e., @valueOf \@42@ +{-# INLINE valueOf #-} +valueOf :: forall n i . (KnownNat n, Num i) => i +valueOf = fromInteger $ natVal (Proxy :: Proxy n) diff --git a/Data/Array/Internal/Dynamic.hs b/Data/Array/Internal/Dynamic.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/Dynamic.hs @@ -0,0 +1,414 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-orphans #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.Dynamic( + Array(..), Vector, ShapeL, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, zipWith4A, zipWith5A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, maximumA, minimumA, + anyA, allA, + broadcast, + update, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Coerce(coerce) +import Data.Data(Data) +import GHC.Generics(Generic) +import qualified Data.Vector as V +import GHC.Stack(HasCallStack) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import qualified Data.Array.DynamicG as G +import Data.Array.Internal(ShapeL, Vector(..), None) + +instance Vector V.Vector where + type VecElem V.Vector = None + {-# INLINE vIndex #-} + vIndex = (V.!) + {-# INLINE vLength #-} + vLength = V.length + {-# INLINE vToList #-} + vToList = V.toList + {-# INLINE vFromList #-} + vFromList = V.fromList + {-# INLINE vSingleton #-} + vSingleton = V.singleton + {-# INLINE vReplicate #-} + vReplicate = V.replicate + {-# INLINE vMap #-} + vMap = fmap + {-# INLINE vZipWith #-} + vZipWith = V.zipWith + {-# INLINE vZipWith3 #-} + vZipWith3 = V.zipWith3 + {-# INLINE vZipWith4 #-} + vZipWith4 = V.zipWith4 + {-# INLINE vZipWith5 #-} + vZipWith5 = V.zipWith5 + {-# INLINE vAppend #-} + vAppend = (V.++) + {-# INLINE vConcat #-} + vConcat = V.concat + {-# INLINE vFold #-} + vFold = V.foldl' + {-# INLINE vSlice #-} + vSlice = V.slice + {-# INLINE vSum #-} + vSum = V.sum + {-# INLINE vProduct #-} + vProduct = V.product + {-# INLINE vMaximum #-} + vMaximum = V.maximum + {-# INLINE vMinimum #-} + vMinimum = V.minimum + {-# INLINE vUpdate #-} + vUpdate = (V.//) + {-# INLINE vGenerate #-} + vGenerate = V.generate + {-# INLINE vAll #-} + vAll = V.all + {-# INLINE vAny #-} + vAny = V.any + +newtype Array a = A { unA :: G.Array V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array a) + +toArrayG :: Array a -> G.Array V.Vector a +toArrayG = unA + +instance (Show a) => Show (Array a) where + showsPrec p = showsPrec p . unA + +instance (Read a) => Read (Array a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array V.Vector a) => Eq (Array a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array V.Vector a) => Ord (Array a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +rank :: Array a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +index :: (HasCallStack) => Array a -> Int -> Array a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: Array a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: (HasCallStack) => ShapeL -> [a] -> Array a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (HasCallStack) => Array a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: (HasCallStack) => ShapeL -> V.Vector a -> Array a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: Array a -> Array a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: (HasCallStack) => ShapeL -> Array a -> Array a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: (HasCallStack) => ShapeL -> Array a -> Array a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack) => Int -> Array a -> Array a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: a -> Array a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (HasCallStack) => Array a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +constant :: ShapeL -> a -> Array a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +mapA :: (a -> b) -> Array a -> Array b +mapA f = A . G.mapA f . unA + +instance Functor Array where + fmap = mapA + +instance Foldable Array where + foldr = foldrA + +instance Traversable Array where + traverse = traverseA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (HasCallStack) => (a -> b -> c) -> Array a -> Array b -> Array c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (HasCallStack) => (a -> b -> c -> d) -> Array a -> Array b -> Array c -> Array d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Map over the array elements. +-- O(n) time. +zipWith4A :: (HasCallStack) => (a -> b -> c -> d -> e) -> Array a -> Array b -> Array c -> Array d -> Array e +zipWith4A f a b c d = A $ G.zipWith4A f (unA a) (unA b) (unA c) (unA d) + +-- | Map over the array elements. +-- O(n) time. +zipWith5A :: (HasCallStack) => (a -> b -> c -> d -> e -> f) -> Array a -> Array b -> Array c -> Array d -> Array e -> Array f +zipWith5A f a b c d e = A $ G.zipWith5A f (unA a) (unA b) (unA c) (unA d) (unA e) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (HasCallStack) => [(Int, Int)] -> a -> Array a -> Array a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (HasCallStack) => [Int] -> Array a -> Array a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (HasCallStack) => Array a -> Array a -> Array a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (HasCallStack) => [Array a] -> Array a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (HasCallStack) => Array (Array a) -> Array a +ravel = A . G.ravel . G.mapA unA . unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (HasCallStack) => Array a -> Array (Array a) +unravel = A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +window :: (HasCallStack) => [Int] -> Array a -> Array a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: (HasCallStack) => [Int] -> Array a -> Array a +stride ws = A . G.stride ws . unA + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: (HasCallStack) => Int -> Int -> Array a -> Array a +rotate d k = A . G.rotate d k . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: (HasCallStack) => [(Int, Int)] -> Array a -> Array a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: (HasCallStack) => Int -> (Array a -> Array b) -> Array a -> Array b +rerank n f = A . G.rerank n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: (HasCallStack) => Int -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c +rerank2 n f ta tb = A $ G.rerank2 n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array a -> Array a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (a -> a -> a) -> a -> Array a -> Array a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +foldrA :: (a -> b -> b) -> b -> Array a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +traverseA :: Applicative f => (a -> f b) -> Array a -> f (Array b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Eq a) => Array a -> Bool +allSameA = G.allSameA . unA + +instance Arbitrary a => Arbitrary (Array a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Num a) => Array a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Num a) => Array a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (HasCallStack, Ord a) => Array a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (HasCallStack, Ord a) => Array a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (a -> Bool) -> Array a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (a -> Bool) -> Array a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: (HasCallStack) => + [Int] -> ShapeL -> Array a -> Array a +broadcast ds sh = A . G.broadcast ds sh . unA + +-- | Update the array at the specified indicies to the associated value. +{-# INLINE update #-} +update :: (HasCallStack) => + Array a -> [([Int], a)] -> Array a +update a = A . G.update (unA a) + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: ShapeL -> ([Int] -> a) -> Array a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: Int -> (a -> a) -> a -> Array a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Enum a, Num a) => Int -> Array a +iota = A . G.iota diff --git a/Data/Array/Internal/DynamicG.hs b/Data/Array/Internal/DynamicG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/DynamicG.hs @@ -0,0 +1,519 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UndecidableInstances #-} +-- | Arrays of dynamic size. The arrays are polymorphic in the underlying +-- linear data structure used to store the actual values. +module Data.Array.Internal.DynamicG( + Array(..), Vector, VecElem, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, zipWith4A, zipWith5A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, maximumA, minimumA, + anyA, allA, + broadcast, + update, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Control.Monad(replicateM) +import Data.Data(Data) +import Data.List(sort) +import GHC.Generics(Generic) +import GHC.Stack +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal + +-- | Arrays stored in a /v/ with values of type /a/. +data Array v a = A !ShapeL !(T v a) + deriving (Generic, Data) + +instance (Vector v, Show a, VecElem v a) => Show (Array v a) where + showsPrec p a@(A s _) = showParen (p > 10) $ + showString "fromList " . showsPrec 11 s . showString " " . showsPrec 11 (toList a) + +instance (Vector v, Read a, VecElem v a) => Read (Array v a) where + readsPrec p = readParen (p > 10) $ \ r1 -> + [(fromList s xs, r4) | ("fromList", r2) <- lex r1, (s, r3) <- readsPrec 11 r2, + (xs, r4) <- readsPrec 11 r3, product s == length xs] + +instance (Vector v, Eq a, VecElem v a, Eq (v a)) => Eq (Array v a) where + (A s v) == (A s' v') = s == s' && equalT s v v' + {-# INLINE (==) #-} + +instance (Vector v, Ord a, Ord (v a), VecElem v a) => Ord (Array v a) where + (A s v) `compare` (A s' v') = compare s s' <> compareT s v v' + {-# INLINE compare #-} + +instance (Vector v, Pretty a, VecElem v a) => Pretty (Array v a) where + pPrintPrec l p (A sh t) = ppT l p sh t + +instance (NFData (v a)) => NFData (Array v a) + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array v a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +{-# INLINE shapeL #-} +shapeL :: Array v a -> ShapeL +shapeL (A s _) = s + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +{-# INLINE rank #-} +rank :: Array v a -> Int +rank (A s _) = length s + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +{-# INLINE index #-} +index :: (HasCallStack, Vector v) => Array v a -> Int -> Array v a +index (A (s:ss) t) i | i < 0 || i >= s = error $ "index: out of bounds " ++ show (i, s) + | otherwise = A ss $ indexT t i +index (A [] _) _ = error "index: scalar" + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +{-# INLINE toList #-} +toList :: (Vector v, VecElem v a) => Array v a -> [a] +toList (A sh t) = toListT sh t + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +{-# INLINE toVector #-} +toVector :: (Vector v, VecElem v a) => Array v a -> v a +toVector (A sh t) = toVectorT sh t + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +{-# INLINE fromList #-} +fromList :: (HasCallStack, Vector v, VecElem v a) => ShapeL -> [a] -> Array v a +fromList ss vs | n /= l = error $ "fromList: size mismatch" ++ show (n, l) + | otherwise = A ss $ T st 0 $ vFromList vs + where n : st = getStridesT ss + l = length vs + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +{-# INLINE fromVector #-} +fromVector :: (HasCallStack, Vector v, VecElem v a) => ShapeL -> v a -> Array v a +fromVector ss v | n /= l = error $ "fromList: size mismatch" ++ show (n, l) + | otherwise = A ss $ T st 0 v + where n : st = getStridesT ss + l = vLength v + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +{-# INLINE normalize #-} +normalize :: (Vector v, VecElem v a) => Array v a -> Array v a +normalize a = fromVector (shapeL a) $ toVector a + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +{-# INLINE reshape #-} +reshape :: (HasCallStack, Vector v, VecElem v a) => ShapeL -> Array v a -> Array v a +reshape sh (A sh' t@(T ost oo v)) + | n /= n' = error $ "reshape: size mismatch " ++ show (sh, sh') + | vLength v == 1 = A sh $ T (map (const 0) sh) 0 v -- Fast special case for singleton vector + | Just nst <- simpleReshape ost sh' sh = A sh $ T nst oo v + | otherwise = A sh $ T st 0 $ toVectorT sh' t + where n : st = getStridesT sh + n' = product sh' + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +{-# INLINE stretch #-} +stretch :: (HasCallStack) => ShapeL -> Array v a -> Array v a +stretch sh (A sh' vs) | Just bs <- str sh sh' = A sh $ stretchT bs vs + | otherwise = error $ "stretch: incompatible " ++ show (sh, sh') + where str [] [] = Just [] + str (x:xs) (y:ys) | x == y = (False :) <$> str xs ys + | y == 1 = (True :) <$> str xs ys + str _ _ = Nothing + +-- | Change the size of the outermost dimension by replication. +{-# INLINE stretchOuter #-} +stretchOuter :: (HasCallStack) => Int -> Array v a -> Array v a +stretchOuter s (A (1:sh) vs) = + A (s:sh) $ stretchT (True : map (const False) (strides vs)) vs +stretchOuter _ _ = error "stretchOuter: needs outermost dimension of size 1" + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +{-# INLINE scalar #-} +scalar :: (Vector v, VecElem v a) => a -> Array v a +scalar = A [] . scalarT + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +{-# INLINE unScalar #-} +unScalar :: (HasCallStack, Vector v, VecElem v a) => Array v a -> a +unScalar (A [] t) = unScalarT t +unScalar _ = error "unScalar: not a scalar" + +-- | Make an array with all elements having the same value. +-- O(1) time +{-# INLINE constant #-} +constant :: (HasCallStack, Vector v, VecElem v a) => ShapeL -> a -> Array v a +constant sh | badShape sh = error "constant: bad shape" + | otherwise = A sh . constantT sh + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE mapA #-} +mapA :: (Vector v, VecElem v a, VecElem v b) => (a -> b) -> Array v a -> Array v b +mapA f (A s t) = A s (mapT s f t) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWithA #-} +zipWithA :: (HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c) => + (a -> b -> c) -> Array v a -> Array v b -> Array v c +zipWithA f (A s t) (A s' t') | s == s' = A s (zipWithT s f t t') + | otherwise = error $ "zipWithA: shape mismatch: " ++ show (s, s') + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith3A #-} +zipWith3A :: (HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d) => + (a -> b -> c -> d) -> Array v a -> Array v b -> Array v c -> Array v d +zipWith3A f (A s t) (A s' t') (A s'' t'') | s == s' && s == s'' = A s (zipWith3T s f t t' t'') + | otherwise = error $ "zipWith3A: shape mismatch: " ++ show (s, s', s'') + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith4A #-} +zipWith4A :: (HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e) => + (a -> b -> c -> d -> e) -> Array v a -> Array v b -> Array v c -> Array v d -> Array v e +zipWith4A f (A s t) (A s' t') (A s'' t'') (A s''' t''') | s == s' && s == s'' && s == s''' = A s (zipWith4T s f t t' t'' t''') + | otherwise = error $ "zipWith4A: shape mismatch: " ++ show (s, s', s'', s''') + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith5A #-} +zipWith5A :: (HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, VecElem v e, VecElem v f) => + (a -> b -> c -> d -> e -> f) -> Array v a -> Array v b -> Array v c -> Array v d -> Array v e -> Array v f +zipWith5A f (A s t) (A s' t') (A s'' t'') (A s''' t''') (A s'''' t'''') | s == s' && s == s'' && s == s''' && s == s'''' = A s (zipWith5T s f t t' t'' t''' t'''') + | otherwise = error $ "zipWith5A: shape mismatch: " ++ show (s, s', s'', s''', s'''') + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +{-# INLINE pad #-} +pad :: forall a v . (Vector v, VecElem v a) => + [(Int, Int)] -> a -> Array v a -> Array v a +pad aps v (A ash at) = uncurry A $ padT v aps ash at + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +{-# INLINE transpose #-} +transpose :: (HasCallStack) => [Int] -> Array v a -> Array v a +transpose is (A sh t) | l > n = error $ "transpose: rank exceeded " ++ show (is, sh) + | sort is /= [0 .. l-1] = + error $ "transpose: not a permutation: " ++ show is + | otherwise = A (permute is' sh) (transposeT is' t) + where l = length is + n = length sh + is' = is ++ [l .. n-1] + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +{-# INLINE append #-} +append :: (HasCallStack, Vector v, VecElem v a) => Array v a -> Array v a -> Array v a +append a@(A (sa:sh) _) b@(A (sb:sh') _) | sh == sh' = + fromVector (sa+sb : sh) (vAppend (toVector a) (toVector b)) +append _ _ = error "append: bad shape" + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +{-# INLINE concatOuter #-} +concatOuter :: (HasCallStack, Vector v, VecElem v a) => [Array v a] -> Array v a +concatOuter [] = error "concatOuter: empty list" +concatOuter as | not $ allSame $ map tail shs = + error $ "concatOuter: non-conforming inner dimensions: " ++ show shs + | otherwise = fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = sum (map head shs) : tail sh + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +{-# INLINE ravel #-} +ravel :: (HasCallStack, Vector v, Vector v', VecElem v a, VecElem v' (Array v a)) => + Array v' (Array v a) -> Array v a +ravel aa | rank aa /= 1 = error "ravel: outermost array does not have rank 1" + | otherwise = + case toList aa of + [] -> error "ravel: empty array" + as | not $ allSame shs -> error $ "ravel: non-conforming inner dimensions: " ++ show shs + | otherwise -> fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = length as : sh + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +{-# INLINE unravel #-} +unravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array v a)) => + Array v a -> Array v' (Array v a) +unravel = rerank 1 scalar + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +{-# INLINE window #-} +window :: (HasCallStack, Vector v) => [Int] -> Array v a -> Array v a +window aws (A ash (T ss o v)) = A (win aws ash) (T (ss' ++ ss) o v) + where ss' = zipWith const ss aws + win (w:ws) (s:sh) | w <= s = s - w + 1 : win ws sh + | otherwise = error $ "window: bad window size : " ++ show (w, s) + win [] sh = aws ++ sh + win _ _ = error $ "window: rank mismatch: " ++ show (aws, ash) + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- The rank of the stride list must not exceed the rank of the array. +-- O(1) time. +{-# INLINE stride #-} +stride :: (HasCallStack, Vector v) => [Int] -> Array v a -> Array v a +stride ats (A ash (T ss o v)) = A (str ats ash) (T (zipWith (*) (ats ++ repeat 1) ss) o v) + where str (t:ts) (s:sh) = (s+t-1) `quot` t : str ts sh + str [] sh = sh + str _ _ = error $ "stride: rank mismatch: " ++ show (ats, ash) + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +{-# INLINE rotate #-} +rotate :: (HasCallStack, Vector v, VecElem v a) => Int -> Int -> Array v a -> Array v a +rotate d k a | d < rank a, k >= 0 = rerank d f a + where + f arr = let h:t = shapeL arr + m = product t + n = h * m + in rev [0] + . reshape (k:h:t) + . stride [n + m] + . window [n] + . reshape [(k + 1) * n] + . stretchOuter (k + 1) + . reshape (1:h:t) $ arr +rotate d k a = error $ "Incorrect arguments to rotate: " ++ show (d, k, rank a) + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +{-# INLINE slice #-} +slice :: (HasCallStack) => [(Int, Int)] -> Array v a -> Array v a +slice asl (A ash (T ats ao v)) = A rsh (T ats o v) + where (o, rsh) = slc asl ash ats + slc ((k,n):sl) (s:sh) (t:ts) | k < 0 || k > s || k+n > s = error $ "slice: out of bounds: slice=" ++ show (k, n) ++ " size=" ++ show s + | otherwise = (i + k*t, n:ns) where (i, ns) = slc sl sh ts + slc (_:_) [] _ = error "slice: slice list too long" + slc [] sh _ = (ao, sh) + slc _ _ _ = error "impossible" + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank #-} +rerank :: (HasCallStack, Vector v, Vector v', VecElem v a, VecElem v' b) => + Int -> (Array v a -> Array v' b) -> Array v a -> Array v' b +rerank n f (A sh t) | n < 0 || n > length sh = error "rerank: rank exceeded" + | otherwise = + ravelOuter osh $ + map (f . A ish) $ + subArraysT osh t + where (osh, ish) = splitAt n sh + +ravelOuter :: (HasCallStack, Vector v, VecElem v a) => ShapeL -> [Array v a] -> Array v a +ravelOuter _ [] = error "ravelOuter: empty list" +ravelOuter osh as | not $ allSame shs = error $ "ravelOuter: non-conforming inner dimensions: " ++ show shs + | otherwise = fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = osh ++ sh + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank2 #-} +rerank2 :: (HasCallStack, Vector v, VecElem v a, VecElem v b, VecElem v c) => + Int -> (Array v a -> Array v b -> Array v c) -> Array v a -> Array v b -> Array v c +rerank2 n f (A sha ta) (A shb tb) | n < 0 || n > length sha || n > length shb = error "rerank: rank exceeded" + | take n sha /= take n shb = error "rerank2: shape mismatch" + | otherwise = + ravelOuter osh $ + zipWith (\ a b -> f (A isha a) (A ishb b)) + (subArraysT osh ta) + (subArraysT osh tb) + where (osh, isha) = splitAt n sha + ishb = drop n shb + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +{-# INLINE rev #-} +rev :: (HasCallStack) => [Int] -> Array v a -> Array v a +rev rs (A sh t) | all (\ r -> r >= 0 && r < n) rs = A sh (reverseT rs sh t) + | otherwise = error "reverse: bad reverse dimension" + where n = length sh + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +{-# INLINE reduce #-} +reduce :: (Vector v, VecElem v a) => + (a -> a -> a) -> a -> Array v a -> Array v a +reduce f z (A sh t) = A [] $ reduceT sh f z t + +-- | Right fold across all elements of an array. +{-# INLINE foldrA #-} +foldrA :: (Vector v, VecElem v a) => (a -> b -> b) -> b -> Array v a -> b +foldrA f z (A sh t) = foldrT sh f z t + +-- | Constrained version of 'traverse' for 'Array's. +{-# INLINE traverseA #-} +traverseA + :: (Vector v, VecElem v a, VecElem v b, Applicative f) + => (a -> f b) -> Array v a -> f (Array v b) +traverseA f (A sh t) = A sh <$> traverseT sh f t + +-- | Check if all elements of the array are equal. +allSameA :: (Vector v, VecElem v a, Eq a) => Array v a -> Bool +allSameA (A sh t) = allSameT sh t + +instance (Vector v, VecElem v a, Arbitrary a) => Arbitrary (Array v a) where + arbitrary = do + r <- choose (0, 5) -- Don't generate huge ranks + -- Don't generate huge number of elements + ss <- replicateM r (getSmall . getPositive <$> arbitrary) `suchThat` ((< 10000) . product) + fromList ss <$> vector (product ss) + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Vector v, VecElem v a, Num a) => Array v a -> a +sumA (A sh t) = sumT sh t + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Vector v, VecElem v a, Num a) => Array v a -> a +productA (A sh t) = productT sh t + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (HasCallStack, Vector v, VecElem v a, Ord a) => Array v a -> a +maximumA a@(A sh t) | size a > 0 = maximumT sh t + | otherwise = error "maximumA called with empty array" + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (HasCallStack, Vector v, VecElem v a, Ord a) => Array v a -> a +minimumA a@(A sh t) | size a > 0 = minimumT sh t + | otherwise = error "minimumA called with empty array" + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Vector v, VecElem v a) => (a -> Bool) -> Array v a -> Bool +anyA p (A sh t) = anyT sh p t + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Vector v, VecElem v a) => (a -> Bool) -> Array v a -> Bool +allA p (A sh t) = anyT sh p t + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: (HasCallStack, Vector v, VecElem v a) => + [Int] -> ShapeL -> Array v a -> Array v a +broadcast ds sh a | length ds /= rank a = error "broadcast: wrong number of broadcasts" + | any (\ d -> d < 0 || d >= r) ds = error "broadcast: bad dimension index" + | not (ascending ds) = error "broadcast: unordered dimensions" + | otherwise = stretch sh $ reshape rsh a + where r = length sh + rsh = [ if i `elem` ds then s else 1 | (i, s) <- zip [0..] sh ] + ascending (x:y:ys) = x < y && ascending (y:ys) + ascending _ = True + +-- | Update the array at the specified indicies to the associated value. +update :: (HasCallStack, Vector v, VecElem v a) => + Array v a -> [([Int], a)] -> Array v a +update (A sh t) us | all (ok . fst) us = A sh $ updateT sh t us + | otherwise = error $ "update: index out of bounds " ++ show (filter (not . ok) $ map fst us) + where ok is = length is == r && and (zipWith (\ i s -> 0 <= i && i < s) is sh) + r = length sh + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Vector v, VecElem v a) => + ShapeL -> ([Int] -> a) -> Array v a +generate sh = A sh . generateT sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall v a . + (Vector v, VecElem v a) => + Int -> (a -> a) -> a -> Array v a +iterateN n f = A [n] . iterateNT n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: forall v a . + (Vector v, VecElem v a, Enum a, Num a) => + Int -> Array v a +iota n = A [n] $ iotaT n diff --git a/Data/Array/Internal/DynamicS.hs b/Data/Array/Internal/DynamicS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/DynamicS.hs @@ -0,0 +1,433 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-orphans #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.DynamicS( + Array(..), Vector, ShapeL, V.Storable, Unbox, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, zipWith4A, zipWith5A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, maximumA, minimumA, + anyA, allA, + broadcast, + update, + generate, iterateN, iota, + bitcast, + ) where +import Control.DeepSeq +import Data.Coerce(coerce) +import Data.Data(Data) +import qualified Data.Vector.Storable as V +import Foreign.Storable(sizeOf) +import GHC.Generics(Generic) +import GHC.Stack(HasCallStack) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import qualified Data.Array.Internal.Dynamic as D +import qualified Data.Array.Internal.DynamicG as G +import Data.Array.Internal(T(..), ShapeL, Vector(..)) + +type Unbox = V.Storable + +instance Vector V.Vector where + type VecElem V.Vector = Unbox + {-# INLINE vIndex #-} + vIndex = (V.!) + {-# INLINE vLength #-} + vLength = V.length + {-# INLINE vToList #-} + vToList = V.toList + {-# INLINE vFromList #-} + vFromList = V.fromList + {-# INLINE vSingleton #-} + vSingleton = V.singleton + {-# INLINE vReplicate #-} + vReplicate = V.replicate + {-# INLINE vMap #-} + vMap = V.map + {-# INLINE vZipWith #-} + vZipWith = V.zipWith + {-# INLINE vZipWith3 #-} + vZipWith3 = V.zipWith3 + {-# INLINE vZipWith4 #-} + vZipWith4 = V.zipWith4 + {-# INLINE vZipWith5 #-} + vZipWith5 = V.zipWith5 + {-# INLINE vAppend #-} + vAppend = (V.++) + {-# INLINE vConcat #-} + vConcat = V.concat + {-# INLINE vFold #-} + vFold = V.foldl' + {-# INLINE vSlice #-} + vSlice = V.slice + {-# INLINE vSum #-} + vSum = V.sum + {-# INLINE vProduct #-} + vProduct = V.product + {-# INLINE vMaximum #-} + vMaximum = V.maximum + {-# INLINE vMinimum #-} + vMinimum = V.minimum + {-# INLINE vUpdate #-} + vUpdate = (V.//) + {-# INLINE vGenerate #-} + vGenerate = V.generate + {-# INLINE vAll #-} + vAll = V.all + {-# INLINE vAny #-} + vAny = V.any + +newtype Array a = A { unA :: G.Array V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array a) + +toArrayG :: Array a -> G.Array V.Vector a +toArrayG = unA + +instance (Show a, Unbox a) => Show (Array a) where + showsPrec p = showsPrec p . unA + +instance (Read a, Unbox a) => Read (Array a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array V.Vector a) => Eq (Array a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array V.Vector a) => Ord (Array a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +rank :: Array a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +index :: (HasCallStack, Unbox a) => Array a -> Int -> Array a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: (HasCallStack, Unbox a) => Array a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: (HasCallStack, Unbox a) => ShapeL -> [a] -> Array a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (HasCallStack, Unbox a) => Array a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: (HasCallStack, Unbox a) => ShapeL -> V.Vector a -> Array a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +{-# INLINE normalize #-} +normalize :: (Unbox a) => Array a -> Array a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: (HasCallStack, Unbox a) => ShapeL -> Array a -> Array a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: (HasCallStack) => ShapeL -> Array a -> Array a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack) => Int -> Array a -> Array a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (HasCallStack, Unbox a) => Array a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +{-# INLINE constant #-} +constant :: (Unbox a) => ShapeL -> a -> Array a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE mapA #-} +mapA :: (Unbox a, Unbox b) => (a -> b) -> Array a -> Array b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWithA #-} +zipWithA :: (HasCallStack, Unbox a, Unbox b, Unbox c) => + (a -> b -> c) -> Array a -> Array b -> Array c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith3A #-} +zipWith3A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d) => + (a -> b -> c -> d) -> Array a -> Array b -> Array c -> Array d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith4A #-} +zipWith4A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e) => + (a -> b -> c -> d -> e) -> Array a -> Array b -> Array c -> Array d -> Array e +zipWith4A f a b c d = A $ G.zipWith4A f (unA a) (unA b) (unA c) (unA d) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith5A #-} +zipWith5A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f) => + (a -> b -> c -> d -> e -> f) -> Array a -> Array b -> Array c -> Array d -> Array e -> Array f +zipWith5A f a b c d e = A $ G.zipWith5A f (unA a) (unA b) (unA c) (unA d) (unA e) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (HasCallStack, Unbox a) => [(Int, Int)] -> a -> Array a -> Array a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (HasCallStack) => [Int] -> Array a -> Array a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (HasCallStack, Unbox a) => Array a -> Array a -> Array a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (HasCallStack, Unbox a) => [Array a] -> Array a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (HasCallStack, Unbox a) => D.Array (Array a) -> Array a +ravel = A . G.ravel . G.mapA unA . D.unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (HasCallStack, Unbox a) => Array a -> D.Array (Array a) +unravel = D.A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: (HasCallStack) => [Int] -> Array a -> Array a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: (HasCallStack) => [Int] -> Array a -> Array a +stride ws = A . G.stride ws . unA + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: (HasCallStack, Unbox a) => Int -> Int -> Array a -> Array a +rotate d k = A . G.rotate d k . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: (HasCallStack) => [(Int, Int)] -> Array a -> Array a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: (HasCallStack, Unbox a, Unbox b) => Int -> (Array a -> Array b) -> Array a -> Array b +rerank n f = A . G.rerank n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: (HasCallStack, Unbox a, Unbox b, Unbox c) => + Int -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c +rerank2 n f ta tb = A $ G.rerank2 n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array a -> Array a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a) => (a -> a -> a) -> a -> Array a -> Array a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a) => (a -> b -> b) -> b -> Array a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f) => (a -> f b) -> Array a -> f (Array b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +{-# INLINE allSameA #-} +allSameA :: (Unbox a, Eq a) => Array a -> Bool +allSameA = G.allSameA . unA + +instance (Arbitrary a, Unbox a) => Arbitrary (Array a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a) => Array a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a) => Array a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: Unbox a => (a -> Bool) -> Array a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: Unbox a => (a -> Bool) -> Array a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +{-# INLINE broadcast #-} +broadcast :: (HasCallStack, Unbox a) => + [Int] -> ShapeL -> Array a -> Array a +broadcast ds sh = A. G.broadcast ds sh . unA + +-- | Update the array at the specified indicies to the associated value. +{-# INLINE update #-} +update :: (HasCallStack, Unbox a) => + Array a -> [([Int], a)] -> Array a +update a = A . G.update (unA a) + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Unbox a) => ShapeL -> ([Int] -> a) -> Array a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: (Unbox a) => + Int -> (a -> a) -> a -> Array a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Unbox a, Enum a, Num a) => Int -> Array a +iota = A . G.iota + +-- | Convert between types by just reinterpreting the bits as another type. +-- For instance the floating point number @(1.5 :: Float)@ will convert to +-- @(0x3fc00000 :: Word32)@ since they have the same bit representation. +{-# INLINE bitcast #-} +bitcast :: forall a b . (HasCallStack, Unbox a, Unbox b) => Array a -> Array b +bitcast (A (G.A sh (T ss o v))) + | sza /= szb + = error $ "bitcast: the types must have the same size. " ++ show (sza, szb) + | otherwise + = A (G.A sh (T ss o (V.unsafeCast v))) + where sza = sizeOf (undefined :: a) + szb = sizeOf (undefined :: b) diff --git a/Data/Array/Internal/DynamicU.hs b/Data/Array/Internal/DynamicU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/DynamicU.hs @@ -0,0 +1,411 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-orphans #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.DynamicU( + Array(..), Vector, ShapeL, Unbox, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, zipWith4A, zipWith5A, + append, concatOuter, + ravel, unravel, + window, stride, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, maximumA, minimumA, + anyA, allA, + broadcast, + update, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Coerce(coerce) +import Data.Data(Data) +import qualified Data.Vector.Unboxed as V +import GHC.Stack(HasCallStack) +import Test.QuickCheck hiding (generate) +import GHC.Generics(Generic) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import qualified Data.Array.Internal.Dynamic as D +import qualified Data.Array.DynamicG as G +import Data.Array.Internal(ShapeL, Vector(..)) + +type Unbox = V.Unbox + +instance Vector V.Vector where + type VecElem V.Vector = Unbox + {-# INLINE vIndex #-} + vIndex = (V.!) + {-# INLINE vLength #-} + vLength = V.length + {-# INLINE vToList #-} + vToList = V.toList + {-# INLINE vFromList #-} + vFromList = V.fromList + {-# INLINE vSingleton #-} + vSingleton = V.singleton + {-# INLINE vReplicate #-} + vReplicate = V.replicate + {-# INLINE vMap #-} + vMap = V.map + {-# INLINE vZipWith #-} + vZipWith = V.zipWith + {-# INLINE vZipWith3 #-} + vZipWith3 = V.zipWith3 + {-# INLINE vZipWith4 #-} + vZipWith4 = V.zipWith4 + {-# INLINE vZipWith5 #-} + vZipWith5 = V.zipWith5 + {-# INLINE vAppend #-} + vAppend = (V.++) + {-# INLINE vConcat #-} + vConcat = V.concat + {-# INLINE vFold #-} + vFold = V.foldl' + {-# INLINE vSlice #-} + vSlice = V.slice + {-# INLINE vSum #-} + vSum = V.sum + {-# INLINE vProduct #-} + vProduct = V.product + {-# INLINE vMaximum #-} + vMaximum = V.maximum + {-# INLINE vMinimum #-} + vMinimum = V.minimum + {-# INLINE vUpdate #-} + vUpdate = (V.//) + {-# INLINE vGenerate #-} + vGenerate = V.generate + {-# INLINE vAll #-} + vAll = V.all + {-# INLINE vAny #-} + vAny = V.any + +newtype Array a = A { unA :: G.Array V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array a) + +toArrayG :: Array a -> G.Array V.Vector a +toArrayG = unA + +instance (Show a, Unbox a) => Show (Array a) where + showsPrec p = showsPrec p . unA + +instance (Read a, Unbox a) => Read (Array a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array V.Vector a) => Eq (Array a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array V.Vector a) => Ord (Array a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +rank :: Array a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +index :: (HasCallStack, Unbox a) => Array a -> Int -> Array a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: (HasCallStack, Unbox a) => Array a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: (HasCallStack, Unbox a) => ShapeL -> [a] -> Array a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (HasCallStack, Unbox a) => Array a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: (HasCallStack, Unbox a) => ShapeL -> V.Vector a -> Array a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +{-# INLINE normalize #-} +normalize :: (Unbox a) => Array a -> Array a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: (HasCallStack, Unbox a) => ShapeL -> Array a -> Array a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: (HasCallStack) => ShapeL -> Array a -> Array a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack) => Int -> Array a -> Array a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (HasCallStack, Unbox a) => Array a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +{-# INLINE constant #-} +constant :: (Unbox a) => ShapeL -> a -> Array a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE mapA #-} +mapA :: (Unbox a, Unbox b) => (a -> b) -> Array a -> Array b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWithA #-} +zipWithA :: (HasCallStack, Unbox a, Unbox b, Unbox c) => + (a -> b -> c) -> Array a -> Array b -> Array c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith3A #-} +zipWith3A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d) => + (a -> b -> c -> d) -> Array a -> Array b -> Array c -> Array d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith4A #-} +zipWith4A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e) => + (a -> b -> c -> d -> e) -> Array a -> Array b -> Array c -> Array d -> Array e +zipWith4A f a b c d = A $ G.zipWith4A f (unA a) (unA b) (unA c) (unA d) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith5A #-} +zipWith5A :: (HasCallStack, Unbox a, Unbox b, Unbox c, Unbox d, Unbox e, Unbox f) => + (a -> b -> c -> d -> e -> f) -> Array a -> Array b -> Array c -> Array d -> Array e -> Array f +zipWith5A f a b c d e = A $ G.zipWith5A f (unA a) (unA b) (unA c) (unA d) (unA e) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (HasCallStack, Unbox a) => [(Int, Int)] -> a -> Array a -> Array a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (HasCallStack) => [Int] -> Array a -> Array a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (HasCallStack, Unbox a) => Array a -> Array a -> Array a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (HasCallStack, Unbox a) => [Array a] -> Array a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (HasCallStack, Unbox a) => D.Array (Array a) -> Array a +ravel = A . G.ravel . G.mapA unA . D.unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (HasCallStack, Unbox a) => Array a -> D.Array (Array a) +unravel = D.A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: (HasCallStack) => [Int] -> Array a -> Array a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: (HasCallStack) => [Int] -> Array a -> Array a +stride ws = A . G.stride ws . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: (HasCallStack) => [(Int, Int)] -> Array a -> Array a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: (HasCallStack, Unbox a, Unbox b) => Int -> (Array a -> Array b) -> Array a -> Array b +rerank n f = A . G.rerank n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: (HasCallStack, Unbox a, Unbox b, Unbox c) => + Int -> (Array a -> Array b -> Array c) -> Array a -> Array b -> Array c +rerank2 n f ta tb = A $ G.rerank2 n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array a -> Array a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a) => (a -> a -> a) -> a -> Array a -> Array a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a) => (a -> b -> b) -> b -> Array a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f) => (a -> f b) -> Array a -> f (Array b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +{-# INLINE allSameA #-} +allSameA :: (Unbox a, Eq a) => Array a -> Bool +allSameA = G.allSameA . unA + +instance (Arbitrary a, Unbox a) => Arbitrary (Array a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a) => Array a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a) => Array a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (HasCallStack, Unbox a, Ord a) => Array a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: Unbox a => (a -> Bool) -> Array a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: Unbox a => (a -> Bool) -> Array a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +{-# INLINE broadcast #-} +broadcast :: (HasCallStack, Unbox a) => + [Int] -> ShapeL -> Array a -> Array a +broadcast ds sh = A. G.broadcast ds sh . unA + +-- | Update the array at the specified indicies to the associated value. +{-# INLINE update #-} +update :: (HasCallStack, Unbox a) => + Array a -> [([Int], a)] -> Array a +update a = A . G.update (unA a) + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Unbox a) => ShapeL -> ([Int] -> a) -> Array a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: (Unbox a) => + Int -> (a -> a) -> a -> Array a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Unbox a, Enum a, Num a) => Int -> Array a +iota = A . G.iota diff --git a/Data/Array/Internal/Ranked.hs b/Data/Array/Internal/Ranked.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/Ranked.hs @@ -0,0 +1,370 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-orphans #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.Ranked( + Array(..), Vector, ShapeL, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Data(Data) +import Data.Coerce(coerce) +import qualified Data.Vector as V +import GHC.Stack +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import Test.QuickCheck hiding (generate) + +import GHC.Generics(Generic) +import Data.Array.Internal.Dynamic() -- Vector instance +import qualified Data.Array.Internal.RankedG as G +import Data.Array.Internal(ShapeL, Vector(..)) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +newtype Array n a = A { unA :: G.Array n V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array n a) + +toArrayG :: Array r a -> G.Array r V.Vector a +toArrayG = unA + +instance (Show a) => Show (Array n a) where + showsPrec p = showsPrec p . unA + +instance (KnownNat n, Read a) => Read (Array n a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array n V.Vector a) => Eq (Array n a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array n V.Vector a) => Ord (Array n a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array n a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array n a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (KnownNat n) => Array n a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: HasCallStack => Array (1+n) a -> Int -> Array n a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: Array n a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: forall n a . (KnownNat n) => ShapeL -> [a] -> Array n a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: Array n a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: forall n a . (KnownNat n) => ShapeL -> V.Vector a -> Array n a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (KnownNat n) => Array n a -> Array n a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: forall n' n a . (KnownNat n, KnownNat n') => ShapeL -> Array n a -> Array n' a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: ShapeL -> Array n a -> Array n a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack, 1 <= n) => Int -> Array n a -> Array n a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: a -> Array 0 a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: Array 0 a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +constant :: forall n a . (KnownNat n) => ShapeL -> a -> Array n a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +mapA :: (a -> b) -> Array n a -> Array n b +mapA f = A . G.mapA f . unA + +instance Functor (Array n) where + fmap = mapA + +instance Foldable (Array n) where + foldr = foldrA + +instance Traversable (Array n) where + traverse = traverseA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (a -> b -> c) -> Array n a -> Array n b -> Array n c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (a -> b -> c -> d) -> Array n a -> Array n b -> Array n c -> Array n d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (KnownNat n) => [(Int, Int)] -> a -> Array n a -> Array n a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (KnownNat n) => [Int] -> Array n a -> Array n a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (KnownNat n) => Array n a -> Array n a -> Array n a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (KnownNat n) => [Array n a] -> Array n a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (KnownNat (1+n)) => + Array 1 (Array n a) -> Array (1+n) a +ravel = A . G.ravel . G.mapA unA . unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: Array (1+n) a -> Array 1 (Array n a) +unravel = A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: (KnownNat n, KnownNat n') => [Int] -> Array n a -> Array n' a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: [Int] -> Array n a -> Array n a +stride ws = A . G.stride ws . unA + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: forall d p a. + (KnownNat p, KnownNat d, + -- Nonsense + (d + (p + 1)) ~ ((p + d) + 1), + (d + p) ~ (p + d), + 1 <= p + 1, + KnownNat ((p + d) + 1), + KnownNat (p + 1), + KnownNat (1 + (p + 1)) + ) => + Int -> Array (p + d) a -> Array (p + d + 1) a +rotate k = A . G.rotate @d @p k . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the array. +-- The extracted slice must fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: [(Int, Int)] -> Array n a -> Array n a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: forall n i o a b . (KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array o b) -> Array (n+i) a -> Array (n+o) b +rerank f = A . G.rerank (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o a b c . + (KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array i b -> Array o c) -> Array (n+i) a -> Array (n+i) b -> Array (n+o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array n a -> Array n a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (a -> a -> a) -> a -> Array n a -> Array 0 a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +foldrA :: (a -> b -> b) -> b -> Array n a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +traverseA :: Applicative f => (a -> f b) -> Array n a -> f (Array n b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Eq a) => Array r a -> Bool +allSameA = G.allSameA . unA + +instance (KnownNat r, Arbitrary a) => Arbitrary (Array r a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Num a) => Array r a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Num a) => Array r a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Ord a) => Array r a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Ord a) => Array r a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (a -> Bool) -> Array r a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (a -> Bool) -> Array r a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: forall r' r a . + (HasCallStack, KnownNat r, KnownNat r') => + [Int] -> ShapeL -> Array r a -> Array r' a +broadcast ds sh = A . G.broadcast ds sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: forall n a . (KnownNat n) => + ShapeL -> ([Int] -> a) -> Array n a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall a . + Int -> (a -> a) -> a -> Array 1 a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Enum a, Num a) => Int -> Array 1 a +iota = A . G.iota diff --git a/Data/Array/Internal/RankedG.hs b/Data/Array/Internal/RankedG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/RankedG.hs @@ -0,0 +1,531 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +-- | Arrays of dynamic size, but static rank. The arrays are polymorphic in the underlying +-- linear data structure used to store the actual values. +module Data.Array.Internal.RankedG( + Array(..), Vector, VecElem, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, maximumA, minimumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.Monad(replicateM) +import Control.DeepSeq +import Data.Data(Data) +import Data.List(sort) +import GHC.Generics(Generic) +import GHC.Stack +import GHC.TypeLits(Nat, type (+), KnownNat, type (<=)) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal + +-- | Arrays stored in a /v/ with values of type /a/. +data Array (n :: Nat) v a = A ShapeL (T v a) + deriving (Generic, Data) + +instance (Vector v, Show a, VecElem v a) => Show (Array n v a) where + showsPrec p a@(A s _) = showParen (p > 10) $ + showString "fromList " . showsPrec 11 s . showString " " . showsPrec 11 (toList a) + +instance (KnownNat n, Vector v, Read a, VecElem v a) => Read (Array n v a) where + readsPrec p = readParen (p > 10) $ \ r1 -> + [(fromList s xs, r4) + | ("fromList", r2) <- lex r1, (s, r3) <- readsPrec 11 r2 + , (xs, r4) <- readsPrec 11 r3, length s == valueOf @n, product s == length xs] + +instance (Vector v, Eq a, VecElem v a, Eq (v a)) => Eq (Array n v a) where + (A s v) == (A s' v') = s == s' && equalT s v v' + {-# INLINE (==) #-} + +instance (Vector v, Ord a, Ord (v a), VecElem v a) => Ord (Array n v a) where + (A s v) `compare` (A s' v') = compare s s' <> compareT s v v' + {-# INLINE compare #-} + +instance (Vector v, Pretty a, VecElem v a) => Pretty (Array n v a) where + pPrintPrec l p (A sh t) = ppT l p sh t + +instance (NFData (v a)) => NFData (Array n v a) where + rnf (A sh v) = rnf sh `seq` rnf v + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array n v a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +{-# INLINE shapeL #-} +shapeL :: Array n v a -> ShapeL +shapeL (A s _) = s + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +{-# INLINE rank #-} +rank :: forall n v a . (KnownNat n) => Array n v a -> Int +rank (A _ _) = valueOf @n + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +{-# INLINE index #-} +index :: (Vector v, HasCallStack) => Array (1+n) v a -> Int -> Array n v a +index (A (s:ss) t) i | i < 0 || i >= s = error $ "index: out of bounds " ++ show (i, s) + | otherwise = A ss $ indexT t i +index (A [] _) _ = error "index: scalar" + +-- | Convert to a list with the elements in the linearization order. +-- O(1) time. +{-# INLINE toList #-} +toList :: (Vector v, VecElem v a) => Array n v a -> [a] +toList (A sh t) = toListT sh t + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +{-# INLINE toVector #-} +toVector :: (Vector v, VecElem v a) => Array n v a -> v a +toVector (A sh t) = toVectorT sh t + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +{-# INLINE fromList #-} +fromList :: forall n v a . (HasCallStack, Vector v, VecElem v a, KnownNat n) => + ShapeL -> [a] -> Array n v a +fromList ss vs | n /= l = error $ "fromList: size mismatch " ++ show (n, l) + | length ss /= valueOf @n = error $ "fromList: rank mismatch " ++ show (length ss, valueOf @n :: Int) + | otherwise = A ss $ T st 0 $ vFromList vs + where n : st = getStridesT ss + l = length vs + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +{-# INLINE fromVector #-} +fromVector :: forall n v a . (HasCallStack, Vector v, VecElem v a, KnownNat n) => + ShapeL -> v a -> Array n v a +fromVector ss v | n /= l = error $ "fromVector: size mismatch" ++ show (n, l) + | length ss /= valueOf @n = error $ "fromVector: rank mismatch " ++ show (length ss, valueOf @n :: Int) + | otherwise = A ss $ T st 0 v + where n : st = getStridesT ss + l = vLength v + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +{-# INLINE normalize #-} +normalize :: (Vector v, VecElem v a, KnownNat n) => Array n v a -> Array n v a +normalize a = fromVector (shapeL a) $ toVector a + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +{-# INLINE reshape #-} +reshape :: forall n n' v a . (HasCallStack,Vector v, VecElem v a, KnownNat n, KnownNat n') => + ShapeL -> Array n v a -> Array n' v a +reshape sh (A sh' t@(T ost oo v)) + | n /= n' = error $ "reshape: size mismatch " ++ show (sh, sh') + | length sh /= valueOf @n' = error $ "reshape: rank mismatch " ++ show (length sh, valueOf @n :: Int) + | vLength v == 1 = A sh $ T (map (const 0) sh) 0 v -- Fast special case for singleton vector + | Just nst <- simpleReshape ost sh' sh = A sh $ T nst oo v + | otherwise = A sh $ T st 0 $ toVectorT sh' t + where n : st = getStridesT sh + n' = product sh' + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +{-# INLINE stretch #-} +stretch :: (HasCallStack) => ShapeL -> Array n v a -> Array n v a +stretch sh (A sh' vs) | Just bs <- str sh sh' = A sh $ stretchT bs vs + | otherwise = error $ "stretch: incompatible " ++ show (sh, sh') + where str [] [] = Just [] + str (x:xs) (y:ys) | x == y = (False :) <$> str xs ys + | y == 1 = (True :) <$> str xs ys + str _ _ = Nothing + +-- | Change the size of the outermost dimension by replication. +{-# INLINE stretchOuter #-} +stretchOuter :: (HasCallStack, 1 <= n) => + Int -> Array n v a -> Array n v a +stretchOuter s (A (1:sh) vs) = + A (s:sh) $ stretchT (True : map (const False) (strides vs)) vs +stretchOuter _ _ = error "stretchOuter: needs outermost dimension of size 1" + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +{-# INLINE scalar #-} +scalar :: (Vector v, VecElem v a) => a -> Array 0 v a +scalar = A [] . scalarT + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +{-# INLINE unScalar #-} +unScalar :: (Vector v, VecElem v a) => Array 0 v a -> a +unScalar (A _ t) = unScalarT t + +-- | Make an array with all elements having the same value. +-- O(1) time +{-# INLINE constant #-} +constant :: forall n v a . (Vector v, VecElem v a, KnownNat n) => + ShapeL -> a -> Array n v a +constant sh | badShape sh = error $ "constant: bad shape: " ++ show sh + | length sh /= valueOf @n = error "constant: rank mismatch" + | otherwise = A sh . constantT sh + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE mapA #-} +mapA :: (Vector v, VecElem v a, VecElem v b) => + (a -> b) -> Array n v a -> Array n v b +mapA f (A s t) = A s (mapT s f t) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWithA #-} +zipWithA :: (Vector v, VecElem v a, VecElem v b, VecElem v c) => + (a -> b -> c) -> Array n v a -> Array n v b -> Array n v c +zipWithA f (A s t) (A s' t') | s == s' = A s (zipWithT s f t t') + | otherwise = error $ "zipWithA: shape mismatch: " ++ show (s, s') + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith3A #-} +zipWith3A :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d) => + (a -> b -> c -> d) -> Array n v a -> Array n v b -> Array n v c -> Array n v d +zipWith3A f (A s t) (A s' t') (A s'' t'') | s == s' && s == s'' = A s (zipWith3T s f t t' t'') + | otherwise = error $ "zipWith3A: shape mismatch: " ++ show (s, s', s'') + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +{-# INLINE pad #-} +pad :: forall n a v . (Vector v, VecElem v a) => + [(Int, Int)] -> a -> Array n v a -> Array n v a +pad aps v (A ash at) = uncurry A $ padT v aps ash at + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +{-# INLINE transpose #-} +transpose :: forall n v a . (KnownNat n) => + [Int] -> Array n v a -> Array n v a +transpose is (A sh t) | l > n = error "transpose: rank exceeded" + | sort is /= [0 .. l-1] = + error $ "transpose: not a permutation: " ++ show is + | otherwise = A (permute is' sh) (transposeT is' t) + where l = length is + n = valueOf @n + is' = is ++ [l .. n-1] + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +{-# INLINE append #-} +append :: (Vector v, VecElem v a, KnownNat n) => + Array n v a -> Array n v a -> Array n v a +append a@(A (sa:sh) _) b@(A (sb:sh') _) | sh == sh' = + fromVector (sa+sb : sh) (vAppend (toVector a) (toVector b)) +append _ _ = error "append: bad shape" + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +{-# INLINE concatOuter #-} +concatOuter :: (Vector v, VecElem v a, KnownNat n) => [Array n v a] -> Array n v a +concatOuter [] = error "concatOuter: empty list" +concatOuter as | not $ allSame $ map tail shs = + error $ "concatOuter: non-conforming inner dimensions: " ++ show shs + | otherwise = fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = sum (map head shs) : tail sh + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +{-# INLINE ravel #-} +ravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array n v a), KnownNat (1+n)) => + Array 1 v' (Array n v a) -> Array (1+n) v a +ravel aa = + case toList aa of + [] -> error "ravel: empty array" + as | not $ allSame shs -> error $ "ravel: non-conforming inner dimensions: " ++ show shs + | otherwise -> fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = length as : sh + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +{-# INLINE unravel #-} +unravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array n v a)) => + Array (1+n) v a -> Array 1 v' (Array n v a) +unravel = rerank @1 scalar + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +{-# INLINE window #-} +window :: forall n n' v a . (Vector v, KnownNat n, KnownNat n') => + [Int] -> Array n v a -> Array n' v a +window aws _ | valueOf @n' /= length aws + valueOf @n = error $ "window: rank mismatch: " ++ show (valueOf @n :: Int, length aws, valueOf @n' :: Int) +window aws (A ash (T ss o v)) = A (win aws ash) (T (ss' ++ ss) o v) + where ss' = zipWith const ss aws + win (w:ws) (s:sh) | w <= s = s - w + 1 : win ws sh + | otherwise = error $ "window: bad window size : " ++ show (w, s) + win [] sh = aws ++ sh + win _ _ = error $ "window: rank mismatch: " ++ show (aws, ash) + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +{-# INLINE stride #-} +stride :: (Vector v) => [Int] -> Array n v a -> Array n v a +stride ats (A ash (T ss o v)) = A (str ats ash) (T (zipWith (*) (ats ++ repeat 1) ss) o v) + where str (t:ts) (s:sh) = (s+t-1) `quot` t : str ts sh + str [] sh = sh + str _ _ = error $ "stride: rank mismatch: " ++ show (ats, ash) + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: forall d p v a. + (KnownNat p, KnownNat d, + Vector v, VecElem v a, + -- Nonsense + (d + (p + 1)) ~ ((p + d) + 1), + (d + p) ~ (p + d), + 1 <= p + 1, + KnownNat ((p + d) + 1), + KnownNat (p + 1), + KnownNat (1 + (p + 1)) + ) => + Int -> Array (p + d) v a -> Array (p + d + 1) v a +rotate k a = rerank @d @p @(p + 1) f a + where + f :: Array p v a -> Array (p + 1) v a + f arr = let h:t = shapeL arr + m = product t + n = h * m + arr' = reshape @p @(p + 1) (1:h:t) arr + repeated = stretchOuter (k + 1) arr' + flattened = reshape @(p + 1) @1 [(k + 1) * n] repeated + batched = window @1 @2 [n] flattened + strided = stride [n + m] batched + in rev [0] (reshape (k:h:t) strided) + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +{-# INLINE slice #-} +slice :: [(Int, Int)] -> Array n v a -> Array n v a +slice asl (A ash (T ats ao v)) = A rsh (T ats o v) + where (o, rsh) = slc asl ash ats + slc ((k,n):sl) (s:sh) (t:ts) | k < 0 || k > s || k+n > s = error "slice: out of bounds" + | otherwise = (i + k*t, n:ns) where (i, ns) = slc sl sh ts + slc [] sh _ = (ao, sh) + slc _ _ _ = error "impossible" + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank #-} +rerank :: forall n i o v v' a b . + (Vector v, Vector v', VecElem v a, VecElem v' b + , KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i v a -> Array o v' b) -> Array (n+i) v a -> Array (n+o) v' b +rerank f (A sh t) = + ravelOuter osh $ + map (f . A ish) $ + subArraysT osh t + where (osh, ish) = splitAt (valueOf @n) sh + +ravelOuter :: (Vector v, VecElem v a, KnownNat m) => ShapeL -> [Array n v a] -> Array m v a +ravelOuter _ [] = error "ravelOuter: empty list" +ravelOuter osh as | not $ allSame shs = error $ "ravelOuter: non-conforming inner dimensions: " ++ show shs + | otherwise = fromVector sh' $ vConcat $ map toVector as + where shs@(sh:_) = map shapeL as + sh' = osh ++ sh + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank2 #-} +rerank2 :: forall n i o a b c v . + (Vector v, VecElem v a, VecElem v b, VecElem v c, + KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i v a -> Array i v b -> Array o v c) -> Array (n+i) v a -> Array (n+i) v b -> Array (n+o) v c +rerank2 f (A sha ta) (A shb tb) | take n sha /= take n shb = error "rerank2: shape mismatch" + | otherwise = + ravelOuter osh $ + zipWith (\ a b -> f (A isha a) (A ishb b)) + (subArraysT osh ta) + (subArraysT osh tb) + where (osh, isha) = splitAt n sha + ishb = drop n shb + n = valueOf @n + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +{-# INLINE rev #-} +rev :: [Int] -> Array n v a -> Array n v a +rev rs (A sh t) | all (\ r -> r >= 0 && r < n) rs = A sh (reverseT rs sh t) + | otherwise = error "reverse: bad reverse dimension" + where n = length sh + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +{-# INLINE reduce #-} +reduce :: (Vector v, VecElem v a) => + (a -> a -> a) -> a -> Array n v a -> Array 0 v a +reduce f z (A sh t) = A [] $ reduceT sh f z t + +-- | Right fold across all elements of an array. +{-# INLINE foldrA #-} +foldrA :: (Vector v, VecElem v a) => (a -> b -> b) -> b -> Array n v a -> b +foldrA f z (A sh t) = foldrT sh f z t + +-- | Constrained version of 'traverse' for 'Array's. +{-# INLINE traverseA #-} +traverseA + :: (Vector v, VecElem v a, VecElem v b, Applicative f) + => (a -> f b) -> Array n v a -> f (Array n v b) +traverseA f (A sh t) = A sh <$> traverseT sh f t + +-- | Check if all elements of the array are equal. +allSameA :: (Vector v, VecElem v a, Eq a) => Array r v a -> Bool +allSameA (A sh t) = allSameT sh t + +instance (KnownNat r, Vector v, VecElem v a, Arbitrary a) => Arbitrary (Array r v a) where + arbitrary = do + -- Don't generate huge number of elements + ss <- replicateM (valueOf @r) (getSmall . getPositive <$> arbitrary) `suchThat` ((< 10000) . product) + fromList ss <$> vector (product ss) + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Vector v, VecElem v a, Num a) => Array r v a -> a +sumA (A sh t) = sumT sh t + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Vector v, VecElem v a, Num a) => Array r v a -> a +productA (A sh t) = productT sh t + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (HasCallStack, Vector v, VecElem v a, Ord a) => Array r v a -> a +maximumA a@(A sh t) | size a > 0 = maximumT sh t + | otherwise = error "maximumA called with empty array" + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (HasCallStack, Vector v, VecElem v a, Ord a) => Array r v a -> a +minimumA a@(A sh t) | size a > 0 = minimumT sh t + | otherwise = error "minimumA called with empty array" + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Vector v, VecElem v a) => (a -> Bool) -> Array r v a -> Bool +anyA p (A sh t) = anyT sh p t + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Vector v, VecElem v a) => (a -> Bool) -> Array r v a -> Bool +allA p (A sh t) = anyT sh p t + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: forall r' r v a . + (HasCallStack, Vector v, VecElem v a, KnownNat r, KnownNat r') => + [Int] -> ShapeL -> Array r v a -> Array r' v a +broadcast ds sh a | length ds /= valueOf @r = error "broadcast: wrong number of broadcasts" + | any (\ d -> d < 0 || d >= r) ds = error "broadcast: bad dimension" + | not (ascending ds) = error "broadcast: unordered dimensions" + | length sh /= r = error "broadcast: wrong rank" + | otherwise = stretch sh $ reshape rsh a + where r = valueOf @r' + rsh = [ if i `elem` ds then s else 1 | (i, s) <- zip [0..] sh ] + ascending (x:y:ys) = x < y && ascending (y:ys) + ascending _ = True + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: forall n v a . + (KnownNat n, Vector v, VecElem v a) => + ShapeL -> ([Int] -> a) -> Array n v a +generate sh | length sh /= valueOf @n = error $ "generate: rank mismatch " ++ show (length sh, valueOf @n :: Int) + | otherwise = A sh . generateT sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall v a . + (Vector v, VecElem v a) => + Int -> (a -> a) -> a -> Array 1 v a +iterateN n f = A [n] . iterateNT n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: forall v a . + (Vector v, VecElem v a, Enum a, Num a) => + Int -> Array 1 v a +iota n = A [n] $ iotaT n diff --git a/Data/Array/Internal/RankedS.hs b/Data/Array/Internal/RankedS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/RankedS.hs @@ -0,0 +1,372 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.RankedS( + Array(..), Vector, ShapeL, Unbox, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Coerce(coerce) +import Data.Data(Data) +import qualified Data.Vector.Storable as V +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import Test.QuickCheck hiding (generate) +import GHC.Generics(Generic) +import GHC.Stack(HasCallStack) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal.DynamicS() -- Vector instance +import qualified Data.Array.Internal.Ranked as R +import qualified Data.Array.Internal.RankedG as G +import Data.Array.Internal(ShapeL, Vector(..)) + +type Unbox = V.Storable + +newtype Array n a = A { unA :: G.Array n V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array n a) + +toArrayG :: Array r a -> G.Array r V.Vector a +toArrayG = unA + +instance (Show a, Unbox a) => Show (Array n a) where + showsPrec p = showsPrec p . unA + +instance (KnownNat n, Read a, Unbox a) => Read (Array n a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array n V.Vector a) => Eq (Array n a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array n V.Vector a) => Ord (Array n a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array n a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array n a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (KnownNat n) => Array n a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: (Unbox a) => Array (1+n) a -> Int -> Array n a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: (Unbox a) => Array n a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: (Unbox a, KnownNat n) => ShapeL -> [a] -> Array n a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (Unbox a) => Array n a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: (Unbox a, KnownNat n) => ShapeL -> V.Vector a -> Array n a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (Unbox a, KnownNat n) => Array n a -> Array n a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: (Unbox a, KnownNat n, KnownNat n') => ShapeL -> Array n a -> Array n' a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: ShapeL -> Array n a -> Array n a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack, 1 <= n) => Int -> Array n a -> Array n a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array 0 a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (Unbox a) => Array 0 a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +constant :: (Unbox a, KnownNat n) => ShapeL -> a -> Array n a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +mapA :: (Unbox a, Unbox b) => + (a -> b) -> Array n a -> Array n b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (Unbox a, Unbox b, Unbox c) => + (a -> b -> c) -> Array n a -> Array n b -> Array n c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (Unbox a, Unbox b, Unbox c, Unbox d) => + (a -> b -> c -> d) -> Array n a -> Array n b -> Array n c -> Array n d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (Unbox a, KnownNat n) => [(Int, Int)] -> a -> Array n a -> Array n a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (KnownNat n) => [Int] -> Array n a -> Array n a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (Unbox a, KnownNat n) => Array n a -> Array n a -> Array n a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (Unbox a, KnownNat n) => [Array n a] -> Array n a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (Unbox a, KnownNat (1+n)) => + R.Array 1 (Array n a) -> Array (1+n) a +ravel = A . G.ravel . G.mapA unA . R.unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (Unbox a) => + Array (1+n) a -> R.Array 1 (Array n a) +unravel = R.A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +window :: (KnownNat n, KnownNat n') => [Int] -> Array n a -> Array n' a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: [Int] -> Array n a -> Array n a +stride ws = A . G.stride ws . unA + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: forall d p a. + (KnownNat p, KnownNat d, Unbox a, + -- Nonsense + (d + (p + 1)) ~ ((p + d) + 1), + (d + p) ~ (p + d), + 1 <= p + 1, + KnownNat ((p + d) + 1), + KnownNat (p + 1), + KnownNat (1 + (p + 1)) + ) => + Int -> Array (p + d) a -> Array (p + d + 1) a +rotate k = A . G.rotate @d @p k . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: [(Int, Int)] -> Array n a -> Array n a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(1) time. +rerank :: forall n i o a b . + (Unbox a, Unbox b, KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array o b) -> Array (n+i) a -> Array (n+o) b +rerank f = A . G.rerank (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o a b c . + (Unbox a, Unbox b, Unbox c, KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array i b -> Array o c) -> Array (n+i) a -> Array (n+i) b -> Array (n+o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array n a -> Array n a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a) => (a -> a -> a) -> a -> Array n a -> Array 0 a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a, Unbox b) => (a -> b -> b) -> b -> Array n a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f) + => (a -> f b) -> Array n a -> f (Array n b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Unbox a, Eq a) => Array n a -> Bool +allSameA = G.allSameA . unA + +instance (KnownNat r, Arbitrary a, Unbox a) => Arbitrary (Array r a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a) => Array r a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a) => Array r a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Unbox a, Ord a) => Array r a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Unbox a, Ord a) => Array r a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: Unbox a => (a -> Bool) -> Array r a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: Unbox a => (a -> Bool) -> Array r a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: forall r' r a . + (HasCallStack, Unbox a, KnownNat r, KnownNat r') => + [Int] -> ShapeL -> Array r a -> Array r' a +broadcast ds sh = A . G.broadcast ds sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: forall n a . (KnownNat n, Unbox a) => + ShapeL -> ([Int] -> a) -> Array n a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall a . (Unbox a) => + Int -> (a -> a) -> a -> Array 1 a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Unbox a, Enum a, Num a) => Int -> Array 1 a +iota = A . G.iota diff --git a/Data/Array/Internal/RankedU.hs b/Data/Array/Internal/RankedU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/RankedU.hs @@ -0,0 +1,371 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.RankedU( + Array(..), Vector, ShapeL, Unbox, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, concatOuter, + ravel, unravel, + window, stride, rotate, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Coerce(coerce) +import Data.Data(Data) +import qualified Data.Vector.Unboxed as V +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import Test.QuickCheck hiding (generate) +import GHC.Generics(Generic) +import GHC.Stack(HasCallStack) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal.DynamicU() -- Vector instance +import qualified Data.Array.Internal.RankedG as G +import Data.Array.Internal(ShapeL, Vector(..)) + +type Unbox = V.Unbox + +newtype Array n a = A { unA :: G.Array n V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array n a) + +toArrayG :: Array r a -> G.Array r V.Vector a +toArrayG = unA + +instance (Show a, Unbox a) => Show (Array n a) where + showsPrec p = showsPrec p . unA + +instance (KnownNat n, Read a, Unbox a) => Read (Array n a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array n V.Vector a) => Eq (Array n a) where + x == y = shapeL x == shapeL y && unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array n V.Vector a) => Ord (Array n a) where + compare x y = compare (shapeL x) (shapeL y) <> compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: Array n a -> Int +size = product . shapeL + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: Array n a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (KnownNat n) => Array n a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: (Unbox a) => Array (1+n) a -> Int -> Array n a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: (Unbox a) => Array n a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: (Unbox a, KnownNat n) => ShapeL -> [a] -> Array n a +fromList ss = A . G.fromList ss + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (Unbox a) => Array n a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: (Unbox a, KnownNat n) => ShapeL -> V.Vector a -> Array n a +fromVector ss = A . G.fromVector ss + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (Unbox a, KnownNat n) => Array n a -> Array n a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: (Unbox a, KnownNat n, KnownNat n') => ShapeL -> Array n a -> Array n' a +reshape s = A . G.reshape s . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: ShapeL -> Array n a -> Array n a +stretch s = A . G.stretch s . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (HasCallStack, 1 <= n) => Int -> Array n a -> Array n a +stretchOuter s = A . G.stretchOuter s . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array 0 a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (Unbox a) => Array 0 a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time +constant :: (Unbox a, KnownNat n) => ShapeL -> a -> Array n a +constant sh = A . G.constant sh + +-- | Map over the array elements. +-- O(n) time. +mapA :: (Unbox a, Unbox b) => + (a -> b) -> Array n a -> Array n b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (Unbox a, Unbox b, Unbox c) => + (a -> b -> c) -> Array n a -> Array n b -> Array n c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (Unbox a, Unbox b, Unbox c, Unbox d) => + (a -> b -> c -> d) -> Array n a -> Array n b -> Array n c -> Array n d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: (Unbox a, KnownNat n) => [(Int, Int)] -> a -> Array n a -> Array n a +pad ps v = A . G.pad ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: (KnownNat n) => [Int] -> Array n a -> Array n a +transpose is = A . G.transpose is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (Unbox a, KnownNat n) => Array n a -> Array n a -> Array n a +append x y = A $ G.append (unA x) (unA y) + +-- | Concatenate a number of arrays into a single array. +-- Fails if any, but the outer, dimensions differ. +-- O(n) time. +concatOuter :: (Unbox a, KnownNat n) => [Array n a] -> Array n a +concatOuter = A . G.concatOuter . coerce + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (Unbox a, Unbox (Array n a), Unbox (G.Array n V.Vector a), KnownNat (1+n)) => + Array 1 (Array n a) -> Array (1+n) a +ravel = A . G.ravel . G.mapA unA . unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (Unbox a, Unbox (Array n a), Unbox (G.Array n V.Vector a)) => + Array (1+n) a -> Array 1 (Array n a) +unravel = A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +window :: (KnownNat n, KnownNat n') => [Int] -> Array n a -> Array n' a +window ws = A . G.window ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: [Int] -> Array n a -> Array n a +stride ws = A . G.stride ws . unA + +-- | Rotate the array k times along the d'th dimension. +-- E.g., if the array shape is @[2, 3, 2]@, d is 1, and k is 4, +-- the resulting shape will be @[2, 4, 3, 2]@. +rotate :: forall d p a. + (KnownNat p, KnownNat d, Unbox a, + -- Nonsense + (d + (p + 1)) ~ ((p + d) + 1), + (d + p) ~ (p + d), + 1 <= p + 1, + KnownNat ((p + d) + 1), + KnownNat (p + 1), + KnownNat (1 + (p + 1)) + ) => + Int -> Array (p + d) a -> Array (p + d + 1) a +rotate k = A . G.rotate @d @p k . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: [(Int, Int)] -> Array n a -> Array n a +slice ss = A . G.slice ss . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(1) time. +rerank :: forall n i o a b . + (Unbox a, Unbox b, KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array o b) -> Array (n+i) a -> Array (n+o) b +rerank f = A . G.rerank (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o a b c . + (Unbox a, Unbox b, Unbox c, KnownNat n, KnownNat o, KnownNat (n+o), KnownNat (1+o)) => + (Array i a -> Array i b -> Array o c) -> Array (n+i) a -> Array (n+i) b -> Array (n+o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: [Int] -> Array n a -> Array n a +rev rs = A . G.rev rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a) => (a -> a -> a) -> a -> Array n a -> Array 0 a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a) => (a -> b -> b) -> b -> Array n a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f) + => (a -> f b) -> Array n a -> f (Array n b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Unbox a, Eq a) => Array n a -> Bool +allSameA = G.allSameA . unA + +instance (KnownNat r, Arbitrary a, Unbox a) => Arbitrary (Array r a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a) => Array r a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a) => Array r a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Unbox a, Ord a) => Array r a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Unbox a, Ord a) => Array r a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: Unbox a => (a -> Bool) -> Array r a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: Unbox a => (a -> Bool) -> Array r a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: forall r' r a . + (HasCallStack, Unbox a, KnownNat r, KnownNat r') => + [Int] -> ShapeL -> Array r a -> Array r' a +broadcast ds sh = A . G.broadcast ds sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: forall n a . (KnownNat n, Unbox a) => + ShapeL -> ([Int] -> a) -> Array n a +generate sh = A . G.generate sh + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall a . (Unbox a) => + Int -> (a -> a) -> a -> Array 1 a +iterateN n f = A . G.iterateN n f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (Unbox a, Enum a, Num a) => Int -> Array 1 a +iota = A . G.iota diff --git a/Data/Array/Internal/Shape.hs b/Data/Array/Internal/Shape.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/Shape.hs @@ -0,0 +1,240 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE FunctionalDependencies #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE NoStarIsType #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.Shape(module Data.Array.Internal.Shape) where +import Data.Proxy +import Type.Reflection +import GHC.TypeLits + +import Data.Array.Internal(valueOf) + +type DivRoundUp n m = Div (n+m-1) m + +----------------- +-- Type level properties. + +-- | Compute the rank, i.e., length of a type level shape. +type family Rank (s :: [Nat]) :: Nat where + Rank '[] = 0 + Rank (n : ns) = 1 + Rank ns + +-- | Compute the size, i.e., total number of elements of a type level shape. +type Size (s :: [Nat]) = Size' 1 s +type family Size' (a :: Nat) (s :: [Nat]) :: Nat where + Size' a '[] = a + Size' a (n : ns) = Size' (a * n) ns +-- Using an accumulating parameter generates fewer constraints. + +---------------------------------- +-- Type level shape operations + +type family (++) (xs :: [Nat]) (ys :: [Nat]) :: [Nat] where + (++) '[] ys = ys + (++) (x ': xs) ys = x ': (xs ++ ys) + +{- +-- XXX O(n^2) +type family Reverse (xs :: [Nat]) :: [Nat] where + Reverse '[] = '[] + Reverse (x ': xs) = Reverse xs ++ '[x] +-} + +type family Take (n :: Nat) (xs :: [Nat]) :: [Nat] where + Take 0 xs = '[] + Take n (x ': xs) = x ': Take (n-1) xs + +type family Drop (n :: Nat) (xs :: [Nat]) :: [Nat] where + Drop 0 xs = xs + Drop n (x ': xs) = Drop (n-1) xs + +type family Last (xs :: [Nat]) where + Last '[x] = x + Last (x ': xs) = Last xs + +type family Init (xs :: [Nat]) where + Init '[x] = '[] + Init (x ': xs) = x ': Init xs + +----------------- + +class ValidStretch (from :: [Nat]) (to :: [Nat]) where + stretching :: Proxy from -> Proxy to -> [Bool] +instance ValidStretch '[] '[] where + stretching _ _ = [] +instance (BoolVal (Stretch s m), ValidStretch ss ms) => ValidStretch (s ': ss) (m ': ms) where + stretching _ _ = boolVal (Proxy :: Proxy (Stretch s m)) : + stretching (Proxy :: Proxy ss) (Proxy :: Proxy ms) +type family Stretch (s::Nat) (m::Nat) :: Bool where + Stretch 1 m = 'True + Stretch m m = 'False + Stretch s m = TypeError ('Text "Cannot stretch " ':<>: 'ShowType s ':<>: 'Text " to " ':<>: 'ShowType m) + +class BoolVal (b :: Bool) where + boolVal :: Proxy b -> Bool +instance BoolVal 'False where + boolVal _ = False +instance BoolVal 'True where + boolVal _ = True + +----------------- + +class Padded (ps :: [(Nat, Nat)]) (sh :: [Nat]) (sh' :: [Nat]) | ps sh -> sh' where + padded :: Proxy ps -> Proxy sh -> [(Int, Int)] +instance Padded '[] sh sh where + padded _ _ = [] +instance (KnownNat l, KnownNat h, (l+s+h) ~ s', Padded ps sh sh') => + Padded ('(l,h) ': ps) (s ': sh) (s' ': sh') where + padded _ _ = (valueOf @l, valueOf @h) : padded (Proxy :: Proxy ps) (Proxy :: Proxy sh) + +----------------- + +class Permutation (is :: [Nat]) +instance (AllElem is (Count 0 is)) => Permutation is + +type family Count (i :: Nat) (xs :: [Nat]) :: [Nat] where + Count i '[] = '[] + Count i (x ': xs) = i ': Count (i+1) xs + +class AllElem (is :: [Nat]) (ns :: [Nat]) +instance AllElem '[] ns +instance (Elem i ns, AllElem is ns) => AllElem (i ': is) ns +class Elem (i :: Nat) (ns :: [Nat]) +instance (Elem' (CmpNat i n) i ns) => Elem i (n : ns) +class Elem' (e :: Ordering) (i :: Nat) (ns :: [Nat]) +instance Elem' 'EQ i ns +instance (Elem i ns) => Elem' 'LT i ns +instance (Elem i ns) => Elem' 'GT i ns + +type Permute (is :: [Nat]) (xs :: [Nat]) = Permute' is (Take (Rank is) xs) ++ Drop (Rank is) xs + +type family Permute' (is :: [Nat]) (xs :: [Nat]) where + Permute' '[] xs = '[] + Permute' (i ': is) xs = Index xs i ': Permute' is xs + +type family Index (xs :: [Nat]) (i :: Nat) where + Index (x : xs) 0 = x + Index (x : xs) i = Index xs (i-1) + +class ValidDims (rs :: [Nat]) (sh :: [Nat]) +instance (AllElem rs (Count 0 sh)) => ValidDims rs sh + +----------------- + +class Window (ws :: [Nat]) (ss :: [Nat]) (rs :: [Nat]) | ws ss -> rs +instance (Window' ws ws ss rs) => Window ws ss rs + +class Window' (ows :: [Nat]) (ws :: [Nat]) (ss :: [Nat]) (rs :: [Nat]) | ows ws ss -> rs +instance ((ows ++ ss) ~ rs) => Window' ows '[] ss rs +instance (Window' ows ws ss rs, w <= s, ((s+1)-w) ~ r) => Window' ows (w ': ws) (s ': ss) (r ': rs) + +----------------- + +class Stride (ts :: [Nat]) (ss :: [Nat]) (rs :: [Nat]) | ts ss -> rs +instance Stride '[] ss ss +instance (Stride ts ss rs, DivRoundUp s t ~ r) => Stride (t ': ts) (s ': ss) (r ': rs) + +----------------- + +class Slice (ls :: [(Nat,Nat)]) (ss :: [Nat]) (rs :: [Nat]) | ls ss -> rs where + sliceOffsets :: Proxy ls -> Proxy ss -> [Int] +instance Slice '[] ss ss where + sliceOffsets _ _ = [] +instance (Slice ls ss rs, (o+n) <= s, KnownNat o) => Slice ('(o,n) ': ls) (s ': ss) (n ': rs) where + sliceOffsets _ _ = valueOf @o : sliceOffsets (Proxy :: Proxy ls) (Proxy :: Proxy ss) + + +----------------- +-- Shape extraction + +class (Typeable s) => Shape (s :: [Nat]) where + shapeP :: Proxy s -> [Int] + sizeP :: Proxy s -> Int + +instance Shape '[] where + {-# INLINE shapeP #-} + shapeP _ = [] + {-# INLINE sizeP #-} + sizeP _ = 1 + +instance forall n s . (Shape s, KnownNat n) => Shape (n ': s) where + {-# INLINE shapeP #-} + shapeP _ = valueOf @n : shapeP (Proxy :: Proxy s) + {-# INLINE sizeP #-} + sizeP _ = valueOf @n * sizeP (Proxy :: Proxy s) + +{-# INLINE shapeT #-} +shapeT :: forall sh . (Shape sh) => [Int] +shapeT = shapeP (Proxy :: Proxy sh) + +{-# INLINE sizeT #-} +sizeT :: forall sh . (Shape sh) => Int +sizeT = sizeP (Proxy :: Proxy sh) + +-- | Turn a dynamic shape back into a type level shape. +-- @withShape sh shapeP == sh@ +withShapeP :: [Int] -> (forall sh . (Shape sh) => Proxy sh -> r) -> r +withShapeP [] f = f (Proxy :: Proxy ('[] :: [Nat])) +withShapeP (n:ns) f = + case someNatVal (toInteger n) of + Just (SomeNat (_ :: Proxy n)) -> withShapeP ns (\ (_ :: Proxy ns) -> f (Proxy :: Proxy (n ': ns))) + _ -> error $ "withShape: bad size " ++ show n + +withShape :: [Int] -> (forall sh . (Shape sh) => r) -> r +withShape sh f = withShapeP sh (\ (_ :: Proxy sh) -> f @sh) + +----------------- + +-- | Using the dimension indices /ds/, can /sh/ be broadcast into shape /sh'/? +class Broadcast (ds :: [Nat]) (sh :: [Nat]) (sh' :: [Nat]) where + broadcasting :: [Bool] +instance (Broadcast' 0 ds sh sh') => Broadcast ds sh sh' where + broadcasting = broadcasting' @0 @ds @sh @sh' + +class Broadcast' (i :: Nat) (ds :: [Nat]) (sh :: [Nat]) (sh' :: [Nat]) where + broadcasting' :: [Bool] +instance Broadcast' i '[] '[] '[] where + broadcasting' = [] +instance (Broadcast' i '[] '[] sh') => Broadcast' i '[] '[] (s : sh') where + broadcasting' = True : broadcasting' @i @'[] @'[] @sh' +instance (TypeError ('Text "Too few dimension indices")) => Broadcast' i '[] (s ': sh) sh' where + broadcasting' = undefined +instance (TypeError ('Text "Too many dimensions indices")) => Broadcast' i (d ': ds) '[] sh' where + broadcasting' = undefined +instance (TypeError ('Text "Too few result dimensions")) => Broadcast' i (d ': ds) (s ': sh) '[] where + broadcasting' = undefined +instance (Broadcast'' (CmpNat i d) i d ds (s ': sh) (s' ': sh')) => Broadcast' i (d ': ds) (s ': sh) (s' ': sh') where + broadcasting' = broadcasting'' @(CmpNat i d) @i @d @ds @(s ': sh) @(s' ': sh') + +class Broadcast'' (o :: Ordering) (i :: Nat) (d :: Nat) (ds :: [Nat]) (sh :: [Nat]) (sh' :: [Nat]) where + broadcasting'' :: [Bool] +instance (Broadcast' (i+1) ds sh rsh) => Broadcast'' 'EQ i d ds (s ': sh) (s ': rsh) where + broadcasting'' = False : broadcasting' @(i+1) @ds @sh @rsh +instance (Broadcast' (i+1) (d ': ds) sh rsh) => Broadcast'' 'LT i d ds sh (s' ': rsh) where + broadcasting'' = True : broadcasting' @(i+1) @(d ': ds) @sh @rsh +instance (TypeError ('Text "unordered dimensions")) => Broadcast'' 'GT i d ds sh rsh where + broadcasting'' = undefined diff --git a/Data/Array/Internal/Shaped.hs b/Data/Array/Internal/Shaped.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/Shaped.hs @@ -0,0 +1,369 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.Shaped( + Array(..), Shape(..), Size, Rank, Vector, ShapeL, + Window, Stride, Permute, Permutation, ValidDims, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, + ravel, unravel, + window, stride, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Data(Data) +import qualified Data.Vector as V +import GHC.Generics(Generic) +import GHC.Stack(HasCallStack) +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal.Dynamic() -- Vector instance +import qualified Data.Array.Internal.ShapedG as G +import Data.Array.Internal(ShapeL, Vector) +import Data.Array.Internal.Shape + +newtype Array sh a = A { unA :: G.Array sh V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array sh a) + +toArrayG :: Array sh a -> G.Array sh V.Vector a +toArrayG = unA + +instance (Show a, Shape sh) => Show (Array sh a) where + showsPrec p = showsPrec p . unA + +instance (Read a, Shape sh) => Read (Array sh a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array sh V.Vector a) => Eq (Array sh a) where + x == y = unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array sh V.Vector a) => Ord (Array sh a) where + compare x y = compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: forall sh a . (Shape sh) => Array sh a -> Int +size = G.size . unA + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: (Shape sh) => Array sh a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (Shape sh, KnownNat (Rank sh)) => Array sh a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: (HasCallStack, KnownNat s) => Array (s:sh) a -> Int -> Array sh a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: forall sh a . (Shape sh) => Array sh a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: forall sh a . (HasCallStack, Shape sh) => [a] -> Array sh a +fromList = A . G.fromList + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (Shape sh) => Array sh a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: forall sh a . (HasCallStack, Shape sh) => V.Vector a -> Array sh a +fromVector = A . G.fromVector + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (Shape sh) => Array sh a -> Array sh a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: forall sh' sh a . (Shape sh, Shape sh', Size sh ~ Size sh') => + Array sh a -> Array sh' a +reshape = A . G.reshape . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: forall sh' sh a . (Shape sh, Shape sh', ValidStretch sh sh') => Array sh a -> Array sh' a +stretch = A . G.stretch . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (KnownNat s, Shape sh) => Array (1 : sh) a -> Array (s : sh) a +stretchOuter = A . G.stretchOuter . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: a -> Array '[] a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: Array '[] a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time. +constant :: forall sh a . (Shape sh) => + a -> Array sh a +constant = A . G.constant + +-- | Map over the array elements. +-- O(n) time. +mapA :: (Shape sh) => (a -> b) -> Array sh a -> Array sh b +mapA f = A . G.mapA f . unA + +instance (Shape sh) => Functor (Array sh) where + fmap = mapA + +instance (Shape sh) => Foldable (Array sh) where + foldr = foldrA + +instance (Shape sh) => Traversable (Array sh) where + traverse = traverseA + +instance (Shape sh) => Applicative (Array sh) where + pure = constant + (<*>) = zipWithA ($) + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (Shape sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (Shape sh) => (a -> b -> c -> d) -> Array sh a -> Array sh b -> Array sh c -> Array sh d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: forall ps sh' sh a . (HasCallStack, Padded ps sh sh', Shape sh) => + a -> Array sh a -> Array sh' a +pad v = A . G.pad @ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: forall is sh a . + (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) => + Array sh a -> Array (Permute is sh) a +transpose = A . G.transpose @is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) => + Array (m ': sh) a -> Array (n ': sh) a -> Array (m+n ': sh) a +append x y = A $ G.append (unA x) (unA y) + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (Shape sh, KnownNat s) => + Array '[s] (Array sh a) -> Array (s:sh) a +ravel = A . G.ravel . G.mapA unA . unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (Shape sh, KnownNat s) => + Array (s:sh) a -> Array '[s] (Array sh a) +unravel = A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: forall ws sh' sh a . + (Window ws sh sh', KnownNat (Rank ws)) => + Array sh a -> Array sh' a +window = A . G.window @ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: forall ts sh' sh a . + (Stride ts sh sh', Shape ts) => + Array sh a -> Array sh' a +stride = A . G.stride @ts . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: forall sl sh' sh a . + (Slice sl sh sh') => + Array sh a -> Array sh' a +slice = A . G.slice @sl . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: forall n i o sh a b . + (Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array o b) -> Array sh a -> Array (Take n sh ++ o) b +rerank f = A . G.rerank @n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o sh a b c . + (Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array i b -> Array o c) -> Array sh a -> Array sh b -> Array (Take n sh ++ o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: forall rs sh a . (ValidDims rs sh, Shape rs, Shape sh) => + Array sh a -> Array sh a +rev = A . G.rev @rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Shape sh) => (a -> a -> a) -> a -> Array sh a -> Array '[] a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +foldrA :: (Shape sh) => (a -> b -> b) -> b -> Array sh a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +-- +-- Note that this 'Array' actually has 'Traversable' anyway. +traverseA + :: (Applicative f, Shape sh) => (a -> f b) -> Array sh a -> f (Array sh b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Shape sh, Eq a) => Array sh a -> Bool +allSameA = G.allSameA . unA + +instance (Shape sh, Arbitrary a) => Arbitrary (Array sh a) where arbitrary = A <$> arbitrary + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Num a, Shape sh) => Array sh a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Num a, Shape sh) => Array sh a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Shape sh) => (a -> Bool) -> Array sh a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Shape sh) => (a -> Bool) -> Array sh a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +{-# INLINE broadcast #-} +broadcast :: forall ds sh' sh a . + (Shape sh, Shape sh', + G.Broadcast ds sh sh') => + Array sh a -> Array sh' a +broadcast = A . G.broadcast @ds @sh' @sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Shape sh) => ([Int] -> a) -> Array sh a +generate = A . G.generate + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall n a . + (KnownNat n) => (a -> a) -> a -> Array '[n] a +iterateN f = A . G.iterateN f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (KnownNat n, Enum a, Num a) => Array '[n] a +iota = A G.iota diff --git a/Data/Array/Internal/ShapedG.hs b/Data/Array/Internal/ShapedG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/ShapedG.hs @@ -0,0 +1,470 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-} +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +-- | Arrays of static size. The arrays are polymorphic in the underlying +-- linear data structure used to store the actual values. +module Data.Array.Internal.ShapedG( + Array(..), Shape(..), Size, Rank, Vector, VecElem, + Window, Stride, Permute, Permutation, ValidDims, + Broadcast, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, + ravel, unravel, + window, stride, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Data(Data) +import Data.Proxy(Proxy(..)) +import GHC.Generics(Generic) +import GHC.Stack(HasCallStack) +import GHC.TypeLits(Nat, type (<=), KnownNat, type (+)) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass + +import Data.Array.Internal +import Data.Array.Internal.Shape + +-- | Arrays stored in a /v/ with values of type /a/. +newtype Array (sh :: [Nat]) v a = A (T v a) + deriving (Generic, Data) + +instance (Vector v, Show a, VecElem v a, Shape sh, Show (v a)) => Show (Array sh v a) where + showsPrec p a@(A _) = showParen (p > 10) $ + showString "fromList @" . showsPrec 11 (shapeL a) . showString" " . showsPrec 11 (toList a) + +instance (Shape sh, Vector v, Read a, VecElem v a) => Read (Array sh v a) where + readsPrec p = readParen (p > 10) $ \ r1 -> + [(fromList xs, r4) + | ("fromList", r2) <- lex r1, ("@", r2') <- lex r2 + , (s, r3) <- readsPrec 11 r2', (xs, r4) <- readsPrec 11 r3 + , s == shapeP (Proxy :: Proxy sh), product s == length xs] + +instance (Vector v, Eq a, VecElem v a, Eq (v a), Shape sh) + => Eq (Array sh v a) where + a@(A v) == (A v') = equalT (shapeL a) v v' + {-# INLINE (==) #-} + +instance (Vector v, Ord a, Ord (v a), VecElem v a, Shape sh) + => Ord (Array sh v a) where + a@(A v) `compare` (A v') = compareT (shapeL a) v v' + {-# INLINE compare #-} + +instance (Vector v, Pretty a, VecElem v a, Shape sh) => Pretty (Array sh v a) where + pPrintPrec l p a@(A t) = ppT l p (shapeL a) t + +instance (NFData (v a)) => NFData (Array sh v a) where + rnf (A t) = rnf t + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: forall sh v a . (Shape sh) => Array sh v a -> Int +size _ = sizeP (Proxy :: Proxy sh) + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +{-# INLINE shapeL #-} +shapeL :: forall sh v a . (Shape sh) => Array sh v a -> ShapeL +shapeL _ = shapeP (Proxy :: Proxy sh) + +-- | The rank of an array, i.e., the number if dimensions it has. +-- O(1) time. +{-# INLINE rank #-} +rank :: forall sh v a . (Shape sh, KnownNat (Rank sh)) => Array sh v a -> Int +rank _ = valueOf @(Rank sh) + +-- | Index into an array. Fails if the array has rank 0 or if the index is out of bounds. +-- O(1) time. +{-# INLINE index #-} +index :: forall s sh v a . (HasCallStack, Vector v, KnownNat s) => + Array (s:sh) v a -> Int -> Array sh v a +index (A t) i | i < 0 || i >= s = error $ "index: out of bounds " ++ show i ++ " >= " ++ show s + | otherwise = A $ indexT t i + where s = valueOf @s + +-- | Convert to a list with the elements in the linearization order. +-- O(1) time. +{-# INLINE toList #-} +toList :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> [a] +toList a@(A t) = toListT (shapeL a) t + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +{-# INLINE toVector #-} +toVector :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> v a +toVector a@(A t) = toVectorT (shapeL a) t + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +{-# INLINE fromList #-} +fromList :: forall sh v a . (HasCallStack, Vector v, VecElem v a, Shape sh) => + [a] -> Array sh v a +fromList vs | n /= l = error $ "fromList: size mismatch " ++ show (n, l) + | otherwise = A $ T st 0 $ vFromList vs + where n : st = getStridesT ss + l = length vs + ss = shapeP (Proxy :: Proxy sh) + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +{-# INLINE fromVector #-} +fromVector :: forall sh v a . (HasCallStack, Vector v, VecElem v a, Shape sh) => + v a -> Array sh v a +fromVector v | n /= l = error $ "fromVector: size mismatch" ++ show (n, l) + | otherwise = A $ T st 0 v + where n : st = getStridesT ss + l = vLength v + ss = shapeP (Proxy :: Proxy sh) + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +{-# INLINE normalize #-} +normalize :: (Vector v, VecElem v a, Shape sh) => Array sh v a -> Array sh v a +normalize = fromVector . toVector + +-- | Change the shape of an array. Type error if the arrays have different number of elements. +-- O(n) or O(1) time. +{-# INLINE reshape #-} +reshape :: forall sh' sh v a . + (Vector v, VecElem v a, Shape sh, Shape sh', Size sh ~ Size sh') => + Array sh v a -> Array sh' v a +reshape a = reshape' (shapeL a) (shapeP (Proxy :: Proxy sh')) a + +reshape' :: (Vector v, VecElem v a) => + ShapeL -> ShapeL -> Array sh v a -> Array sh' v a +reshape' sh sh' (A t@(T ost oo v)) + | vLength v == 1 = A $ T (map (const 0) sh) 0 v -- Fast special case for singleton vector + | Just nst <- simpleReshape ost sh sh' = A $ T nst oo v + | otherwise = A $ fromVectorT sh' $ toVectorT sh t + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +{-# INLINE stretch #-} +stretch :: forall sh' sh v a . (Shape sh, Shape sh', ValidStretch sh sh') => + Array sh v a -> Array sh' v a +stretch = stretch' (stretching (Proxy :: Proxy sh) (Proxy :: Proxy sh')) + +stretch' :: [Bool] -> Array sh v a -> Array sh' v a +stretch' str (A vs) = A $ stretchT str vs + +-- | Change the size of the outermost dimension by replication. +{-# INLINE stretchOuter #-} +stretchOuter :: forall s sh v a . (Shape sh) => + Array (1 : sh) v a -> Array (s : sh) v a +stretchOuter (A vs) = A $ stretchT (True : map (const False) (strides vs)) vs + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +{-# INLINE scalar #-} +scalar :: (Vector v, VecElem v a) => a -> Array '[] v a +scalar = A . scalarT + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +{-# INLINE unScalar #-} +unScalar :: (Vector v, VecElem v a) => Array '[] v a -> a +unScalar (A t) = unScalarT t + +-- | Make an array with all elements having the same value. +-- O(1) time. +{-# INLINE constant #-} +constant :: forall sh v a . (Vector v, VecElem v a, Shape sh) => + a -> Array sh v a +constant = A . constantT (shapeP (Proxy :: Proxy sh)) + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE mapA #-} +mapA :: (Vector v, VecElem v a, VecElem v b, Shape sh) => + (a -> b) -> Array sh v a -> Array sh v b +mapA f a@(A t) = A $ mapT (shapeL a) f t + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWithA #-} +zipWithA :: (Vector v, VecElem v a, VecElem v b, VecElem v c, Shape sh) => + (a -> b -> c) -> Array sh v a -> Array sh v b -> Array sh v c +zipWithA f a@(A t) (A t') = A $ zipWithT (shapeL a) f t t' + +-- | Map over the array elements. +-- O(n) time. +{-# INLINE zipWith3A #-} +zipWith3A :: (Vector v, VecElem v a, VecElem v b, VecElem v c, VecElem v d, Shape sh) => + (a -> b -> c -> d) -> Array sh v a -> Array sh v b -> Array sh v c -> Array sh v d +zipWith3A f a@(A t) (A t') (A t'') = A $ zipWith3T (shapeL a) f t t' t'' + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +{-# INLINE pad #-} +pad :: forall ps sh' sh a v . (HasCallStack, Vector v, VecElem v a, Padded ps sh sh', Shape sh) => + a -> Array sh v a -> Array sh' v a +pad v a@(A at) = A $ snd $ padT v aps ash at + where ash = shapeL a + aps = padded (Proxy :: Proxy ps) (Proxy :: Proxy sh) + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +{-# INLINE transpose #-} +transpose :: forall is sh v a . + (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) => + Array sh v a -> Array (Permute is sh) v a +transpose (A t) = A (transposeT is' t) + where l = length is + n = valueOf @(Rank sh) + is' = is ++ [l .. n-1] + is = shapeP (Proxy :: Proxy is) + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +{-# INLINE append #-} +append :: (Vector v, VecElem v a, Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) => + Array (m ': sh) v a -> Array (n ': sh) v a -> Array (m+n ': sh) v a +append a b = fromVector (vAppend (toVector a) (toVector b)) + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +{-# INLINE ravel #-} +ravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array sh v a) + , Shape sh, KnownNat s) => + Array '[s] v' (Array sh v a) -> Array (s:sh) v a +ravel = fromVector . vConcat . map toVector . toList + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +-- O(n) time. +{-# INLINE unravel #-} +unravel :: (Vector v, Vector v', VecElem v a, VecElem v' (Array sh v a) + , Shape sh, KnownNat s) => + Array (s:sh) v a -> Array '[s] v' (Array sh v a) +unravel = rerank @1 scalar + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- +-- E.g., @window [2] (fromList [4] [1,2,3,4]) == fromList [3,2] [1,2, 2,3, 3,4]@ +-- O(1) time. +-- +-- If the window parameter @ws = [w1,...,wk]@ and @wa = window ws a@ then +-- @wa `index` i1 ... `index` ik == slice [(i1,w1),...,(ik,wk)] a@. +{-# INLINE window #-} +window :: forall ws sh' sh v a . + (Window ws sh sh', Vector v, KnownNat (Rank ws)) => + Array sh v a -> Array sh' v a +window (A (T ss o v)) = A (T (ss' ++ ss) o v) + where ss' = take (valueOf @(Rank ws)) ss + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +{-# INLINE stride #-} +stride :: forall ts sh' sh v a . + (Stride ts sh sh', Vector v, Shape ts) => + Array sh v a -> Array sh' v a +stride (A (T ss o v)) = A (T (zipWith (*) (ats ++ repeat 1) ss) o v) + where ats = shapeP (Proxy :: Proxy ts) + +-- | Extract a slice of an array. +-- The first type argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the array. +-- The extracted slice must fall within the array dimensions. +-- E.g. @slice @'[ '(1,2)] (fromList @'[4] [1,2,3,4]) == fromList @'[2] [2,3]@. +-- O(1) time. +{-# INLINE slice #-} +slice :: forall sl sh' sh v a . + (Slice sl sh sh') => + Array sh v a -> Array sh' v a +slice (A (T ts o v)) = A (T ts (o+i) v) + where i = sum $ zipWith (*) ts $ sliceOffsets (Proxy :: Proxy sl) (Proxy :: Proxy sh) + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank #-} +rerank :: forall n i o sh v v' a b . + (Vector v, Vector v', VecElem v a, VecElem v' b, + Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i v a -> Array o v' b) -> Array sh v a -> Array (Take n sh ++ o) v' b +rerank f a@(A t) = + fromVector $ + vConcat $ + map (toVector . f . A) $ + subArraysT osh t + where osh = take (valueOf @n) (shapeL a) + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +{-# INLINE rerank2 #-} +rerank2 :: forall n i1 i2 o sh1 sh2 r v a b c . + (Vector v, VecElem v a, VecElem v b, VecElem v c, + Drop n sh1 ~ i1, Drop n sh2 ~ i2, Shape sh1, Shape sh2, + Take n sh1 ~ r, Take n sh2 ~ r, + KnownNat n, Shape o, Shape (r ++ o)) => + (Array i1 v a -> Array i2 v b -> Array o v c) -> Array sh1 v a -> Array sh2 v b -> Array (r ++ o) v c +rerank2 f aa@(A ta) (A tb) = + fromVector $ + vConcat $ + zipWith (\ a b -> toVector $ f (A a) (A b)) + (subArraysT osh ta) + (subArraysT osh tb) + where osh = take (valueOf @n) (shapeL aa) + + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +{-# INLINE rev #-} +rev :: forall rs sh v a . (ValidDims rs sh, Shape rs, Shape sh) => Array sh v a -> Array sh v a +rev a@(A t) = A (reverseT rs sh t) + where rs = shapeP (Proxy :: Proxy rs) + sh = shapeL a + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +{-# INLINE reduce #-} +reduce :: (Vector v, VecElem v a, Shape sh) => + (a -> a -> a) -> a -> Array sh v a -> Array '[] v a +reduce f z a@(A t) = A $ reduceT (shapeL a) f z t + +-- | Right fold across all elements of an array. +{-# INLINE foldrA #-} +foldrA + :: (Vector v, VecElem v a, Shape sh) + => (a -> b -> b) -> b -> Array sh v a -> b +foldrA f z a@(A t) = foldrT (shapeL a) f z t + +-- | Constrained version of 'traverse' for 'Array's. +{-# INLINE traverseA #-} +traverseA + :: (Vector v, VecElem v a, VecElem v b, Applicative f, Shape sh) + => (a -> f b) -> Array sh v a -> f (Array sh v b) +traverseA f a@(A t) = A <$> traverseT (shapeL a) f t + +-- | Check if all elements of the array are equal. +allSameA :: (Shape sh, Vector v, VecElem v a, Eq a) => Array sh v a -> Bool +allSameA a@(A t) = allSameT (shapeL a) t + +instance (Shape sh, Vector v, VecElem v a, Arbitrary a) => Arbitrary (Array sh v a) where + arbitrary = fromList <$> vector (sizeP (Proxy :: Proxy sh)) + +-- | Sum of all elements. +{-# INLINE sumA #-} +sumA :: (Vector v, VecElem v a, Num a, Shape sh) => Array sh v a -> a +sumA a@(A t) = sumT (shapeL a) t + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Vector v, VecElem v a, Num a, Shape sh) => Array sh v a -> a +productA a@(A t) = productT (shapeL a) t + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) => Array sh v a -> a +maximumA a@(A t) = maximumT (shapeL a) t + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Vector v, VecElem v a, Ord a, Shape sh, 1 <= Size sh) => Array sh v a -> a +minimumA a@(A t) = minimumT (shapeL a) t + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Vector v, VecElem v a, Shape sh) => (a -> Bool) -> Array sh v a -> Bool +anyA p a@(A t) = anyT (shapeL a) p t + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Vector v, VecElem v a, Shape sh) => (a -> Bool) -> Array sh v a -> Bool +allA p a@(A t) = anyT (shapeL a) p t + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +broadcast :: forall ds sh' sh v a . + (Shape sh, Shape sh', + Broadcast ds sh sh', + Vector v, VecElem v a) => + Array sh v a -> Array sh' v a +broadcast a = stretch' bc $ + reshape' sh rsh a + where sh' = shapeP (Proxy :: Proxy sh') + sh = shapeP (Proxy :: Proxy sh) + rsh = [ if b then 1 else s | (s, b) <- zip sh' bc ] + bc = broadcasting @ds @sh @sh' + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: forall sh v a . + (Vector v, VecElem v a, Shape sh) => + ([Int] -> a) -> Array sh v a +generate = A . generateT (shapeP (Proxy :: Proxy sh)) + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall n v a . + (Vector v, VecElem v a, KnownNat n) => + (a -> a) -> a -> Array '[n] v a +iterateN f = A . iterateNT (valueOf @n) f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: forall n v a . + (Vector v, VecElem v a, KnownNat n, Enum a, Num a) => + Array '[n] v a +iota = A $ iotaT (valueOf @n) diff --git a/Data/Array/Internal/ShapedS.hs b/Data/Array/Internal/ShapedS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/ShapedS.hs @@ -0,0 +1,358 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.ShapedS( + Array(..), Shape(..), Size, Rank, Vector, ShapeL, Unbox, + Window, Stride, Permute, Permutation, ValidDims, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, + ravel, unravel, + window, stride, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + anyA, allA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Data(Data) +import qualified Data.Vector.Storable as V +import GHC.Generics(Generic) +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import GHC.Stack(HasCallStack) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal.DynamicS() -- Vector instance +import qualified Data.Array.Internal.Shaped as S +import qualified Data.Array.Internal.ShapedG as G +import Data.Array.Internal(ShapeL, Vector) +import Data.Array.Internal.Shape + +type Unbox = V.Storable + +newtype Array sh a = A { unA :: G.Array sh V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array sh a) + +toArrayG :: Array sh a -> G.Array sh V.Vector a +toArrayG = unA + +instance (Unbox a, Show a, Shape sh) => Show (Array sh a) where + showsPrec p = showsPrec p . unA + +instance (Read a, Unbox a, Shape sh) => Read (Array sh a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array sh V.Vector a) => Eq (Array sh a) where + x == y = unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array sh V.Vector a) => Ord (Array sh a) where + compare x y = compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: forall sh a . (Shape sh) => Array sh a -> Int +size = G.size . unA + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: (Shape sh) => Array sh a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (Shape sh, KnownNat (Rank sh)) => Array sh a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: (Unbox a, KnownNat s) => Array (s:sh) a -> Int -> Array sh a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: forall sh a . (Unbox a, Shape sh) => Array sh a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: forall sh a . (HasCallStack, Unbox a, Shape sh) => [a] -> Array sh a +fromList = A . G.fromList + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (Unbox a, Shape sh) => Array sh a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: forall sh a . (HasCallStack, Unbox a, Shape sh) => V.Vector a -> Array sh a +fromVector = A . G.fromVector + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (Unbox a, Shape sh) => Array sh a -> Array sh a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: forall sh' sh a . (Unbox a, Shape sh, Shape sh', Size sh ~ Size sh') => + Array sh a -> Array sh' a +reshape = A . G.reshape . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: forall sh' sh a . (Shape sh, Shape sh', ValidStretch sh sh') => Array sh a -> Array sh' a +stretch = A . G.stretch . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (KnownNat s, Shape sh) => Array (1 : sh) a -> Array (s : sh) a +stretchOuter = A . G.stretchOuter . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array '[] a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (Unbox a) => Array '[] a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time. +constant :: forall sh a . (Unbox a, Shape sh) => + a -> Array sh a +constant = A . G.constant + +-- | Map over the array elements. +-- O(n) time. +mapA :: (Unbox a, Unbox b, Shape sh) => (a -> b) -> Array sh a -> Array sh b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (Unbox a, Unbox b, Unbox c, Shape sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (Unbox a, Unbox b, Unbox c, Unbox d, Shape sh) => (a -> b -> c -> d) -> Array sh a -> Array sh b -> Array sh c -> Array sh d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: forall ps sh' sh a . (Unbox a, Padded ps sh sh', Shape sh) => + a -> Array sh a -> Array sh' a +pad v = A . G.pad @ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: forall is sh a . + (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) => + Array sh a -> Array (Permute is sh) a +transpose = A . G.transpose @is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (Unbox a, Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) => + Array (m ': sh) a -> Array (n ': sh) a -> Array (m+n ': sh) a +append x y = A $ G.append (unA x) (unA y) + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (Unbox a, Shape sh, KnownNat s) => + S.Array '[s] (Array sh a) -> Array (s:sh) a +ravel = A . G.ravel . G.mapA unA . S.unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (Unbox a, Shape sh, KnownNat s) => + Array (s:sh) a -> S.Array '[s] (Array sh a) +unravel = S.A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: forall ws sh' sh a . + (Window ws sh sh', KnownNat (Rank ws)) => + Array sh a -> Array sh' a +window = A . G.window @ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: forall ts sh' sh a . + (Stride ts sh sh', Shape ts) => + Array sh a -> Array sh' a +stride = A . G.stride @ts . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: forall sl sh' sh a . + (Slice sl sh sh') => + Array sh a -> Array sh' a +slice = A . G.slice @sl . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: forall n i o sh a b . + (Unbox a, Unbox b, + Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array o b) -> Array sh a -> Array (Take n sh ++ o) b +rerank f = A . G.rerank @n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o sh a b c . + (Unbox a, Unbox b, Unbox c, + Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array i b -> Array o c) -> Array sh a -> Array sh b -> Array (Take n sh ++ o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: forall rs sh a . (ValidDims rs sh, Shape rs, Shape sh) => + Array sh a -> Array sh a +rev = A . G.rev @rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a, Shape sh) => (a -> a -> a) -> a -> Array sh a -> Array '[] a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a, Shape sh) => (a -> b -> b) -> b -> Array sh a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f, Shape sh) + => (a -> f b) -> Array sh a -> f (Array sh b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Shape sh, Unbox a, Eq a) => Array sh a -> Bool +allSameA = G.allSameA . unA + +instance (Shape sh, Arbitrary a, Unbox a) => Arbitrary (Array sh a) where arbitrary = A <$> arbitrary + +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a, Shape sh) => Array sh a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a, Shape sh) => Array sh a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +{-# INLINE broadcast #-} +broadcast :: forall ds sh' sh a . + (Unbox a, Shape sh, Shape sh', + G.Broadcast ds sh sh') => + Array sh a -> Array sh' a +broadcast = A . G.broadcast @ds @sh' @sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Unbox a, Shape sh) => ([Int] -> a) -> Array sh a +generate = A . G.generate + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall n a . + (Unbox a, KnownNat n) => (a -> a) -> a -> Array '[n] a +iterateN f = A . G.iterateN f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (KnownNat n, Unbox a, Enum a, Num a) => Array '[n] a +iota = A G.iota diff --git a/Data/Array/Internal/ShapedU.hs b/Data/Array/Internal/ShapedU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Internal/ShapedU.hs @@ -0,0 +1,357 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE ConstraintKinds #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE ExplicitNamespaces #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} +module Data.Array.Internal.ShapedU( + Array(..), Shape(..), Size, Rank, Vector, ShapeL, Unbox, + Window, Stride, Permute, Permutation, ValidDims, + toArrayG, + size, shapeL, rank, + toList, fromList, toVector, fromVector, + normalize, + scalar, unScalar, constant, + reshape, stretch, stretchOuter, transpose, + index, pad, + mapA, zipWithA, zipWith3A, + append, + ravel, unravel, + window, stride, + slice, rerank, rerank2, rev, + reduce, foldrA, traverseA, + allSameA, + sumA, productA, minimumA, maximumA, + allA, anyA, + broadcast, + generate, iterateN, iota, + ) where +import Control.DeepSeq +import Data.Data(Data) +import qualified Data.Vector.Unboxed as V +import GHC.Generics(Generic) +import GHC.TypeLits(KnownNat, type (+), type (<=)) +import Test.QuickCheck hiding (generate) +import Text.PrettyPrint.HughesPJClass hiding ((<>)) + +import Data.Array.Internal.DynamicU() -- Vector instance +import qualified Data.Array.Internal.ShapedG as G +import Data.Array.Internal(ShapeL, Vector) +import Data.Array.Internal.Shape + +type Unbox = V.Unbox + +newtype Array sh a = A { unA :: G.Array sh V.Vector a } + deriving (Pretty, Generic, Data) + +instance NFData a => NFData (Array sh a) + +toArrayG :: Array sh a -> G.Array sh V.Vector a +toArrayG = unA + +instance (Unbox a, Show a, Shape sh) => Show (Array sh a) where + showsPrec p = showsPrec p . unA + +instance (Read a, Unbox a, Shape sh) => Read (Array sh a) where + readsPrec p s = [(A a, r) | (a, r) <- readsPrec p s] + +instance Eq (G.Array sh V.Vector a) => Eq (Array sh a) where + x == y = unA x == unA y + {-# INLINE (==) #-} + +instance Ord (G.Array sh V.Vector a) => Ord (Array sh a) where + compare x y = compare (unA x) (unA y) + {-# INLINE compare #-} + +-- | The number of elements in the array. +{-# INLINE size #-} +size :: forall sh a . (Shape sh) => Array sh a -> Int +size = G.size . unA + +-- | The shape of an array, i.e., a list of the sizes of its dimensions. +-- In the linearization of the array the outermost (i.e. first list element) +-- varies most slowly. +-- O(1) time. +shapeL :: (Shape sh) => Array sh a -> ShapeL +shapeL = G.shapeL . unA + +-- | The rank of an array, i.e., the number if dimensions it has, +-- which is the @n@ in @Array n a@. +-- O(1) time. +rank :: (Shape sh, KnownNat (Rank sh)) => Array sh a -> Int +rank = G.rank . unA + +-- | Index into an array. Fails if the index is out of bounds. +-- O(1) time. +index :: (Unbox a, KnownNat s) => Array (s:sh) a -> Int -> Array sh a +index a = A . G.index (unA a) + +-- | Convert to a list with the elements in the linearization order. +-- O(n) time. +toList :: forall sh a . (Unbox a, Shape sh) => Array sh a -> [a] +toList = G.toList . unA + +-- | Convert from a list with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(n) time. +fromList :: forall sh a . (Unbox a, Shape sh) => [a] -> Array sh a +fromList = A . G.fromList + +-- | Convert to a vector with the elements in the linearization order. +-- O(n) or O(1) time (the latter if the vector is already in the linearization order). +toVector :: (Unbox a, Shape sh) => Array sh a -> V.Vector a +toVector = G.toVector . unA + +-- | Convert from a vector with the elements given in the linearization order. +-- Fails if the given shape does not have the same number of elements as the list. +-- O(1) time. +fromVector :: forall sh a . (Unbox a, Shape sh) => V.Vector a -> Array sh a +fromVector = A . G.fromVector + +-- | Make sure the underlying vector is in the linearization order. +-- This is semantically an identity function, but can have big performance +-- implications. +-- O(n) or O(1) time. +normalize :: (Unbox a, Shape sh) => Array sh a -> Array sh a +normalize = A . G.normalize . unA + +-- | Change the shape of an array. Fails if the arrays have different number of elements. +-- O(n) or O(1) time. +reshape :: forall sh' sh a . (Unbox a, Shape sh, Shape sh', Size sh ~ Size sh') => + Array sh a -> Array sh' a +reshape = A . G.reshape . unA + +-- | Change the size of dimensions with size 1. These dimension can be changed to any size. +-- All other dimensions must remain the same. +-- O(1) time. +stretch :: forall sh' sh a . (Shape sh, Shape sh', ValidStretch sh sh') => Array sh a -> Array sh' a +stretch = A . G.stretch . unA + +-- | Change the size of the outermost dimension by replication. +stretchOuter :: (KnownNat s, Shape sh) => Array (1 : sh) a -> Array (s : sh) a +stretchOuter = A . G.stretchOuter . unA + +-- | Convert a value to a scalar (rank 0) array. +-- O(1) time. +scalar :: (Unbox a) => a -> Array '[] a +scalar = A . G.scalar + +-- | Convert a scalar (rank 0) array to a value. +-- O(1) time. +unScalar :: (Unbox a) => Array '[] a -> a +unScalar = G.unScalar . unA + +-- | Make an array with all elements having the same value. +-- O(1) time. +constant :: forall sh a . (Unbox a, Shape sh) => + a -> Array sh a +constant = A . G.constant + +-- | Map over the array elements. +-- O(n) time. +mapA :: (Unbox a, Unbox b, Shape sh) => (a -> b) -> Array sh a -> Array sh b +mapA f = A . G.mapA f . unA + +-- | Map over the array elements. +-- O(n) time. +zipWithA :: (Unbox a, Unbox b, Unbox c, Shape sh) => (a -> b -> c) -> Array sh a -> Array sh b -> Array sh c +zipWithA f a b = A $ G.zipWithA f (unA a) (unA b) + +-- | Map over the array elements. +-- O(n) time. +zipWith3A :: (Unbox a, Unbox b, Unbox c, Unbox d, Shape sh) => (a -> b -> c -> d) -> Array sh a -> Array sh b -> Array sh c -> Array sh d +zipWith3A f a b c = A $ G.zipWith3A f (unA a) (unA b) (unA c) + +-- | Pad each dimension on the low and high side with the given value. +-- O(n) time. +pad :: forall ps sh' sh a . (Unbox a, Padded ps sh sh', Shape sh) => + a -> Array sh a -> Array sh' a +pad v = A . G.pad @ps v . unA + +-- | Do an arbitrary array transposition. +-- Fails if the transposition argument is not a permutation of the numbers +-- [0..r-1], where r is the rank of the array. +-- O(1) time. +transpose :: forall is sh a . + (Permutation is, Rank is <= Rank sh, Shape sh, Shape is, KnownNat (Rank sh)) => + Array sh a -> Array (Permute is sh) a +transpose = A . G.transpose @is . unA + +-- | Append two arrays along the outermost dimension. +-- All dimensions, except the outermost, must be the same. +-- O(n) time. +append :: (Unbox a, Shape sh, KnownNat m, KnownNat n, KnownNat (m+n)) => + Array (m ': sh) a -> Array (n ': sh) a -> Array (m+n ': sh) a +append x y = A $ G.append (unA x) (unA y) + +-- | Turn a rank-1 array of arrays into a single array by making the outer array into the outermost +-- dimension of the result array. All the arrays must have the same shape. +-- O(n) time. +ravel :: (Unbox a, Unbox (Array sh a), Unbox (G.Array sh V.Vector a), Shape sh, KnownNat s) => + Array '[s] (Array sh a) -> Array (s:sh) a +ravel = A . G.ravel . G.mapA unA . unA + +-- | Turn an array into a nested array, this is the inverse of 'ravel'. +-- I.e., @ravel . unravel == id@. +unravel :: (Unbox a, Unbox (Array sh a), Unbox (G.Array sh V.Vector a), Shape sh, KnownNat s) => + Array (s:sh) a -> Array '[s] (Array sh a) +unravel = A . G.mapA A . G.unravel . unA + +-- | Make a window of the outermost dimensions. +-- The rank increases with the length of the window list. +-- E.g., if the shape of the array is @[10,12,8]@ and +-- the window size is @[3,3]@ then the resulting array will have shape +-- @[8,10,3,3,8]@. +-- O(1) time. +window :: forall ws sh' sh a . + (Window ws sh sh', KnownNat (Rank ws)) => + Array sh a -> Array sh' a +window = A . G.window @ws . unA + +-- | Stride the outermost dimensions. +-- E.g., if the array shape is @[10,12,8]@ and the strides are +-- @[2,2]@ then the resulting shape will be @[5,6,8]@. +-- O(1) time. +stride :: forall ts sh' sh a . + (Stride ts sh sh', Shape ts) => + Array sh a -> Array sh' a +stride = A . G.stride @ts . unA + +-- | Extract a slice of an array. +-- The first argument is a list of (offset, length) pairs. +-- The length of the slicing argument must not exceed the rank of the arrar. +-- The extracted slice mul fall within the array dimensions. +-- E.g. @slice [1,2] (fromList [4] [1,2,3,4]) == [2,3]@. +-- O(1) time. +slice :: forall sl sh' sh a . + (Slice sl sh sh') => + Array sh a -> Array sh' a +slice = A . G.slice @sl . unA + +-- | Apply a function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank :: forall n i o sh a b . + (Unbox a, Unbox b, + Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array o b) -> Array sh a -> Array (Take n sh ++ o) b +rerank f = A . G.rerank @n (unA . f . A) . unA + +-- | Apply a two-argument function to the subarrays /n/ levels down and make +-- the results into an array with the same /n/ outermost dimensions. +-- The /n/ must not exceed the rank of the array. +-- O(n) time. +rerank2 :: forall n i o sh a b c . + (Unbox a, Unbox b, Unbox c, + Drop n sh ~ i, Shape sh, KnownNat n, Shape o, Shape (Take n sh ++ o)) => + (Array i a -> Array i b -> Array o c) -> Array sh a -> Array sh b -> Array (Take n sh ++ o) c +rerank2 f ta tb = A $ G.rerank2 @n (\ a b -> unA $ f (A a) (A b)) (unA ta) (unA tb) + +-- | Reverse the given dimensions, with the outermost being dimension 0. +-- O(1) time. +rev :: forall rs sh a . (ValidDims rs sh, Shape rs, Shape sh) => + Array sh a -> Array sh a +rev = A . G.rev @rs . unA + +-- | Reduce all elements of an array into a rank 0 array. +-- To reduce parts use 'rerank' and 'transpose' together with 'reduce'. +-- O(n) time. +reduce :: (Unbox a, Shape sh) => (a -> a -> a) -> a -> Array sh a -> Array '[] a +reduce f z = A . G.reduce f z . unA + +-- | Constrained version of 'foldr' for Arrays. +foldrA :: (Unbox a, Shape sh) => (a -> b -> b) -> b -> Array sh a -> b +foldrA f z = G.foldrA f z . unA + +-- | Constrained version of 'traverse' for Arrays. +traverseA + :: (Unbox a, Unbox b, Applicative f, Shape sh) + => (a -> f b) -> Array sh a -> f (Array sh b) +traverseA f = fmap A . G.traverseA f . unA + +-- | Check if all elements of the array are equal. +allSameA :: (Shape sh, Unbox a, Eq a) => Array sh a -> Bool +allSameA = G.allSameA . unA + +instance (Shape sh, Arbitrary a, Unbox a) => Arbitrary (Array sh a) where arbitrary = A <$> arbitrary +-- | Sum of all elements. + +{-# INLINE sumA #-} +sumA :: (Unbox a, Num a, Shape sh) => Array sh a -> a +sumA = G.sumA . unA + +-- | Product of all elements. +{-# INLINE productA #-} +productA :: (Unbox a, Num a, Shape sh) => Array sh a -> a +productA = G.productA . unA + +-- | Maximum of all elements. +{-# INLINE maximumA #-} +maximumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +maximumA = G.maximumA . unA + +-- | Minimum of all elements. +{-# INLINE minimumA #-} +minimumA :: (Unbox a, Ord a, Shape sh, 1 <= Size sh) => Array sh a -> a +minimumA = G.minimumA . unA + +-- | Test if the predicate holds for any element. +{-# INLINE anyA #-} +anyA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool +anyA p = G.anyA p . unA + +-- | Test if the predicate holds for all elements. +{-# INLINE allA #-} +allA :: (Shape sh, Unbox a) => (a -> Bool) -> Array sh a -> Bool +allA p = G.allA p . unA + +-- | Put the dimensions of the argument into the specified dimensions, +-- and just replicate the data along all other dimensions. +-- The list of dimensions indicies must have the same rank as the argument array +-- and it must be strictly ascending. +{-# INLINE broadcast #-} +broadcast :: forall ds sh' sh a . + (Unbox a, Shape sh, Shape sh', + G.Broadcast ds sh sh') => + Array sh a -> Array sh' a +broadcast = A . G.broadcast @ds @sh' @sh . unA + +-- | Generate an array with a function that computes the value for each index. +{-# INLINE generate #-} +generate :: (Unbox a, Shape sh) => ([Int] -> a) -> Array sh a +generate = A . G.generate + +-- | Iterate a function n times. +{-# INLINE iterateN #-} +iterateN :: forall n a . + (Unbox a, KnownNat n) => (a -> a) -> a -> Array '[n] a +iterateN f = A . G.iterateN f + +-- | Generate a vector from 0 to n-1. +{-# INLINE iota #-} +iota :: (KnownNat n, Unbox a, Enum a, Num a) => Array '[n] a +iota = A G.iota diff --git a/Data/Array/Ranked.hs b/Data/Array/Ranked.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Ranked.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.Ranked(module A) where +import Data.Array.Internal.Ranked as A hiding(Array(..)) +import Data.Array.Internal.Ranked as A(Array) diff --git a/Data/Array/RankedG.hs b/Data/Array/RankedG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/RankedG.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.RankedG(module A) where +import Data.Array.Internal.RankedG as A hiding(Array(..)) +import Data.Array.Internal.RankedG as A(Array) diff --git a/Data/Array/RankedS.hs b/Data/Array/RankedS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/RankedS.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.RankedS(module A) where +import Data.Array.Internal.RankedS as A hiding(Array(..)) +import Data.Array.Internal.RankedS as A(Array) diff --git a/Data/Array/RankedU.hs b/Data/Array/RankedU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/RankedU.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.RankedU(module A) where +import Data.Array.Internal.RankedU as A hiding(Array(..)) +import Data.Array.Internal.RankedU as A(Array) diff --git a/Data/Array/Shape.hs b/Data/Array/Shape.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Shape.hs @@ -0,0 +1,16 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.Shape(module Data.Array.Internal.Shape) where +import Data.Array.Internal.Shape diff --git a/Data/Array/Shaped.hs b/Data/Array/Shaped.hs new file mode 100644 --- /dev/null +++ b/Data/Array/Shaped.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.Shaped(module A) where +import Data.Array.Internal.Shaped as A hiding(Array(..)) +import Data.Array.Internal.Shaped as A(Array) diff --git a/Data/Array/ShapedG.hs b/Data/Array/ShapedG.hs new file mode 100644 --- /dev/null +++ b/Data/Array/ShapedG.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.ShapedG(module A) where +import Data.Array.Internal.ShapedG as A hiding(Array(..)) +import Data.Array.Internal.ShapedG as A(Array) diff --git a/Data/Array/ShapedS.hs b/Data/Array/ShapedS.hs new file mode 100644 --- /dev/null +++ b/Data/Array/ShapedS.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.ShapedS(module A) where +import Data.Array.Internal.ShapedS as A hiding(Array(..)) +import Data.Array.Internal.ShapedS as A(Array) diff --git a/Data/Array/ShapedU.hs b/Data/Array/ShapedU.hs new file mode 100644 --- /dev/null +++ b/Data/Array/ShapedU.hs @@ -0,0 +1,17 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +module Data.Array.ShapedU(module A) where +import Data.Array.Internal.ShapedU as A hiding(Array(..)) +import Data.Array.Internal.ShapedU as A(Array) diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,13 @@ +Copyright 2018-2020 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,116 @@ +# Orthotope + +## Disclaimer + +This is not an officially supported Google product. + +## Summary + +This is a library for multi-dimensional arrays inspired by APL. + +## Multi-dimensional arrays + +Each array has a number of elements of the same type, and a *shape*. The shape +can be described by a list of integers that gives the size for each of the +dimensions. E.g. the array shape `[2,3]` is a 2x3 matrix (2 rows, 3 +columns), and the shape `[]` is a single value (a scalar). +The number of dimensions is called the *rank* of the array. + +The shape may or may not be part of the type, depending on which version of the +API you use. + +## API variants + +The API comes in many variants, depending on how strongly typed it is and what +the underlying storage is. + +### Types + +* `Dynamic`, the shape is not part of the type, but is checked at runtime. + E.g., `Array Float` is an array of `Float` which can have any shape. + +* `Ranked`, the rank of the array is part of the type, but the actual sizes of + the dimensions are checked at runtime. E.g., `Array 2 Float` is the type of + 2-dimensional arrays (i.e., matrices) of `Float`. + +* `Shaped`, the shape of the array is part of the type and is checked + statically. E.g., `Array [2,3] Float` is the type of 2x3 arrays of `Float`. + +Converting between these types is cheap since they all share the same underlying +trepresentation. + +### Storage + +Each of the type variants has several storage variants, indicated by a suffix of +the module names. + +* `G` The generic array type where you can provide your own storage. + +* `S` Uses `Data.Vector.Storable` for storage. + +* `U` Uses `Data.Vector.Unboxed` for storage. + +* ` ` (empty suffix) Uses `Data.Vector` for storage. + +Conversion between different storage types requires copying the data, so it is +not a cheap operation. + +## API + +The library API is mostly structural operations, i.e., operations that +treat the elements in a uniform way. For more algorithmic operations, +e.g., matrix multiplication, we suggest using a different library, +like `hmatrix`. + +### Examples using `Dynamic` + +Some preliminaries: + +``` +> import Data.Array.Dynamic +> import Text.PrettyPrint.HughesPJClass +> pp = putStrLn . prettyShow +``` + +An easy way to create an array from a list is to use `fromList`; +the first argument is the shape of the array. + +``` +> m = fromList [2,3] [1..6] +> m +fromList [2,3] [1,2,3,4,5,6] +> shapeL m +[2,3] +> size m +6 +``` + +Arrays can be pretty printed. They are shown in the APL way: +The innermost dimension on a line, the next dimension vertically, +the next dimension vertically with an empty line in betwee, and so on. + +``` +> pp m +1 2 3 +4 5 6 +``` + +We can have an arbitrary number of dimensions. + +``` +> s = fromList [] [42] +> v = fromList [3] [7,8,9] +> a = fromList [2,3,4] [1..24] +> pp s +42 +> pp v +7 8 9 +> pp a + 1 2 3 4 + 5 6 7 8 + 9 10 11 12 + +13 14 15 16 +17 18 19 20 +21 22 23 24 +``` 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/orthotope.cabal b/orthotope.cabal new file mode 100644 --- /dev/null +++ b/orthotope.cabal @@ -0,0 +1,92 @@ +name: orthotope +version: 0.1.0.0 +synopsis: Multidimensional arrays inspired by APL +license: Apache +license-file: LICENSE +copyright: 2018 Google Inc +category: array +maintainer: lennart@augustsson.net +description: Multidimensional arrays inspired by APL. + The library contains a wide variety of structurual + operations on arrays, but not actual algorithms. +build-type: Simple +cabal-version: >=1.10 + +extra-source-files: + CHANGELOG.md + LICENSE + README.md + +source-repository head + type: git + location: https://github.com/augustss/orthotope + +library + hs-source-dirs: . + ghc-options: -Wall + exposed-modules: Data.Array.Convert + , Data.Array.Dynamic + , Data.Array.DynamicG + , Data.Array.DynamicS + , Data.Array.DynamicU + , Data.Array.Ranked + , Data.Array.RankedG + , Data.Array.RankedS + , Data.Array.RankedU + , Data.Array.Shape + , Data.Array.Shaped + , Data.Array.ShapedG + , Data.Array.ShapedS + , Data.Array.ShapedU + + , Data.Array.Internal + , Data.Array.Internal.Dynamic + , Data.Array.Internal.DynamicG + , Data.Array.Internal.DynamicS + , Data.Array.Internal.DynamicU + , Data.Array.Internal.Ranked + , Data.Array.Internal.RankedG + , Data.Array.Internal.RankedS + , Data.Array.Internal.RankedU + , Data.Array.Internal.Shape + , Data.Array.Internal.Shaped + , Data.Array.Internal.ShapedG + , Data.Array.Internal.ShapedS + , Data.Array.Internal.ShapedU + + build-depends: base >= 4.14 && < 4.15 + , deepseq + , dlist + , pretty + , QuickCheck + , vector + + default-language: Haskell2010 + +test-suite tests + type: exitcode-stdio-1.0 + hs-source-dirs: tests + main-is: Tests.hs + + default-language: Haskell2010 + + other-modules: DynamicTest + , DynamicUTest + , RankedTest + , RankedUTest + , ShapedTest + , ShapedUTest + + ghc-options: + -Wall -rtsopts + + build-depends: + base, + deepseq, + orthotope, + test-framework >= 0.3.3, + test-framework-quickcheck2, + test-framework-hunit, + HUnit, + QuickCheck >= 2.4.0.1, + vector diff --git a/tests/DynamicTest.hs b/tests/DynamicTest.hs new file mode 100644 --- /dev/null +++ b/tests/DynamicTest.hs @@ -0,0 +1,361 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE ScopedTypeVariables #-} +module DynamicTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.Dynamic +import qualified Data.Vector as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "Dynamic" $ + let a1, a2 :: Array Int + a1 = fromList [2,3] [1..6] + a2 = transpose [1,0] a1 + show_1 = assertEqual "1" "fromList [2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList [3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a2) + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a2) + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [3] [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [2] [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [] [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [] [1,2::Int]) + fromList_2 = assertThrows "sh" (fromList [4,5] [1,2::Int]) + fromVector_1 = assertEqual "1" a1 (fromVector [2,3] $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [6] [1..6]) (reshape [6] a1) + reshape_2 = assertEqual "1" (fromList [1,2,3,1] [1,4,2,5,3,6]) (reshape [1,2,3,1] a2) + a3, a4 :: Array Int + a3 = fromList [1] [5] + a4 = fromList [] [5] + stretch_1 = assertEqual "1" (fromList [3] [5,5,5]) (stretch [3] a3) + stretch_2 = assertEqual "2" (fromList [2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch [2,2,3,2] $ reshape [1,2,3,1] a1) + stretch_3 = assertThrows "3" (stretch [1,2] a3) + stretch_4 = assertThrows "4" (stretch [4,3] a1) + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + unScalar_2 = assertThrows "2" (unScalar a3) + constant_1 = assertEqual "1" (fromList [2,3] [1,1,1,1,1,1]) (constant [2,3] (1::Int)) + mapA_1 = assertEqual "1" (fromList [2,3] [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [3,2] [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,3] [2,4..12]) (zipWithA (+) a1 a1) + zipWithA_2 = assertThrows "2" (zipWithA (+) a1 a2) + zipWith3A_1 = assertEqual "1" (fromList [2,3] [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [5,10] [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad [(1,2),(3,4)] 9 a1) + pad_2 = assertThrows "2" (pad [(1,1),(1,1),(1,1)] 0 a1) + a5 :: Array Int + a5 = fromList [2,3,4] [1..24] + transpose_1 = assertEqual "1" (fromList [2,3,4] [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose [0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [2,4,3] [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose [0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [3,4,2] [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose [1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [4,2,3] [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose [2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [4,3,2] [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose [2,1,0] a5) + transpose_7 = assertThrows "7" (transpose [0,1,2,3] a5) + transpose_8 = assertThrows "7" (transpose [0,1,3] a5) + transpose_9 = assertEqual "9" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0] a5) + append_1 = assertEqual "1" (fromList [3,3] [1..9]) + (append a1 (fromList [1,3] [7,8,9])) + concatOuter_1 = assertEqual "1" (fromList [6,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (concatOuter [a1, concatOuter [a1,a1]]) + concatOuter_2 = assertThrows "2" (concatOuter [a1, a2]) + ravel_1 = assertEqual "1" (fromList [3,2,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (ravel $ fromList [3] [a1,a1,a1]) + ravel_2 = assertThrows "2" (ravel $ fromList [2] [a1, concatOuter [a1,a1]]) + unravel_1 = assertEqual "1" [a1,a1,a1] + (toList $ unravel $ ravel $ fromList [3] [a1,a1,a1]) + a6 :: Array Int + a6 = fromList [4,5] [1..20] + window_1 = assertEqual "1" (fromList [2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window [3,3] a6) + window_2 = assertThrows "2" (window [3,6] a6) + window_3 = assertThrows "3" (window [3,3,3] a6) + stride_1 = assertEqual "1" (fromList [2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride [1,2,2] a5) + stride_2 = assertThrows "2" (stride [1,2,2] a1) + rotate_1 = assertEqual "1" (fromList [2, 4, 3, 2] + [1, 2, 3, 4, 5, 6, + 5, 6, 1, 2, 3, 4, + 3, 4, 5, 6, 1, 2, + 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, + 11, 12, 7, 8, 9, 10, + 9, 10, 11, 12, 7, 8, + 7, 8, 9, 10, 11, 12]) + (rotate 1 4 $ fromList [2, 3, 2] [1 .. 12::Int]) + slice_1 = assertEqual "1" (fromList [2,2,1] [8,12,20,24]) + (slice [(0,2),(1,2),(3,1)] a5) + slice_2 = assertThrows "2" (slice [(0,0)] a4) + slice_3 = assertThrows "3" (slice [(-1,1)] a5) + slice_4 = assertThrows "4" (slice [(10,0)] a5) + slice_5 = assertThrows "5" (slice [(0,3)] a5) + box = scalar . Just + rerank_1 = assertEqual "1" (box a5) + (rerank 0 box a5) + rerank_2 = assertEqual "2" (fromList [2] [Just $ fromList [3,4] [1,2,3,4,5,6,7,8,9,10,11,12], + Just $ fromList [3,4] [13,14,15,16,17,18,19,20,21,22,23,24]]) + (rerank 1 box a5) + rerank_3 = assertEqual "3" (fromList [2,3] [Just $ fromList [4] [1,2,3,4], + Just $ fromList [4] [5,6,7,8], + Just $ fromList [4] [9,10,11,12], + Just $ fromList [4] [13,14,15,16], + Just $ fromList [4] [17,18,19,20], + Just $ fromList [4] [21,22,23,24]]) + (rerank 2 box a5) + rerank_4 = assertEqual "4" (mapA (Just . scalar) a5) + (rerank 3 box a5) + rerank_5 = assertThrows "4" (rerank 4 box a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [2,3] [40,200,488,904,1448,2120]) + (rerank2 2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [2,3] [3,2,1,6,5,4]) + (rev [1] a1) + rev_2 = assertEqual "2" (fromList [2,3] [6,5,4,3,2,1]) + (rev [0,1] a1) + rev_3 = assertThrows "3" (rev [2] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList [2] [6,120]) (rerank 1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList [3] [4,10,18]) (rerank 1 (reduce (*) 1) a2) + + -- Test fast toVector + toVector_10 = + assertEqual "10" (V.fromList [1..24]) (toVector a5) -- full vector + toVector_11 = + assertEqual "11" (V.fromList [1..12]) (toVector $ slice [(0,1)] a5) -- reduced length + toVector_12 = + assertEqual "12" (V.fromList [13..24]) (toVector $ slice [(1,1)] a5) -- offset + toVector_13 = + assertEqual "13" (V.fromList $ [13..24] ++ [1..12]) + (toVector $ rev [0] a5) -- non-normal dim 0 + toVector_14 = + assertEqual "14" (V.fromList [9,10,11,12,5,6,7,8,1,2,3,4,21,22,23,24,17,18,19,20,13,14,15,16]) + (toVector $ rev [1] a5) -- non-normal dim 1 + toVector_15 = + assertEqual "15" (V.fromList [4,3,2,1,8,7,6,5,12,11,10,9,16,15,14,13,20,19,18,17,24,23,22,21]) + (toVector $ rev [2] a5) -- non-normal dim 2 + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "stretch_3" stretch_3 + , testCase "stretch_4" stretch_4 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "unScalar_2" unScalar_2 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWithA_2" zipWithA_2 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "pad_2" pad_2 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_7" transpose_7 + , testCase "transpose_8" transpose_8 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "concatOuter_1" concatOuter_1 + , testCase "concatOuter_2" concatOuter_2 + , testCase "ravel_1" ravel_1 + , testCase "ravel_2" ravel_2 + , testCase "unravel_1" unravel_1 + , testCase "window_1" window_1 + , testCase "window_2" window_2 + , testCase "window_3" window_3 + , testCase "stride_1" stride_1 + , testCase "stride_2" stride_2 + , testCase "rotate_1" rotate_1 + , testCase "slice_1" slice_1 + , testCase "slice_2" slice_2 + , testCase "slice_3" slice_3 + , testCase "slice_4" slice_4 + , testCase "slice_5" slice_5 + , testCase "rerank_1" rerank_1 + , testCase "rerank_2" rerank_2 + , testCase "rerank_3" rerank_3 + , testCase "rerank_4" rerank_4 + , testCase "rerank_5" rerank_5 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "rev_3" rev_3 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + , testCase "toVector_10" toVector_10 + , testCase "toVector_11" toVector_11 + , testCase "toVector_12" toVector_12 + , testCase "toVector_13" toVector_13 + , testCase "toVector_14" toVector_14 + , testCase "toVector_15" toVector_15 + ] + in tests diff --git a/tests/DynamicUTest.hs b/tests/DynamicUTest.hs new file mode 100644 --- /dev/null +++ b/tests/DynamicUTest.hs @@ -0,0 +1,298 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE ScopedTypeVariables #-} +module DynamicUTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.DynamicU +import qualified Data.Vector.Unboxed as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "DynamicU" $ + let a1, a2 :: Array Int + a1 = fromList [2,3] [1..6] + a2 = transpose [1,0] a1 + show_1 = assertEqual "1" "fromList [2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList [3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a2) + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a2) + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [3] [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [2] [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [] [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [] [1,2::Int]) + fromList_2 = assertThrows "sh" (fromList [4,5] [1,2::Int]) + fromVector_1 = assertEqual "1" a1 (fromVector [2,3] $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [6] [1..6]) (reshape [6] a1) + reshape_2 = assertEqual "1" (fromList [1,2,3,1] [1,4,2,5,3,6]) (reshape [1,2,3,1] a2) + a3, a4 :: Array Int + a3 = fromList [1] [5] + a4 = fromList [] [5] + stretch_1 = assertEqual "1" (fromList [3] [5,5,5]) (stretch [3] a3) + stretch_2 = assertEqual "2" (fromList [2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch [2,2,3,2] $ reshape [1,2,3,1] a1) + stretch_3 = assertThrows "3" (stretch [1,2] a3) + stretch_4 = assertThrows "4" (stretch [4,3] a1) + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + unScalar_2 = assertThrows "2" (unScalar a3) + constant_1 = assertEqual "1" (fromList [2,3] [1,1,1,1,1,1]) (constant [2,3] (1::Int)) + mapA_1 = assertEqual "1" (fromList [2,3] [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [3,2] [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,3] [2,4..12]) (zipWithA (+) a1 a1) + zipWithA_2 = assertThrows "2" (zipWithA (+) a1 a2) + zipWith3A_1 = assertEqual "1" (fromList [2,3] [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [5,10] [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad [(1,2),(3,4)] 9 a1) + pad_2 = assertThrows "2" (pad [(1,1),(1,1),(1,1)] 0 a1) + a5 :: Array Int + a5 = fromList [2,3,4] [1..24] + transpose_1 = assertEqual "1" (fromList [2,3,4] [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose [0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [2,4,3] [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose [0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [3,4,2] [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose [1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [4,2,3] [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose [2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [4,3,2] [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose [2,1,0] a5) + transpose_7 = assertThrows "7" (transpose [0,1,2,3] a5) + transpose_8 = assertThrows "7" (transpose [0,1,3] a5) + transpose_9 = assertEqual "9" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0] a5) + append_1 = assertEqual "1" (fromList [3,3] [1..9]) + (append a1 (fromList [1,3] [7,8,9])) + concatOuter_1 = assertEqual "1" (fromList [6,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (concatOuter [a1, concatOuter [a1,a1]]) + concatOuter_2 = assertThrows "2" (concatOuter [a1, a2]) + a6 :: Array Int + a6 = fromList [4,5] [1..20] + window_1 = assertEqual "1" (fromList [2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window [3,3] a6) + window_2 = assertThrows "2" (window [3,6] a6) + window_3 = assertThrows "3" (window [3,3,3] a6) + stride_1 = assertEqual "1" (fromList [2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride [1,2,2] a5) + stride_2 = assertThrows "2" (stride [1,2,2] a1) + slice_1 = assertEqual "1" (fromList [2,2,1] [8,12,20,24]) + (slice [(0,2),(1,2),(3,1)] a5) + slice_2 = assertThrows "2" (slice [(0,0)] a4) + slice_3 = assertThrows "3" (slice [(-1,1)] a5) + slice_4 = assertThrows "4" (slice [(10,0)] a5) + slice_5 = assertThrows "5" (slice [(0,3)] a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [2,3] [40,200,488,904,1448,2120]) + (rerank2 2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [2,3] [3,2,1,6,5,4]) + (rev [1] a1) + rev_2 = assertEqual "2" (fromList [2,3] [6,5,4,3,2,1]) + (rev [0,1] a1) + rev_3 = assertThrows "3" (rev [2] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList [2] [6,120]) (rerank 1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList [3] [4,10,18]) (rerank 1 (reduce (*) 1) a2) + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "stretch_3" stretch_3 + , testCase "stretch_4" stretch_4 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "unScalar_2" unScalar_2 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWithA_2" zipWithA_2 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "pad_2" pad_2 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_7" transpose_7 + , testCase "transpose_8" transpose_8 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "concatOuter_1" concatOuter_1 + , testCase "concatOuter_2" concatOuter_2 + , testCase "window_1" window_1 + , testCase "window_2" window_2 + , testCase "window_3" window_3 + , testCase "stride_1" stride_1 + , testCase "stride_2" stride_2 + , testCase "slice_1" slice_1 + , testCase "slice_2" slice_2 + , testCase "slice_3" slice_3 + , testCase "slice_4" slice_4 + , testCase "slice_5" slice_5 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "rev_3" rev_3 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + ] + in tests diff --git a/tests/RankedTest.hs b/tests/RankedTest.hs new file mode 100644 --- /dev/null +++ b/tests/RankedTest.hs @@ -0,0 +1,338 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ScopedTypeVariables #-} +module RankedTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.Ranked +import qualified Data.Vector as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "Ranked" $ + let a1, a2 :: Array 2 Int + a1 = fromList [2,3] [1..6] + a2 = transpose [1,0] a1 + show_1 = assertEqual "1" "fromList [2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList [3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a2) + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a2) + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [3] [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [2] [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [] [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [] [1,2] :: Array 0 Int) + fromList_2 = assertThrows "sh" (fromList [4,5] [1,2] :: Array 2 Int) + fromVector_1 = assertEqual "1" a1 (fromVector [2,3] $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [6] [1..6] :: Array 1 Int) (reshape [6] a1) + reshape_2 = assertEqual "1" (fromList [1,2,3,1] [1,4,2,5,3,6]) (reshape [1,2,3,1] a2 :: Array 4 Int) + a3 :: Array 1 Int + a3 = fromList [1] [5] + a4 :: Array 0 Int + a4 = fromList [] [5] + stretch_1 = assertEqual "1" (fromList [3] [5,5,5]) (stretch [3] a3) + stretch_2 = assertEqual "2" (fromList [2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch [2,2,3,2] (reshape [1,2,3,1] a1 :: Array 4 Int)) + stretch_3 = assertThrows "3" (stretch [1,2] a3) + stretch_4 = assertThrows "4" (stretch [4,3] a1) + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + constant_1 = assertEqual "1" (fromList [2,3] [1,1,1,1,1,1]) (constant [2,3] 1 :: Array 2 Int) + mapA_1 = assertEqual "1" (fromList [2,3] [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [3,2] [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,3] [2,4..12]) (zipWithA (+) a1 a1) + zipWithA_2 = assertThrows "2" (zipWithA (+) a1 a2) + zipWith3A_1 = assertEqual "1" (fromList [2,3] [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [5,10] [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad [(1,2),(3,4)] 9 a1) + pad_2 = assertThrows "2" (pad [(1,1),(1,1),(1,1)] 0 a1) + a5 :: Array 3 Int + a5 = fromList [2,3,4] [1..24] + transpose_1 = assertEqual "1" (fromList [2,3,4] [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose [0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [2,4,3] [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose [0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [3,4,2] [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose [1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [4,2,3] [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose [2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [4,3,2] [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose [2,1,0] a5) + transpose_7 = assertThrows "7" (transpose [0,1,2,3] a5) + transpose_8 = assertThrows "7" (transpose [0,1,3] a5) + transpose_9 = assertEqual "9" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0] a5) + append_1 = assertEqual "1" (fromList [3,3] [1..9]) + (append a1 (fromList [1,3] [7,8,9])) + concatOuter_1 = assertEqual "1" (fromList [6,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (concatOuter [a1, concatOuter [a1,a1]]) + concatOuter_2 = assertThrows "2" (concatOuter [a1, a2]) + ravel_1 = assertEqual "1" (fromList [3,2,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (ravel $ fromList [3] [a1,a1,a1]) + ravel_2 = assertThrows "2" (ravel $ fromList [2] [a1, concatOuter [a1,a1]]) + unravel_1 = assertEqual "1" [a1,a1,a1] + (toList $ unravel $ ravel $ fromList [3] [a1,a1,a1]) + a6 :: Array 2 Int + a6 = fromList [4,5] [1..20] + window_1 = assertEqual "1" (fromList [2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window [3,3] a6 :: Array 4 Int) + window_2 = assertThrows "2" (window [3,6] a6 :: Array 4 Int) + window_3 = assertThrows "3" (window [3,3,3] a6 :: Array 5 Int) + stride_1 = assertEqual "1" (fromList [2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride [1,2,2] a5) + stride_2 = assertThrows "2" (stride [1,2,2] a1) + rotate_1 = assertEqual "1" (fromList [2, 4, 3, 2] + [1, 2, 3, 4, 5, 6, + 5, 6, 1, 2, 3, 4, + 3, 4, 5, 6, 1, 2, + 1, 2, 3, 4, 5, 6, + 7, 8, 9, 10, 11, 12, + 11, 12, 7, 8, 9, 10, + 9, 10, 11, 12, 7, 8, + 7, 8, 9, 10, 11, 12] :: Array 4 Int) + (rotate @1 4 (fromList [2, 3, 2] [1 .. 12] :: Array 3 Int)) + slice_1 = assertEqual "1" (fromList [2,2,1] [8,12,20,24]) + (slice [(0,2),(1,2),(3,1)] a5) + slice_2 = assertThrows "2" (slice [(0,0)] a4) + slice_3 = assertThrows "3" (slice [(-1,1)] a5) + slice_4 = assertThrows "4" (slice [(10,0)] a5) + slice_5 = assertThrows "5" (slice [(0,3)] a5) + + box = scalar . Just + rerank_1 = assertEqual "1" (box a5) + (rerank @0 box a5) + rerank_2 = assertEqual "2" (fromList [2] [Just $ fromList [3,4] [1,2,3,4,5,6,7,8,9,10,11,12], + Just $ fromList [3,4] [13,14,15,16,17,18,19,20,21,22,23,24]]) + (rerank @1 box a5) + rerank_3 = assertEqual "3" (fromList [2,3] [Just $ fromList [4] [1,2,3,4], + Just $ fromList [4] [5,6,7,8], + Just $ fromList [4] [9,10,11,12], + Just $ fromList [4] [13,14,15,16], + Just $ fromList [4] [17,18,19,20], + Just $ fromList [4] [21,22,23,24]]) + (rerank @2 box a5) + rerank_4 = assertEqual "4" (mapA (Just . scalar) a5) + (rerank @3 box a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [2,3] [40,200,488,904,1448,2120]) + (rerank2 @2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [2,3] [3,2,1,6,5,4]) + (rev [1] a1) + rev_2 = assertEqual "2" (fromList [2,3] [6,5,4,3,2,1]) + (rev [0,1] a1) + rev_3 = assertThrows "3" (rev [2] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList [2] [6,120]) (rerank @1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList [3] [4,10,18]) (rerank @1 (reduce (*) 1) a2) + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "stretch_3" stretch_3 + , testCase "stretch_4" stretch_4 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWithA_2" zipWithA_2 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "pad_2" pad_2 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_7" transpose_7 + , testCase "transpose_8" transpose_8 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "concatOuter_1" concatOuter_1 + , testCase "concatOuter_2" concatOuter_2 + , testCase "ravel_1" ravel_1 + , testCase "ravel_2" ravel_2 + , testCase "unravel_1" unravel_1 + , testCase "window_1" window_1 + , testCase "window_2" window_2 + , testCase "window_3" window_3 + , testCase "stride_1" stride_1 + , testCase "stride_2" stride_2 + , testCase "rotate_1" rotate_1 + , testCase "slice_1" slice_1 + , testCase "slice_2" slice_2 + , testCase "slice_3" slice_3 + , testCase "slice_4" slice_4 + , testCase "slice_5" slice_5 + , testCase "rerank_1" rerank_1 + , testCase "rerank_2" rerank_2 + , testCase "rerank_3" rerank_3 + , testCase "rerank_4" rerank_4 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "rev_3" rev_3 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + ] + in tests diff --git a/tests/RankedUTest.hs b/tests/RankedUTest.hs new file mode 100644 --- /dev/null +++ b/tests/RankedUTest.hs @@ -0,0 +1,299 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ScopedTypeVariables #-} +module RankedUTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.RankedU +import qualified Data.Vector.Unboxed as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "RankedU" $ + let a1, a2 :: Array 2 Int + a1 = fromList [2,3] [1..6] + a2 = transpose [1,0] a1 + show_1 = assertEqual "1" "fromList [2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList [3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a2) + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a2) + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [3] [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [2] [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [] [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [] [1,2] :: Array 0 Int) + fromList_2 = assertThrows "sh" (fromList [4,5] [1,2] :: Array 2 Int) + fromVector_1 = assertEqual "1" a1 (fromVector [2,3] $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [6] [1..6] :: Array 1 Int) (reshape [6] a1) + reshape_2 = assertEqual "1" (fromList [1,2,3,1] [1,4,2,5,3,6]) (reshape [1,2,3,1] a2 :: Array 4 Int) + a3 :: Array 1 Int + a3 = fromList [1] [5] + a4 :: Array 0 Int + a4 = fromList [] [5] + stretch_1 = assertEqual "1" (fromList [3] [5,5,5]) (stretch [3] a3) + stretch_2 = assertEqual "2" (fromList [2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch [2,2,3,2] (reshape [1,2,3,1] a1 :: Array 4 Int)) + stretch_3 = assertThrows "3" (stretch [1,2] a3) + stretch_4 = assertThrows "4" (stretch [4,3] a1) + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + constant_1 = assertEqual "1" (fromList [2,3] [1,1,1,1,1,1]) (constant [2,3] 1 :: Array 2 Int) + mapA_1 = assertEqual "1" (fromList [2,3] [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [3,2] [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,3] [2,4..12]) (zipWithA (+) a1 a1) + zipWithA_2 = assertThrows "2" (zipWithA (+) a1 a2) + zipWith3A_1 = assertEqual "1" (fromList [2,3] [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [5,10] [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad [(1,2),(3,4)] 9 a1) + pad_2 = assertThrows "2" (pad [(1,1),(1,1),(1,1)] 0 a1) + a5 :: Array 3 Int + a5 = fromList [2,3,4] [1..24] + transpose_1 = assertEqual "1" (fromList [2,3,4] [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose [0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [2,4,3] [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose [0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [3,4,2] [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose [1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [4,2,3] [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose [2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [4,3,2] [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose [2,1,0] a5) + transpose_7 = assertThrows "7" (transpose [0,1,2,3] a5) + transpose_8 = assertThrows "7" (transpose [0,1,3] a5) + transpose_9 = assertEqual "9" (fromList [3,2,4] [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose [1,0] a5) + append_1 = assertEqual "1" (fromList [3,3] [1..9]) + (append a1 (fromList [1,3] [7,8,9])) + concatOuter_1 = assertEqual "1" (fromList [6,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (concatOuter [a1, concatOuter [a1,a1]]) + concatOuter_2 = assertThrows "2" (concatOuter [a1, a2]) + a6 :: Array 2 Int + a6 = fromList [4,5] [1..20] + window_1 = assertEqual "1" (fromList [2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window [3,3] a6 :: Array 4 Int) + window_2 = assertThrows "2" (window [3,6] a6 :: Array 4 Int) + window_3 = assertThrows "3" (window [3,3,3] a6 :: Array 5 Int) + stride_1 = assertEqual "1" (fromList [2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride [1,2,2] a5) + stride_2 = assertThrows "2" (stride [1,2,2] a1) + slice_1 = assertEqual "1" (fromList [2,2,1] [8,12,20,24]) + (slice [(0,2),(1,2),(3,1)] a5) + slice_2 = assertThrows "2" (slice [(0,0)] a4) + slice_3 = assertThrows "3" (slice [(-1,1)] a5) + slice_4 = assertThrows "4" (slice [(10,0)] a5) + slice_5 = assertThrows "5" (slice [(0,3)] a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [2,3] [40,200,488,904,1448,2120]) + (rerank2 @2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [2,3] [3,2,1,6,5,4]) + (rev [1] a1) + rev_2 = assertEqual "2" (fromList [2,3] [6,5,4,3,2,1]) + (rev [0,1] a1) + rev_3 = assertThrows "3" (rev [2] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList [2] [6,120]) (rerank @1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList [3] [4,10,18]) (rerank @1 (reduce (*) 1) a2) + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "stretch_3" stretch_3 + , testCase "stretch_4" stretch_4 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWithA_2" zipWithA_2 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "pad_2" pad_2 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_7" transpose_7 + , testCase "transpose_8" transpose_8 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "concatOuter_1" concatOuter_1 + , testCase "concatOuter_2" concatOuter_2 + , testCase "window_1" window_1 + , testCase "window_2" window_2 + , testCase "window_3" window_3 + , testCase "stride_1" stride_1 + , testCase "stride_2" stride_2 + , testCase "slice_1" slice_1 + , testCase "slice_2" slice_2 + , testCase "slice_3" slice_3 + , testCase "slice_4" slice_4 + , testCase "slice_5" slice_5 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "rev_3" rev_3 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + ] + in tests diff --git a/tests/ShapedTest.hs b/tests/ShapedTest.hs new file mode 100644 --- /dev/null +++ b/tests/ShapedTest.hs @@ -0,0 +1,294 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ScopedTypeVariables #-} +module ShapedTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.Shaped +import qualified Data.Vector as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "Shaped" $ + let a1, a1' :: Array [2,3] Int + a1 = fromList [1..6] + a1' = fromList [2..7] + a2 :: Array [3,2] Int + a2 = fromList [1,4,2,5,3,6] + show_1 = assertEqual "1" "fromList @[2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList @[3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a1') + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a1') + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [1,2] :: Array '[] Int) + fromList_2 = assertThrows "sh" (fromList [1,2] :: Array [4,5] Int) + fromVector_1 = assertEqual "1" a1 (fromVector $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [1..6] :: Array '[6] Int) (reshape a1) + reshape_2 = assertEqual "1" (fromList [1,4,2,5,3,6]) (reshape a2 :: Array [1,2,3,1] Int) + a3 :: Array '[1] Int + a3 = fromList [5] + a4 :: Array '[] Int + a4 = fromList [5] + stretch_1 = assertEqual "1" (fromList @'[3] [5,5,5]) (stretch @'[3] a3) + stretch_2 = assertEqual "2" (fromList @[2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch @[2,2,3,2] (reshape @[1,2,3,1] a1)) + + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + constant_1 = assertEqual "1" (fromList [1,1,1,1,1,1]) (constant 1 :: Array [2,3] Int) + mapA_1 = assertEqual "1" (fromList [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,4..12]) (zipWithA (+) a1 a1) + zipWith3A_1 = assertEqual "1" (fromList [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad @['(1,2), '(3,4)] 9 a1) + a5 :: Array '[2,3,4] Int + a5 = fromList [1..24] + transpose_1 = assertEqual "1" (fromList [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose @'[0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose @'[0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose @'[1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose @'[1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose @'[2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose @'[2,1,0] a5) + transpose_9 = assertEqual "9" (fromList [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose @'[1,0] a5) + append_1 = assertEqual "1" (fromList [1..9]) + (append a1 (fromList @[1,3] [7,8,9])) + ravel_1 = assertEqual "1" (fromList @[3,2,3] [1,2,3,4,5,6,1,2,3,4,5,6,1,2,3,4,5,6]) + (ravel $ fromList @'[3] [a1,a1,a1]) + unravel_1 = assertEqual "1" [a1,a1,a1] + (toList $ unravel $ ravel $ fromList @'[3] [a1,a1,a1]) + a6 :: Array [4,5] Int + a6 = fromList [1..20] + window_1 = assertEqual "1" (fromList @[2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window @[3,3] a6) + stride_1 = assertEqual "1" (fromList @[2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride @[1,2,2] a5) + slice_1 = assertEqual "1" (fromList @[2,2,1] [8,12,20,24]) + (slice @['(0,2), '(1,2), '(3,1)] a5) + box = scalar . Just + rerank_1 = assertEqual "1" (box a5) + (rerank @0 box a5) + rerank_2 = assertEqual "2" (fromList [Just $ fromList @[3,4] [1,2,3,4,5,6,7,8,9,10,11,12], + Just $ fromList @[3,4] [13,14,15,16,17,18,19,20,21,22,23,24]]) + (rerank @1 box a5) + rerank_3 = assertEqual "3" (fromList @[2,3] [Just $ fromList [1,2,3,4], + Just $ fromList [5,6,7,8], + Just $ fromList [9,10,11,12], + Just $ fromList [13,14,15,16], + Just $ fromList [17,18,19,20], + Just $ fromList [21,22,23,24]]) + (rerank @2 box a5) + rerank_4 = assertEqual "4" (mapA (Just . scalar) a5) + (rerank @3 box a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [40,200,488,904,1448,2120]) + (rerank2 @2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [3,2,1,6,5,4]) + (rev @'[1] a1) + rev_2 = assertEqual "2" (fromList [6,5,4,3,2,1]) + (rev @[0,1] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList @'[2] [6,120]) (rerank @1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList @'[3] [4,10,18]) (rerank @1 (reduce (*) 1) a2) + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "ravel_1" ravel_1 + , testCase "unravel_1" unravel_1 + , testCase "window_1" window_1 + , testCase "stride_1" stride_1 + , testCase "slice_1" slice_1 + , testCase "rerank_1" rerank_1 + , testCase "rerank_2" rerank_2 + , testCase "rerank_3" rerank_3 + , testCase "rerank_4" rerank_4 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + ] + in tests diff --git a/tests/ShapedUTest.hs b/tests/ShapedUTest.hs new file mode 100644 --- /dev/null +++ b/tests/ShapedUTest.hs @@ -0,0 +1,269 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE ScopedTypeVariables #-} +module ShapedUTest(test) where + +import Control.DeepSeq +import Control.Exception +import Data.Array.ShapedU +import qualified Data.Vector.Unboxed as V +import Test.Framework (Test, testGroup) +import Test.Framework.Providers.HUnit (testCase) +import Test.HUnit (assertEqual, assertFailure, Assertion) + +assertThrows :: (NFData a) => String -> a -> Assertion +assertThrows s a = catch (deepseq a $ assertFailure s) (\ (_ :: ErrorCall) -> return ()) + +test :: Test +test = testGroup "ShapedU" $ + let a1, a1' :: Array [2,3] Int + a1 = fromList [1..6] + a1' = fromList [2..7] + a2 :: Array [3,2] Int + a2 = fromList [1,4,2,5,3,6] + show_1 = assertEqual "1" "fromList @[2,3] [1,2,3,4,5,6]" (show a1) + show_2 = assertEqual "2" "fromList @[3,2] [1,4,2,5,3,6]" (show a2) + eq_1 = assertEqual "1" True (a1 == a1) + eq_2 = assertEqual "2" False (a1 == a1') + ord_1 = assertEqual "1" EQ (a1 `compare` a1) + ord_2 = assertEqual "2" LT (a1 `compare` a1') + shapeL_1 = assertEqual "1" [2,3] (shapeL a1) + shapeL_2 = assertEqual "2" [3,2] (shapeL a2) + rank_1 = assertEqual "1" 2 (rank a1) + rank_2 = assertEqual "2" 2 (rank a2) + index_1 = assertEqual "1" (fromList [1,2,3]) (index a1 0) + index_2 = assertEqual "2" (fromList [1,4]) (index a2 0) + index_3 = assertEqual "3" (fromList [4]) (a2 `index` 0 `index` 1) + index_4 = assertThrows "<0" (index a1 (-1)) + index_5 = assertThrows ">" (index a1 2) + toList_1 = assertEqual "1" [1,2,3,4,5,6] (toList a1) + toList_2 = assertEqual "2" [1,4,2,5,3,6] (toList a2) + toVector_1 = assertEqual "1" (V.fromList [1,2,3,4,5,6]) (toVector a1) + toVector_2 = assertEqual "2" (V.fromList [1,4,2,5,3,6]) (toVector a2) + fromList_1 = assertThrows "sh" (fromList [1,2] :: Array '[] Int) + fromList_2 = assertThrows "sh" (fromList [1,2] :: Array [4,5] Int) + fromVector_1 = assertEqual "1" a1 (fromVector $ V.fromList [1..6]) + normalize_1 = assertEqual "1" a1 (normalize a1) + reshape_1 = assertEqual "1" (fromList [1..6] :: Array '[6] Int) (reshape a1) + reshape_2 = assertEqual "1" (fromList [1,4,2,5,3,6]) (reshape a2 :: Array [1,2,3,1] Int) + a3 :: Array '[1] Int + a3 = fromList [5] + a4 :: Array '[] Int + a4 = fromList [5] + stretch_1 = assertEqual "1" (fromList @'[3] [5,5,5]) (stretch @'[3] a3) + stretch_2 = assertEqual "2" (fromList @[2,2,3,2] [1,1,2,2,3,3,4,4,5,5,6,6,1,1,2,2,3,3,4,4,5,5,6,6]) + (stretch @[2,2,3,2] (reshape @[1,2,3,1] a1)) + + scalar_1 = assertEqual "1" a4 (scalar 5) + unScalar_1 = assertEqual "1" 5 (unScalar a4) + constant_1 = assertEqual "1" (fromList [1,1,1,1,1,1]) (constant 1 :: Array [2,3] Int) + mapA_1 = assertEqual "1" (fromList [2..7]) (mapA succ a1) + mapA_2 = assertEqual "1" (fromList [2,5,3,6,4,7]) (mapA succ a2) + zipWithA_1 = assertEqual "1" (fromList [2,4..12]) (zipWithA (+) a1 a1) + zipWith3A_1 = assertEqual "1" (fromList [2,6,12,20,30,42]) (zipWith3A (\ x y z -> x*y+z) a1 a1 a1) + pad_1 = assertEqual "1" (fromList [9,9,9,9,9,9,9,9,9,9, + 9,9,9,1,2,3,9,9,9,9, + 9,9,9,4,5,6,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9, + 9,9,9,9,9,9,9,9,9,9]) + (pad @['(1,2), '(3,4)] 9 a1) + a5 :: Array '[2,3,4] Int + a5 = fromList [1..24] + transpose_1 = assertEqual "1" (fromList [1,2,3,4, + 5,6,7,8, + 9,10,11,12, + + 13,14,15,16, + 17,18,19,20, + 21,22,23,24]) + (transpose @'[0,1,2] a5) + transpose_2 = assertEqual "2" (fromList [1,5,9, + 2,6,10, + 3,7,11, + 4,8,12, + + 13,17,21, + 14,18,22, + 15,19,23, + 16,20,24]) + (transpose @'[0,2,1] a5) + transpose_3 = assertEqual "3" (fromList [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose @'[1,0,2] a5) + transpose_4 = assertEqual "4" (fromList [1,13, + 2,14, + 3,15, + 4,16, + + 5,17, + 6,18, + 7,19, + 8,20, + + 9,21, + 10,22, + 11,23, + 12,24]) + (transpose @'[1,2,0] a5) + transpose_5 = assertEqual "5" (fromList [1,5,9, + 13,17,21, + + 2,6,10, + + 14,18,22, + + 3,7,11, + 15,19,23, + + 4,8,12, + 16,20,24]) + (transpose @'[2,0,1] a5) + transpose_6 = assertEqual "6" (fromList [1,13, + 5,17, + 9,21, + + 2,14, + 6,18, + 10,22, + + 3,15, + 7,19, + 11,23, + + 4,16, + 8,20, + 12,24]) + (transpose @'[2,1,0] a5) + transpose_9 = assertEqual "9" (fromList [1,2,3,4, + 13,14,15,16, + + 5,6,7,8, + 17,18,19,20, + + 9,10,11,12, + 21,22,23,24]) + (transpose @'[1,0] a5) + append_1 = assertEqual "1" (fromList [1..9]) + (append a1 (fromList @[1,3] [7,8,9])) + a6 :: Array [4,5] Int + a6 = fromList [1..20] + window_1 = assertEqual "1" (fromList @[2,3,3,3] [1,2,3, + 6,7,8, + 11,12,13, + + 2,3,4, + 7,8,9, + 12,13,14, + + 3,4,5, + 8,9,10, + 13,14,15, + + + 6,7,8, + 11,12,13, + 16,17,18, + + 7,8,9, + 12,13,14, + 17,18,19, + + 8,9,10, + 13,14,15, + 18,19,20]) + (window @[3,3] a6) + stride_1 = assertEqual "1" (fromList @[2,2,2] [1,3, + 9,11, + + 13,15, + 21,23]) + (stride @[1,2,2] a5) + slice_1 = assertEqual "1" (fromList @[2,2,1] [8,12,20,24]) + (slice @['(0,2), '(1,2), '(3,1)] a5) + a7 = mapA succ a5 + dot x y = reduce (+) 0 $ zipWithA (*) x y + rerank2_1 = assertEqual "1" (fromList [40,200,488,904,1448,2120]) + (rerank2 @2 dot a5 a7) + rev_1 = assertEqual "1" (fromList [3,2,1,6,5,4]) + (rev @'[1] a1) + rev_2 = assertEqual "2" (fromList [6,5,4,3,2,1]) + (rev @[0,1] a1) + reduce_1 = assertEqual "1" (scalar 720) (reduce (*) 1 a1) + reduce_2 = assertEqual "2" (fromList @'[2] [6,120]) (rerank @1 (reduce (*) 1) a1) + reduce_3 = assertEqual "3" (fromList @'[3] [4,10,18]) (rerank @1 (reduce (*) 1) a2) + + tests = + [ testCase "show_1" show_1 + , testCase "show_2" show_2 + , testCase "eq_1" eq_1 + , testCase "eq_2" eq_2 + , testCase "ord_1" ord_1 + , testCase "ord_2" ord_2 + , testCase "shapeL_1" shapeL_1 + , testCase "shapeL_2" shapeL_2 + , testCase "rank_1" rank_1 + , testCase "rank_2" rank_2 + , testCase "index_1" index_1 + , testCase "index_2" index_2 + , testCase "index_3" index_3 + , testCase "index_4" index_4 + , testCase "index_5" index_5 + , testCase "toList_1" toList_1 + , testCase "toList_2" toList_2 + , testCase "toVector_1" toVector_1 + , testCase "toVector_2" toVector_2 + , testCase "fromList_1" fromList_1 + , testCase "fromList_2" fromList_2 + , testCase "fromVector_1" fromVector_1 + , testCase "normalize_1" normalize_1 + , testCase "reshape_1" reshape_1 + , testCase "reshape_2" reshape_2 + , testCase "stretch_1" stretch_1 + , testCase "stretch_2" stretch_2 + , testCase "scalar_1" scalar_1 + , testCase "unScalar_1" unScalar_1 + , testCase "constant_1" constant_1 + , testCase "mapA_1" mapA_1 + , testCase "mapA_2" mapA_2 + , testCase "zipWithA_1" zipWithA_1 + , testCase "zipWith3A_1" zipWith3A_1 + , testCase "pad_1" pad_1 + , testCase "transpose_1" transpose_1 + , testCase "transpose_2" transpose_2 + , testCase "transpose_3" transpose_3 + , testCase "transpose_4" transpose_4 + , testCase "transpose_5" transpose_5 + , testCase "transpose_6" transpose_6 + , testCase "transpose_9" transpose_9 + , testCase "append_1" append_1 + , testCase "window_1" window_1 + , testCase "stride_1" stride_1 + , testCase "slice_1" slice_1 + , testCase "rerank2_1" rerank2_1 + , testCase "rev_1" rev_1 + , testCase "rev_2" rev_2 + , testCase "reduce_1" reduce_1 + , testCase "reduce_2" reduce_2 + , testCase "reduce_3" reduce_3 + ] + in tests diff --git a/tests/Tests.hs b/tests/Tests.hs new file mode 100644 --- /dev/null +++ b/tests/Tests.hs @@ -0,0 +1,32 @@ +-- Copyright 2020 Google LLC +-- +-- Licensed under the Apache License, Version 2.0 (the "License"); +-- you may not use this file except in compliance with the License. +-- You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, software +-- distributed under the License is distributed on an "AS IS" BASIS, +-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +-- See the License for the specific language governing permissions and +-- limitations under the License. + +import Test.Framework (defaultMain) + +import qualified DynamicTest +import qualified DynamicUTest +import qualified RankedTest +import qualified RankedUTest +import qualified ShapedTest +import qualified ShapedUTest + +main :: IO () +main = defaultMain + [ DynamicTest.test + , DynamicUTest.test + , RankedTest.test + , RankedUTest.test + , ShapedTest.test + , ShapedUTest.test + ]