diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,3 @@ +# 0.1.0.0 + +Initial version. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,202 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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/short-vec.cabal b/short-vec.cabal new file mode 100644 --- /dev/null +++ b/short-vec.cabal @@ -0,0 +1,106 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.33.0. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: d250fd4c14d88f72bb1abb89d02d30c1f6f6cbcbab3b4edfb356e404089e770e + +name: short-vec +version: 0.1.0.0 +synopsis: A length-indexed vector type build on 'SmallArray#' +description: This provides performant length-indexed vectors with a suite of rewrite rules + implementing fusion of intermediate structures, so that expressions involving + many operations can usually be compiled to a single pass of computing and + writing each index of the final 'Vec'. +category: Data +homepage: https://github.com/google/hs-fin-vec#readme +bug-reports: https://github.com/google/hs-fin-vec/issues +author: Lennart Augustsson , Wren Romano, Andrew Pritchard + +maintainer: Andrew Pritchard +copyright: 2018-2021 Google LLC +license: Apache-2.0 +license-file: LICENSE +build-type: Simple +extra-source-files: + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/google/hs-fin-vec + subdir: short-vec + +library + exposed-modules: + Data.Vec.Short + Data.Vec.Short.Explicit + Data.Vec.Short.Internal + other-modules: + Paths_short_vec + hs-source-dirs: + src + build-depends: + QuickCheck >=2.5 && <2.15 + , adjunctions >=4.4 && <4.5 + , base >=4.12 && <4.16 + , data-default-class >=0.0 && <0.2 + , deepseq >=1.1 && <1.5 + , distributive >=0.1 && <0.7 + , fin-int >=0.1 && <0.2 + , indexed-traversable >=0.1 && <0.2 + , integer-gmp >=0.5 && <1.2 + , portray >=0.1 && <0.2 + , portray-diff >=0.1 && <0.2 + , semigroupoids >=1.0 && <5.4 + , sint >=0.1 && <0.2 + default-language: Haskell2010 + +test-suite Vec-test + type: exitcode-stdio-1.0 + main-is: Tests.hs + hs-source-dirs: + test + build-depends: + HUnit + , QuickCheck >=2.4.0.1 + , adjunctions >=4.4 && <4.5 + , base + , data-default-class >=0.0 && <0.2 + , deepseq >=1.1 && <1.5 + , distributive >=0.1 && <0.7 + , fin-int >=0.1 && <0.2 + , indexed-traversable >=0.1 && <0.2 + , integer-gmp >=0.5 && <1.2 + , portray >=0.1 && <0.2 + , portray-diff >=0.1 && <0.2 + , semigroupoids >=1.0 && <5.4 + , short-vec + , sint >=0.1 && <0.2 + , test-framework >=0.3.3 + , test-framework-hunit + , test-framework-quickcheck2 + default-language: Haskell2010 + +benchmark Vec-benchmark + type: exitcode-stdio-1.0 + main-is: Benchmark.hs + hs-source-dirs: + test + build-depends: + QuickCheck >=2.5 && <2.15 + , adjunctions >=4.4 && <4.5 + , base >=4.12 && <4.16 + , data-default-class >=0.0 && <0.2 + , deepseq + , distributive >=0.1 && <0.7 + , fin-int + , gauge + , indexed-traversable >=0.1 && <0.2 + , integer-gmp >=0.5 && <1.2 + , portray >=0.1 && <0.2 + , portray-diff >=0.1 && <0.2 + , semigroupoids >=1.0 && <5.4 + , short-vec + , sint >=0.1 && <0.2 + default-language: Haskell2010 diff --git a/src/Data/Vec/Short.hs b/src/Data/Vec/Short.hs new file mode 100644 --- /dev/null +++ b/src/Data/Vec/Short.hs @@ -0,0 +1,176 @@ +-- Copyright 2018-2021 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 BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE NoStarIsType #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeOperators #-} + +-- | An implementation of short vectors. +-- +-- The underlying implementation uses the 'GHC.Exts.SmallArray#' primitive, +-- which is best-suited for short vectors (less than a few hundred elements). +-- +-- In contrast to "Data.Vec.Short.Explicit", this module provides an API where +-- bounds parameters are passed implicitly by 'KnownNat'. This can be more +-- convenient in cases where the bounds are obvious and type-level arithmetic +-- is not involved, but it comes at the cost of some runtime +-- 'Numeric.Natural.Natural'-to-'Int' conversions. +-- +-- When type-level arithmetic is involved, the +-- [ghc-typelits-knownnat](https://hackage.haskell.org/package/ghc-typelits-knownnat) +-- plugin may be useful to derive 'KnownNat' instances for bounds automatically. + +module Data.Vec.Short + ( Vec + -- * Core constructors\/generators + -- ** 'Fin'-based constructors + , mkVec, mkVec' + , backpermute + -- ** List-based constructors + , fromList, withVec + -- ** Arity-based constructors + , vec1, vec2, vec3, vec4, vec6, vec8 + -- ** 'Enum'-based constructors + , viota + + -- * Core operators + -- ** Size\/length + , svSize, vLength, vSize, withSize + -- ** Indexing + , (!), indexK + -- ** Add/remove element + , insert, remove, overIx + + -- * List-like operators + -- ** Constructor views + -- *** The nil view + , nil + -- ** Operator views + -- *** The append view + , (++), split + -- *** The concat view + , concat, concatMap, reshape + -- *** The reverse view + , rev, rot + -- *** The transposition view + , vtranspose + -- ** Misc list-like operators + , iterate, iterate' + , vsort, vsortBy, vsortOn + , vfindIndex + + -- * Additional zips, maps, folds, etc. + , map', imap', withPos + , cross + , vscanl + , liftA2Lazy + ) where + +import Prelude hiding (concatMap, concat, iterate, (++)) + +import GHC.TypeNats (KnownNat, type (+), type (*)) +import GHC.Stack (HasCallStack) + +import Data.Fin.Int (Fin) +import Data.SInt (sintVal, unSInt, reifySInt) + +import Data.Vec.Short.Internal hiding + ( backpermute, mkVec, mkVec', split, reshape, vtranspose + , iterate, iterate', fromList, viota, liftA2Lazy + ) +import qualified Data.Vec.Short.Internal as V + +-- | Create a known-length vector using a pure function. +-- +-- Note if you don't have the 'KnownNat' instance at hand, but you already have +-- a 'Vec' of the desired length, you can use 'withSize' to get the 'KnownNat' +-- instance. +mkVec :: KnownNat n => (Fin n -> a) -> Vec n a +mkVec = V.mkVec sintVal +{-# INLINE mkVec #-} + +-- | Create a known-length vector using a pure function, strictly. +mkVec' :: KnownNat n => (Fin n -> a) -> Vec n a +mkVec' = V.mkVec' sintVal +{-# INLINE mkVec' #-} + +-- | Create a 'Vec' by selecting indices of another 'Vec'. +backpermute :: KnownNat m => (Fin m -> Fin n) -> Vec n a -> Vec m a +backpermute = V.backpermute sintVal +{-# INLINE backpermute #-} + +-- | Convert a list to a vector, throwing an error if the list has the +-- wrong length. +-- Note: Because this walks @xs@ to check its length, this cannot be +-- used with the list fusion optimization rules. +fromList :: (HasCallStack, KnownNat n) => [a] -> Vec n a +fromList = V.fromList sintVal +{-# INLINE fromList #-} + +-- | Return a vector with all elements of the type in ascending order. +viota :: KnownNat n => Vec n (Fin n) +viota = V.viota sintVal +{-# INLINE viota #-} + +-- | Split a 'Vec' into two at a given offset. +split :: KnownNat m => Vec (m + n) a -> (Vec m a, Vec n a) +split = V.split sintVal +{-# INLINE split #-} + +-- | Split a 'Vec' into a 'Vec' of equal-sized chunks. +reshape :: KnownNat m => Vec (n * m) a -> Vec n (Vec m a) +reshape = V.reshape sintVal +{-# INLINE reshape #-} + +-- | Transpose a vector of vectors. +vtranspose :: KnownNat m => Vec n (Vec m a) -> Vec m (Vec n a) +vtranspose = V.vtranspose sintVal +{-# INLINE vtranspose #-} + +-- | Statically determine the (purported) size\/length of the vector. +-- If you'd rather not require the 'KnownNat' constraint, see 'vSize'. +vLength :: forall n a. KnownNat n => Vec n a -> Int +vLength _ = unSInt @n sintVal +{-# INLINE vLength #-} + +-- | Generate a Vec by repeated application of a function. +-- +-- > toList (Vec.iterate @n f z) === take (valueOf @n) (Prelude.iterate f z) +iterate :: KnownNat n => (a -> a) -> a -> Vec n a +iterate = V.iterate sintVal +{-# INLINE iterate #-} + +-- | A strict version of 'iterate'. +iterate' :: KnownNat n => (a -> a) -> a -> Vec n a +iterate' = V.iterate' sintVal +{-# INLINE iterate' #-} + +-- | A truly lazy version of @liftA2@. +-- +-- As opposed to the actual @liftA2@ it does not inspect the arguments which +-- makes it possible it to use in code that has lazy knot-tying. +liftA2Lazy :: KnownNat n => (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c +liftA2Lazy = V.liftA2Lazy sintVal +{-# INLINE liftA2Lazy #-} + +-- | Dynamically determine the (actual) size\/length of the vector, +-- returning evidence that @n@ is \"known\". If you'd rather obtain @n@ +-- as a standard 'Int', see 'vSize'. +withSize :: forall n a r . Vec n a -> (KnownNat n => r) -> r +withSize !xs f = reifySInt (svSize xs) f +{-# INLINE withSize #-} diff --git a/src/Data/Vec/Short/Explicit.hs b/src/Data/Vec/Short/Explicit.hs new file mode 100644 --- /dev/null +++ b/src/Data/Vec/Short/Explicit.hs @@ -0,0 +1,73 @@ +-- Copyright 2021 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. + +-- | An API for "Data.Vec.Short" with 'Data.SInt.SInt's for all size parameters. +-- +-- In contrast to "Data.Vec.Short", this module provides an API where runtime +-- values of bound parameters are provided explicitly by 'Data.SInt.SInt's, +-- which can be more intuitive than passing implicitly via +-- 'GHC.TypeLits.KnownNat', and can avoid some runtime +-- 'Numeric.Natural.Natural'-to-'Int' conversions and bounds checks resulting +-- from @KnownNat@, at the cost of making some code more tedious where the +-- bounds "should" be obvious. + +module Data.Vec.Short.Explicit + ( Vec + -- * Core constructors\/generators + -- ** 'Data.Fin.Int.Fin'-based constructors + , mkVec, mkVec' + , backpermute + -- ** List-based constructors + , fromList, withVec + -- ** Arity-based constructors + , vec1, vec2, vec3, vec4, vec6, vec8 + -- ** 'Enum'-based constructors + , viota + + -- * Core operators + -- ** Size\/length + , svSize, vSize + -- ** Indexing + , (!), indexK + -- ** Add/remove element + , insert, remove, overIx + + -- * List-like operators + -- ** Constructor views + -- *** The nil view + , nil + -- ** Operator views + -- *** The append view + , (++), split + -- *** The concat view + , concat, concatMap, reshape + -- *** The reverse view + , rev, rot + -- *** The transposition view + , vtranspose + -- ** Misc list-like operators + , iterate, iterate' + , vsort, vsortBy, vsortOn + , vfindIndex + + -- * Additional zips, maps, folds, etc. + , map', imap', withPos + , cross + , vscanl + , liftA2Lazy + ) where + +import Prelude hiding (concatMap, concat, iterate, (++)) + +import Data.Vec.Short.Internal diff --git a/src/Data/Vec/Short/Internal.hs b/src/Data/Vec/Short/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Vec/Short/Internal.hs @@ -0,0 +1,1417 @@ +-- Copyright 2018-2021 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. + +-- Work around +{-# OPTIONS_GHC -fno-float-in #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +-- Make Haddock prefer to link to Data.Vec.Short rather than here, and not +-- complain about missing docs for package-internal functions. +{-# OPTIONS_HADDOCK not-home, prune #-} + +{-# LANGUAGE AllowAmbiguousTypes #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RoleAnnotations #-} +{-# LANGUAGE NoStarIsType #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ViewPatterns #-} + +-- | An implementation of short vectors. +-- +-- The underlying implementation uses the 'SmallArray#' primitive, +-- which lacks the \"card marking\" of 'GHC.Exts.Array#': the upside being +-- that it avoids the overhead of maintaining the card state, the downside +-- being that the garbage collector must scan through the entire array +-- rather than just the parts marked as having changed since the last GC. +-- Using 'SmallArray#' is typically a win for arrays with fewer than 128 +-- elements. + +-- TODO(b/109667526): add rewrite rules, and maybe builder and view +-- interfaces along the way. +-- +-- TODO(b/109668556): revisit all the inline pragmas. + +module Data.Vec.Short.Internal where + +import Prelude hiding ((++), concat, iterate) + +import Control.Applicative (Applicative(..)) +import Control.DeepSeq (NFData(rnf)) +import Control.Exception (assert) +import qualified Data.Data as D +import qualified Data.Foldable as F +import Data.Function (on) +import Data.Functor ((<&>)) +import Data.Kind (Type) +import qualified Data.List as L (sort, sortBy, sortOn, findIndex) +import Data.Semigroup (All(..), Any(..), Sum(..), Product(..)) +import GHC.Exts + ( Int(I#), Proxy#, State#, SmallMutableArray#, SmallArray# + , cloneSmallArray#, copySmallArray#, indexSmallArray#, newSmallArray# + , sizeofSmallArray#, thawSmallArray#, unsafeFreezeSmallArray# + , writeSmallArray#, proxy#, coerce + ) +import qualified GHC.Exts as GHC (IsList(..)) +import GHC.Stack (HasCallStack) +import GHC.ST (ST(..), runST) +import GHC.TypeNats (Nat, KnownNat, type (+), type (*), natVal') + +import Data.Default.Class (Default(..)) +import Data.Distributive (Distributive(..)) +import Data.Foldable.WithIndex (FoldableWithIndex(..)) +import Data.Functor.Apply (Apply(..)) +import Data.Functor.Bind (Bind(..)) +import Data.Functor.Rep (Representable(..)) +import Data.Functor.WithIndex (FunctorWithIndex) +import qualified Data.Functor.WithIndex as X (imap) +import Data.Portray (Portray(..), Portrayal(..), strAtom) +import Data.Portray.Diff (Diff(..)) +import Data.Traversable.WithIndex (TraversableWithIndex(..)) +import qualified Test.QuickCheck as QC + +import Data.Fin.Int (Fin, finToInt, unsafeFin) +import Data.SInt (SInt(SI#, unSInt), sintVal, subSIntL, divSIntR, withSInt, addSInt) + +#if !MIN_VERSION_base(4,15,0) +import GHC.Natural (naturalToInteger, naturalToInt) +import GHC.Integer (integerToInt) +#endif + +-------------------------------------------------------------------------------- +-------------------------------------------------------------------------------- + +foldrEnumFin :: SInt n -> (Fin n -> a -> a) -> a -> a +foldrEnumFin sn c n = go 0 + where + go i + | i == unSInt sn = n + | otherwise = c (unsafeFin i) (go (i + 1)) +{-# INLINE [0] foldrEnumFin #-} + +forMFin_ :: Applicative f => SInt n -> (Fin n -> f a) -> f () +forMFin_ n f = foldrEnumFin n (\i rest -> f i *> rest) (pure ()) +{-# INLINE forMFin_ #-} + +foldMapFin :: Monoid m => SInt n -> (Fin n -> (# m #)) -> m +foldMapFin sn f = go 0 mempty + where + go i acc + | i == unSInt sn = acc + | otherwise = case f (unsafeFin i) of (# x #) -> go (i + 1) (acc <> x) +{-# INLINE foldMapFin #-} + +foldMFin_ + :: Monad m + => SInt n -> (a -> Fin n -> m a) -> a -> m () +foldMFin_ n f z = foldrEnumFin n (\i rest a -> f a i >>= rest) (\_ -> pure ()) z +{-# INLINE foldMFin_ #-} + +forMFin + :: Applicative f => SInt n -> (Fin n -> (# f a #)) -> f [a] +forMFin n f = foldrEnumFin n + (\i -> case f i of (# fx #) -> liftA2 (:) fx) + (pure []) +{-# INLINE forMFin #-} + +-- | Wrap stateful primops which don't return a value. +prim_ :: (State# s -> State# s) -> ST s () +prim_ f = ST $ \s0 -> case f s0 of s1 -> (# s1, () #) +{-# INLINE prim_ #-} + +-- Alas, due to a confluence of problems, we cannot define the combinator: +-- +-- > ($#) :: forall (a :: #) (b :: Type) (s :: Type) +-- > . (a -> b) -> (State# s -> (# State# s, a #)) -> ST s b +-- > ($#) f p = ST $ \s0 -> case p s0 of (# s1, x #) -> (# s1, f x #) +-- +-- While we can hack around most of those problems, and get a version +-- that works for all our use cases, for both ghc-8.0.2 and ghc-8.2.1; +-- the final straw is if we want to keep the code lint clean, HLint-2.1 +-- cannot parse the necessary combination of DataKinds and KindSignatures. +-- Thus, we're forced to inline this combinator everywhere we want it. + + +-- [Note TypeUnsafe]: If the type-unsafe primitives are used without +-- validating their implicit premises, then the 'Nat' type-indices +-- of 'Vec'\/'MutableVec' will become out of sync with the actual +-- 'sizeofSmallArray#'; which in turn invalidates all the safety and +-- correctness guarantees we assumed we could rely on those type-indices +-- to provide. + +-- [Note MemoryUnsafe]: There are two sorts of memory unsafety introduced +-- by GHC's primops. First is the usual index-out-of-bounds unsafety. +-- In the functions defined here, this sort of unsafety only leaks out as +-- a result of type-safety having been violated. Second is the impurity +-- introduced by 'unsafeFreezeSmallArray#' and 'unsafeThawSmallArray#' +-- if references are used non-linearly; that is, because these primops +-- freeze\/thaw arrays in place, they allow a term which holds the +-- 'MutableVec' view to make mutations which are then visible to terms +-- holding the 'Vec' view, thereby violating the purity of Haskell. + +-- Without these annotations GHC will infer that the @n@ parameter is +-- phantom, which opens the door to bugs by allowing folks to coerce it +-- to a different 'Nat'. +type role Vec nominal representational + +-- | @'Vec' n a@ is an element-lazy array of @n@ values of type @a@. +-- +-- This comes with a fusion framework, so many intermediate vectors are +-- optimized away, and generally the only Vecs that correspond to actual arrays +-- are those stored in data constructors, accessed multiple times, or appearing +-- as inputs or outputs of non-inlinable functions. Additionally, some +-- operations that rely on building up a vector incrementally (as opposed to +-- computing each index independently of the others) cannot be fused; notably +-- 'fromList', 'traverse', 'iterate', and 'vscanl'; these will always construct +-- real arrays for their results. +-- +-- Most operations are access-strict unless otherwise noted, which means that +-- forcing the result (usually a Vec, but possibly something else, like with +-- 'foldMap') eagerly performs all indexing and drops references to any input +-- arrays. +data Vec (n :: Nat) (a :: Type) = V# (SmallArray# a) + -- Alas, cannot derive a 'Generic' instance: + -- \"\"V# must not have exotic unlifted or polymorphic arguments\"\" + -- Nor can we derive 'Data', but at least that one we can give our + -- own instance for. + +type role MutableVec nominal nominal representational +data MutableVec (s :: Type) (n :: Nat) (a :: Type) + = MV# (SmallMutableArray# s a) + +newMV :: SInt n -> a -> ST s (MutableVec s n a) +newMV (SI# n) = unsafeNewMV n +{-# INLINE newMV #-} + +-- TODO(b/109668129): We should be able to replace most of the remaining +-- calls to this function with 'newMV' if we made use of the various +-- combinators in "Utils.NatMath". Would be nice to get that extra +-- type-safety, supposing it doesn't introduce significant performance +-- regressions. + +-- | This function is /type-unsafe/: because it assumes the 'Int' +-- argument is in fact the reflection of @n@. +unsafeNewMV :: Int -> a -> ST s (MutableVec s n a) +unsafeNewMV (I# n) x = + ST $ \s0 -> + case newSmallArray# n x s0 of { (# s1, sma #) -> + (# s1, MV# sma #) } +{-# INLINE unsafeNewMV #-} + +-- Unsafe version of writeMV, using Int. +-- Each use should be vetted for being in bounds. +unsafeWriteMV :: MutableVec s n a -> Int -> a -> ST s () +unsafeWriteMV (MV# sma) (I# i) x = prim_ (writeSmallArray# sma i x) +{-# INLINE unsafeWriteMV #-} + +writeMV :: MutableVec s n a -> Fin n -> a -> ST s () +writeMV mv i = unsafeWriteMV mv (finToInt i) +{-# INLINE writeMV #-} + +-- | This function is /memory-unsafe/: because it freezes the 'MutableVec' +-- in place. See [Note MemoryUnsafe]. +unsafeFreezeMV :: MutableVec s n a -> ST s (Vec n a) +unsafeFreezeMV (MV# sma) = + ST $ \s0 -> + case unsafeFreezeSmallArray# sma s0 of { (# s1, sa #) -> + (# s1, V# sa #) } +{-# INLINE unsafeFreezeMV #-} + +-- | Safely thaw a vector, by allocating a new array and copying the +-- elements over. This is both type-safe and memory-safe. +safeThawMV :: Vec n a -> ST s (MutableVec s n a) +safeThawMV (V# sa) = + ST $ \s0 -> + case thawSmallArray# sa 0# (sizeofSmallArray# sa) s0 of { (# s1, sma #) -> + (# s1, MV# sma #) } +{-# INLINE safeThawMV #-} + +-- | This function is /type-unsafe/: because it assumes all the integers +-- are in bounds for their respective arrays. It is also /memory-unsafe/, +-- because we don't do any dynamic checks on those integers. We could +-- add such, but have avoided doing so for performance reasons. +-- See [Note TypeUnsafe] and [Note MemoryUnsafe]. +-- +-- TODO(b/109671227): would assertions /really/ affect the performance +-- significantly? +unsafeCopyVec :: Vec n a -> Int -> MutableVec s m a -> Int -> Int -> ST s () +unsafeCopyVec (V# src) (I# srcOff) (MV# dst) (I# dstOff) (I# len) = + prim_ (copySmallArray# src srcOff dst dstOff len) +{-# INLINE[1] unsafeCopyVec #-} + +-- Avoid 0 length copies. +{-# RULES "unsafeCopyVec/0" forall v s m d . unsafeCopyVec v s m d 0 = return () #-} + +-- | Return a known-length slice of a given vector. +-- +-- Since the type is insufficiently specific to ensure memory-safety on its own +-- (because the offset argument is just 'Int'), this needs to perform runtime +-- bounds checks to ensure memory safety. +sliceVec :: Vec n a -> Int -> SInt m -> Vec m a +sliceVec xs@(V# sa) off@(I# o) (SI# len@(I# l)) = + assert (0 <= off && 0 <= len && len <= vSize xs - off) $ + V# (cloneSmallArray# sa o l) +{-# INLINE sliceVec #-} + +{- +-- If we define a @i :<: n@ type whose witnesses are isomorphic to @i@ +-- itself, then we can implement these safely by rephrasing the Pi-type +-- @(i :: Fin n) -> t@ as the non-Pi @forall i. i :<: n -> t@. Otherwise +-- we can only do unsafe implementations, like the 'unsafeCopyVec' and +-- 'sliceVec' above. All the ones that use 'Min' will want to add +-- {-# OPTIONS_GHC -fplugin=GHC.TypeLits.Extra.Solver #-} to infer well. + +sliceVec + :: Vec n a + -> (o :: Fin n) + -> (m :: Fin (n - o)) + -> ST s (Vec m a) +sliceVec (V# sa) (finToInt -> I# off) (finToInt -> I# len) = + V# (cloneSmallArray# sa off len) + +copyVec + :: Vec n a + -> (srcOff :: Fin n) + -> MutableVec s m a + -> (dstOff :: Fin m) + -> (len :: Fin (m - dstOff `Min` n - srcOff)) + -> ST s () +copyVec (V# src) (finToInt -> I# srcOff) (MV# dst) (finToInt -> I# dstOff) (finToInt -> I# len) = + prim_ (copySmallArray# src srcOff dst dstOff len) + +-- And similarly for 'copySmallMutableArray#', 'cloneSmallArray#', +-- 'cloneSmallMutableArray#', 'freezeSmallArray#', 'thawSmallArray#'. +-} + + +-------------------------------------------------------------------------------- +-------------------------------------------------------------------------------- + +fetch :: Vec n a -> Fin n -> (# a #) +fetch (V# arr) (finToInt -> I# i) = indexSmallArray# arr i +{-# INLINE fetch #-} + +fusibleFetch :: Vec n a -> Fin n -> (# a #) +fusibleFetch = _aIndex . access +{-# INLINE fusibleFetch #-} + +-- | Extract the given index from a 'Vec'. +-- +-- This is subject to fusion if this is the only use of its input, so code like +-- @fmap f v ! i@ (which might arise due to inlining) will optimize to +-- @f (v ! i)@. +(!) :: Vec n a -> Fin n -> a +(!) xs i = case fusibleFetch xs i of (# x #) -> x +{-# INLINE (!) #-} + +-- | Eagerly look up the value at a given position, without forcing the +-- value itself. +-- +-- One of the problems with @(!)@ is that it will hold onto the underlying +-- array until the @xs!i@ expression is forced; which is a source of +-- space leaks. However, forcing the @xs!i@ expression will force +-- not only the array lookup but also the value itself; which can be +-- undesirably strict, thereby ruining the compositionality benefits +-- of laziness. The 'indexK' function is designed to overcome those +-- limitations. That is, forcing the expression @indexK xs i k@ will +-- force the array lookup and the @r@ value; thereby leaving it up to +-- @k@ to decide whether or not to force the @a@ before returning @r@. +-- So, for example, if we force @indexK xs i Just@ this will force the +-- array lookup, and wrap the unforced element in the 'Just' constructor. +{- HLINT ignore indexK "Eta reduce" -} +indexK :: Vec n a -> Fin n -> (a -> r) -> r +indexK v i k = case fetch v i of (# x #) -> k x +{-# INLINE indexK #-} + +-- | Return the size of a vector as 'SInt'. +svSize :: Vec n a -> SInt n +-- Note this strongly relies on @n@ matching the actual size of the array: if +-- it didn't, we'd be constructing an invalid 'SInt', which manifests +-- unsafety. So, it's unsafe for a Vec to have the wrong length. +svSize (V# sa) = SI# (I# (sizeofSmallArray# sa)) +{-# INLINE svSize #-} + +-- | Dynamically determine the (actual) size\/length of the vector, +-- as a standard term-level 'Int'. If you'd rather obtain @n@ at the +-- type-level, and\/or to prove that the returned value is indeed the +-- @n@ of the input, see 'svSize' and 'Data.Vec.Short.withSize' instead. If +-- you'd rather obtain @n@ statically, see 'Data.Vec.Short.vLength'. +vSize :: Vec n a -> Int +vSize = unSInt . svSize +{-# INLINE vSize #-} + + +-------------------------------------------------------------------------------- +uninitialized :: a +uninitialized = error "Vec: uninitialized" +{-# NOINLINE uninitialized #-} + +-- Unsafe version of createVec, with Int instead of SInt. Each use +-- should be vetted for size == valueOf @n. Using this rather than writing out +-- 'SI# at each call site means we have a place to insert assertions more +-- easily. +unsafeCreateVec :: Int -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a +unsafeCreateVec n = createVec (SI# n) +{-# INLINE unsafeCreateVec #-} + +createVec + :: SInt n + -> (forall s. MutableVec s n a -> ST s ()) -> Vec n a +createVec n action = runST $ do + mv <- newMV n uninitialized + action mv + unsafeFreezeMV mv +{-# INLINE createVec #-} + +-------------------------------------------------------------------------------- +-- Fusion Internals +-------------------------------------------------------------------------------- + +-- Fusion framework overview: +-- +-- Like with list fusion, the goal is to replace actual intermediate Vec +-- objects with a non-materialized representation. In this case, that's a +-- function for accessing vector elements by their Fin index, paired with a +-- runtime representation of the Vec length; this representation is called +-- 'Accessor'. Also like with list fusion, we arrange to rewrite the +-- user-facing API functions to a "fusion form", which converts any input Vecs +-- into Accessors (with 'access'), implements the actual logic in terms of the +-- non-materialized representation, and converts any results into actual Vecs +-- with 'materialize'. Then a rewrite rule deletes opposing +-- 'access'/'materialize' pairs, eliminating the intermediate allocations +-- wherever this happens. +-- +-- To avoid duplicating work computing the Vec elements, we have a soft +-- requirement that a particular call to 'access' must not be used to fetch any +-- index more than once. Violating this would mean the access/materialize rule +-- can reduce sharing. No current functions break this rule, but we could +-- imagine adding cases where it's the client's responsibility to make sure +-- elements aren't used multiple times, like a "linear" variant of +-- 'backpermute'. +-- +-- Since some Vec functions admit more efficient implementations than you'd get +-- by materializing an Accessor of their contents (e.g. implementing 'take_' by +-- 'unsafeCopyVec'), we adapt another trick from the list fusion library: keep +-- relevant operations on 'Accessor's in symbolic form for one extra simplifier +-- phase, and detect when these operations are still left alone and un-fused, +-- to rewrite them back to specialized implementations. When we need to do +-- this, we have the specialized implementation under a different name from the +-- user-facing function, write the fusion form as the implementation of the +-- user-facing one with an INLINE pragma, and have an extra rule to bring back +-- the specialized form starting in phase 1. When we don't need to rewrite +-- back to a specialized implementation (e.g. with 'fmap'), there's simply no +-- specialized implementation provided. +-- +-- Why the difference from how list fusion does it (namely, by writing the +-- specialized implementation as the user-facing function and rewriting fusion +-- forms back to the original with phase-controlled RULES)? I've seen GHC's +-- specialization pass do bad things with that approach, seemingly re-applying +-- rules in the wrong order to the output of specialization, resulting in using +-- element-by-element implementations instead of specialized ones. By having +-- totally different functions for things like 'append_' and 'pureVec_', we +-- can't accidentally mess them up post-specialization with RULES. +-- +-- Here's what happens in each of the GHC simplifier phases: +-- +-- Phase [gentle]: GHC isn't willing to inline worker-wrapper'd function bodies +-- yet, so if we tried to use this phase to make real progress, we'd miss +-- things that didn't get exposed to RULES until phase 2. So, we just bide our +-- time. Some of our rules are enabled, but we don't change anything after +-- [gentle]. +-- +-- Phase 2: expand ops into their fusion form with RULES or INLINEs, and do the +-- actual fusion of adjacent ops with the "materialize/access" rule. Also, do +-- some limited single-op fusion by explicitly merging the fusible form of +-- map-of-map with RULES. This allows more cases to get rewritten back to +-- specialized implementations in phase 1. +-- +-- Phase 1: detect cases that are still just a single op with a specialized +-- implementation available, and rewrite them to use it. That is, when no +-- fusion actually happened, go back into non-fusion land. +-- +-- Phase 0: inline everything and let GHC optimize the things that did get +-- subjected to nontrivial fusion. + +data Accessor n a = Accessor + { _aSize :: SInt n + , _aIndex :: Fin n -> (# a #) + } + +-- | Convert a 'Vec' into its size and an indexing function. +access :: Vec n a -> Accessor n a +access v = Accessor (svSize v) (fetch v) +{-# INLINE CONLIKE [0] access #-} + +-- | Construct an actual 'Vec' from an 'Accessor'. +-- +-- Strictness properties: forcing the resulting Vec forces all of the unboxed +-- tuple accesses, so you can make Vecs that are strict in whatever you want by +-- controlling what work goes inside/outside the unboxed tuple construction. +-- Generally this is used to force all of the array accessing so that input +-- 'Vec's are no longer retained after the result is forced; but it's also used +-- to implement element-strict variants of some functions. +materialize :: Accessor n a -> Vec n a +materialize (Accessor n get) = createVec n $ \mv -> forMFin_ n $ \i -> + case get i of (# x #) -> writeMV mv i x +{-# INLINE [0] materialize #-} + +{-# RULES + +-- Fuses adjacent Vec ops, keeping everything in Accessor form. +"access/materialize" forall va. access (materialize va) = va + +-- Transports coercions out from between access/materialize pairs so they can +-- fuse. +"access/coerce/materialize" + forall v. access (coerce v) = mapVA coerce (access v) + +-- Eliminates no-op copies of a vector. +"materialize/access" forall v. materialize (access v) = v + + #-} + +pureVA :: SInt n -> a -> Accessor n a +pureVA n x = Accessor n $ \_ -> (# x #) +{-# INLINE [0] pureVA #-} + +mapVA :: (a -> b) -> Accessor n a -> Accessor n b +mapVA f (Accessor n get) = Accessor n $ \i -> case get i of (# x #) -> (# f x #) +{-# INLINE [0] mapVA #-} + +mapWithPosVA :: (Fin n -> a -> b) -> Accessor n a -> Accessor n b +mapWithPosVA f (Accessor n get) = Accessor n $ + \i -> case get i of (# x #) -> (# f i x #) +{-# INLINE [0] mapWithPosVA #-} + +-- Make an 'Accessor' force its elements before returning them. +seqVA :: Accessor n a -> Accessor n a +seqVA (Accessor n get) = Accessor n $ + \i -> case get i of (# x #) -> x `seq` (# x #) +{-# INLINE [0] seqVA #-} + +takeVA + :: SInt m -> Accessor (m + n) a -> Accessor m a +takeVA m (Accessor _ get) = Accessor m (\i -> get (embedPlus i)) + where + embedPlus :: Fin m -> Fin (m + n) + embedPlus = unsafeFin . finToInt +{-# INLINE [0] takeVA #-} + +dropVA :: SInt m -> Accessor (m + n) a -> Accessor n a +dropVA m (Accessor mn get) = Accessor (SI# (unSInt mn - unSInt m)) $ + \i -> get (unsafeFin (finToInt i + unSInt m)) +{-# INLINE [0] dropVA #-} + +revVA :: Accessor n a -> Accessor n a +revVA (Accessor n get) = Accessor n $ \i -> get (complementIt i) + where + !nMinus1 = unSInt n - 1 + + complementIt :: Fin n -> Fin n + complementIt = unsafeFin . (nMinus1 -) . finToInt +{-# INLINE [0] revVA #-} + +rotVA :: Fin n -> Accessor n a -> Accessor n a +rotVA (finToInt -> !o) (Accessor n get) = Accessor n $ + \(finToInt -> !i) -> get $ unsafeFin $ if i >= o then i - o else nmo + i + where + !nmo = unSInt n - o +{-# INLINE [0] rotVA #-} + +liftA2VA :: (a -> b -> c) -> Accessor n a -> Accessor n b -> Accessor n c +liftA2VA f (Accessor n getA) (Accessor _ getB) = Accessor n $ + \i -> case getA i of (# a #) -> case getB i of (# b #) -> (# f a b #) +{-# INLINE [0] liftA2VA #-} + +foldMapVA :: Monoid m => (a -> m) -> Accessor n a -> m +foldMapVA f (Accessor n get) = + foldMapFin n (\i -> case get i of (# x #) -> (# f x #)) +{-# INLINE [0] foldMapVA #-} + +sequenceVA :: Applicative f => Accessor n (f a) -> f (Vec n a) +sequenceVA (Accessor n get) = listVec n <$> forMFin n get +{-# INLINE [0] sequenceVA #-} + +-- SInt version of 'splitFin'. Maybe I'll change the Fin library to provide +-- an SInt API at some point? +splitFinS :: SInt n -> Fin (n + m) -> Either (Fin n) (Fin m) +splitFinS (SI# n) (finToInt -> i) + | i < n = Left (unsafeFin i) + | otherwise = Right (unsafeFin (i - n)) + +addPosSInt :: SInt n -> SInt m -> SInt (n + m) +addPosSInt (SI# n) (SI# m) = + let nm = n + m + in if nm < 0 + then error "addPosSInt: Int overflow" + else SI# (n + m) +{-# INLINE addPosSInt #-} + +appendVA :: Accessor n a -> Accessor m a -> Accessor (n + m) a +appendVA (Accessor n getN) (Accessor m getM) = Accessor + (addPosSInt n m) + (\i -> case splitFinS n i of + Left i' -> getN i' + Right i' -> getM i') +{-# INLINE [0] appendVA #-} + +-------------------------------------------------------------------------------- +-- User-facing API with fusion rules +-------------------------------------------------------------------------------- + +-- Unsafe version of mkVec, with Int instead of SInt. Each use should be +-- vetted for s == valueOf @n. Using this rather than writing out 'SI# and +-- 'unsafeFin' at each call site means we have a place to insert assertions +-- more easily. +unsafeMkVec :: Int -> (Int -> a) -> Vec n a +unsafeMkVec n f = mkVec (SI# n) $ \i -> f (finToInt i) +{-# INLINE unsafeMkVec #-} + +-- | Create a known-length vector using a pure function. +-- +-- Note if you already have a 'Vec' of the desired length, you can use 'svSize' +-- to get the 'SInt' parameter. +tabulateVec, mkVec :: SInt n -> (Fin n -> a) -> Vec n a +tabulateVec n f = materialize $ Accessor n $ \i -> (# f i #) +mkVec = tabulateVec +{-# INLINE tabulateVec #-} +{-# INLINE mkVec #-} + +-- | Element-strict form of 'mkVec': elements are forced when forcing the Vec. +tabulateVec', mkVec' :: SInt n -> (Fin n -> a) -> Vec n a +tabulateVec' n f = materialize $ + Accessor n $ \i -> let x = f i in x `seq` (# x #) +mkVec' = tabulateVec' +{-# INLINE tabulateVec' #-} +{-# INLINE mkVec' #-} + +-- | Construct a vector by choosing an element of another vector for each index. +-- +-- @ +-- backpermute n f v ! i === v ! f i +-- @ +backpermute :: SInt m -> (Fin m -> Fin n) -> Vec n a -> Vec m a +-- Take care: backpermute can reference the same index of the input vector +-- multiple times, so if we subjected the input side to fusion, we'd +-- potentially duplicate work. It might make sense to make a variant of +-- 'backpermute' that assumes either indices are not duplicated or the +-- computation behind each value is cheap enough to duplicate, but we can't +-- just blindly fuse things into all 'backpermute's. +backpermute m f xs = materialize $ Accessor m $ \i -> fetch xs (f i) +{-# INLINE backpermute #-} + +-------------------------------------------------------------------------------- +-- N.B., since the same @KnownNat n@ instance is passed to both 'createVecP' +-- and 'enumFinP', this will be memory-safe even if the @KnownNat n@ +-- instance is illegitimate (e.g., by using 'blackMagic' unsafely). +-- An illegitimate 'KnownNat' instance would only compromise type-safety: +-- since it'd mean that the actual length of the resulting 'SmallArray#' +-- differs from the @n@ the 'Vec' claims it has. + +-- | Create a vector of the specified length from a list. If @n < length xs@ +-- then the suffix of the vector will be filled with 'uninitialized' +-- values. If @n > length xs@ then the suffix of @xs@ won't be included +-- in the vector. Either way, this function is both type-safe and memory-safe. +listVec :: SInt n -> [a] -> Vec n a +listVec n xs = createVec n $ \mv -> ($ xs) $ foldrEnumFin n + (\i rest xs' -> case xs' of + [] -> writeMV mv i uninitialized >> rest [] + (x:xs'') -> writeMV mv i x >> rest xs'') + (\_ -> return ()) +{-# INLINE listVec #-} + + +-- | Convert a list to a vector of the same length. +withVec :: [a] -> (forall n. Vec n a -> r) -> r +withVec xs f = withSInt (length xs) $ \sn -> f (listVec sn xs) +{-# INLINABLE withVec #-} + +-- | Convert a list to a vector, given a hint for the length of the list. +-- If the hint does not match the actual length of the list, then the +-- behavior of this function is left unspecified. If the hint does not +-- match the desired @n@, then we throw an error just like 'fromList'. +-- For a non-errorful version, see 'withVec' instead. +fromListN :: HasCallStack => SInt n -> Int -> [a] -> Vec n a +fromListN sn l xs + | l == n = listVec sn xs + | otherwise = error $ "Vec.fromListN: " <> show l <> " /= " <> show n + where + !n = unSInt sn +{-# INLINABLE fromListN #-} + + +-- | Convert a list to a vector, throwing an error if the list has the +-- wrong length. +-- Note: Because this walks @xs@ to check its length, this cannot be +-- used with the list fusion optimization rules. +fromList :: HasCallStack => SInt n -> [a] -> Vec n a +fromList sn xs + | n `eqLength` xs = listVec sn xs + | otherwise = error $ "Vec.fromList: length /= " <> show n + where + !n = unSInt sn +{-# INLINABLE fromList #-} + + +-- TODO(b/109757264): move this out of library into @ListUtils.hs@ +-- | An implementation of @n == length xs@ which short-circuits +-- once it can determine the answer, rather than necessarily recursing +-- through the entire list to compute its length. +eqLength :: Int -> [a] -> Bool +eqLength 0 [] = True +eqLength 0 (_:_) = False -- too long +eqLength _ [] = False -- too short +eqLength n (_:xs) = eqLength (n - 1) xs + + +-- To support -XOverloadedLists +instance KnownNat n => GHC.IsList (Vec n a) where + type Item (Vec n a) = a + fromListN = fromListN sintVal + fromList = fromList sintVal -- Not subject to list fusion optimizations. + toList = F.toList + +-------------------------------------------------------------------------------- +-- | Claim that 'vecConstr' is the only data-constructor of 'Vec'. +vecDataType :: D.DataType +vecDataType = D.mkDataType "Vec.Vec" [vecConstr] + +-- | Treat the 'fromList' function as a data-constructor for 'Vec'. +vecConstr :: D.Constr +vecConstr = D.mkConstr vecDataType "fromList" [] D.Prefix + +-- The 'KnownNat' constraint is necessary for 'fromList'. +instance (KnownNat n, D.Data a) => D.Data (Vec n a) where + toConstr _ = vecConstr + dataTypeOf _ = vecDataType + gfoldl app pur v = pur (fromList sintVal) `app` F.toList v + gunfold app pur c + | D.constrIndex c == 1 = app (pur (fromList sintVal)) + | otherwise = error "gunfold@Vec: invalid constrIndex" + +-------------------------------------------------------------------------------- + +instance Show a => Show (Vec n a) where + showsPrec p xs = showParen (p >= precedence) + $ showString "fromListN " + . shows (vSize xs) + . showString " " + . shows (F.toList xs) + +instance (KnownNat n, Read a) => Read (Vec n a) where + readsPrec p = readParen (p >= precedence) $ \s -> + [ assertSize (length ls) (fromListN n m ls, s''') + | ("fromListN", s') <- lex s + , (m, s'') <- readsPrec precedence s' + , (ls, s''') <- assertSize m readsPrec precedence s'' + ] + where + n = sintVal @n + + assertSize :: Int -> b -> b + assertSize m x = if m /= unSInt n + then error $ "Can't read a Vec with " <> show m <> + " elements into a type `Vec " <> show n <> + "`" + else x + +-- Vec is being shown as a function application which has precedence 10. Thus, +-- if we are already in function application context or one that binds even +-- tightlier (i.e. with higher precedence) we need to wrap the expression in +-- parantheses. +precedence :: Int +precedence = 10 + +instance Portray a => Portray (Vec n a) where + portray xs = Apply (strAtom "fromListN") + [portray (vSize xs), portray $ F.toList xs] + +instance (Portray a, Diff a) => Diff (Vec n a) where + diff x y = (diff `on` F.toList) x y <&> + \d -> Apply (strAtom "fromListN") [portray (vSize x), d] + +instance NFData a => NFData (Vec n a) where + rnf !xs = foldMapFin (svSize xs) $ + \i -> case indexK xs i rnf of () -> (# () #) + {-# INLINE rnf #-} + +-- | Point-wise @(<>)@. +instance Semigroup a => Semigroup (Vec n a) where + (<>) = liftF2 (<>) + +-- | Point-wise @mempty@. +instance (KnownNat n, Monoid a) => Monoid (Vec n a) where + mempty = pure mempty + +instance forall a n. (QC.Arbitrary a, KnownNat n) => QC.Arbitrary (Vec n a) + where + -- While we could get rid of the intermediate list by digging into + -- how 'Gen' works under the hood, the benefit doesn't seem worth it. + arbitrary = listVec sn <$> QC.vectorOf n QC.arbitrary + where + !sn@(SI# n) = sintVal + + -- If @a@ admits too many ways to shrink, we might prefer to + -- interleave the @shrink(xs!i)@ lists, rather than concatenating + -- them as list comprehension will. + shrink xs = + [ upd i xs x' + | i <- foldrEnumFin sn (:) [], x' <- indexK xs i QC.shrink + ] + where + !sn = svSize xs + +-- | Safely construct a new vector that differs only in one element. +-- This is inefficient, and only intended for internal use. +upd :: Fin n -> Vec n a -> a -> Vec n a +upd i xs x = runST $ do + mv <- safeThawMV xs + writeMV mv i x + unsafeFreezeMV mv +{-# INLINE upd #-} + +instance (Show a) => QC.CoArbitrary (Vec n a) where + coarbitrary = QC.coarbitraryShow + +instance (KnownNat n, Num a) => Num (Vec n a) where + (+) = liftA2 (+) + (*) = liftA2 (*) + (-) = liftA2 (-) + abs = fmap abs + signum = fmap signum + negate = fmap negate + fromInteger = pure . fromInteger + +instance (KnownNat n, Default a) => Default (Vec n a) where + def = pure def + +instance Eq a => Eq (Vec n a) where + xs == ys = getAll $ foldMap All $ liftF2 (==) xs ys + xs /= ys = getAny $ foldMap Any $ liftF2 (/=) xs ys + +instance Ord a => Ord (Vec n a) where + compare xs ys = F.fold $ liftF2 compare xs ys + +mapVec :: (a -> b) -> Vec n a -> Vec n b +mapVec f = materialize . mapVA f . access +{-# INLINE mapVec #-} + +{-# RULES + +"mapVec/merge" forall f g va. mapVA f (mapVA g va) = mapVA (f . g) va +"mapVec/coerce" [1] forall v. materialize (mapVA coerce (access v)) = coerce v +"mapVec/id" mapVA (\x -> x) = id + + #-} + +instance Functor (Vec n) where + fmap = mapVec + + -- This is just 'pureVec' getting its size from an existing 'Vec'. Since + -- the output is independent of the values in the input, we can tie into + -- the fusion framework to get the size of the input Vec, and still + -- potentially use the specialized form of 'pureVec'. + x <$ v = pureVec (_aSize (access v)) x + {-# INLINE (<$) #-} + +instance FunctorWithIndex (Fin n) (Vec n) where imap = imap + +instance KnownNat n => FoldableWithIndex (Fin n) (Vec n) where + ifoldMap f = F.fold . imap f + +instance KnownNat n => TraversableWithIndex (Fin n) (Vec n) where + itraverse f = sequenceA . imap f + +-- | An element-strict version of 'fmap'. +map' :: (a -> b) -> Vec n a -> Vec n b +map' f = materialize . seqVA . mapVA f . access +{-# INLINE map' #-} + +-- | A variant of 'fmap' that provides the index in addition to the element. +imap :: (Fin n -> a -> b) -> Vec n a -> Vec n b +imap f = materialize . mapWithPosVA f . access +{-# INLINE imap #-} + +-- | Pair each element of a 'Vec' with its index. +-- +-- You can also use 'imap', but there should be no harm in using this +-- because the fusion framework should optimize away the intermediate Vec. +withPos :: Vec n a -> Vec n (Fin n, a) +withPos = imap (,) +{-# INLINE withPos #-} + +-- | An element-strict version of 'imap'. +imap' :: (Fin n -> a -> b) -> Vec n a -> Vec n b +imap' f = materialize . seqVA . mapWithPosVA f . access +{-# INLINE imap' #-} + +pureVec_, pureVec :: SInt n -> a -> Vec n a +pureVec_ n x = runST $ newMV n x >>= unsafeFreezeMV +{-# NOINLINE pureVec_ #-} + +pureVec = \n x -> materialize (pureVA n x) +{-# INLINE pureVec #-} + +{-# RULES + +"pureVec/spec" [1] forall n x. materialize (pureVA n x) = pureVec_ n x + +"mapVA/pureVA" forall f n x. mapVA f (pureVA n x) = pureVA n (f x) + + #-} + +liftA2Vec :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c +liftA2Vec f as bs = materialize (liftA2VA f (access as) (access bs)) +{-# INLINE liftA2Vec #-} + +instance Apply (Vec n) where + liftF2 = liftA2Vec + {-# INLINE liftF2 #-} + +instance KnownNat n => Applicative (Vec n) where + pure = pureVec sintVal + {-# INLINE pure #-} + + liftA2 = liftA2Vec + {-# INLINE liftA2 #-} + + (<*>) = liftA2Vec ($) + {-# INLINE (<*>) #-} + + _ *> ys = ys + {-# INLINE (*>) #-} + + xs <* _ = xs + {-# INLINE (<*) #-} + +instance Bind (Vec n) where + xs >>- f = materialize (case access xs of + Accessor n get -> Accessor n (\i -> case get i of + (# x #) -> f x `fusibleFetch` i)) + {-# INLINE (>>-) #-} + +instance KnownNat n => Monad (Vec n) where (>>=) = (>>-) + +-- | A truly lazy version of @liftA2@. +-- +-- As opposed to the actual @liftA2@ it does not inspect the arguments which +-- makes it possible it to use in code that has lazy knot-tying. +liftA2Lazy :: SInt n -> (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c +liftA2Lazy sn f xs ys = tabulateVec sn $ \i -> + indexK xs i $ \x -> + indexK ys i $ \y -> + f x y + +-------------------------------------------------------------------------------- +-- | > unsafeIndexK xs i === indexK xs (unsafeFin i) +-- +-- TODO(b/109672429): try to get rid of all the uses of this function, +-- and other uses of 'unsafeFin' as well. +unsafeIndexK :: Vec n a -> Int -> (a -> r) -> r +unsafeIndexK (V# sa) (I# i) k = case indexSmallArray# sa i of (# x #) -> k x +{-# INLINE unsafeIndexK #-} + +instance Foldable (Vec n) where + length = vSize + {-# INLINE length #-} + + null = (0 ==) . length + {-# INLINE null #-} + + foldMap f = foldMapVA f . access + {-# INLINE foldMap #-} + + fold = foldMapVA id . access + {-# INLINE fold #-} + + foldr f acc0 = \v -> + let Accessor n get = access v + in foldrEnumFin n + (\i acc -> case get i of (# x #) -> f x acc) + acc0 + {-# INLINE foldr #-} + + foldr' f acc0 = \v -> + let !(Accessor n get) = access v + go !i !acc + | i < 0 = acc + | otherwise = + case get (unsafeFin i) of (# x #) -> go (i - 1) (f x acc) + in go (unSInt n - 1) acc0 + {-# INLINE foldr' #-} + + foldl f acc0 = \v -> + let !(Accessor n get) = access v + go !i + | i < 0 = acc0 + | otherwise = case get (unsafeFin i) of (# x #) -> f (go (i - 1)) x + in go (unSInt n - 1) + {-# INLINE foldl #-} + + foldl' f acc0 = \v -> + let !(Accessor (unSInt -> !n) get) = access v + go !i !acc + | i >= n = acc + | otherwise = + case get (unsafeFin i) of (# a #) -> go (i+1) (f acc a) + in go 0 acc0 + {-# INLINE foldl' #-} + + foldr1 f = \v -> + case access v of + Accessor (subtract 1 . unSInt -> lMinus1) get + | lMinus1 < 0 -> error "foldr1@Vec: empty list" + | otherwise -> + let !(# z #) = get (unsafeFin lMinus1) + go !i | i >= lMinus1 = z + | otherwise = + let !(# x #) = get (unsafeFin i) + in f x (go (i + 1)) + in go 0 + {-# INLINE foldr1 #-} + + foldl1 f = \v -> + case access v of + Accessor (subtract 1 . unSInt -> lMinus1) get + | lMinus1 < 0 -> error "foldl1@Vec: empty list" + | otherwise -> + let !(# z #) = get (unsafeFin (0 :: Int)) + go !i | i <= 0 = z + | otherwise = + let !(# x #) = get (unsafeFin i) + in f (go (i - 1)) x + in go lMinus1 + {-# INLINE foldl1 #-} + + -- The INLINE declarations here are important to fusion: without it, GHC + -- fully optimizes the default implementation (which is identical to this + -- one) when compiling this module, and exposes the post-optimized + -- unfolding. Then when we use 'sum', the thing that gets inlined is + -- already post-fusion, and no fusion can happen. + -- + -- With the INLINE, GHC exposes (the Core desugaring of) this exact term as + -- the unfolding, and 'sum' can participate in fusion. + + sum = coerce . foldMap Sum + {-# INLINE sum #-} + + product = coerce . foldMap Product + {-# INLINE product #-} + + minimum = foldr1 min + {-# INLINE minimum #-} + + maximum = foldr1 max + {-# INLINE maximum #-} + + elem x = coerce . foldMap (Any . (== x)) + {-# INLINE elem #-} + +traverseVec :: Applicative f => (a -> f b) -> Vec n a -> f (Vec n b) +traverseVec f = sequenceVA . mapVA f . access +{-# INLINE traverseVec #-} + +-- TODO(b/109674103): for linear-use applicatives ('IO', 'Maybe',...) we +-- can give a more efficient definition. +instance Traversable (Vec n) where + traverse = traverseVec + {-# INLINE traverse #-} + +-- We don't bother defining 'collect', since anything other than the +-- default instance would be egregiously inefficient if the function +-- actually allocates vectors, since we'd end up allocating @n@ identical +-- copies of every vector in the @f@ structure! If, however, we were +-- to have a bunch of fusion rules for things like @mkVec f ! i = f i@; +-- then, it might be worth defining 'collect' directly, since we could +-- avoid allocating any of the intermediate arrays! +-- +-- TODO(b/109674103): for linear-use functors ('IO', 'Maybe',...) we +-- can give a more efficient definition. +instance KnownNat n => Distributive (Vec n) where + distribute fxs = tabulate $ \i -> fmap (! i) fxs + {-# INLINE distribute #-} + +instance KnownNat n => Representable (Vec n) where + type Rep (Vec n) = Fin n + + tabulate = tabulateVec sintVal + {-# INLINE tabulate #-} + + index = (!) + {-# INLINE index #-} + +-- | 'Prelude.scanl', for 'Vec'. +vscanl :: (b -> a -> b) -> b -> Vec n a -> Vec (1 + n) b +-- TODO(awpr): we can probably subject the input Vec to fusion here. +vscanl f b v = listVec (sintVal `addSInt` svSize v) . scanl f b $ F.toList v + + +-------------------------------------------------------------------------------- +-------------------------------------------------------------------------------- +-- List like ops and their isomorphisms. + +-- | A zero-length 'Vec' of any element type. +nil :: Vec 0 a +-- Note: in the C-- code, this is a single global thunk with a polymorphic +-- type; we won't re-create it separately for different types. The NOINLINE +-- serves to ensure we just reference the thunk rather than inlining the code +-- that builds it. This blocks fusion, but we have another trick: rewrite any +-- materialize @0 to nil. Then anything that would've fused with this will no +-- longer reference it at all. +nil = mkVec sintVal (\i -> i `seq` error "Vec.nil: the impossible happened") +{-# NOINLINE nil #-} + +{-# RULES + +-- Note this matches only zero-length vectors because the LHS type is unified +-- with the RHS type. +"materialize@0" forall va. materialize va = nil + + #-} + +---------------- + +-- | Concatenate two 'Vec's. +infixr 5 ++ +append_, (++) :: Vec n a -> Vec m a -> Vec (n + m) a +append_ xs ys = runST $ do + let !n = vSize xs + !m = vSize ys + zs <- unsafeNewMV (n + m) uninitialized + unsafeCopyVec xs 0 zs 0 n + unsafeCopyVec ys 0 zs n m + unsafeFreezeMV zs +{-# NOINLINE append_ #-} + +(++) = \xs ys -> materialize (appendVA (access xs) (access ys)) +{-# INLINE (++) #-} + +{-# RULES + +"++/spec" [1] + forall xs ys. materialize (appendVA (access xs) (access ys)) = xs `append_` ys + + #-} + +-- TODO(b/109675695): may want other variants of this operational-function +-- with different types. +-- TODO: might as well simply call the underlying +-- 'thawSmallArray#' directly, instead of passing through +-- 'sliceVec'... ne? (i.e., to avoid the dynamic bounds checks) +take_ :: SInt m -> Vec (m + n) a -> Vec m a +take_ m xs = sliceVec xs 0 m +{-# NOINLINE take_ #-} +-- TODO(awpr): fusion behaves poorly here because of Core-level casts arising +-- from unifying @m + n@ with some other Nat: we get +-- @materialize (takeVA m (access ((materialize _) `cast` )))@, which +-- can't match the access/materialize rule. The @cast@ comes from converting +-- the equal-but-not-syntactically-identical types @Vec (m + n) a@ and +-- @Vec _something a@ using the coercion @_something ~N# (m + n)@. If we +-- define this with inequality constraints instead, +-- @take_ :: (m <= n) => SInt m -> Vec n a -> Vec m a@, then there's no need +-- for a cast there; @n@ and @m@ are just unified away to be identical to the +-- sizes of the input and output vectors. + +{-# RULES + +"take_/spec" [1] forall m xs. materialize (takeVA m (access xs)) = take_ m xs + + #-} + +-- TODO(b/109675695): may want other variants of this operational-function +-- with different types. +drop_ :: forall m n a. SInt m -> Vec (m + n) a -> Vec n a +drop_ m xs = + sliceVec xs (unSInt m) $ + svSize xs `subSIntL` m +{-# NOINLINE drop_ #-} +-- TODO(awpr): as with 'take_', casts are causing trouble here. Consider +-- messing with the type signature to avoid them. + +{-# RULES + +"drop_/spec" [1] forall m xs. materialize (dropVA m (access xs)) = drop_ m xs + + #-} + +-- | Split a vector into two shorter vectors at the given index. +split + :: forall m n a. SInt m -> Vec (m + n) a -> (Vec m a, Vec n a) +-- TODO(b/109675695): may want other variants of this operational-function +-- with different types. +split m xs = + let va = access xs + in (materialize (takeVA m va), materialize (dropVA m va)) +{-# INLINE split #-} + +-- TODO(awpr): fusion for 'concat' and 'reshape'? + +-- | Concatenate a nested 'Vec' into one longer 'Vec'. +concat :: forall m n a. Vec n (Vec m a) -> Vec (n * m) a +concat xs = + let !n = vSize xs + in if n == 0 then + -- Outer size is 0, return the empty array + unsafeCreateVec 0 $ \ _ -> return () + else + let !m = unsafeIndexK xs 0 vSize -- We know the size is > 0. + in unsafeCreateVec (n * m) $ \ marr -> + F.forM_ [0..n-1] $ \ i -> + unsafeIndexK xs i $ \ ys -> -- i is [0..n-1], so in bounds. + unsafeCopyVec ys 0 marr (i * m) m +{-# INLINE concat #-} + +-- | Turn a vector into a vector of vector by chunking it. +reshape :: SInt m -> Vec (n * m) a -> Vec n (Vec m a) +reshape m = + let !m' = unSInt m + in \xs -> + mkVec (svSize xs `divSIntR` m) (\i -> sliceVec xs (finToInt i * m') m) +{-# INLINE reshape #-} + +-- | Map each element of a 'Vec' to a (same-sized) sub-'Vec' of the result. +concatMap :: forall m n a b. (a -> Vec m b) -> Vec n a -> Vec (n * m) b +concatMap f = concat . fmap f +{-# INLINE concatMap #-} + + +-- | Generate a Vec by repeated application of a function. +-- +-- > toList (Vec.iterate @n f z) === take (valueOf @n) (Prelude.iterate f z) +iterate :: SInt n -> (a -> a) -> a -> Vec n a +iterate sn f z = + createVec sn $ \mv -> + foldMFin_ sn (\x i -> f x <$ writeMV mv i x) z +{-# INLINE iterate #-} + + +-- | A strict version of 'iterate'. +iterate' :: SInt n -> (a -> a) -> a -> Vec n a +iterate' sn f !z = + createVec sn $ \mv -> + foldMFin_ sn (\x i -> f x <$ (writeMV mv i $! x)) z +{-# INLINE iterate' #-} + + +-- | Return a copy of the array with elements in the reverse order. +rev :: Vec n a -> Vec n a +rev = materialize . revVA . access +{-# INLINE rev #-} + + +-- | Rotate a vector right by @i@ positions. +-- +-- @rot 1 [x, y, z] = [z, x, y]@ +rot, rot_ :: Fin n -> Vec n a -> Vec n a +rot_ o' xs = createVec (svSize xs) $ \mv -> do + let o = finToInt o' + nmo = vSize xs - o + unsafeCopyVec xs nmo mv 0 o + unsafeCopyVec xs 0 mv o nmo +{-# NOINLINE rot_ #-} + +rot o = \v -> materialize (rotVA o (access v)) +{-# INLINE rot #-} + +{-# RULES + +"rot/spec" [1] forall o v. materialize (rotVA o (access v)) = rot_ o v + + #-} + +-- | Return a vector with all elements of the type in ascending order. +viota :: SInt n -> Vec n (Fin n) +viota sn = mkVec sn id +{-# INLINE viota #-} + +-- | One variant of the cross product of two vectors. +cross :: Vec m a -> Vec n b -> Vec n (Vec m (a, b)) +cross xs = fmap (\y -> fmap (, y) xs) +{-# INLINE cross #-} + +-- TODO: the concatenated version. +-- cross :: Vec m a -> Vec n b -> Vec (m * n) (a, b) +-- +-- TODO: the transposed nested version. +-- cross :: Vec m a -> Vec n b -> Vec m (Vec n (a, b)) + +-------------------------------- + +-- Element insertion and removal. + +-- TODO(awpr): fusion implementations of insert and remove? + +-- Unsafe version of insert. Assumes i < valueOf @(n+1) +-- Statically determined 0 length copies are removed by a RULE. +unsafeInsert :: Int -> a -> Vec n a -> Vec (n+1) a +unsafeInsert i xi xs = + let !n = vSize xs + in unsafeCreateVec (n+1) $ \ mv -> do + -- Explicitly: @mv[0..i-1] := xs[0..i-1]@ + unsafeCopyVec xs 0 mv 0 i + -- Explicitly: @mv[i] := xi@ + unsafeWriteMV mv i xi + -- Explicitly: @mv[i+1..n] := xs[i..n-1]@ + unsafeCopyVec xs i mv (i + 1) (n - i) + +-- | Insert an element at the given position in a vector. +-- O(n) +insert :: Fin (n+1) -> a -> Vec n a -> Vec (n+1) a +insert i = unsafeInsert (finToInt i) + +-- Unsafe version of remove. Assumes i < valueOf @(n+1) +-- Statically determined 0 length copies are removed by a RULE. +unsafeRemove :: Int -> Vec (n+1) a -> Vec n a +unsafeRemove i xs = + let !np1 = vSize xs + in unsafeCreateVec (np1 - 1) $ \ mv -> do + -- Explicitly: @mv[0..i-1] := xs[0..i-1]@ + unsafeCopyVec xs 0 mv 0 i + -- Explicitly: @mv[i..n-1] := xs[i+1..n]@ + unsafeCopyVec xs (i+1) mv i (np1 - 1 - i) + +-- | Remove an element at the given position in a vector. +-- O(n) +remove :: Fin (n+1) -> Vec (n+1) a -> Vec n a +remove i = unsafeRemove (finToInt i) + +-------------------------------- + +-- | Sort a 'Vec' according to its 'Ord' instance. +vsort :: Ord a => Vec n a -> Vec n a +vsort xs = listVec (svSize xs) . L.sort . F.toList $ xs + +-- | Sort a 'Vec' with a given comparison function. +vsortBy :: (a -> a -> Ordering) -> Vec n a -> Vec n a +vsortBy f xs = listVec (svSize xs). L.sortBy f . F.toList $ xs + +-- | Sort a 'Vec' with a given sort-key function. +vsortOn :: Ord b => (a -> b) -> Vec n a -> Vec n a +vsortOn f xs = listVec (svSize xs). L.sortOn f . F.toList $ xs + + +-------------------------------- + +-- | Transpose a vector of vectors. +vtranspose :: SInt m -> Vec n (Vec m a) -> Vec m (Vec n a) +vtranspose sm xs = + let !s = vSize xs + !t = unSInt sm + in unsafeMkVec t $ \ i -> + -- s is the size of the outer vector, i.e. valueOf @n + unsafeMkVec s $ \ j -> + unsafeIndexK xs j $ \ v -> unsafeIndexK v i id + +-------------------------------- + +-- | Find the index of the first element, if any, that satisfies a predicate. +vfindIndex :: (a -> Bool) -> Vec n a -> Maybe (Fin n) +vfindIndex p = fmap unsafeFin . L.findIndex p . F.toList + +-------------------------------- + +-- | Create a singleton 'Vec'. +vec1 :: a -> Vec 1 a +vec1 = pure +{-# INLINE vec1 #-} + +-- | Create a 'Vec' from two elements. +vec2 :: a -> a -> Vec 2 a +vec2 x0 x1 = mkVec sintVal $ \i -> case finToInt i of + 0 -> x0 + 1 -> x1 + _ -> error "Impossible: Fin out of range" +{-# INLINE vec2 #-} + +-- | Create a 'Vec' from three elements. +vec3 :: a -> a -> a -> Vec 3 a +vec3 x0 x1 x2 = mkVec sintVal $ \i -> case finToInt i of + 0 -> x0 + 1 -> x1 + 2 -> x2 + _ -> error "Impossible: Fin out of range" +{-# INLINE vec3 #-} + +-- | Create a 'Vec' from four elements. +vec4 :: a -> a -> a -> a -> Vec 4 a +vec4 x0 x1 x2 x3 = mkVec sintVal $ \i -> case finToInt i of + 0 -> x0 + 1 -> x1 + 2 -> x2 + 3 -> x3 + _ -> error "Impossible: Fin out of range" +{-# INLINE vec4 #-} + +-- | Create a 'Vec' from six elements. +vec6 :: a -> a -> a -> a -> a -> a ->Vec 6 a +vec6 x0 x1 x2 x3 x4 x5 = mkVec sintVal $ \i -> case finToInt i of + 0 -> x0 + 1 -> x1 + 2 -> x2 + 3 -> x3 + 4 -> x4 + 5 -> x5 + _ -> error "Impossible: Fin out of range" +{-# INLINE vec6 #-} + +-- | Create a 'Vec' from eight elements. +vec8 :: a -> a -> a -> a -> a -> a -> a -> a -> Vec 8 a +vec8 x0 x1 x2 x3 x4 x5 x6 x7 = mkVec sintVal $ \i -> case finToInt i of + 0 -> x0 + 1 -> x1 + 2 -> x2 + 3 -> x3 + 4 -> x4 + 5 -> x5 + 6 -> x6 + 7 -> x7 + _ -> error "Impossible: Fin out of range" +{-# INLINE vec8 #-} + +--------------------------- + +-- | Get the value of a statically known natural number. +{-# INLINE valueOf #-} +valueOf :: forall (n :: Nat) (i :: Type) . (KnownNat n, Num i) => i +valueOf = fromIntegral $ natVal' (proxy# :: Proxy# n) + +#if !MIN_VERSION_base(4,15,0) +-- base-4.15.0.0 removed naturalToInt. +{-# RULES "integerToInt . naturalToInteger => naturalToInt" + forall a. integerToInt (naturalToInteger a) = + let !(I# i) = naturalToInt a + in i + #-} +#endif + +-- | Modify the given index of a 'Vec'. +overIx :: Fin n -> (a -> a) -> Vec n a -> Vec n a +overIx i f v = runST $ do + mv <- safeThawMV v + indexK v i (writeMV mv i . f) + unsafeFreezeMV mv +{-# INLINE overIx #-} diff --git a/test/Benchmark.hs b/test/Benchmark.hs new file mode 100644 --- /dev/null +++ b/test/Benchmark.hs @@ -0,0 +1,129 @@ +-- Copyright 2018-2021 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 TypeFamilies #-} + +module Main where + +import Prelude hiding ((++)) + +import Control.Applicative (liftA2, liftA3) +import Control.Exception (evaluate) +import Data.Foldable (foldl', foldr', toList) +import Data.Semigroup (Sum(..)) + +import Control.DeepSeq (force) +import Data.Vec.Short +import Data.Fin.Int (fin, finToInt, modNegate) +import Data.Functor.WithIndex (FunctorWithIndex(..)) + +import qualified Gauge as G + +theAnswer :: Vec 64 Int +theAnswer = pure 42 + +benchForce_DeepSeq + , benchPure, benchRot + , benchMap, benchMapId, benchMapCoerce, benchMap2 + , benchSum, benchSumMap + , benchFoldr, benchFoldr1, benchFoldr', benchFoldl, benchFoldl1, benchFoldl' + , benchAppend, benchAppendMap + , benchTraverseIO + , benchLiftA2 + , benchConvolve + , benchMapWithPos + , benchMapToList, benchMapToListMap, benchSumToListMap + , benchEq_true, benchEq_false + , benchTake, benchMapTake, benchDrop, benchMapDrop + , benchSplit + :: G.Benchmark + +benchForce_DeepSeq = G.bench "force_DeepSeq" $ G.whnf force theAnswer +benchPure = G.bench "pure" $ G.whnf (pure :: Int -> Vec 64 Int) 42 +benchMap = G.bench "map" $ G.whnf (map' (+2)) theAnswer +benchMapId = G.bench "mapId" $ G.whnf (fmap id) theAnswer +benchMapCoerce = G.bench "mapCoerce" $ G.whnf (fmap Sum) theAnswer +benchMap2 = G.bench "map2" $ G.whnf (map' (+2) . fmap (+2)) theAnswer +benchRot = G.bench "rot" $ G.whnf (rot (fin 7)) theAnswer +benchSum = G.bench "sum" $ G.whnf sum theAnswer +benchSumMap = G.bench "sumMap" $ G.whnf (sum . fmap (+7)) theAnswer +benchFoldr = G.bench "foldr" $ G.whnf (foldr (+) 0) theAnswer +benchFoldr1 = G.bench "foldr1" $ G.whnf (foldr1 (+)) theAnswer +benchFoldr' = G.bench "foldr'" $ G.whnf (foldr' (+) 0) theAnswer +benchFoldl = G.bench "foldl" $ G.whnf (foldl (+) 0) theAnswer +benchFoldl1 = G.bench "foldl1" $ G.whnf (foldl1 (+)) theAnswer +benchFoldl' = G.bench "foldl'" $ G.whnf (foldl' (+) 0) theAnswer +benchAppend = G.bench "append" $ G.whnf (\x -> x ++ x) theAnswer +benchAppendMap = G.bench "appendMap" $ G.whnf (\x -> x ++ fmap (+2) x) theAnswer + +benchTraverseIO = G.bench "traverseIO" $ + G.whnfAppIO (traverse evaluate) theAnswer + +benchLiftA2 = G.bench "liftA2" $ + G.whnf (liftA2 (+) (pure 2 :: Vec 64 Int)) theAnswer + +benchConvolve = G.bench "convolve" $ G.nf + (\x -> liftA3 + (\a b c -> a + b + c) + (rot (fin 1) x) + x + (rot (modNegate (fin 1)) x)) + theAnswer + +benchMapWithPos = G.bench "mapWithPos'" $ + G.whnf (imap (\i -> (+) (finToInt i))) theAnswer + +benchMapToList = G.bench "mapToList" $ G.nf (map (+2) . toList) theAnswer +benchMapToListMap = G.bench "mapToListMap" $ + G.nf (map (+2) . toList . fmap (+2)) theAnswer + +benchSumToListMap = G.bench "sumToListMap" $ + G.nf (sum . toList . fmap (+2)) theAnswer + +benchEq_true = G.bench "eq_true" $ G.whnf (\x -> x == x) theAnswer + +benchEq_false = G.bench "eq_false" $ + G.whnf (\ (x, y) -> x == y) (theAnswer, (+1) <$> theAnswer) + +benchTake = G.bench "take" $ G.whnf (fst . split @50) theAnswer +benchMapTake = G.bench "mapTake" $ + G.whnf (fst . split @50 . fmap (+2)) theAnswer + +benchDrop = G.bench "drop" $ G.whnf (snd . split @50) theAnswer +benchMapDrop = G.bench "mapDrop" $ + G.whnf (snd . split @50 . fmap (+2)) theAnswer + +benchSplit = G.bench "split" $ + G.whnf (\x -> let (l, r) = split @50 x in l `seq` r `seq` (l, r)) theAnswer + +main :: IO () +main = G.defaultMainWith (G.defaultConfig { G.minSamples = Just 5 }) $ + [ benchForce_DeepSeq + , benchPure + , benchMap, benchMap2, benchMapCoerce, benchMapId + , benchTraverseIO + , benchLiftA2 + , benchConvolve + , benchMapWithPos + , benchRot + , benchSum, benchSumMap + , benchFoldr, benchFoldr1, benchFoldr', benchFoldl, benchFoldl1, benchFoldl' + , benchAppend, benchAppendMap + , benchMapToList, benchMapToListMap, benchSumToListMap + , benchEq_true, benchEq_false + , benchTake, benchMapTake, benchDrop, benchMapDrop + , benchSplit + ] diff --git a/test/Tests.hs b/test/Tests.hs new file mode 100644 --- /dev/null +++ b/test/Tests.hs @@ -0,0 +1,43 @@ +-- Copyright 2018-2021 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 Main where + +import Prelude hiding ((++)) + +import Test.Framework (defaultMain) +import Test.Framework.Providers.QuickCheck2 (testProperty) +import Test.QuickCheck (Property, (===), counterexample) + +import Data.Vec.Short ((++), Vec, split, vec2, vec3) + +main :: IO () +main = defaultMain + [ testProperty "vec" $ readInvertsShow $ + vec2 (3 :: Int) 5 + , testProperty "vec of vec" $ readInvertsShow $ + vec3 (vec2 (2 :: Int) 4) (vec2 3 5) (vec2 4 6) + , testProperty "append of split" $ \ (v :: Vec 8 Int) -> + let (l, r) = split @5 v + in v === l ++ r + ] + where + readInvertsShow :: (Eq a, Read a, Show a) => a -> Property + readInvertsShow a = + let astr = show a + in counterexample astr $ a === read astr