diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,9 @@ +Changelog +========= + +Version 0.1.0.0 +--------------- + + + +* Initial release diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,30 @@ +Copyright Justin Le (c) 2018 + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Justin Le nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,100 @@ +# nonempty-containers + +Efficient and optimized non-empty (by construction) versions of types from +*[containers][]*. Inspired by *[non-empty-containers][]* library, except +attempting a more faithful port (with under-the-hood optimizations) of the full +*containers* API. Also contains a convenient typeclass abstraction for +converting betwewen non-empty and possibly-empty variants, as well as pattern +synonym-based conversion methods. + +[containers]: http://hackage.haskell.org/package/containers +[non-empty-containers]: http://hackage.haskell.org/package/non-empty-containers + +Non-empty *by construction* means that the data type is implemented using a +data structure where it is structurally impossible to represent an empty +collection. + +Unlike similar packages (see below), this package is defined to be a +*drop-in replacement* for the *containers* API in most situations. More or +less every single function is implemented with the same asymptotics and +typeclass constraints. An extensive test suite (with 457 total tests) is +provided to ensure that the behavior of functions are identical to their +original *containers* counterparts. + +Care is also taken to modify the interface of specific functions to reflect +non-emptiness and emptiness as concepts, including: + +1. Functions that might return empty results (like `delete`, `filter`) return + possibly-empty variants instead. + +2. Functions that totally partition a non-empty collection (like `partition`, + `splitAt`, `span`) would previously return a tuple of either halves: + + ```haskell + mapEither :: (a -> Either b c) -> Map k a -> (Map k b, Map k c) + ``` + + The final result is always a total partition (every item in the original + map is represented in the result), so, to reflect this, [`These`][these] is + returned instead: + + ```haskell + data These a b = This a + | That b + | These a b + + mapEither :: (a -> Either b c) -> NEMap k a -> These (NEMap k c) (NEMap k c) + ``` + + This preserves the invariance of non-emptiness: either we have a non-empty + map in the first camp (containing all original values), a non-empty map in + the second camp (containing all original values), or a split between two + non-empty maps in either camp. + + [these]: https://hackage.haskell.org/package/these + +3. Typeclass-polymorphic functions are made more general (or have more general + variants provided) whenever possible. This means that functions like + `foldMapWithKey` are written for all `Semigroup m` instead of only `Monoid + m`, and `traverseWithKey1` is provided to work for all `Apply f` instances + (instead of only `Applicative f` instances). + + `Foldable1` and `Traversable1` instances are also provided, to provide + `foldMap1` and `traverse1`. + +4. Functions that can "potentially delete" (like `alter` and `updateAt`) + return possibly-empty variants. However, alternatives are offered + (whenever not already present) with variants that disallow deletion, + allowing for guaranteed non-empty maps to be returned. + +Contains non-empty versions for: + +* `Map` +* `IntMap` +* `Set` +* `IntSet` +* `Sequence` + +A typeclass abstraction (in *Data.Containers.NonEmpty*) is provided to allow +for easy conversions between non-empty and possibly-empty variants. Note that +`Tree`, from *Data.Tree*, is already non-empty by construction. + +Similar packages include: + +* [non-empty-containers][]: Similar approach with similar data types, but API + is limited to a few choice functions. +* [non-empty][]: Similar approach with similar data types, but is meant to be + more general and work for a variety of more data types. +* [nonempty-alternative][]: Similar approach, but is instead a generalized + data type for all `Alternative` instances. + +[non-empty]: http://hackage.haskell.org/package/non-empty +[nonempty-alternative]: http://hackage.haskell.org/package/nonempty-alternative + +Currently not implemented: + +* Extended merging functions. However, there aren't too many benefits to be + gained from lifting extended merging functions, because their + emptiness/non-emptiness guarantees are difficult to statically conclude. +* Strict variants of Map functions. This is something that I wouldn't mind, + and might add in the future. PR's are welcomed! 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/nonempty-containers.cabal b/nonempty-containers.cabal new file mode 100644 --- /dev/null +++ b/nonempty-containers.cabal @@ -0,0 +1,89 @@ +cabal-version: 1.12 + +-- This file has been generated from package.yaml by hpack version 0.31.0. +-- +-- see: https://github.com/sol/hpack +-- +-- hash: bc9e6411940584bc108b18b524d5f1ef9a293724bd59219e1b637fc2d2566ea3 + +name: nonempty-containers +version: 0.1.0.0 +synopsis: Non-empty variants of containers data types, with full API +description: Efficient and optimized non-empty versions of types from /containers/. + Inspired by /non-empty-containers/ library, except attempting a more + faithful port (with under-the-hood optimizations) of the full /containers/ + API. Also contains a convenient typeclass abstraction for converting + betwewen non-empty and possibly-empty variants. See README.md for more + information. +category: Data Structures +homepage: https://github.com/mstksg/nonempty-containers#readme +bug-reports: https://github.com/mstksg/nonempty-containers/issues +author: Justin Le +maintainer: justin@jle.im +copyright: (c) Justin Le 2018 +license: BSD3 +license-file: LICENSE +tested-with: GHC >= 8.2 && < 8.8 +build-type: Simple +extra-source-files: + README.md + CHANGELOG.md + +source-repository head + type: git + location: https://github.com/mstksg/nonempty-containers + +library + exposed-modules: + Data.Containers.NonEmpty + Data.IntMap.NonEmpty + Data.IntMap.NonEmpty.Internal + Data.IntSet.NonEmpty + Data.IntSet.NonEmpty.Internal + Data.Map.NonEmpty + Data.Map.NonEmpty.Internal + Data.Sequence.NonEmpty + Data.Sequence.NonEmpty.Internal + Data.Set.NonEmpty + Data.Set.NonEmpty.Internal + other-modules: + Paths_nonempty_containers + hs-source-dirs: + src + ghc-options: -Wall -Wcompat -Werror=incomplete-patterns -Wredundant-constraints + build-depends: + base >=4.9 && <5 + , comonad + , containers >=0.5.9 + , deepseq + , semigroupoids + , these + default-language: Haskell2010 + +test-suite nonempty-containers-test + type: exitcode-stdio-1.0 + main-is: Spec.hs + other-modules: + Tests.IntMap + Tests.IntSet + Tests.Map + Tests.Sequence + Tests.Set + Tests.Util + Paths_nonempty_containers + hs-source-dirs: + test + ghc-options: -Wall -Wcompat -Werror=incomplete-patterns -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N + build-depends: + base >=4.9 && <5 + , comonad + , containers >=0.5.9 + , hedgehog + , hedgehog-fn + , nonempty-containers + , semigroupoids + , tasty + , tasty-hedgehog + , text + , these + default-language: Haskell2010 diff --git a/src/Data/Containers/NonEmpty.hs b/src/Data/Containers/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/Containers/NonEmpty.hs @@ -0,0 +1,217 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeFamilyDependencies #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.Containers.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Typeclass +-- +-- Provides the typeclass 'HasNonEmpty', which abstracts over different +-- types which have a "non-empty" variant. +-- +-- Used to convert between and in between possibly-empty and non-empty +-- types. Instances are provided for all modules in this package, as well +-- as for 'NonEmpty' in /base/. +module Data.Containers.NonEmpty ( + HasNonEmpty(..) + , pattern IsNonEmpty, pattern IsEmpty + ) where + +import Data.IntMap (IntMap) +import Data.IntMap.NonEmpty (NEIntMap) +import Data.IntSet (IntSet) +import Data.IntSet.NonEmpty (NEIntSet) +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Map (Map) +import Data.Map.NonEmpty (NEMap) +import Data.Maybe +import Data.Sequence (Seq(..)) +import Data.Sequence.NonEmpty (NESeq(..)) +import Data.Set (Set) +import Data.Set.NonEmpty (NESet) +import qualified Data.IntMap as IM +import qualified Data.IntMap.NonEmpty as NEIM +import qualified Data.IntSet as IS +import qualified Data.IntSet.NonEmpty as NEIS +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import qualified Data.Map.NonEmpty as NEM +import qualified Data.Sequence as Seq +import qualified Data.Sequence.NonEmpty as NESeq +import qualified Data.Set as S +import qualified Data.Set.NonEmpty as NES + +-- | If @s@ is an instance of @HasNonEmpty@, it means that there is +-- a corresponding "non-empty" version of @s@, @'NE' s@. +-- +-- In order for things to be well-behaved, we expect that 'nonEmpty' and +-- @maybe 'empty' 'fromNonEmpty'@ should form an isomorphism (or that +-- @'withNonEmpty' 'empty' 'fromNonEmpty' == id@. In addition, +-- the following properties should hold for most exectations: +-- +-- * @(x == empty) ==> isEmpty x@ +-- * @(x == empty) ==> isNothing (nonEmpty x)@ +-- * @isEmpty x ==> isNothing (nonEmpty x)@ +-- * @unsafeToNonEmpty x == fromJust (nonEmpty x)@ +-- * Usually, @not (isEmpty x) ==> isJust (nonEmpty x)@, but this isn't +-- necessary. +class HasNonEmpty s where + {-# MINIMAL (nonEmpty | withNonEmpty), fromNonEmpty, empty #-} + + -- | @'NE' s@ is the "non-empty" version of @s@. + type NE s = t | t -> s + + -- | "Smart constructor" for @'NE' s@ given a (potentailly empty) @s@. + -- Will return 'Nothing' if the @s@ was empty, and @'Just' n@ if the + -- @s@ was not empty, with @n :: 'NE' s@. + -- + -- Should form an isomorphism with @'maybe' 'empty' 'fromNonEmpty'@. + nonEmpty :: s -> Maybe (NE s) + nonEmpty = withNonEmpty Nothing Just + + -- | Convert a @'NE' s@ (non-empty @s@) back into an @s@, "obscuring" + -- its non-emptiness from its type. + fromNonEmpty :: NE s -> s + + -- | Continuation-based version of 'nonEmpty', which can be more + -- efficient in certain situations. + -- + -- @'withNonEmpty' 'empty' 'fromNonEmpty'@ should be @id@. + withNonEmpty :: r -> (NE s -> r) -> s -> r + withNonEmpty def f = maybe def f . nonEmpty + + -- | An empty @s@. + empty :: s + + -- | Check if an @s@ is empty. + isEmpty :: s -> Bool + isEmpty = isNothing . nonEmpty + + -- | Unsafely coerce an @s@ into an @'NE' s@ (non-empty @s@). Is + -- undefined (throws a runtime exception when evaluation is attempted) + -- when the @s@ is empty. + unsafeToNonEmpty :: s -> NE s + unsafeToNonEmpty = fromMaybe e . nonEmpty + where + e = errorWithoutStackTrace "unsafeToNonEmpty: empty input provided" + +instance HasNonEmpty [a] where + type NE [a] = NonEmpty a + nonEmpty = NE.nonEmpty + fromNonEmpty = NE.toList + withNonEmpty def f = \case + [] -> def + x:xs -> f (x :| xs) + empty = [] + isEmpty = null + unsafeToNonEmpty = NE.fromList + +instance HasNonEmpty (Map k a) where + type NE (Map k a) = NEMap k a + nonEmpty = NEM.nonEmptyMap + fromNonEmpty = NEM.toMap + withNonEmpty = NEM.withNonEmpty + empty = M.empty + isEmpty = M.null + unsafeToNonEmpty = NEM.unsafeFromMap + +instance HasNonEmpty (IntMap a) where + type NE (IntMap a) = NEIntMap a + nonEmpty = NEIM.nonEmptyMap + fromNonEmpty = NEIM.toMap + withNonEmpty = NEIM.withNonEmpty + empty = IM.empty + isEmpty = IM.null + unsafeToNonEmpty = NEIM.unsafeFromMap + +instance HasNonEmpty (Set a) where + type NE (Set a) = NESet a + nonEmpty = NES.nonEmptySet + fromNonEmpty = NES.toSet + withNonEmpty = NES.withNonEmpty + empty = S.empty + isEmpty = S.null + unsafeToNonEmpty = NES.unsafeFromSet + +instance HasNonEmpty IntSet where + type NE IntSet = NEIntSet + nonEmpty = NEIS.nonEmptySet + fromNonEmpty = NEIS.toSet + withNonEmpty = NEIS.withNonEmpty + empty = IS.empty + isEmpty = IS.null + unsafeToNonEmpty = NEIS.unsafeFromSet + +instance HasNonEmpty (Seq a) where + type NE (Seq a) = NESeq a + nonEmpty = NESeq.nonEmptySeq + fromNonEmpty = NESeq.toSeq + withNonEmpty = NESeq.withNonEmpty + empty = Seq.empty + isEmpty = Seq.null + unsafeToNonEmpty = NESeq.unsafeFromSeq + + +-- | The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat a @s@ as +-- if it were either a @'IsNonEmpty' n@ (where @n@ is a non-empty version +-- of @s@, type @'NE' s@) or an 'IsEmpty'. +-- +-- For example, you can pattern match on a list to get a 'NonEmpty' +-- (non-empty list): +-- +-- @ +-- safeHead :: [Int] -> Int +-- safeHead ('IsNonEmpty' (x :| _)) = x -- here, the list was not empty +-- safehead 'IsEmpty' = 0 -- here, the list was empty +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original input was /not/ +-- empty, and you have a verified-non-empty @n :: 'NE' s@ to use. +-- +-- Note that because of the way coverage checking works for polymorphic +-- pattern synonyms, you will unfortunatelly still get incomplete pattern +-- match warnings if you match on both 'IsNonEmpty' and 'NonEmpty', even +-- though the two are meant to provide complete coverage. However, many +-- instances of 'HasNonEmpty' (like 'NEMap', 'NEIntMap', 'NESet', +-- 'NEIntSet') will provide their own monomorphic versions of these +-- patterns that can be verified as complete covers by GHC. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a @'NE' s@ back into an @s@, "obscuring" its non-emptiness (see +-- 'fromNonEmpty'). +pattern IsNonEmpty :: HasNonEmpty s => NE s -> s +pattern IsNonEmpty n <- (nonEmpty->Just n) + where + IsNonEmpty n = fromNonEmpty n + +-- | The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat a @s@ as +-- if it were either a @'IsNonEmpty' n@ (where @n@ is a non-empty version +-- of @s@, type @'NE' s@) or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original item was empty. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'empty'. +-- +-- Note that because of the way coverage checking works for polymorphic +-- pattern synonyms, you will unfortunatelly still get incomplete pattern +-- match warnings if you match on both 'IsNonEmpty' and 'NonEmpty', even +-- though the two are meant to provide complete coverage. However, many +-- instances of 'HasNonEmpty' (like 'NEMap', 'NEIntMap', 'NESet', +-- 'NEIntSet') will provide their own monomorphic versions of these +-- patterns that can be verified as complete covers by GHC. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: HasNonEmpty s => s +pattern IsEmpty <- (isEmpty->True) + where + IsEmpty = empty diff --git a/src/Data/IntMap/NonEmpty.hs b/src/Data/IntMap/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/IntMap/NonEmpty.hs @@ -0,0 +1,1971 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.IntMap.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Finite Integer-Indexed Maps (lazy interface) +-- +-- The @'NEIntMap' v@ type represents a non-empty finite map (sometimes +-- called a dictionary) from integer keys to values of type @v@. +-- An 'NEIntMap' is strict in its keys but lazy in its values. +-- +-- See documentation for 'NEIntMap' for information on how to convert and +-- manipulate such non-empty maps. +-- +-- This module essentially re-imports the API of "Data.IntMap.Lazy" and its +-- 'IntMap' type, along with semantics and asymptotics. In most +-- situations, asymptotics are different only by a constant factor. In +-- some situations, asmyptotics are even better (constant-time instead of +-- log-time). +-- +-- Because 'NEIntMap' is implemented using 'IntMap', all of the caveats of using +-- 'IntMap' apply (such as the limitation of the maximum size of maps). +-- +-- All functions take non-empty maps as inputs. In situations where their +-- results can be guarunteed to also be non-empty, they also return +-- non-empty maps. In situations where their results could potentially be +-- empty, 'IntMap' is returned instead. +-- +-- Some variants of functions (like 'alter'', 'alterF'', 'adjustMin', +-- 'adjustMax', 'adjustMinWithKey', 'adjustMaxWithKey') are provided in +-- a way restructured to preserve guaruntees of non-empty maps being +-- returned. +-- +-- Some functions (like 'mapEither', 'partition', 'split') +-- have modified return types to account for possible configurations of +-- non-emptiness. +-- +-- This module is intended to be imported qualified, to avoid name clashes with +-- "Prelude" and "Data.IntMap" functions: +-- +-- > import qualified Data.IntMap.NonEmpty as NEIM +-- +-- Note that all asmyptotics /O(f(n))/ in this module are actually +-- /O(min(W, f(n)))/, where @W@ is the number of bits in an 'Int' (32 or +-- 64). That is, if @f(n)@ is greater than @W@, all operations are +-- constant-time. +-- +-- At the moment, this package does not provide a variant strict on values +-- for these functions, like /containers/ does. This is a planned future +-- implementation (PR's are appreciated). For now, you can simulate +-- a strict interface by manually forcing values before returning results. +module Data.IntMap.NonEmpty ( + -- * Non-Empty IntMap Type + NEIntMap + , Key + + -- ** Conversions between empty and non-empty maps + , pattern IsNonEmpty + , pattern IsEmpty + , nonEmptyMap + , toMap + , withNonEmpty + , insertMap + , insertMapWith + , insertMapWithKey + , insertMapMin + , insertMapMax + , unsafeFromMap + + -- * Construction + , singleton + , fromSet + + -- ** From Unordered Lists + , fromList + , fromListWith + , fromListWithKey + + -- ** From Ascending Lists + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- * Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- * Deletion\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + , alter' + , alterF' + + -- * Query + -- ** Lookup + , lookup + , (!?) + , (!) + , findWithDefault + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + + -- ** Size + , size + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , (\\) + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- -- ** Universal combining function + -- , mergeWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey1 + , traverseWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldr1 + , foldl1 + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldr1' + , foldl' + , foldl1' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + + -- ** Lists + , toList + + -- ** Ordered lists + , toAscList + , toDescList + + -- * Filter + , filter + , filterWithKey + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , adjustMin + , adjustMax + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , minView + , maxView + + -- * Debugging + , valid + ) where + +import Control.Applicative +import Data.Bifunctor +import Data.Functor.Identity +import Data.IntMap.Internal (IntMap(..), Key) +import Data.IntMap.NonEmpty.Internal +import Data.IntSet (IntSet) +import Data.IntSet.NonEmpty.Internal (NEIntSet(..)) +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Maybe hiding (mapMaybe) +import Data.Semigroup.Foldable (Foldable1) +import Data.These +import Prelude hiding (map, filter, lookup, foldl, foldr, foldl1, foldr1) +import qualified Data.Foldable as F +import qualified Data.IntMap as M +import qualified Data.IntSet as S +import qualified Data.List.NonEmpty as NE +import qualified Data.Maybe as Maybe +import qualified Data.Semigroup.Foldable as F1 + +-- | /O(1)/ match, /O(log n)/ usage of contents. The 'IsNonEmpty' and +-- 'IsEmpty' patterns allow you to treat a 'IntMap' as if it were either +-- a @'IsNonEmpty' n@ (where @n@ is a 'NEIntMap') or an 'IsEmpty'. +-- +-- For example, you can pattern match on a 'IntMap': +-- +-- @ +-- myFunc :: 'IntMap' K X -> Y +-- myFunc ('IsNonEmpty' n) = -- here, the user provided a non-empty map, and @n@ is the 'NEIntMap' +-- myFunc 'IsEmpty' = -- here, the user provided an empty map. +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original 'IntMap' was /not/ +-- empty, and you have a verified-non-empty 'NEIntMap' @n@ to use. +-- +-- Note that patching on this pattern is /O(1)/. However, using the +-- contents requires a /O(log n)/ cost that is deferred until after the +-- pattern is matched on (and is not incurred at all if the contents are +-- never used). +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a 'NEIntMap' back into a 'IntMap', obscuring its non-emptiness (see 'toMap'). +pattern IsNonEmpty :: NEIntMap a -> IntMap a +pattern IsNonEmpty n <- (nonEmptyMap->Just n) + where + IsNonEmpty n = toMap n + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'IntMap' as if it were either a @'IsNonEmpty' n@ (where @n@ is +-- a 'NEIntMap') or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original 'IntMap' was empty. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'Data.IntMap.empty'. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: IntMap a +pattern IsEmpty <- (M.null->True) + where + IsEmpty = M.empty + +{-# COMPLETE IsNonEmpty, IsEmpty #-} + +-- | /O(log n)/. Unsafe version of 'nonEmptyMap'. Coerces a 'IntMap' into an +-- 'NEIntMap', but is undefined (throws a runtime exception when evaluation is +-- attempted) for an empty 'IntMap'. +unsafeFromMap + :: IntMap a + -> NEIntMap a +unsafeFromMap = withNonEmpty e id + where + e = errorWithoutStackTrace "NEIntMap.unsafeFromMap: empty map" +{-# INLINE unsafeFromMap #-} + +-- | /O(log n)/. Convert a 'IntMap' into an 'NEIntMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. If key is already present, +-- will overwrite the original value. +-- +-- See 'insertMapMin' for a version that is constant-time if the new key is +-- /strictly smaller than/ all keys in the original map. +-- +-- > insertMap 4 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(4,"c"), (5,"a")]) +-- > insertMap 4 "c" Data.IntMap.empty == singleton 4 "c" +insertMap :: Key -> a -> IntMap a -> NEIntMap a +insertMap k v = withNonEmpty (singleton k v) (insert k v) +{-# INLINE insertMap #-} + +-- | /O(log n)/. Convert a 'IntMap' into an 'NEIntMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. Uses a combining function +-- with the new value as the first argument if the key is already present. +-- +-- > insertMapWith (++) 4 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(4,"c"), (5,"a")]) +-- > insertMapWith (++) 5 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(5,"ca")]) +insertMapWith + :: (a -> a -> a) + -> Key + -> a + -> IntMap a + -> NEIntMap a +insertMapWith f k v = withNonEmpty (singleton k v) (insertWith f k v) +{-# INLINE insertMapWith #-} + +-- | /O(log n)/. Convert a 'IntMap' into an 'NEIntMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. Uses a combining function +-- with the key and new value as the first and second arguments if the key +-- is already present. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3, "b") :| [(5, "5:xxx|a")]) +-- > insertWithKey f 7 "xxx" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +-- > insertWithKey f 5 "xxx" Data.IntMap.empty == singleton 5 "xxx" +insertMapWithKey + :: (Key -> a -> a -> a) + -> Key + -> a + -> IntMap a + -> NEIntMap a +insertMapWithKey f k v = withNonEmpty (singleton k v) (insertWithKey f k v) +{-# INLINE insertMapWithKey #-} + +-- | /O(1)/ Convert a 'IntMap' into an 'NEIntMap' by adding a key-value pair +-- where the key is /strictly less than/ all keys in the input map. The +-- keys in the original map must all be /strictly greater than/ the new +-- key. /The precondition is not checked./ +-- +-- > insertMapMin 2 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((2,"c") :| [(3,"b"), (5,"a")]) +-- > valid (insertMapMin 2 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == True +-- > valid (insertMapMin 7 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == False +-- > valid (insertMapMin 3 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == False +insertMapMin + :: Key + -> a + -> IntMap a + -> NEIntMap a +insertMapMin = NEIntMap +{-# INLINE insertMapMin #-} + +-- | /O(log n)/ Convert a 'IntMap' into an 'NEIntMap' by adding a key-value pair +-- where the key is /strictly greater than/ all keys in the input map. The +-- keys in the original map must all be /strictly less than/ the new +-- key. /The precondition is not checked./ +-- +-- At the current moment, this is identical simply 'insertMap'; however, +-- it is left both for consistency and as a placeholder for a future +-- version where optimizations are implemented to allow for a faster +-- implementation. +-- +-- > insertMap 7 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(5,"a"), (7,"c")]) + +-- these currently are all valid, but shouldn't be +-- > valid (insertMap 7 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == True +-- > valid (insertMap 2 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == False +-- > valid (insertMap 5 "c" (Data.IntMap.fromList [(5,"a"), (3,"b")])) == False +insertMapMax + :: Key + -> a + -> IntMap a + -> NEIntMap a +insertMapMax k v = withNonEmpty (singleton k v) go + where + go (NEIntMap k0 v0 m0) = NEIntMap k0 v0 . insertMaxMap k v $ m0 +{-# INLINE insertMapMax #-} + +-- | /O(n)/. Build a non-empty map from a non-empty set of keys and +-- a function which for each key computes its value. +-- +-- > fromSet (\k -> replicate k 'a') (Data.Set.NonEmpty.fromList (3 :| [5])) == fromList ((5,"aaaaa") :| [(3,"aaa")]) +fromSet + :: (Key -> a) + -> NEIntSet + -> NEIntMap a +fromSet f (NEIntSet k ks) = NEIntMap k (f k) (M.fromSet f ks) +{-# INLINE fromSet #-} + +-- | /O(n*log n)/. Build a map from a non-empty list of key\/value pairs +-- with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) ((5,"a") :| [(5,"b"), (3,"b"), (3,"a"), (5,"a")]) == fromList ((3, "ab") :| [(5, "aba")]) +fromListWith + :: (a -> a -> a) + -> NonEmpty (Key, a) + -> NEIntMap a +fromListWith f = fromListWithKey (const f) +{-# INLINE fromListWith #-} + +-- | /O(n*log n)/. Build a map from a non-empty list of key\/value pairs +-- with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f ((5,"a") :| [(5,"b"), (3,"b"), (3,"a"), (5,"a")]) == fromList ((3, "3ab") :| [(5, "5a5ba")]) +fromListWithKey + :: (Key -> a -> a -> a) + -> NonEmpty (Key, a) + -> NEIntMap a +fromListWithKey f ((k0, v0) :| xs) = F.foldl' go (singleton k0 v0) xs + where + go m (k, v) = insertWithKey f k v m + {-# INLINE go #-} +{-# INLINE fromListWithKey #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList ((3,"b") :| [(5,"a")]) == fromList ((3, "b") :| [(5, "a")]) +-- > fromAscList ((3,"b") :| [(5,"a"), (5,"b")]) == fromList ((3, "b") :| [(5, "b")]) +-- > valid (fromAscList ((3,"b") :| [(5,"a"), (5,"b")])) == True +-- > valid (fromAscList ((5,"a") :| [(3,"b"), (5,"b")])) == False +fromAscList + :: NonEmpty (Key, a) + -> NEIntMap a +fromAscList = fromDistinctAscList . combineEq +{-# INLINE fromAscList #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is ascending) is not checked./ +-- +-- > fromAscListWith (++) ((3,"b") :| [(5,"a"), (5,"b")]) == fromList ((3, "b") :| [(5, "ba")]) +-- > valid (fromAscListWith (++) ((3,"b") :| [(5,"a"), (5,"b"))]) == True +-- > valid (fromAscListWith (++) ((5,"a") :| [(3,"b"), (5,"b"))]) == False +fromAscListWith + :: (a -> a -> a) + -> NonEmpty (Key, a) + -> NEIntMap a +fromAscListWith f = fromAscListWithKey (const f) +{-# INLINE fromAscListWith #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f ((3,"b") :| [(5,"a"), (5,"b"), (5,"b")]) == fromList ((3, "b") :| [(5, "5:b5:ba")]) +-- > valid (fromAscListWithKey f ((3,"b") :| [(5,"a"), (5,"b"), (5,"b")])) == True +-- > valid (fromAscListWithKey f ((5,"a") :| [(3,"b"), (5,"b"), (5,"b")])) == False +fromAscListWithKey + :: (Key -> a -> a -> a) + -> NonEmpty (Key, a) + -> NEIntMap a +fromAscListWithKey f = fromDistinctAscList . combineEqWith f +{-# INLINE fromAscListWithKey #-} + +-- | /O(n)/. Build a map from an ascending non-empty list of distinct +-- elements in linear time. /The precondition is not checked./ +-- +-- > fromDistinctAscList ((3,"b") :| [(5,"a")]) == fromList ((3, "b") :| [(5, "a")]) +-- > valid (fromDistinctAscList ((3,"b") :| [(5,"a")])) == True +-- > valid (fromDistinctAscList ((3,"b") :| [(5,"a"), (5,"b")])) == False +fromDistinctAscList :: NonEmpty (Key, a) -> NEIntMap a +fromDistinctAscList ((k, v) :| xs) = insertMapMin k v + . M.fromDistinctAscList + $ xs +{-# INLINE fromDistinctAscList #-} + +-- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- See 'insertMap' for a version where the first argument is a 'IntMap'. +-- +-- > insert 5 'x' (fromList ((5,'a') :| [(3,'b')])) == fromList ((3, 'b') :| [(5, 'x')]) +-- > insert 7 'x' (fromList ((5,'a') :| [(3,'b')])) == fromList ((3, 'b') :| [(5, 'a'), (7, 'x')]) +insert + :: Key + -> a + -> NEIntMap a + -> NEIntMap a +insert k v n@(NEIntMap k0 v0 m) = case compare k k0 of + LT -> NEIntMap k v . toMap $ n + EQ -> NEIntMap k v m + GT -> NEIntMap k0 v0 . M.insert k v $ m +{-# INLINE insert #-} + +-- | /O(log n)/. Insert with a function, combining key, new value and old +-- value. @'insertWithKey' f key value mp@ will insert the pair (key, +-- value) into @mp@ if key does not exist in the map. If the key does +-- exist, the function will insert the pair @(key,f key new_value +-- old_value)@. Note that the key passed to f is the same key passed to +-- 'insertWithKey'. +-- +-- See 'insertMapWithKey' for a version where the first argument is a 'IntMap'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "5:xxx|a")]) +-- > insertWithKey f 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +insertWithKey + :: (Key -> a -> a -> a) + -> Key + -> a + -> NEIntMap a + -> NEIntMap a +insertWithKey f k v n@(NEIntMap k0 v0 m) = case compare k k0 of + LT -> NEIntMap k v . toMap $ n + EQ -> NEIntMap k (f k v v0) m + GT -> NEIntMap k0 v0 $ M.insertWithKey f k v m +{-# INLINE insertWithKey #-} + +-- | /O(log n)/. Combines insert operation with old value retrieval. The +-- expression (@'insertLookupWithKey' f k x map@) is a pair where the first +-- element is equal to (@'lookup' k map@) and the second element equal to +-- (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == (Just "a", fromList ((3, "b") :| [(5, "5:xxx|a")])) +-- > insertLookupWithKey f 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == (Nothing, fromList ((3, "b") :| [(5, "a"), (7, "xxx")])) +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList ((5,"a") :| [(3,"b")])) == (Just "a", fromList ((3, "b") :| [(5, "x")])) +-- > insertLookup 7 "x" (fromList ((5,"a") :| [(3,"b")])) == (Nothing, fromList ((3, "b") :| [(5, "a"), (7, "x")])) +insertLookupWithKey + :: (Key -> a -> a -> a) + -> Key + -> a + -> NEIntMap a + -> (Maybe a, NEIntMap a) +insertLookupWithKey f k v n@(NEIntMap k0 v0 m) = case compare k k0 of + LT -> (Nothing, NEIntMap k v . toMap $ n ) + EQ -> (Just v , NEIntMap k (f k v v0) m ) + GT -> NEIntMap k0 v0 <$> M.insertLookupWithKey f k v m +{-# INLINE insertLookupWithKey #-} + +-- | /O(log n)/. Delete a key and its value from the non-empty map. +-- A potentially empty map ('IntMap') is returned, since this might delete the +-- last item in the 'NEIntMap'. When the key is not a member of the map, is +-- equivalent to 'toMap'. +-- +-- > delete 5 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 3 "b" +-- > delete 7 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.Singleton [(3, "b"), (5, "a")] +delete :: Key -> NEIntMap a -> IntMap a +delete k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> toMap n + EQ -> m + GT -> insertMinMap k0 v . M.delete k $ m +{-# INLINE delete #-} + +-- | /O(log n)/. Update a value at a specific key with the result of the +-- provided function. When the key is not a member of the map, the original +-- map is returned. +-- +-- > adjust ("new " ++) 5 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "new a")]) +-- > adjust ("new " ++) 7 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a")]) +adjust + :: (a -> a) + -> Key + -> NEIntMap a + -> NEIntMap a +adjust f = adjustWithKey (const f) +{-# INLINE adjust #-} + +-- | /O(log n)/. Adjust a value at a specific key. When the key is not +-- a member of the map, the original map is returned. +-- +-- > let f key x = (show key) ++ ":new " ++ x +-- > adjustWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "5:new a")]) +-- > adjustWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a")]) +adjustWithKey + :: (Key -> a -> a) + -> Key + -> NEIntMap a + -> NEIntMap a +adjustWithKey f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> n + EQ -> NEIntMap k0 (f k0 v) m + GT -> NEIntMap k0 v . M.adjustWithKey f k $ m +{-# INLINE adjustWithKey #-} + +-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- Returns a potentially empty map ('IntMap'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEIntMap'. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +update + :: (a -> Maybe a) + -> Key + -> NEIntMap a + -> IntMap a +update f = updateWithKey (const f) +{-# INLINE update #-} + +-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- Returns a potentially empty map ('IntMap'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEIntMap'. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +updateWithKey + :: (Key -> a -> Maybe a) + -> Key + -> NEIntMap a + -> IntMap a +updateWithKey f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> toMap n + EQ -> maybe m (flip (insertMinMap k0) m) . f k0 $ v + GT -> insertMinMap k0 v . M.updateWithKey f k $ m +{-# INLINE updateWithKey #-} + +-- | /O(min(n,W))/. Lookup and update. +-- The function returns original value, if it is updated. +-- This is different behavior than @Data.Map.NonEmpty.updateLookupWithKey@. +-- Returns the original key value if the map entry is deleted. +-- +-- Returns a potentially empty map ('IntMap') in the case that we delete +-- the final key of a singleton map. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == (Just "5:new a", Data.IntMap.fromList ((3, "b") :| [(5, "5:new a")])) +-- > updateLookupWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == (Nothing, Data.IntMap.fromList ((3, "b") :| [(5, "a")])) +-- > updateLookupWithKey f 3 (fromList ((5,"a") :| [(3,"b")])) == (Just "b", Data.IntMap.singleton 5 "a") +updateLookupWithKey + :: (Key -> a -> Maybe a) + -> Key + -> NEIntMap a + -> (Maybe a, IntMap a) +updateLookupWithKey f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> (Nothing, toMap n) + EQ -> let u = f k0 v + in (Just v, maybe m (flip (insertMinMap k0) m) u) + GT -> fmap (insertMinMap k0 v) . M.updateLookupWithKey f k $ m +{-# INLINE updateLookupWithKey #-} + +-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alter' can be used to insert, delete, or +-- update a value in a 'IntMap'. In short : @Data.IntMap.lookup k ('alter' +-- f k m) = f ('lookup' k m)@. +-- +-- Returns a potentially empty map ('IntMap'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEIntMap'. +-- +-- See 'alterF'' for a version that disallows deletion, and so therefore +-- can return 'NEIntMap'. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "c")] +alter + :: (Maybe a -> Maybe a) + -> Key + -> NEIntMap a + -> IntMap a +alter f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> ($ toMap n) . maybe id (insertMinMap k ) $ f Nothing + EQ -> ($ m ) . maybe id (insertMinMap k0) $ f (Just v) + GT -> insertMinMap k0 v . M.alter f k $ m +{-# INLINE alter #-} + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ +-- at @k@, or absence thereof. 'alterF' can be used to inspect, insert, +-- delete, or update a value in a 'IntMap'. In short: @Data.IntMap.lookup +-- k \<$\> 'alterF' f k m = f ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> NEIntMap Int String -> IO (IntMap Int String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing = do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) = do +-- putStrLn $ "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserResponse2 :: IO (Maybe String) +-- @ +-- +-- Like @Data.IntMap.alterF@ for 'IntMap', 'alterF' can be considered +-- to be a unifying generalization of 'lookup' and 'delete'; however, as +-- a constrast, it cannot be used to implement 'insert', because it must +-- return a 'IntMap' instead of an 'NEIntMap' (because the function might delete +-- the final item in the 'NEIntMap'). When used with trivial functors like +-- 'Identity' and 'Const', it is often slightly slower than +-- specialized 'lookup' and 'delete'. However, when the functor is +-- non-trivial and key comparison is not particularly cheap, it is the +-- fastest way. +-- +-- See 'alterF'' for a version that disallows deletion, and so therefore +-- can return 'NEIntMap' and be used to implement 'insert' +-- +-- Note on rewrite rules: +-- +-- This module includes GHC rewrite rules to optimize 'alterF' for +-- the 'Const' and 'Identity' functors. In general, these rules +-- improve performance. The sole exception is that when using +-- 'Identity', deleting a key that is already absent takes longer +-- than it would without the rules. If you expect this to occur +-- a very large fraction of the time, you might consider using a +-- private copy of the 'Identity' type. +-- +-- Note: Unlike @Data.IntMap.alterF@ for 'IntMap', 'alterF' is /not/ a flipped +-- version of the 'Control.Lens.At.at' combinator from "Control.Lens.At". +-- However, it match the shape expected from most functions expecting +-- lenses, getters, and setters, so can be thought of as a "psuedo-lens", +-- with virtually the same practical applications as a legitimate lens. +alterF + :: Functor f + => (Maybe a -> f (Maybe a)) + -> Key + -> NEIntMap a + -> f (IntMap a) +alterF f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> ($ toMap n) . maybe id (insertMinMap k ) <$> f Nothing + EQ -> ($ m ) . maybe id (insertMinMap k0) <$> f (Just v) + GT -> insertMinMap k0 v <$> M.alterF f k m +{-# INLINABLE [2] alterF #-} + +-- if f ~ Const b, it's a lookup +{-# RULES +"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m + #-} +-- if f ~ Identity, it's an 'alter' +{-# RULES +"alterF/Identity" forall k (f :: Maybe a -> Identity (Maybe a)) . alterF f k = Identity . alter (runIdentity . f) k + #-} + +-- | /O(log n)/. Variant of 'alter' that disallows deletion. Allows us to +-- guarantee that the result is also a non-empty IntMap. +alter' + :: (Maybe a -> a) + -> Key + -> NEIntMap a + -> NEIntMap a +alter' f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> NEIntMap k (f Nothing) . toMap $ n + EQ -> NEIntMap k0 (f (Just v)) $ m + GT -> NEIntMap k0 v . M.alter (Just . f) k $ m +{-# INLINE alter' #-} + +-- | /O(log n)/. Variant of 'alterF' that disallows deletion. Allows us to +-- guarantee that the result is also a non-empty IntMap. +-- +-- Like @Data.IntMap.alterF@ for 'IntMap', can be used to generalize and unify +-- 'lookup' and 'insert'. However, because it disallows deletion, it +-- cannot be used to implement 'delete'. +-- +-- See 'alterF' for usage information and caveats. +-- +-- Note: Neither 'alterF' nor 'alterF'' can be considered flipped versions +-- of the 'Control.Lens.At.at' combinator from "Control.Lens.At". However, +-- this can match the shape expected from most functions expecting lenses, +-- getters, and setters, so can be thought of as a "psuedo-lens", with +-- virtually the same practical applications as a legitimate lens. +-- +-- __WARNING__: The rewrite rule for 'Identity' exposes an inconsistency in +-- undefined behavior for "Data.IntMap". @Data.IntMap.alterF@ will actually +-- /maintain/ the original key in the map when used with 'Identity'; +-- however, @Data.IntMap.insertWith@ will /replace/ the orginal key in the +-- map. The rewrite rule for 'alterF'' has chosen to be faithful to +-- @Data.IntMap.insertWith@, and /not/ @Data.IntMap.alterF@, for the sake of +-- a cleaner implementation. +alterF' + :: Functor f + => (Maybe a -> f a) + -> Key + -> NEIntMap a + -> f (NEIntMap a) +alterF' f k n@(NEIntMap k0 v m) = case compare k k0 of + LT -> flip (NEIntMap k ) (toMap n) <$> f Nothing + EQ -> flip (NEIntMap k0) m <$> f (Just v) + GT -> NEIntMap k0 v <$> M.alterF (fmap Just . f) k m +{-# INLINABLE [2] alterF' #-} + +-- if f ~ Const b, it's a lookup +{-# RULES +"alterF'/Const" forall k (f :: Maybe a -> Const b a) . alterF' f k = \m -> Const . getConst . f $ lookup k m + #-} +-- if f ~ Identity, it's an insertWith +{-# RULES +"alterF'/Identity" forall k (f :: Maybe a -> Identity a) . alterF' f k = Identity . insertWith (\_ -> runIdentity . f . Just) k (runIdentity (f Nothing)) + #-} + +-- | /O(log n)/. Lookup the value at a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +-- +-- An example of using @lookup@: +-- +-- > import Prelude hiding (lookup) +-- > import Data.Map.NonEmpty +-- > +-- > employeeDept = fromList (("John","Sales") :| [("Bob","IT")]) +-- > deptCountry = fromList (("IT","USA") :| [("Sales","France")]) +-- > countryCurrency = fromList (("USA", "Dollar") :| [("France", "Euro")]) +-- > +-- > employeeCurrency :: String -> Maybe String +-- > employeeCurrency name = do +-- > dept <- lookup name employeeDept +-- > country <- lookup dept deptCountry +-- > lookup country countryCurrency +-- > +-- > main = do +-- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) +-- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) +-- +-- The output of this program: +-- +-- > John's currency: Just "Euro" +-- > Pete's currency: Nothing +lookup + :: Key + -> NEIntMap a + -> Maybe a +lookup k (NEIntMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Just v + GT -> M.lookup k m +{-# INLINE lookup #-} + +-- | /O(log n)/. Find the value at a key. Returns 'Nothing' when the +-- element can not be found. +-- +-- prop> fromList ((5, 'a') :| [(3, 'b')]) !? 1 == Nothing +-- prop> fromList ((5, 'a') :| [(3, 'b')]) !? 5 == Just 'a' +(!?) :: NEIntMap a -> Key -> Maybe a +(!?) = flip lookup +{-# INLINE (!?) #-} + +-- | /O(log n)/. Find the value at a key. Calls 'error' when the element +-- can not be found. +-- +-- > fromList ((5,'a') :| [(3,'b')]) ! 1 Error: element not in the map +-- > fromList ((5,'a') :| [(3,'b')]) ! 5 == 'a' +(!) :: NEIntMap a -> Key -> a +(!) m k = fromMaybe e $ m !? k + where + e = error "NEIntMap.!: given key is not an element in the map" +{-# INLINE (!) #-} + +infixl 9 !? +infixl 9 ! + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList ((5,'a') :| [(3,'b')])) == 'x' +-- > findWithDefault 'x' 5 (fromList ((5,'a') :| [(3,'b')])) == 'a' +findWithDefault + :: a + -> Key + -> NEIntMap a + -> a +findWithDefault def k (NEIntMap k0 v m) = case compare k k0 of + LT -> def + EQ -> v + GT -> M.findWithDefault def k m +{-# INLINE findWithDefault #-} + +-- | /O(log n)/. Is the key a member of the map? See also 'notMember'. +-- +-- > member 5 (fromList ((5,'a') :| [(3,'b')])) == True +-- > member 1 (fromList ((5,'a') :| [(3,'b')])) == False +member :: Key -> NEIntMap a -> Bool +member k (NEIntMap k0 _ m) = case compare k k0 of + LT -> False + EQ -> True + GT -> M.member k m +{-# INLINE member #-} + +-- | /O(log n)/. Is the key not a member of the map? See also 'member'. +-- +-- > notMember 5 (fromList ((5,'a') :| [(3,'b')])) == False +-- > notMember 1 (fromList ((5,'a') :| [(3,'b')])) == True +notMember :: Key -> NEIntMap a -> Bool +notMember k (NEIntMap k0 _ m) = case compare k k0 of + LT -> True + EQ -> False + GT -> M.notMember k m +{-# INLINE notMember #-} + +-- | /O(log n)/. Find largest key smaller than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupLT 3 (fromList ((3,'a') :| [(5,'b')])) == Nothing +-- > lookupLT 4 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +lookupLT :: Key -> NEIntMap a -> Maybe (Key, a) +lookupLT k (NEIntMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Nothing + GT -> M.lookupLT k m <|> Just (k0, v) +{-# INLINE lookupLT #-} + +-- | /O(log n)/. Find smallest key greater than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupGT 4 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +-- > lookupGT 5 (fromList ((3,'a') :| [(5,'b')])) == Nothing +lookupGT :: Key -> NEIntMap a -> Maybe (Key, a) +lookupGT k (NEIntMap k0 v m) = case compare k k0 of + LT -> Just (k0, v) + EQ -> lookupMinMap m + GT -> M.lookupGT k m +{-# INLINE lookupGT #-} + +-- | /O(log n)/. Find largest key smaller or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupLE 2 (fromList ((3,'a') :| [(5,'b')])) == Nothing +-- > lookupLE 4 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +-- > lookupLE 5 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +lookupLE :: Key -> NEIntMap a -> Maybe (Key, a) +lookupLE k (NEIntMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Just (k0, v) + GT -> M.lookupLE k m <|> Just (k0, v) +{-# INLINE lookupLE #-} + +-- | /O(log n)/. Find smallest key greater or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupGE 3 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +-- > lookupGE 4 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +-- > lookupGE 6 (fromList ((3,'a') :| [(5,'b')])) == Nothing +lookupGE :: Key -> NEIntMap a -> Maybe (Key, a) +lookupGE k (NEIntMap k0 v m) = case compare k k0 of + LT -> Just (k0, v) + EQ -> Just (k0, v) + GT -> M.lookupGE k m +{-# INLINE lookupGE #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Union with a combining function. +-- +-- > unionWith (++) (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "aA"), (7, "C")]) +unionWith + :: (a -> a -> a) + -> NEIntMap a + -> NEIntMap a + -> NEIntMap a +unionWith f n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 v2 m2) = case compare k1 k2 of + LT -> NEIntMap k1 v1 . M.unionWith f m1 . toMap $ n2 + EQ -> NEIntMap k1 (f v1 v2) . M.unionWith f m1 $ m2 + GT -> NEIntMap k2 v2 . M.unionWith f (toMap n1) $ m2 +{-# INLINE unionWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- Union with a combining function, given the matching key. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "5:a|A"), (7, "C")]) +unionWithKey + :: (Key -> a -> a -> a) + -> NEIntMap a + -> NEIntMap a + -> NEIntMap a +unionWithKey f n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 v2 m2) = case compare k1 k2 of + LT -> NEIntMap k1 v1 . M.unionWithKey f m1 . toMap $ n2 + EQ -> NEIntMap k1 (f k1 v1 v2) . M.unionWithKey f m1 $ m2 + GT -> NEIntMap k2 v2 . M.unionWithKey f (toMap n1) $ m2 +{-# INLINE unionWithKey #-} + +-- | The union of a non-empty list of maps, with a combining operation: +-- (@'unionsWith' f == 'Data.Foldable.foldl1' ('unionWith' f)@). +-- +-- > unionsWith (++) (fromList ((5, "a") :| [(3, "b")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "A3") :| [(3, "B3")])]) +-- > == fromList ((3, "bB3") :| [(5, "aAA3"), (7, "C")]) +unionsWith + :: Foldable1 f + => (a -> a -> a) + -> f (NEIntMap a) + -> NEIntMap a +unionsWith f (F1.toNonEmpty->(m :| ms)) = F.foldl' (unionWith f) m ms +{-# INLINE unionsWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +-- +-- Returns a potentially empty map ('IntMap'), in case the first map is +-- a subset of the second map. +-- +-- > difference (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.IntMap.singleton 3 "b" +difference + :: NEIntMap a + -> NEIntMap b + -> IntMap a +difference n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 _ m2) = case compare k1 k2 of + -- k1 is not in n2, so cannot be deleted + LT -> insertMinMap k1 v1 $ m1 `M.difference` toMap n2 + -- k2 deletes k1, and only k1 + EQ -> m1 `M.difference` m2 + -- k2 is not in n1, so cannot delete anything, so we can just difference n1 // m2. + GT -> toMap n1 `M.difference` m2 +{-# INLINE difference #-} + +-- | Same as 'difference'. +(\\) + :: NEIntMap a + -> NEIntMap b + -> IntMap a +(\\) = difference +{-# INLINE (\\) #-} + +-- | /O(n+m)/. Difference with a combining function. +-- When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- Returns a potentially empty map ('IntMap'), in case the first map is +-- a subset of the second map and the function returns 'Nothing' for every +-- pair. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(3, "B"), (7, "C")])) +-- > == Data.IntMap.singleton 3 "b:B" +differenceWith + :: (a -> b -> Maybe a) + -> NEIntMap a + -> NEIntMap b + -> IntMap a +differenceWith f = differenceWithKey (const f) +{-# INLINE differenceWith #-} + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- Returns a potentially empty map ('IntMap'), in case the first map is +-- a subset of the second map and the function returns 'Nothing' for every +-- pair. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(3, "B"), (10, "C")])) +-- > == Data.IntMap.singleton 3 "3:b|B" +differenceWithKey + :: (Key -> a -> b -> Maybe a) + -> NEIntMap a + -> NEIntMap b + -> IntMap a +differenceWithKey f n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 v2 m2) = case compare k1 k2 of + -- k1 is not in n2, so cannot be deleted + LT -> insertMinMap k1 v1 $ M.differenceWithKey f m1 (toMap n2) + -- k2 deletes k1, and only k1 + EQ -> ($ M.differenceWithKey f m1 m2) . maybe id (insertMinMap k1) $ f k1 v1 v2 + -- k2 is not in n1, so cannot delete anything, so we can just difference n1 // m2. + GT -> M.differenceWithKey f (toMap n1) m2 +{-# INLINE differenceWithKey #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection of two maps. +-- Return data in the first map for the keys existing in both maps. +-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). +-- +-- Returns a potentially empty map ('IntMap'), in case the two maps share no +-- keys in common. +-- +-- > intersection (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.IntMap.singleton 5 "a" +intersection + :: NEIntMap a + -> NEIntMap b + -> IntMap a +intersection n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 _ m2) = case compare k1 k2 of + -- k1 is not in n2 + LT -> m1 `M.intersection` toMap n2 + -- k1 and k2 are a part of the result + EQ -> insertMinMap k1 v1 $ m1 `M.intersection` m2 + -- k2 is not in n1 + GT -> toMap n1 `M.intersection` m2 +{-# INLINE intersection #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- Returns a potentially empty map ('IntMap'), in case the two maps share no +-- keys in common. +-- +-- > intersectionWith (++) (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.IntMap.singleton 5 "aA" +intersectionWith + :: (a -> b -> c) + -> NEIntMap a + -> NEIntMap b + -> IntMap c +intersectionWith f = intersectionWithKey (const f) +{-# INLINE intersectionWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- Returns a potentially empty map ('IntMap'), in case the two maps share no +-- keys in common. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.IntMap.singleton 5 "5:a|A" +intersectionWithKey + :: (Key -> a -> b -> c) + -> NEIntMap a + -> NEIntMap b + -> IntMap c +intersectionWithKey f n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 v2 m2) = case compare k1 k2 of + -- k1 is not in n2 + LT -> M.intersectionWithKey f m1 (toMap n2) + -- k1 and k2 are a part of the result + EQ -> insertMinMap k1 (f k1 v1 v2) $ M.intersectionWithKey f m1 m2 + -- k2 is not in n1 + GT -> M.intersectionWithKey f (toMap n1) m2 +{-# INLINE intersectionWithKey #-} + +-- | /O(n)/. IntMap a function over all values in the map. +-- +-- > let f key x = (show key) ++ ":" ++ x +-- > mapWithKey f (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "3:b") :| [(5, "5:a")]) +mapWithKey :: (Key -> a -> b) -> NEIntMap a -> NEIntMap b +mapWithKey f (NEIntMap k v m) = NEIntMap k (f k v) (M.mapWithKey f m) +{-# NOINLINE [1] mapWithKey #-} +{-# RULES +"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = + mapWithKey (\k a -> f k (g k a)) xs +"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = + mapWithKey (\k a -> f k (g a)) xs +"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = + mapWithKey (\k a -> f (g k a)) xs + #-} + +-- | /O(n)/. The function 'mapAccum' threads an accumulating argument +-- through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList ((5,"a") :| [(3,"b")])) == ("Everything: ba", fromList ((3, "bX") :| [(5, "aX")])) +mapAccum + :: (a -> b -> (a, c)) + -> a + -> NEIntMap b + -> (a, NEIntMap c) +mapAccum f = mapAccumWithKey (\x _ -> f x) +{-# INLINE mapAccum #-} + +-- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") +-- > mapAccumWithKey f "Everything:" (fromList ((5,"a") :| [(3,"b")])) == ("Everything: 3-b 5-a", fromList ((3, "bX") :| [(5, "aX")])) +mapAccumWithKey + :: (a -> Key -> b -> (a, c)) + -> a + -> NEIntMap b + -> (a, NEIntMap c) +mapAccumWithKey f z0 (NEIntMap k v m) = (z2, NEIntMap k v' m') + where + ~(z1, v') = f z0 k v + ~(z2, m') = M.mapAccumWithKey f z1 m +{-# INLINE mapAccumWithKey #-} + +-- | /O(n)/. The function 'mapAccumRWithKey' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey + :: (a -> Key -> b -> (a, c)) + -> a + -> NEIntMap b + -> (a, NEIntMap c) +mapAccumRWithKey f z0 (NEIntMap k v m) = (z2, NEIntMap k v' m') + where + ~(z1, m') = M.mapAccumRWithKey f z0 m + ~(z2, v') = f z1 k v +{-# INLINE mapAccumRWithKey #-} + +-- | /O(n*log n)/. +-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the value at the greatest of the +-- original keys is retained. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeys (+ 1) (fromList ((5,"a") :| [(3,"b")])) == fromList ((4, "b") :| [(6, "a")]) +-- > mapKeys (\ _ -> 1) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 1 "c" +-- > mapKeys (\ _ -> 3) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 3 "c" +mapKeys + :: (Key -> Key) + -> NEIntMap a + -> NEIntMap a +mapKeys f (NEIntMap k0 v0 m) = fromListWith const + . ((f k0, v0) :|) + . M.foldrWithKey (\k v kvs -> (f k, v) : kvs) [] + $ m +{-# INLINABLE mapKeys #-} + +-- | /O(n*log n)/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. The value at the greater of the two original keys +-- is used as the first argument to @c@. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 3 "cdab" +mapKeysWith + :: (a -> a -> a) + -> (Key -> Key) + -> NEIntMap a + -> NEIntMap a +mapKeysWith c f (NEIntMap k0 v0 m) = fromListWith c + . ((f k0, v0) :|) + . M.foldrWithKey (\k v kvs -> (f k, v) : kvs) [] + $ m +{-# INLINABLE mapKeysWith #-} + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList ((5,"a") :| [(3,"b")])) == fromList ((6, "b") :| [(10, "a")]) +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList ((5,"a") :| [(3,"b")]))) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList ((5,"a") :| [(3,"b")]))) == False +mapKeysMonotonic + :: (Key -> Key) + -> NEIntMap a + -> NEIntMap a +mapKeysMonotonic f (NEIntMap k v m) = NEIntMap (f k) v + . M.mapKeysMonotonic f + $ m +{-# INLINE mapKeysMonotonic #-} + +-- | /O(n)/. Fold the keys and values in the map using the given right-associative +-- binary operator, such that +-- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- +-- For example, +-- +-- > keysList map = foldrWithKey (\k x ks -> k:ks) [] map +foldrWithKey :: (Key -> a -> b -> b) -> b -> NEIntMap a -> b +foldrWithKey f z (NEIntMap k v m) = f k v . M.foldrWithKey f z $ m +{-# INLINE foldrWithKey #-} + +-- | /O(n)/. Fold the keys and values in the map using the given left-associative +-- binary operator, such that +-- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@. +-- +-- For example, +-- +-- > keysList = reverse . foldlWithKey (\ks k x -> k:ks) [] +foldlWithKey :: (a -> Key -> b -> a) -> a -> NEIntMap b -> a +foldlWithKey f z (NEIntMap k v m) = M.foldlWithKey f (f z k v) m +{-# INLINE foldlWithKey #-} + +-- | /O(n)/. A strict version of 'foldr1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr1' :: (a -> a -> a) -> NEIntMap a -> a +foldr1' f (NEIntMap _ v m) = case M.maxView m of + Nothing -> v + Just (y, m') -> let !z = M.foldr' f y m' in v `f` z +{-# INLINE foldr1' #-} + +-- | /O(n)/. A strict version of 'foldl1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl1' :: (a -> a -> a) -> NEIntMap a -> a +foldl1' f (NEIntMap _ v m) = M.foldl' f v m +{-# INLINE foldl1' #-} + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldrWithKey' :: (Key -> a -> b -> b) -> b -> NEIntMap a -> b +foldrWithKey' f z (NEIntMap k v m) = f k v y + where + !y = M.foldrWithKey f z m +{-# INLINE foldrWithKey' #-} + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldlWithKey' :: (a -> Key -> b -> a) -> a -> NEIntMap b -> a +foldlWithKey' f z (NEIntMap k v m) = M.foldlWithKey' f x m + where + !x = f z k v +{-# INLINE foldlWithKey' #-} + +-- | /O(n)/. Return all keys of the map in ascending order. +-- +-- > keys (fromList ((5,"a") :| [(3,"b")])) == (3 :| [5]) +keys :: NEIntMap a -> NonEmpty Key +keys (NEIntMap k _ m) = k :| M.keys m +{-# INLINE keys #-} + +-- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map +-- in ascending key order. +-- +-- > assocs (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +assocs :: NEIntMap a -> NonEmpty (Key, a) +assocs = toList +{-# INLINE assocs #-} + +-- | /O(n)/. The non-empty set of all keys of the map. +-- +-- > keysSet (fromList ((5,"a") :| [(3,"b")])) == Data.Set.NonEmpty.fromList (3 :| [5]) +keysSet :: NEIntMap a -> NEIntSet +keysSet (NEIntMap k _ m) = NEIntSet k (M.keysSet m) +{-# INLINE keysSet #-} + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are +-- in ascending order. +-- +-- > toAscList (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +toAscList :: NEIntMap a -> NonEmpty (Key, a) +toAscList = toList +{-# INLINE toAscList #-} + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys +-- are in descending order. +-- +-- > toDescList (fromList ((5,"a") :| [(3,"b")])) == ((5,"a") :| [(3,"b")]) +toDescList :: NEIntMap a -> NonEmpty (Key, a) +toDescList (NEIntMap k0 v0 m) = M.foldlWithKey' go ((k0, v0) :| []) m + where + go xs k v = (k, v) NE.<| xs +{-# INLINE toDescList #-} + +-- | /O(n)/. Filter all values that satisfy the predicate. +-- +-- Returns a potentially empty map ('IntMap'), because we could +-- potentailly filter out all items in the original 'NEIntMap'. +-- +-- > filter (> "a") (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 3 "b" +-- > filter (> "x") (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.empty +-- > filter (< "a") (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.empty +filter + :: (a -> Bool) + -> NEIntMap a + -> IntMap a +filter f (NEIntMap k v m) + | f v = insertMinMap k v . M.filter f $ m + | otherwise = M.filter f m +{-# INLINE filter #-} + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- Returns a potentially empty map ('IntMap'), because we could +-- potentailly filter out all items in the original 'NEIntMap'. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +filterWithKey + :: (Key -> a -> Bool) + -> NEIntMap a + -> IntMap a +filterWithKey f (NEIntMap k v m) + | f k v = insertMinMap k v . M.filterWithKey f $ m + | otherwise = M.filterWithKey f m +{-# INLINE filterWithKey #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Restrict an 'NEIntMap' to only those keys +-- found in a 'Data.Set.Set'. +-- +-- @ +-- m \`restrictKeys\` s = 'filterWithKey' (\k _ -> k ``Set.member`` s) m +-- m \`restrictKeys\` s = m ``intersection`` 'fromSet' (const ()) s +-- @ +restrictKeys + :: NEIntMap a + -> IntSet + -> IntMap a +restrictKeys n@(NEIntMap k v m) xs = case S.minView xs of + Nothing -> M.empty + Just (y, ys) -> case compare k y of + -- k is not in xs + LT -> m `M.restrictKeys` xs + -- k and y are a part of the result + EQ -> insertMinMap k v $ m `M.restrictKeys` ys + -- y is not in m + GT -> toMap n `M.restrictKeys` ys +{-# INLINE restrictKeys #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Remove all keys in a 'Data.Set.Set' from +-- an 'NEIntMap'. +-- +-- @ +-- m \`withoutKeys\` s = 'filterWithKey' (\k _ -> k ``Set.notMember`` s) m +-- m \`withoutKeys\` s = m ``difference`` 'fromSet' (const ()) s +-- @ +withoutKeys + :: NEIntMap a + -> IntSet + -> IntMap a +withoutKeys n@(NEIntMap k v m) xs = case S.minView xs of + Nothing -> toMap n + Just (y, ys) -> case compare k y of + -- k is not in xs, so cannot be deleted + LT -> insertMinMap k v $ m `M.withoutKeys` xs + -- y deletes k, and only k + EQ -> m `M.withoutKeys` ys + -- y is not in n, so cannot delete anything, so we can just difference n and ys + GT -> toMap n `M.withoutKeys` ys +{-# INLINE withoutKeys #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the predicate was true for all items. +-- * @'That' n2@ means that the predicate was false for all items. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partition (> "a") (fromList ((5,"a") :| [(3,"b")])) == These (singleton 3 "b") (singleton 5 "a") +-- > partition (< "x") (fromList ((5,"a") :| [(3,"b")])) == This (fromList ((3, "b") :| [(5, "a")])) +-- > partition (> "x") (fromList ((5,"a") :| [(3,"b")])) == That (fromList ((3, "b") :| [(5, "a")])) +partition + :: (a -> Bool) + -> NEIntMap a + -> These (NEIntMap a) (NEIntMap a) +partition f = partitionWithKey (const f) +{-# INLINE partition #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the predicate was true for all items, +-- returning the original map. +-- * @'That' n2@ means that the predicate was false for all items, +-- returning the original map. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList ((5,"a") :| [(3,"b")])) == These (singleton 5 "a") (singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList ((5,"a") :| [(3,"b")])) == This (fromList ((3, "b") :| [(5, "a")])) +-- > partitionWithKey (\ k _ -> k > 7) (fromList ((5,"a") :| [(3,"b")])) == That (fromList ((3, "b") :| [(5, "a")])) +partitionWithKey + :: (Key -> a -> Bool) + -> NEIntMap a + -> These (NEIntMap a) (NEIntMap a) +partitionWithKey f n@(NEIntMap k v m0) = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) + | f k v -> This n + | otherwise -> That n + (Just n1, Nothing) + | f k v -> This n + | otherwise -> These n1 (singleton k v) + (Nothing, Just n2) + | f k v -> These (singleton k v) n2 + | otherwise -> That n + (Just n1, Just n2) + | f k v -> These (insertMapMin k v m1) n2 + | otherwise -> These n1 (insertMapMin k v m2) + where + (m1, m2) = M.partitionWithKey f m0 +{-# INLINABLE partitionWithKey #-} + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('IntMap'), because the function could +-- potentially return 'Nothing' on all items in the 'NEIntMap'. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "new a" +mapMaybe + :: (a -> Maybe b) + -> NEIntMap a + -> IntMap b +mapMaybe f = mapMaybeWithKey (const f) +{-# INLINE mapMaybe #-} + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('IntMap'), because the function could +-- potentially return 'Nothing' on all items in the 'NEIntMap'. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 3 "key : 3" +mapMaybeWithKey + :: (Key -> a -> Maybe b) + -> NEIntMap a + -> IntMap b +mapMaybeWithKey f (NEIntMap k v m) = ($ M.mapMaybeWithKey f m) + . maybe id (insertMinMap k) + $ f k v +{-# INLINE mapMaybeWithKey #-} + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the results were all 'Left'. +-- * @'That' n2@ means that the results were all 'Right'. +-- * @'These' n1 n2@ gives @n1@ (the map where the results were 'Left') +-- and @n2@ (the map where the results were 'Right') +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == These (fromList ((3,"b") :| [(5,"a")])) (fromList ((1,"x") :| [(7,"z")])) +-- > +-- > mapEither (\ a -> Right a) (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == That (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +mapEither + :: (a -> Either b c) + -> NEIntMap a + -> These (NEIntMap b) (NEIntMap c) +mapEither f = mapEitherWithKey (const f) +{-# INLINE mapEither #-} + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the results were all 'Left'. +-- * @'That' n2@ means that the results were all 'Right'. +-- * @'These' n1 n2@ gives @n1@ (the map where the results were 'Left') +-- and @n2@ (the map where the results were 'Right') +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == These (fromList ((1,2) :| [(3,6)])) (fromList ((5,"aa") :| [(7,"zz")])) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == That (fromList ((1,"x") :| [(3,"b"), (5,"a"), (7,"z")])) +mapEitherWithKey + :: (Key -> a -> Either b c) + -> NEIntMap a + -> These (NEIntMap b) (NEIntMap c) +mapEitherWithKey f (NEIntMap k v m0) = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> case f k v of + Left v' -> This (singleton k v') + Right v' -> That (singleton k v') + (Just n1, Nothing) -> case f k v of + Left v' -> This (insertMapMin k v' m1) + Right v' -> These n1 (singleton k v') + (Nothing, Just n2) -> case f k v of + Left v' -> These (singleton k v') n2 + Right v' -> That (insertMapMin k v' m2) + (Just n1, Just n2) -> case f k v of + Left v' -> These (insertMapMin k v' m1) n2 + Right v' -> These n1 (insertMapMin k v' m2) + where + (m1, m2) = M.mapEitherWithKey f m0 +{-# INLINABLE mapEitherWithKey #-} + +-- | /O(log n)/. The expression (@'split' k map@) is potentially a 'These' +-- containing up to two 'NEIntMap's based on splitting the map into maps +-- containing items before and after the given key @k@. It will never +-- return a map that contains @k@ itself. +-- +-- * 'Nothing' means that @k@ was the only key in the the original map, +-- and so there are no items before or after it. +-- * @'Just' ('This' n1)@ means @k@ was larger than or equal to all items +-- in the map, and @n1@ is the entire original map (minus @k@, if it was +-- present) +-- * @'Just' ('That' n2)@ means @k@ was smaller than or equal to all +-- items in the map, and @n2@ is the entire original map (minus @k@, if +-- it was present) +-- * @'Just' ('These' n1 n2)@ gives @n1@ (the map of all keys from the +-- original map less than @k@) and @n2@ (the map of all keys from the +-- original map greater than @k@) +-- +-- > split 2 (fromList ((5,"a") :| [(3,"b")])) == Just (That (fromList ((3,"b") :| [(5,"a")])) ) +-- > split 3 (fromList ((5,"a") :| [(3,"b")])) == Just (That (singleton 5 "a") ) +-- > split 4 (fromList ((5,"a") :| [(3,"b")])) == Just (These (singleton 3 "b") (singleton 5 "a")) +-- > split 5 (fromList ((5,"a") :| [(3,"b")])) == Just (This (singleton 3 "b") ) +-- > split 6 (fromList ((5,"a") :| [(3,"b")])) == Just (This (fromList ((3,"b") :| [(5,"a")])) ) +-- > split 5 (singleton 5 "a") == Nothing +split + :: Key + -> NEIntMap a + -> Maybe (These (NEIntMap a) (NEIntMap a)) +split k n@(NEIntMap k0 v m0) = case compare k k0 of + LT -> Just $ That n + EQ -> That <$> nonEmptyMap m0 + GT -> case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> Just $ This (singleton k0 v) + (Just _ , Nothing) -> Just $ This (insertMapMin k0 v m1) + (Nothing, Just n2) -> Just $ These (singleton k0 v) n2 + (Just _ , Just n2) -> Just $ These (insertMapMin k0 v m1) n2 + where + (m1, m2) = M.split k m0 +{-# INLINABLE split #-} + +-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@, as a @'Maybe' a@. +-- +-- > splitLookup 2 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (That (fromList ((3,"b") :| [(5,"a")])))) +-- > splitLookup 3 (fromList ((5,"a") :| [(3,"b")])) == (Just "b", Just (That (singleton 5 "a"))) +-- > splitLookup 4 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (These (singleton 3 "b") (singleton 5 "a"))) +-- > splitLookup 5 (fromList ((5,"a") :| [(3,"b")])) == (Just "a", Just (This (singleton 3 "b")) +-- > splitLookup 6 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (This (fromList ((3,"b") :| [(5,"a")]))) +-- > splitLookup 5 (singleton 5 "a") == (Just "a", Nothing) +splitLookup + :: Key + -> NEIntMap a + -> (Maybe a, Maybe (These (NEIntMap a) (NEIntMap a))) +splitLookup k n@(NEIntMap k0 v0 m0) = case compare k k0 of + LT -> (Nothing, Just $ That n) + EQ -> (Just v0, That <$> nonEmptyMap m0) + GT -> (v ,) $ case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> Just $ This (singleton k0 v0) + (Just _ , Nothing) -> Just $ This (insertMapMin k0 v0 m1) + (Nothing, Just n2) -> Just $ These (singleton k0 v0) n2 + (Just _ , Just n2) -> Just $ These (insertMapMin k0 v0 m1) n2 + where + (m1, v, m2) = M.splitLookup k m0 +{-# INLINABLE splitLookup #-} + +-- | /O(1)/. Decompose a map into pieces based on the structure of the +-- underlying tree. This function is useful for consuming a map in +-- parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that +-- the pieces returned will be in ascending order (all elements in the +-- first submap less than all elements in the second, and so on). +-- +-- Note that the current implementation does not return more than four +-- submaps, but you should not depend on this behaviour because it can +-- change in the future without notice. +splitRoot + :: NEIntMap a + -> NonEmpty (NEIntMap a) +splitRoot (NEIntMap k v m) = singleton k v + :| Maybe.mapMaybe nonEmptyMap (M.splitRoot m) +{-# INLINE splitRoot #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +isSubmapOf :: Eq a => NEIntMap a -> NEIntMap a -> Bool +isSubmapOf = isSubmapOfBy (==) +{-# INLINE isSubmapOf #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if +-- all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when +-- applied to their respective values. For example, the following +-- expressions are all 'True': +-- +-- > isSubmapOfBy (==) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (<=) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (==) (fromList (('a',1) :| [('b',2)])) (fromList (('a',1) :| [('b',2)])) +-- +-- But the following are all 'False': +-- +-- > isSubmapOfBy (==) (singleton 'a' 2) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (<) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (==) (fromList (('a',1) :| [('b',2)])) (singleton 'a' 1) +isSubmapOfBy + :: (a -> b -> Bool) + -> NEIntMap a + -> NEIntMap b + -> Bool +isSubmapOfBy f (NEIntMap k v m0) (toMap->m1) = kvSub + && M.isSubmapOfBy f m0 m1 + where + kvSub = case M.lookup k m1 of + Just v0 -> f v v0 + Nothing -> False +{-# INLINE isSubmapOfBy #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap +-- but not equal). Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' +-- (==)@). +isProperSubmapOf :: Eq a => NEIntMap a -> NEIntMap a -> Bool +isProperSubmapOf = isProperSubmapOfBy (==) +{-# INLINE isProperSubmapOf #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap +-- but not equal). The expression (@'isProperSubmapOfBy' f m1 m2@) returns +-- 'True' when @m1@ and @m2@ are not equal, all keys in @m1@ are in @m2@, +-- and when @f@ returns 'True' when applied to their respective values. For +-- example, the following expressions are all 'True': +-- +-- > isProperSubmapOfBy (==) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +-- > isProperSubmapOfBy (<=) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +-- +-- But the following are all 'False': +-- +-- > isProperSubmapOfBy (==) (fromList ((1,1) :| [(2,2)])) (fromList ((1,1) :| [(2,2)])) +-- > isProperSubmapOfBy (==) (fromList ((1,1) :| [(2,2)])) (singleton 1 1)) +-- > isProperSubmapOfBy (<) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +isProperSubmapOfBy + :: (a -> b -> Bool) + -> NEIntMap a + -> NEIntMap b + -> Bool +isProperSubmapOfBy f m1 m2 = M.size (neimIntMap m1) < M.size (neimIntMap m2) + && isSubmapOfBy f m1 m2 +{-# INLINE isProperSubmapOfBy #-} + +-- | /O(1)/. The minimal key of the map. Note that this is total, making +-- 'Data.IntMap.lookupMin' obsolete. It is constant-time, so has better +-- asymptotics than @Data.IntMap.lookupMin@ and @Data.IntMap.findMin@, as well. +-- +-- > findMin (fromList ((5,"a") :| [(3,"b")])) == (3,"b") +findMin :: NEIntMap a -> (Key, a) +findMin (NEIntMap k v _) = (k, v) +{-# INLINE findMin #-} + +-- | /O(log n)/. The maximal key of the map. Note that this is total, making +-- 'Data.IntMap.lookupMin' obsolete. +-- +-- > findMax (fromList ((5,"a") :| [(3,"b")])) == (5,"a") +findMax :: NEIntMap a -> (Key, a) +findMax (NEIntMap k v m) = fromMaybe (k, v) . lookupMaxMap $ m +{-# INLINE findMax #-} + +-- | /O(1)/. Delete the minimal key. Returns a potentially empty map +-- ('IntMap'), because we might end up deleting the final key in a singleton +-- map. It is constant-time, so has better asymptotics than +-- 'Data.IntMap.deleteMin'. +-- +-- > deleteMin (fromList ((5,"a") :| [(3,"b"), (7,"c")])) == Data.IntMap.fromList [(5,"a"), (7,"c")] +-- > deleteMin (singleton 5 "a") == Data.IntMap.empty +deleteMin :: NEIntMap a -> IntMap a +deleteMin (NEIntMap _ _ m) = m +{-# INLINE deleteMin #-} + +-- | /O(log n)/. Delete the maximal key. Returns a potentially empty map +-- ('IntMap'), because we might end up deleting the final key in a singleton +-- map. +-- +-- > deleteMax (fromList ((5,"a") :| [(3,"b"), (7,"c")])) == Data.IntMap.fromList [(3,"b"), (5,"a")] +-- > deleteMax (singleton 5 "a") == Data.IntMap.empty +deleteMax :: NEIntMap a -> IntMap a +deleteMax (NEIntMap k v m) = insertMinMap k v . M.deleteMax $ m +{-# INLINE deleteMax #-} + +-- | /O(1)/ if delete, /O(log n)/ otherwise. Update the value at the +-- minimal key. Returns a potentially empty map ('IntMap'), because we might +-- end up deleting the final key in the map if the function returns +-- 'Nothing'. See 'adjustMin' for a version that can guaruntee that we +-- return a non-empty map. +-- +-- > updateMin (\ a -> Just ("X" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "Xb"), (5, "a")] +-- > updateMin (\ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +updateMin :: (a -> Maybe a) -> NEIntMap a -> IntMap a +updateMin f = updateMinWithKey (const f) +{-# INLINE updateMin #-} + +-- | /O(1)/. A version of 'updateMin' that disallows deletion, allowing us +-- to guarantee that the result is also non-empty. +adjustMin :: (a -> a) -> NEIntMap a -> NEIntMap a +adjustMin f = adjustMinWithKey (const f) +{-# INLINE adjustMin #-} + +-- | /O(1)/ if delete, /O(log n)/ otherwise. Update the value at the +-- minimal key. Returns a potentially empty map ('IntMap'), because we might +-- end up deleting the final key in the map if the function returns +-- 'Nothing'. See 'adjustMinWithKey' for a version that guaruntees +-- a non-empty map. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +updateMinWithKey :: (Key -> a -> Maybe a) -> NEIntMap a -> IntMap a +updateMinWithKey f (NEIntMap k v m) = ($ m) . maybe id (insertMinMap k) $ f k v +{-# INLINE updateMinWithKey #-} + +-- | /O(1)/. A version of 'adjustMaxWithKey' that disallows deletion, +-- allowing us to guarantee that the result is also non-empty. Note that +-- it also is able to have better asymptotics than 'updateMinWithKey' in +-- general. +adjustMinWithKey :: (Key -> a -> a) -> NEIntMap a -> NEIntMap a +adjustMinWithKey f (NEIntMap k v m) = NEIntMap k (f k v) m +{-# INLINE adjustMinWithKey #-} + +-- | /O(log n)/. Update the value at the maximal key. Returns +-- a potentially empty map ('IntMap'), because we might end up deleting the +-- final key in the map if the function returns 'Nothing'. See 'adjustMax' +-- for a version that can guarantee that we return a non-empty map. +-- +-- > updateMax (\ a -> Just ("X" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3, "b"), (5, "Xa")] +-- > updateMax (\ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 3 "b" +updateMax :: (a -> Maybe a) -> NEIntMap a -> IntMap a +updateMax f = updateMaxWithKey (const f) +{-# INLINE updateMax #-} + +-- | /O(log n)/. A version of 'updateMax' that disallows deletion, allowing +-- us to guarantee that the result is also non-empty. +adjustMax :: (a -> a) -> NEIntMap a -> NEIntMap a +adjustMax f = adjustMaxWithKey (const f) +{-# INLINE adjustMax #-} + +-- | /O(log n)/. Update the value at the maximal key. Returns +-- a potentially empty map ('IntMap'), because we might end up deleting the +-- final key in the map if the function returns 'Nothing'. See +-- 'adjustMaxWithKey' for a version that guaruntees a non-empty map. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.IntMap.singleton 5 "a" +updateMaxWithKey :: (Key -> a -> Maybe a) -> NEIntMap a -> IntMap a +updateMaxWithKey f (NEIntMap k v m) + | M.null m = maybe m (M.singleton k) $ f k v + | otherwise = insertMinMap k v + . M.updateMaxWithKey f + $ m +{-# INLINE updateMaxWithKey #-} + +-- | /O(log n)/. A version of 'updateMaxWithKey' that disallows deletion, +-- allowing us to guarantee that the result is also non-empty. +adjustMaxWithKey :: (Key -> a -> a) -> NEIntMap a -> NEIntMap a +adjustMaxWithKey f (NEIntMap k0 v m) + | M.null m = NEIntMap k0 (f k0 v) m + | otherwise = insertMapMin k0 v + . M.updateMaxWithKey (\k -> Just . f k) + $ m +{-# INLINE adjustMaxWithKey #-} + +-- | /O(1)/. Retrieves the value associated with minimal key of the +-- map, and the map stripped of that element. It is constant-time, so has +-- better asymptotics than @Data.IntMap.minView@ for 'IntMap'. +-- +-- Note that unlike @Data.IntMap.minView@ for 'IntMap', this cannot ever fail, +-- so doesn't need to return in a 'Maybe'. However, the result 'IntMap' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > minView (fromList ((5,"a") :| [(3,"b")])) == ("b", Data.IntMap.singleton 5 "a") +minView :: NEIntMap a -> (a, IntMap a) +minView = first snd . deleteFindMin +{-# INLINE minView #-} + +-- | /O(1)/. Delete and find the minimal key-value pair. It is +-- constant-time, so has better asymptotics that @Data.IntMap.minView@ for +-- 'IntMap'. +-- +-- Note that unlike @Data.IntMap.deleteFindMin@ for 'IntMap', this cannot ever +-- fail, and so is a total function. However, the result 'IntMap' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > deleteFindMin (fromList ((5,"a") :| [(3,"b"), (10,"c")])) == ((3,"b"), Data.IntMap.fromList [(5,"a"), (10,"c")]) +deleteFindMin :: NEIntMap a -> ((Key, a), IntMap a) +deleteFindMin (NEIntMap k v m) = ((k, v), m) +{-# INLINE deleteFindMin #-} + +-- | /O(log n)/. Retrieves the value associated with maximal key of the +-- map, and the map stripped of that element. +-- +-- Note that unlike @Data.IntMap.maxView@ from 'IntMap', this cannot ever fail, +-- so doesn't need to return in a 'Maybe'. However, the result 'IntMap' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > maxView (fromList ((5,"a") :| [(3,"b")])) == ("a", Data.IntMap.singleton 3 "b") +maxView :: NEIntMap a -> (a, IntMap a) +maxView = first snd . deleteFindMax +{-# INLINE maxView #-} + +-- | /O(log n)/. Delete and find the minimal key-value pair. +-- +-- Note that unlike @Data.IntMap.deleteFindMax@ for 'IntMap', this cannot ever +-- fail, and so is a total function. However, the result 'IntMap' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > deleteFindMax (fromList ((5,"a") :| [(3,"b"), (10,"c")])) == ((10,"c"), Data.IntMap.fromList [(3,"b"), (5,"a")]) +deleteFindMax :: NEIntMap a -> ((Key, a), IntMap a) +deleteFindMax (NEIntMap k v m) = maybe ((k, v), M.empty) (second (insertMinMap k v)) + . M.maxViewWithKey + $ m +{-# INLINE deleteFindMax #-} + +-- --------------------------- +-- Combining functions +-- --------------------------- +-- +-- Code comes from "Data.Map.Internal" from containers, modified slightly +-- to work with NonEmpty +-- +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 + +combineEq :: NonEmpty (Key, b) -> NonEmpty (Key, b) +combineEq = \case + x :| [] -> x :| [] + x :| xx@(_:_) -> go x xx + where + go z [] = z :| [] + go z@(kz,_) (x@(kx,xx):xs') + | kx==kz = go (kx,xx) xs' + | otherwise = z NE.<| go x xs' + +combineEqWith + :: (Key -> b -> b -> b) + -> NonEmpty (Key, b) + -> NonEmpty (Key, b) +combineEqWith f = \case + x :| [] -> x :| [] + x :| xx@(_:_) -> go x xx + where + go z [] = z :| [] + go z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in go (kx,yy) xs' + | otherwise = z NE.<| go x xs' diff --git a/src/Data/IntMap/NonEmpty/Internal.hs b/src/Data/IntMap/NonEmpty/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/IntMap/NonEmpty/Internal.hs @@ -0,0 +1,639 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Data.IntMap.NonEmpty.Internal +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Unsafe internal-use functions used in the implementation of +-- "Data.IntMap.NonEmpty". These functions can potentially be used to +-- break the abstraction of 'NEIntMap' and produce unsound maps, so be +-- wary! +module Data.IntMap.NonEmpty.Internal ( + -- * Non-Empty IntMap type + NEIntMap(..) + , Key + , singleton + , nonEmptyMap + , withNonEmpty + , fromList + , toList + , map + , insertWith + , union + , unions + , elems + , size + , toMap + -- * Folds + , foldr + , foldr' + , foldr1 + , foldl + , foldl' + , foldl1 + -- * Traversals + , traverseWithKey + , traverseWithKey1 + , foldMapWithKey + , traverseMapWithKey + -- * Unsafe IntMap Functions + , insertMinMap + , insertMaxMap + -- * Debug + , valid + -- * CPP compatibility + , lookupMinMap + , lookupMaxMap + ) where + +import Control.Applicative +import Control.DeepSeq +import Data.Coerce +import Data.Data +import Data.Function +import Data.Functor.Apply +import Data.Functor.Classes +import Data.IntMap.Internal (IntMap(..), Key) +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Maybe +import Data.Semigroup +import Data.Semigroup.Foldable (Foldable1(fold1)) +import Data.Semigroup.Traversable (Traversable1(..)) +import Data.Typeable (Typeable) +import Prelude hiding (foldr1, foldl1, foldr, foldl, map) +import Text.Read +import qualified Data.Foldable as F +import qualified Data.IntMap as M +import qualified Data.Semigroup.Foldable as F1 + +-- | A non-empty (by construction) map from integer keys to values @a@. At +-- least one key-value pair exists in an @'NEIntMap' v@ at all times. +-- +-- Functions that /take/ an 'NEIntMap' can safely operate on it with the +-- assumption that it has at least one key-value pair. +-- +-- Functions that /return/ an 'NEIntMap' provide an assurance that the result +-- has at least one key-value pair. +-- +-- "Data.IntMap.NonEmpty" re-exports the API of "Data.IntMap", faithfully +-- reproducing asymptotics, typeclass constraints, and semantics. +-- Functions that ensure that input and output maps are both non-empty +-- (like 'Data.IntMap.NonEmpty.insert') return 'NEIntMap', but functions that +-- might potentially return an empty map (like 'Data.IntMap.NonEmpty.delete') +-- return a 'IntMap' instead. +-- +-- You can directly construct an 'NEIntMap' with the API from +-- "Data.IntMap.NonEmpty"; it's more or less the same as constructing a normal +-- 'IntMap', except you don't have access to 'Data.IntMap.empty'. There are also +-- a few ways to construct an 'NEIntMap' from a 'IntMap': +-- +-- 1. The 'nonEmptyMap' smart constructor will convert a @'IntMap' k a@ into +-- a @'Maybe' ('NEIntMap' k a)@, returning 'Nothing' if the original 'IntMap' +-- was empty. +-- 2. You can use the 'Data.IntMap.NonEmpty.insertIntMap' family of functions to +-- insert a value into a 'IntMap' to create a guaranteed 'NEIntMap'. +-- 3. You can use the 'Data.IntMap.NonEmpty.IsNonEmpty' and +-- 'Data.IntMap.NonEmpty.IsEmpty' patterns to "pattern match" on a 'IntMap' +-- to reveal it as either containing a 'NEIntMap' or an empty map. +-- 4. 'withNonEmpty' offers a continuation-based interface for +-- deconstructing a 'IntMap' and treating it as if it were an +-- 'NEIntMap'. +-- +-- You can convert an 'NEIntMap' into a 'IntMap' with 'toMap' or +-- 'Data.IntMap.NonEmpty.IsNonEmpty', essentially "obscuring" the non-empty +-- property from the type. +data NEIntMap a = + NEIntMap { neimK0 :: !Key -- ^ invariant: must be smaller than smallest key in map + , neimV0 :: a + , neimIntMap :: !(IntMap a) + } + deriving (Typeable) + +instance Eq a => Eq (NEIntMap a) where + t1 == t2 = M.size (neimIntMap t1) == M.size (neimIntMap t2) + && toList t1 == toList t2 + +instance Ord a => Ord (NEIntMap a) where + compare = compare `on` toList + (<) = (<) `on` toList + (>) = (>) `on` toList + (<=) = (<=) `on` toList + (>=) = (>=) `on` toList + +instance Eq1 NEIntMap where + liftEq eq m1 m2 = M.size (neimIntMap m1) == M.size (neimIntMap m2) + && liftEq (liftEq eq) (toList m1) (toList m2) + +instance Ord1 NEIntMap where + liftCompare cmp m n = + liftCompare (liftCompare cmp) (toList m) (toList n) + +instance Show1 NEIntMap where + liftShowsPrec sp sl d m = + showsUnaryWith (liftShowsPrec sp' sl') "fromList" d (toList m) + where + sp' = liftShowsPrec sp sl + sl' = liftShowList sp sl + +instance Read1 NEIntMap where + liftReadsPrec rp rl = readsData $ + readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList + where + rp' = liftReadsPrec rp rl + rl' = liftReadList rp rl + +instance Read e => Read (NEIntMap e) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- parens . prec 10 $ readPrec + return (fromList xs) + readListPrec = readListPrecDefault + +instance Show a => Show (NEIntMap a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList (" . shows (toList m) . showString ")" + +instance NFData a => NFData (NEIntMap a) where + rnf (NEIntMap k v a) = rnf k `seq` rnf v `seq` rnf a + +-- Data instance code from Data.IntMap.Internal +-- +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +-- (c) wren romano 2016 +instance Data a => Data (NEIntMap a) where + gfoldl f z im = z fromList `f` toList im + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = intMapDataType + dataCast1 = gcast1 + +fromListConstr :: Constr +fromListConstr = mkConstr intMapDataType "fromList" [] Prefix + +intMapDataType :: DataType +intMapDataType = mkDataType "Data.IntMap.NonEmpty.Internal.NEIntMap" [fromListConstr] + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@. +-- +-- > elemsList map = foldr (:) [] map +-- +-- > let f a len = len + (length a) +-- > foldr f 0 (fromList ((5,"a") :| [(3,"bbb")])) == 4 +foldr :: (a -> b -> b) -> b -> NEIntMap a -> b +foldr f z (NEIntMap _ v m) = v `f` M.foldr f z m +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (a -> b -> b) -> b -> NEIntMap a -> b +foldr' f z (NEIntMap _ v m) = v `f` y + where + !y = M.foldr' f z m +{-# INLINE foldr' #-} + +-- | /O(n)/. A version of 'foldr' that uses the value at the maximal key in +-- the map as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldr1' for 'IntMap', this function is +-- total if the input function is total. +foldr1 :: (a -> a -> a) -> NEIntMap a -> a +foldr1 f (NEIntMap _ v m) = maybe v (f v . uncurry (M.foldr f)) + . M.maxView + $ m +{-# INLINE foldr1 #-} + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@. +-- +-- > elemsList = reverse . foldl (flip (:)) [] +-- +-- > let f len a = len + (length a) +-- > foldl f 0 (fromList ((5,"a") :| [(3,"bbb")])) == 4 +foldl :: (a -> b -> a) -> a -> NEIntMap b -> a +foldl f z (NEIntMap _ v m) = M.foldl f (f z v) m +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> NEIntMap b -> a +foldl' f z (NEIntMap _ v m) = M.foldl' f x m + where + !x = f z v +{-# INLINE foldl' #-} + +-- | /O(n)/. A version of 'foldl' that uses the value at the minimal key in +-- the map as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldl1' for 'IntMap', this function is +-- total if the input function is total. +foldl1 :: (a -> a -> a) -> NEIntMap a -> a +foldl1 f (NEIntMap _ v m) = M.foldl f v m +{-# INLINE foldl1 #-} + +-- | /O(n)/. Fold the keys and values in the map using the given semigroup, +-- such that +-- +-- @'foldMapWithKey' f = 'Data.Semigroup.Foldable.fold1' . 'Data.IntMap.NonEmpty.mapWithKey' f@ +-- +-- __WARNING__: Differs from @Data.IntMap.foldMapWithKey@, which traverses +-- positive items first, then negative items. +-- +-- This can be an asymptotically faster than +-- 'Data.IntMap.NonEmpty.foldrWithKey' or 'Data.IntMap.NonEmpty.foldlWithKey' for +-- some monoids. + +-- TODO: benchmark against maxView method +foldMapWithKey + :: Semigroup m + => (Key -> a -> m) + -> NEIntMap a + -> m +foldMapWithKey f = F1.foldMap1 (uncurry f) . toList +{-# INLINE foldMapWithKey #-} + +-- | /O(n)/. IntMap a function over all values in the map. +-- +-- > map (++ "x") (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "bx") :| [(5, "ax")]) +map :: (a -> b) -> NEIntMap a -> NEIntMap b +map f (NEIntMap k0 v m) = NEIntMap k0 (f v) (M.map f m) +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs + #-} +{-# RULES +"map/coerce" map coerce = coerce + #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and +-- @t2@. It prefers @t1@ when duplicate keys are encountered, i.e. +-- (@'union' == 'Data.IntMap.NonEmpty.unionWith' 'const'@). +-- +-- > union (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "a"), (7, "C")]) +union + :: NEIntMap a + -> NEIntMap a + -> NEIntMap a +union n1@(NEIntMap k1 v1 m1) n2@(NEIntMap k2 v2 m2) = case compare k1 k2 of + LT -> NEIntMap k1 v1 . M.union m1 . toMap $ n2 + EQ -> NEIntMap k1 v1 . M.union m1 $ m2 + GT -> NEIntMap k2 v2 . M.union (toMap n1) $ m2 +{-# INLINE union #-} + +-- | The left-biased union of a non-empty list of maps. +-- +-- > unions (fromList ((5, "a") :| [(3, "b")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "A3") :| [(3, "B3")])]) +-- > == fromList [(3, "b"), (5, "a"), (7, "C")] +-- > unions (fromList ((5, "A3") :| [(3, "B3")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "a") :| [(3, "b")])]) +-- > == fromList ((3, "B3") :| [(5, "A3"), (7, "C")]) +unions + :: Foldable1 f + => f (NEIntMap a) + -> NEIntMap a +unions (F1.toNonEmpty->(m :| ms)) = F.foldl' union m ms +{-# INLINE unions #-} + +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- +-- > elems (fromList ((5,"a") :| [(3,"b")])) == ("b" :| ["a"]) +elems :: NEIntMap a -> NonEmpty a +elems (NEIntMap _ v m) = v :| M.elems m +{-# INLINE elems #-} + +-- | /O(1)/. The number of elements in the map. Guaranteed to be greater +-- than zero. +-- +-- > size (singleton 1 'a') == 1 +-- > size (fromList ((1,'a') :| [(2,'c'), (3,'b')])) == 3 +size :: NEIntMap a -> Int +size (NEIntMap _ _ m) = 1 + M.size m +{-# INLINE size #-} + +-- | /O(log n)/. +-- Convert a non-empty map back into a normal possibly-empty map, for usage +-- with functions that expect 'IntMap'. +-- +-- Can be thought of as "obscuring" the non-emptiness of the map in its +-- type. See the 'Data.IntMap.NonEmpty.IsNotEmpty' pattern. +-- +-- 'nonEmptyMap' and @'maybe' 'Data.IntMap.empty' 'toMap'@ form an isomorphism: they +-- are perfect structure-preserving inverses of eachother. +-- +-- > toMap (fromList ((3,"a") :| [(5,"b")])) == Data.IntMap.fromList [(3,"a"), (5,"b")] +toMap :: NEIntMap a -> IntMap a +toMap (NEIntMap k v m) = insertMinMap k v m +{-# INLINE toMap #-} + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- That is, behaves exactly like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value. +-- +-- /Use 'traverseWithKey1'/ whenever possible (if your 'Applicative' +-- also has 'Apply' instance). This version is provided only for types +-- that do not have 'Apply' instance, since 'Apply' is not at the moment +-- (and might not ever be) an official superclass of 'Applicative'. +-- +-- __WARNING__: Differs from @Data.IntMap.traverseWithKey@, which traverses +-- positive items first, then negative items. +-- +-- @ +-- 'traverseWithKey' f = 'unwrapApplicative' . 'traverseWithKey1' (\\k -> WrapApplicative . f k) +-- @ +traverseWithKey + :: Applicative t + => (Key -> a -> t b) + -> NEIntMap a + -> t (NEIntMap b) +traverseWithKey f (NEIntMap k v m0) = + NEIntMap k <$> f k v + <*> traverseMapWithKey f m0 +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. +-- @'traverseWithKey1' f m == 'fromList' <$> 'traverse1' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- +-- That is, behaves exactly like a regular 'traverse1' except that the traversing +-- function also has access to the key associated with a value. +-- +-- __WARNING__: Differs from @Data.IntMap.traverseWithKey@, which traverses +-- positive items first, then negative items. +-- +-- Is more general than 'traverseWithKey', since works with all 'Apply', +-- and not just 'Applicative'. + +-- TODO: benchmark against maxView-based methods +traverseWithKey1 + :: Apply t + => (Key -> a -> t b) + -> NEIntMap a + -> t (NEIntMap b) +traverseWithKey1 f (NEIntMap k0 v m0) = case runMaybeApply m1 of + Left m2 -> NEIntMap k0 <$> f k0 v <.> m2 + Right m2 -> flip (NEIntMap k0) m2 <$> f k0 v + where + m1 = traverseMapWithKey (\k -> MaybeApply . Left . f k) m0 +{-# INLINABLE traverseWithKey1 #-} + +-- | /O(n)/. Convert the map to a non-empty list of key\/value pairs. +-- +-- > toList (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +toList :: NEIntMap a -> NonEmpty (Key, a) +toList (NEIntMap k v m) = (k,v) :| M.toList m +{-# INLINE toList #-} + +-- | /O(log n)/. Smart constructor for an 'NEIntMap' from a 'IntMap'. Returns +-- 'Nothing' if the 'IntMap' was originally actually empty, and @'Just' n@ +-- with an 'NEIntMap', if the 'IntMap' was not empty. +-- +-- 'nonEmptyMap' and @'maybe' 'Data.IntMap.empty' 'toMap'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- See 'Data.IntMap.NonEmpty.IsNonEmpty' for a pattern synonym that lets you +-- "match on" the possiblity of a 'IntMap' being an 'NEIntMap'. +-- +-- > nonEmptyMap (Data.IntMap.fromList [(3,"a"), (5,"b")]) == Just (fromList ((3,"a") :| [(5,"b")])) +nonEmptyMap :: IntMap a -> Maybe (NEIntMap a) +nonEmptyMap = (fmap . uncurry . uncurry) NEIntMap . M.minViewWithKey +{-# INLINE nonEmptyMap #-} + +-- | /O(log n)/. A general continuation-based way to consume a 'IntMap' as if +-- it were an 'NEIntMap'. @'withNonEmpty' def f@ will take a 'IntMap'. If map is +-- empty, it will evaluate to @def@. Otherwise, a non-empty map 'NEIntMap' +-- will be fed to the function @f@ instead. +-- +-- @'nonEmptyMap' == 'withNonEmpty' 'Nothing' 'Just'@ +withNonEmpty + :: r -- ^ value to return if map is empty + -> (NEIntMap a -> r) -- ^ function to apply if map is not empty + -> IntMap a + -> r +withNonEmpty def f = maybe def f . nonEmptyMap +{-# INLINE withNonEmpty #-} + +-- | /O(n*log n)/. Build a non-empty map from a non-empty list of +-- key\/value pairs. See also 'Data.IntMap.NonEmpty.fromAscList'. If the list +-- contains more than one value for the same key, the last value for the +-- key is retained. +-- +-- > fromList ((5,"a") :| [(3,"b"), (5, "c")]) == fromList ((5,"c") :| [(3,"b")]) +-- > fromList ((5,"c") :| [(3,"b"), (5, "a")]) == fromList ((5,"a") :| [(3,"b")]) + +-- TODO: write manually and optimize to be equivalent to +-- 'fromDistinctAscList' if items are ordered, just like the actual +-- 'M.fromList'. +fromList :: NonEmpty (Key, a) -> NEIntMap a +fromList ((k, v) :| xs) = withNonEmpty (singleton k v) (insertWith (const id) k v) + . M.fromList + $ xs +{-# INLINE fromList #-} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList ((1, 'a') :| []) +-- > size (singleton 1 'a') == 1 +singleton :: Key -> a -> NEIntMap a +singleton k v = NEIntMap k v M.empty +{-# INLINE singleton #-} + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ will insert the pair (key, value) into +-- @mp@ if key does not exist in the map. If the key does exist, the +-- function will insert the pair @(key, f new_value old_value)@. +-- +-- See 'Data.IntMap.NonEmpty.insertIntMapWith' for a version where the first +-- argument is a 'IntMap'. +-- +-- > insertWith (++) 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "xxxa")]) +-- > insertWith (++) 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +insertWith + :: (a -> a -> a) + -> Key + -> a + -> NEIntMap a + -> NEIntMap a +insertWith f k v n@(NEIntMap k0 v0 m) = case compare k k0 of + LT -> NEIntMap k v . toMap $ n + EQ -> NEIntMap k (f v v0) m + GT -> NEIntMap k0 v0 $ M.insertWith f k v m +{-# INLINE insertWith #-} + + +-- | Left-biased union +instance Semigroup (NEIntMap a) where + (<>) = union + {-# INLINE (<>) #-} + sconcat = unions + {-# INLINE sconcat #-} + +instance Functor NEIntMap where + fmap = map + {-# INLINE fmap #-} + x <$ NEIntMap k _ m = NEIntMap k x (x <$ m) + {-# INLINE (<$) #-} + +-- | Traverses elements in order of ascending keys. +-- +-- __WARNING:__ 'F.fold' and 'F.foldMap' are different than for the +-- 'IntMap' instance. They traverse elements in order of ascending keys, +-- while 'IntMap' traverses positive keys first, then negative keys. +-- +-- 'Data.Foldable.foldr1', 'Data.Foldable.foldl1', 'Data.Foldable.minimum', +-- 'Data.Foldable.maximum' are all total. +instance Foldable NEIntMap where +#if MIN_VERSION_base(4,11,0) + fold (NEIntMap _ v m) = v <> F.fold (M.elems m) + {-# INLINE fold #-} + foldMap f (NEIntMap _ v m) = f v <> foldMap f (M.elems m) + {-# INLINE foldMap #-} +#else + fold (NEIntMap _ v m) = v `mappend` F.fold (M.elems m) + {-# INLINE fold #-} + foldMap f (NEIntMap _ v m) = f v `mappend` foldMap f (M.elems m) + {-# INLINE foldMap #-} +#endif + foldr = foldr + {-# INLINE foldr #-} + foldr' = foldr' + {-# INLINE foldr' #-} + foldr1 = foldr1 + {-# INLINE foldr1 #-} + foldl = foldl + {-# INLINE foldl #-} + foldl' = foldl' + {-# INLINE foldl' #-} + foldl1 = foldl1 + {-# INLINE foldl1 #-} + null _ = False + {-# INLINE null #-} + length = size + {-# INLINE length #-} + elem x (NEIntMap _ v m) = F.elem x m + || x == v + {-# INLINE elem #-} + -- TODO: use build + toList = F.toList . elems + {-# INLINE toList #-} + +-- | Traverses elements in order of ascending keys +-- +-- __WARNING:__ Different than for the 'IntMap' instance. They traverse +-- elements in order of ascending keys, while 'IntMap' traverses positive +-- keys first, then negative keys. +instance Traversable NEIntMap where + traverse f = traverseWithKey (const f) + {-# INLINE traverse #-} + +-- | Traverses elements in order of ascending keys +-- +-- __WARNING:__ 'F1.fold1' and 'F1.foldMap1' are different than 'F.fold' and +-- 'F.foldMap' for the 'IntMap' instance of 'Foldable'. They traverse +-- elements in order of ascending keys, while 'IntMap' traverses positive +-- keys first, then negative keys. +instance Foldable1 NEIntMap where + fold1 (NEIntMap _ v m) = maybe v (v <>) + . getOption + . F.foldMap (Option . Just) + . M.elems + $ m + {-# INLINE fold1 #-} + foldMap1 f = foldMapWithKey (const f) + {-# INLINE foldMap1 #-} + toNonEmpty = elems + {-# INLINE toNonEmpty #-} + +-- | Traverses elements in order of ascending keys +-- +-- __WARNING:__ 'traverse1' and 'sequence1' are different 'traverse' and +-- 'sequence' for the 'IntMap' instance of 'Traversable'. They traverse +-- elements in order of ascending keys, while 'IntMap' traverses positive +-- keys first, then negative keys. +instance Traversable1 NEIntMap where + traverse1 f = traverseWithKey1 (const f) + {-# INLINE traverse1 #-} + +-- | /O(n)/. Test if the internal map structure is valid. +valid :: NEIntMap a -> Bool +valid (NEIntMap k _ m) = all ((k <) . fst . fst) (M.minViewWithKey m) + + + + + +-- | /O(log n)/. Insert new key and value into a map where keys are +-- /strictly greater than/ the new key. That is, the new key must be +-- /strictly less than/ all keys present in the 'IntMap'. /The precondition +-- is not checked./ +-- +-- At the moment this is simply an alias for @Data.IntSet.insert@, but it's +-- left here as a placeholder in case this eventually gets implemented in +-- a more efficient way. + +-- TODO: implementation +insertMinMap :: Key -> a -> IntMap a -> IntMap a +insertMinMap = M.insert +{-# INLINABLE insertMinMap #-} + +-- | /O(log n)/. Insert new key and value into a map where keys are +-- /strictly less than/ the new key. That is, the new key must be +-- /strictly greater than/ all keys present in the 'IntMap'. /The +-- precondition is not checked./ +-- +-- At the moment this is simply an alias for @Data.IntSet.insert@, but it's +-- left here as a placeholder in case this eventually gets implemented in +-- a more efficient way. + +-- TODO: implementation +insertMaxMap :: Key -> a -> IntMap a -> IntMap a +insertMaxMap = M.insert +{-# INLINABLE insertMaxMap #-} + +-- | /O(n)/. A fixed version of 'Data.IntMap.traverseWithKey' that +-- traverses items in ascending order of keys. +traverseMapWithKey :: Applicative t => (Key -> a -> t b) -> IntMap a -> t (IntMap b) +traverseMapWithKey f = go + where + go Nil = pure Nil + go (Tip k v) = Tip k <$> f k v + go (Bin p m l r) = liftA2 (flip (Bin p m)) (go r) (go l) +{-# INLINE traverseMapWithKey #-} + +-- --------------------------------------------- +-- | CPP for new functions not in old containers +-- --------------------------------------------- + +-- | Compatibility layer for 'Data.IntMap.Lazy.lookupMinMap'. +lookupMinMap :: IntMap a -> Maybe (Key, a) +#if MIN_VERSION_containers(0,5,11) +lookupMinMap = M.lookupMin +#else +lookupMinMap = fmap fst . M.minViewWithKey +#endif +{-# INLINE lookupMinMap #-} + +-- | Compatibility layer for 'Data.IntMap.Lazy.lookupMaxMap'. +lookupMaxMap :: IntMap a -> Maybe (Key, a) +#if MIN_VERSION_containers(0,5,11) +lookupMaxMap = M.lookupMax +#else +lookupMaxMap = fmap fst . M.maxViewWithKey +#endif +{-# INLINE lookupMaxMap #-} + diff --git a/src/Data/IntSet/NonEmpty.hs b/src/Data/IntSet/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/IntSet/NonEmpty.hs @@ -0,0 +1,755 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.IntSet.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Finite Integer Sets +-- +-- The 'NEIntSet' type represents a non-empty set of integers. +-- +-- See documentation for 'NEIntSet' for information on how to convert and +-- manipulate such non-empty set. +-- +-- This module essentially re-imports the API of "Data.IntSet" and its 'IntSet' +-- type, along with semantics and asymptotics. In most situations, +-- asymptotics are different only by a constant factor. In some +-- situations, asmyptotics are even better (constant-time instead of +-- log-time). +-- +-- Because 'NEIntSet' is implemented using 'IntSet', all of the caveats of +-- using 'IntSet' apply (such as the limitation of the maximum size of +-- sets). +-- +-- All functions take non-empty sets as inputs. In situations where their +-- results can be guarunteed to also be non-empty, they also return +-- non-empty sets. In situations where their results could potentially be +-- empty, 'IntSet' is returned instead. +-- +-- Some functions ('partition', 'split') have modified return types to +-- account for possible configurations of non-emptiness. +-- +-- This module is intended to be imported qualified, to avoid name clashes +-- with "Prelude" and "Data.IntSet" functions: +-- +-- > import qualified Data.IntSet.NonEmpty as NEIS +-- +-- Note that all asmyptotics /O(f(n))/ in this module are actually +-- /O(min(W, f(n)))/, where @W@ is the number of bits in an 'Int' (32 or +-- 64). That is, if @f(n)@ is greater than @W@, all operations are +-- constant-time. +module Data.IntSet.NonEmpty ( + -- * Non-Empty Set Type + NEIntSet + , Key + + -- ** Conversions between empty and non-empty sets + , pattern IsNonEmpty + , pattern IsEmpty + , nonEmptySet + , toSet + , withNonEmpty + , insertSet + , insertSetMin + , insertSetMax + , unsafeFromSet + + -- * Construction + , singleton + , fromList + , fromAscList + , fromDistinctAscList + + -- * Insertion + , insert + + -- * Deletion + , delete + + -- * Query + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , size + , isSubsetOf + , isProperSubsetOf + , disjoint + + -- * Combine + , union + , unions + , difference + , (\\) + , intersection + + -- * Filter + , filter + , partition + , split + , splitMember + , splitRoot + + -- * Map + , map + + -- * Folds + , foldr + , foldl + , foldr1 + , foldl1 + -- ** Strict folds + , foldr' + , foldl' + , foldr1' + , foldl1' + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + + -- * Conversion + + -- ** List + , elems + , toList + , toAscList + , toDescList + + -- * Debugging + , valid + ) where + + +import Control.Applicative +import Data.Bifunctor +import Data.IntSet (IntSet) +import Data.IntSet.NonEmpty.Internal +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Maybe +import Data.These +import Prelude hiding (foldr, foldl, foldr1, foldl1, filter, map) +import qualified Data.IntSet as S +import qualified Data.List.NonEmpty as NE + +-- | /O(1)/ match, /O(log n)/ usage of contents. The 'IsNonEmpty' and +-- 'IsEmpty' patterns allow you to treat a 'IntSet' as if it were either +-- a @'IsNonEmpty' n@ (where @n@ is a 'NEIntSet') or an 'IsEmpty'. +-- +-- For example, you can pattern match on a 'IntSet': +-- +-- @ +-- myFunc :: 'IntSet' X -> Y +-- myFunc ('IsNonEmpty' n) = -- here, the user provided a non-empty set, and @n@ is the 'NEIntSet' +-- myFunc 'IsEmpty' = -- here, the user provided an empty set +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original 'IntSet' was /not/ +-- empty, and you have a verified-non-empty 'NEIntSet' @n@ to use. +-- +-- Note that patching on this pattern is /O(1)/. However, using the +-- contents requires a /O(log n)/ cost that is deferred until after the +-- pattern is matched on (and is not incurred at all if the contents are +-- never used). +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a 'NEIntSet' back into a 'IntSet', obscuring its non-emptiness (see 'toSet'). +pattern IsNonEmpty :: NEIntSet -> IntSet +pattern IsNonEmpty n <- (nonEmptySet->Just n) + where + IsNonEmpty n = toSet n + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'IntSet' as if it were either a @'IsNonEmpty' n@ (where @n@ is +-- a 'NEIntSet') or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original 'IntSet' was empty. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'Data.IntSet.empty'. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: IntSet +pattern IsEmpty <- (S.null->True) + where + IsEmpty = S.empty + +{-# COMPLETE IsNonEmpty, IsEmpty #-} + +-- | /O(log n)/. Convert a 'IntSet' into an 'NEIntSet' by adding a value. +-- Because of this, we know that the set must have at least one +-- element, and so therefore cannot be empty. +-- +-- See 'insertSetMin' for a version that is constant-time if the new +-- value is /strictly smaller than/ all values in the original set +-- +-- > insertSet 4 (Data.IntSet.fromList [5, 3]) == fromList (3 :| [4, 5]) +-- > insertSet 4 Data.IntSet.empty == singleton 4 "c" +insertSet :: Key -> IntSet -> NEIntSet +insertSet x = withNonEmpty (singleton x) (insert x) +{-# INLINE insertSet #-} + +-- | /O(1)/ Convert a 'IntSet' into an 'NEIntSet' by adding a value where the +-- value is /strictly less than/ all values in the input set The values in +-- the original map must all be /strictly greater than/ the new value. +-- /The precondition is not checked./ +-- +-- > insertSetMin 2 (Data.IntSet.fromList [5, 3]) == fromList (2 :| [3, 5]) +-- > valid (insertSetMin 2 (Data.IntSet.fromList [5, 3])) == True +-- > valid (insertSetMin 7 (Data.IntSet.fromList [5, 3])) == False +-- > valid (insertSetMin 3 (Data.IntSet.fromList [5, 3])) == False +insertSetMin :: Key -> IntSet -> NEIntSet +insertSetMin = NEIntSet +{-# INLINE insertSetMin #-} + +-- | /O(log n)/ Convert a 'IntSet' into an 'NEIntSet' by adding a value +-- where the value is /strictly less than/ all values in the input set The +-- values in the original map must all be /strictly greater than/ the new +-- value. /The precondition is not checked./ +-- +-- At the current moment, this is identical simply 'insertSet'; however, +-- it is left both for consistency and as a placeholder for a future +-- version where optimizations are implemented to allow for a faster +-- implementation. +-- +-- > insertSetMin 7 (Data.IntSet.fromList [5, 3]) == fromList (3 :| [5, 7]) + +-- these currently are all valid, but shouldn't be +-- > valid (insertSetMin 7 (Data.IntSet.fromList [5, 3])) == True +-- > valid (insertSetMin 2 (Data.IntSet.fromList [5, 3])) == False +-- > valid (insertSetMin 5 (Data.IntSet.fromList [5, 3])) == False +insertSetMax :: Key -> IntSet -> NEIntSet +insertSetMax x = withNonEmpty (singleton x) go + where + go (NEIntSet x0 s0) = NEIntSet x0 . insertMaxSet x $ s0 +{-# INLINE insertSetMax #-} + +-- | /O(log n)/. Unsafe version of 'nonEmptySet'. Coerces a 'IntSet' +-- into an 'NEIntSet', but is undefined (throws a runtime exception when +-- evaluation is attempted) for an empty 'IntSet'. +unsafeFromSet + :: IntSet + -> NEIntSet +unsafeFromSet = withNonEmpty e id + where + e = errorWithoutStackTrace "NEIntSet.unsafeFromSet: empty set" +{-# INLINE unsafeFromSet #-} + +-- | /O(n)/. Build a set from an ascending list in linear time. /The +-- precondition (input list is ascending) is not checked./ +fromAscList :: NonEmpty Key -> NEIntSet +fromAscList = fromDistinctAscList . combineEq +{-# INLINE fromAscList #-} + +-- | /O(n)/. Build a set from an ascending list of distinct elements in linear time. +-- /The precondition (input list is strictly ascending) is not checked./ +fromDistinctAscList :: NonEmpty Key -> NEIntSet +fromDistinctAscList (x :| xs) = insertSetMin x + . S.fromDistinctAscList + $ xs +{-# INLINE fromDistinctAscList #-} + +-- | /O(log n)/. Insert an element in a set. +-- If the set already contains an element equal to the given value, +-- it is replaced with the new value. +insert :: Key -> NEIntSet -> NEIntSet +insert x n@(NEIntSet x0 s) = case compare x x0 of + LT -> NEIntSet x $ toSet n + EQ -> NEIntSet x s + GT -> NEIntSet x0 $ S.insert x s +{-# INLINE insert #-} + +-- | /O(log n)/. Delete an element from a set. +delete :: Key -> NEIntSet -> IntSet +delete x n@(NEIntSet x0 s) = case compare x x0 of + LT -> toSet n + EQ -> s + GT -> insertMinSet x0 . S.delete x $ s +{-# INLINE delete #-} + +-- | /O(log n)/. Is the element in the set? +member :: Key -> NEIntSet -> Bool +member x (NEIntSet x0 s) = case compare x x0 of + LT -> False + EQ -> True + GT -> S.member x s +{-# INLINE member #-} + +-- | /O(log n)/. Is the element not in the set? +notMember :: Key -> NEIntSet -> Bool +notMember x (NEIntSet x0 s) = case compare x x0 of + LT -> True + EQ -> False + GT -> S.notMember x s +{-# INLINE notMember #-} + +-- | /O(log n)/. Find largest element smaller than the given one. +-- +-- > lookupLT 3 (fromList (3 :| [5])) == Nothing +-- > lookupLT 5 (fromList (3 :| [5])) == Just 3 +lookupLT :: Key -> NEIntSet -> Maybe Key +lookupLT x (NEIntSet x0 s) = case compare x x0 of + LT -> Nothing + EQ -> Nothing + GT -> S.lookupLT x s <|> Just x0 +{-# INLINE lookupLT #-} + +-- | /O(log n)/. Find smallest element greater than the given one. +-- +-- > lookupLT 4 (fromList (3 :| [5])) == Just 5 +-- > lookupLT 5 (fromList (3 :| [5])) == Nothing +lookupGT :: Key -> NEIntSet -> Maybe Key +lookupGT x (NEIntSet x0 s) = case compare x x0 of + LT -> Just x0 + EQ -> fst <$> S.minView s + GT -> S.lookupGT x s +{-# INLINE lookupGT #-} + +-- | /O(log n)/. Find largest element smaller or equal to the given one. +-- +-- > lookupLT 2 (fromList (3 :| [5])) == Nothing +-- > lookupLT 4 (fromList (3 :| [5])) == Just 3 +-- > lookupLT 5 (fromList (3 :| [5])) == Just 5 +lookupLE :: Key -> NEIntSet -> Maybe Key +lookupLE x (NEIntSet x0 s) = case compare x x0 of + LT -> Nothing + EQ -> Just x0 + GT -> S.lookupLE x s <|> Just x0 +{-# INLINE lookupLE #-} + +-- | /O(log n)/. Find smallest element greater or equal to the given one. +-- +-- > lookupLT 3 (fromList (3 :| [5])) == Just 3 +-- > lookupLT 4 (fromList (3 :| [5])) == Just 5 +-- > lookupLT 6 (fromList (3 :| [5])) == Nothing +lookupGE :: Key -> NEIntSet -> Maybe Key +lookupGE x (NEIntSet x0 s) = case compare x x0 of + LT -> Just x0 + EQ -> Just x0 + GT -> S.lookupGE x s +{-# INLINE lookupGE #-} + +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'Data.IntSet.NonEmpty.toAscList'@. +-- +-- For example, +-- +-- > elemsList set = foldr (:) [] set +foldr :: (Key -> b -> b) -> b -> NEIntSet -> b +foldr f z (NEIntSet x s) = x `f` S.foldr f z s +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (Key -> b -> b) -> b -> NEIntSet -> b +foldr' f z (NEIntSet x s) = x `f` y + where + !y = S.foldr' f z s +{-# INLINE foldr' #-} + +-- | /O(n)/. A version of 'foldr' that uses the value at the maximal value +-- in the set as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldr1' for 'IntSet', this function is +-- total if the input function is total. +foldr1 :: (Key -> Key -> Key) -> NEIntSet -> Key +foldr1 f (NEIntSet x s) = maybe x (f x . uncurry (S.foldr f)) + . S.maxView + $ s +{-# INLINE foldr1 #-} + +-- | /O(n)/. Fold the elements in the set using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'Data.IntSet.NonEmpty.toAscList'@. +-- +-- For example, +-- +-- > descElemsList set = foldl (flip (:)) [] set +foldl :: (a -> Key -> a) -> a -> NEIntSet -> a +foldl f z (NEIntSet x s) = S.foldl f (f z x) s +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> Key -> a) -> a -> NEIntSet -> a +foldl' f z (NEIntSet x s) = S.foldl' f y s + where + !y = f z x +{-# INLINE foldl' #-} + +-- | /O(n)/. A version of 'foldl' that uses the value at the minimal value +-- in the set as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldl1' for 'IntSet', this function is +-- total if the input function is total. +foldl1 :: (Key -> Key -> Key) -> NEIntSet -> Key +foldl1 f (NEIntSet x s) = S.foldl f x s +{-# INLINE foldl1 #-} + +-- | /O(n)/. A strict version of 'foldr1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr1' :: (Key -> Key -> Key) -> NEIntSet -> Key +foldr1' f (NEIntSet x s) = case S.maxView s of + Nothing -> x + Just (y, s') -> let !z = S.foldr' f y s' in x `f` z +{-# INLINE foldr1' #-} + +-- | /O(n)/. A strict version of 'foldl1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl1' :: (Key -> Key -> Key) -> NEIntSet -> Key +foldl1' f (NEIntSet x s) = S.foldl' f x s +{-# INLINE foldl1' #-} + +-- | /O(1)/. The number of elements in the set. Guaranteed to be greater +-- than zero. +size :: NEIntSet -> Int +size (NEIntSet _ s) = 1 + S.size s +{-# INLINE size #-} + +-- | /O(n+m)/. Is this a subset? +-- @(s1 \`isSubsetOf\` s2)@ tells whether @s1@ is a subset of @s2@. +isSubsetOf + :: NEIntSet + -> NEIntSet + -> Bool +isSubsetOf (NEIntSet x s0) (toSet->s1) = x `S.member` s1 + && s0 `S.isSubsetOf` s1 +{-# INLINE isSubsetOf #-} + +-- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). +isProperSubsetOf + :: NEIntSet + -> NEIntSet + -> Bool +isProperSubsetOf s0 s1 = S.size (neisIntSet s0) < S.size (neisIntSet s1) + && s0 `isSubsetOf` s1 +{-# INLINE isProperSubsetOf #-} + +-- | /O(n+m)/. Check whether two sets are disjoint (i.e. their intersection +-- is empty). +-- +-- > disjoint (fromList (2:|[4,6])) (fromList (1:|[3])) == True +-- > disjoint (fromList (2:|[4,6,8])) (fromList (2:|[3,5,7])) == False +-- > disjoint (fromList (1:|[2])) (fromList (1:|[2,3,4])) == False +disjoint + :: NEIntSet + -> NEIntSet + -> Bool +disjoint n1@(NEIntSet x1 s1) n2@(NEIntSet x2 s2) = case compare x1 x2 of + -- x1 is not in n2 + LT -> s1 `disjointSet` toSet n2 + -- k1 and k2 are a part of the result + EQ -> False + -- k2 is not in n1 + GT -> toSet n1 `disjointSet` s2 +{-# INLINE disjoint #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two sets. +-- +-- Returns a potentially empty set ('IntSet') because the first set might be +-- a subset of the second set, and therefore have all of its elements +-- removed. +difference + :: NEIntSet + -> NEIntSet + -> IntSet +difference n1@(NEIntSet x1 s1) n2@(NEIntSet x2 s2) = case compare x1 x2 of + -- x1 is not in n2, so cannot be deleted + LT -> insertMinSet x1 $ s1 `S.difference` toSet n2 + -- x2 deletes x1, and only x1 + EQ -> s1 `S.difference` s2 + -- x2 is not in n1, so cannot delete anything, so we can just difference n1 // s2. + GT -> toSet n1 `S.difference` s2 +{-# INLINE difference #-} + +-- | Same as 'difference'. +(\\) + :: NEIntSet + -> NEIntSet + -> IntSet +(\\) = difference +{-# INLINE (\\) #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. The intersection of two sets. +-- +-- Returns a potentially empty set ('IntSet'), because the two sets might have +-- an empty intersection. +-- +-- Elements of the result come from the first set, so for example +-- +-- > import qualified Data.IntSet.NonEmpty as NES +-- > data AB = A | B deriving Show +-- > instance Ord AB where compare _ _ = EQ +-- > instance Eq AB where _ == _ = True +-- > main = print (NES.singleton A `NES.intersection` NES.singleton B, +-- > NES.singleton B `NES.intersection` NES.singleton A) +-- +-- prints @(fromList (A:|[]),fromList (B:|[]))@. +intersection + :: NEIntSet + -> NEIntSet + -> IntSet +intersection n1@(NEIntSet x1 s1) n2@(NEIntSet x2 s2) = case compare x1 x2 of + -- x1 is not in n2 + LT -> s1 `S.intersection` toSet n2 + -- x1 and x2 are a part of the result + EQ -> insertMinSet x1 $ s1 `S.intersection` s2 + -- x2 is not in n1 + GT -> toSet n1 `S.intersection` s2 +{-# INLINE intersection #-} + +-- | /O(n)/. Filter all elements that satisfy the predicate. +-- +-- Returns a potentially empty set ('IntSet') because the predicate might +-- filter out all items in the original non-empty set. +filter + :: (Key -> Bool) + -> NEIntSet + -> IntSet +filter f (NEIntSet x s1) + | f x = insertMinSet x . S.filter f $ s1 + | otherwise = S.filter f s1 +{-# INLINE filter #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty sets: +-- +-- * @'This' n1@ means that the predicate was true for all items. +-- * @'That' n2@ means that the predicate was false for all items. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partition (> 3) (fromList (5 :| [3])) == These (singleton 5) (singleton 3) +-- > partition (< 7) (fromList (5 :| [3])) == This (fromList (3 :| [5])) +-- > partition (> 7) (fromList (5 :| [3])) == That (fromList (3 :| [5])) +partition + :: (Key -> Bool) + -> NEIntSet + -> These NEIntSet NEIntSet +partition f n@(NEIntSet x s0) = case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) + | f x -> This n + | otherwise -> That n + (Just n1, Nothing) + | f x -> This n + | otherwise -> These n1 (singleton x) + (Nothing, Just n2) + | f x -> These (singleton x) n2 + | otherwise -> That n + (Just n1, Just n2) + | f x -> These (insertSetMin x s1) n2 + | otherwise -> These n1 (insertSetMin x s2) + where + (s1, s2) = S.partition f s0 +{-# INLINABLE partition #-} + +-- | /O(log n)/. The expression (@'split' x set@) is potentially a 'These' +-- containing up to two 'NEIntSet's based on splitting the set into sets +-- containing items before and after the value @x@. It will never return +-- a set that contains @x@ itself. +-- +-- * 'Nothing' means that @x@ was the only value in the the original set, +-- and so there are no items before or after it. +-- * @'Just' ('This' n1)@ means @x@ was larger than or equal to all items +-- in the set, and @n1@ is the entire original set (minus @x@, if it +-- was present) +-- * @'Just' ('That' n2)@ means @x@ was smaller than or equal to all +-- items in the set, and @n2@ is the entire original set (minus @x@, if +-- it was present) +-- * @'Just' ('These' n1 n2)@ gives @n1@ (the set of all values from the +-- original set less than @x@) and @n2@ (the set of all values from the +-- original set greater than @x@). +-- +-- > split 2 (fromList (5 :| [3])) == Just (That (fromList (3 :| [5])) ) +-- > split 3 (fromList (5 :| [3])) == Just (That (singleton 5) ) +-- > split 4 (fromList (5 :| [3])) == Just (These (singleton 3) (singleton 5)) +-- > split 5 (fromList (5 :| [3])) == Just (This (singleton 3) ) +-- > split 6 (fromList (5 :| [3])) == Just (This (fromList (3 :| [5])) ) +-- > split 5 (singleton 5) == Nothing +split + :: Key + -> NEIntSet + -> Maybe (These NEIntSet NEIntSet) +split x n@(NEIntSet x0 s0) = case compare x x0 of + LT -> Just $ That n + EQ -> That <$> nonEmptySet s0 + GT -> case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> Just $ This (singleton x0) + (Just _ , Nothing) -> Just $ This (insertSetMin x0 s1) + (Nothing, Just n2) -> Just $ These (singleton x0) n2 + (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2 + where + (s1, s2) = S.split x s0 +{-# INLINABLE split #-} + +-- | /O(log n)/. The expression (@'splitMember' x set@) splits a set just +-- like 'split' but also returns @'member' x set@ (whether or not @x@ was +-- in @set@) +-- +-- > splitMember 2 (fromList (5 :| [3])) == (False, Just (That (fromList (3 :| [5)])))) +-- > splitMember 3 (fromList (5 :| [3])) == (True , Just (That (singleton 5))) +-- > splitMember 4 (fromList (5 :| [3])) == (False, Just (These (singleton 3) (singleton 5))) +-- > splitMember 5 (fromList (5 :| [3])) == (True , Just (This (singleton 3)) +-- > splitMember 6 (fromList (5 :| [3])) == (False, Just (This (fromList (3 :| [5]))) +-- > splitMember 5 (singleton 5) == (True , Nothing) +splitMember + :: Key + -> NEIntSet + -> (Bool, Maybe (These NEIntSet NEIntSet)) +splitMember x n@(NEIntSet x0 s0) = case compare x x0 of + LT -> (False, Just $ That n) + EQ -> (True , That <$> nonEmptySet s0) + GT -> (mem ,) $ case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> Just $ This (singleton x0) + (Just _ , Nothing) -> Just $ This (insertSetMin x0 s1) + (Nothing, Just n2) -> Just $ These (singleton x0) n2 + (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2 + where + (s1, mem, s2) = S.splitMember x s0 +{-# INLINABLE splitMember #-} + +-- | /O(1)/. Decompose a set into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a set in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that +-- the pieces returned will be in ascending order (all elements in the +-- first subset less than all elements in the second, and so on). +-- +-- Note that the current implementation does not return more than four +-- subsets, but you should not depend on this behaviour because it can +-- change in the future without notice. +splitRoot + :: NEIntSet + -> NonEmpty NEIntSet +splitRoot (NEIntSet x s) = singleton x + :| mapMaybe nonEmptySet (S.splitRoot s) +{-# INLINE splitRoot #-} + +-- | /O(n*log n)/. +-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. +-- +-- It's worth noting that the size of the result may be smaller if, +-- for some @(x,y)@, @x \/= y && f x == f y@ +map :: (Key -> Key) + -> NEIntSet + -> NEIntSet +map f (NEIntSet x0 s) = fromList + . (f x0 :|) + . S.foldr (\x xs -> f x : xs) [] + $ s +{-# INLINE map #-} + +-- | /O(1)/. The minimal element of a set. Note that this is total, making +-- 'Data.IntSet.lookupMin' obsolete. It is constant-time, so has better +-- asymptotics than @Data.IntSet.lookupMin@ and @Data.Map.findMin@ as well. +-- +-- > findMin (fromList (5 :| [3])) == 3 +findMin :: NEIntSet -> Key +findMin (NEIntSet x _) = x +{-# INLINE findMin #-} + +-- | /O(log n)/. The maximal key of a set Note that this is total, +-- making 'Data.IntSet.lookupMin' obsolete. +-- +-- > findMax (fromList (5 :| [3])) == 5 +findMax :: NEIntSet -> Key +findMax (NEIntSet x s) = maybe x fst . S.maxView $ s +{-# INLINE findMax #-} + +-- | /O(1)/. Delete the minimal element. Returns a potentially empty set +-- ('IntSet'), because we might delete the final item in a singleton set. It +-- is constant-time, so has better asymptotics than @Data.IntSet.deleteMin@. +-- +-- > deleteMin (fromList (5 :| [3, 7])) == Data.IntSet.fromList [5, 7] +-- > deleteMin (singleton 5) == Data.IntSet.empty +deleteMin :: NEIntSet -> IntSet +deleteMin (NEIntSet _ s) = s +{-# INLINE deleteMin #-} + +-- | /O(log n)/. Delete the maximal element. Returns a potentially empty +-- set ('IntSet'), because we might delete the final item in a singleton set. +-- +-- > deleteMax (fromList (5 :| [3, 7])) == Data.IntSet.fromList [3, 5] +-- > deleteMax (singleton 5) == Data.IntSet.empty +deleteMax :: NEIntSet -> IntSet +deleteMax (NEIntSet x s) = insertMinSet x . S.deleteMax $ s +{-# INLINE deleteMax #-} + +-- | /O(1)/. Delete and find the minimal element. It is constant-time, so +-- has better asymptotics that @Data.IntSet.minView@ for 'IntSet'. +-- +-- Note that unlike @Data.IntSet.deleteFindMin@ for 'IntSet', this cannot ever +-- fail, and so is a total function. However, the result 'IntSet' is +-- potentially empty, since the original set might have contained just +-- a single item. +-- +-- > deleteFindMin (fromList (5 :| [3, 10])) == (3, Data.IntSet.fromList [5, 10]) +deleteFindMin :: NEIntSet -> (Key, IntSet) +deleteFindMin (NEIntSet x s) = (x, s) +{-# INLINE deleteFindMin #-} + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- Note that unlike @Data.IntSet.deleteFindMax@ for 'IntSet', this cannot ever +-- fail, and so is a total function. However, the result 'IntSet' is +-- potentially empty, since the original set might have contained just +-- a single item. +-- +-- > deleteFindMax (fromList (5 :| [3, 10])) == (10, Data.IntSet.fromList [3, 5]) +deleteFindMax :: NEIntSet -> (Key, IntSet) +deleteFindMax (NEIntSet x s) = maybe (x, S.empty) (second (insertMinSet x)) + . S.maxView + $ s +{-# INLINE deleteFindMax #-} + +-- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending +-- order. +elems :: NEIntSet -> NonEmpty Key +elems = toList +{-# INLINE elems #-} + +-- | /O(n)/. Convert the set to an ascending non-empty list of elements. +toAscList :: NEIntSet -> NonEmpty Key +toAscList = toList +{-# INLINE toAscList #-} + +-- | /O(n)/. Convert the set to a descending non-empty list of elements. +toDescList :: NEIntSet -> NonEmpty Key +toDescList (NEIntSet x s) = S.foldl' (flip (NE.<|)) (x :| []) s +{-# INLINE toDescList #-} + +combineEq :: NonEmpty Key -> NonEmpty Key +combineEq (x :| xs) = go x xs + where + go z [] = z :| [] + go z (y:ys) + | z == y = go z ys + | otherwise = z NE.<| go y ys diff --git a/src/Data/IntSet/NonEmpty/Internal.hs b/src/Data/IntSet/NonEmpty/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/IntSet/NonEmpty/Internal.hs @@ -0,0 +1,281 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Data.IntSet.NonEmpty.Internal +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Unsafe internal-use functions used in the implementation of +-- "Data.IntSet.NonEmpty". These functions can potentially be used to break +-- the abstraction of 'NEIntSet' and produce unsound sets, so be wary! +module Data.IntSet.NonEmpty.Internal ( + NEIntSet(..) + , Key + , nonEmptySet + , withNonEmpty + , toSet + , singleton + , fromList + , toList + , union + , unions + , valid + , insertMinSet + , insertMaxSet + , disjointSet + ) where + +import Control.DeepSeq +import Data.Data +import Data.Function +import Data.IntSet.Internal (IntSet(..), Key) +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup +import Data.Semigroup.Foldable (Foldable1) +import Data.Typeable (Typeable) +import Text.Read +import qualified Data.Foldable as F +import qualified Data.IntSet as S +import qualified Data.Semigroup.Foldable as F1 + +-- | A non-empty (by construction) set of integers. At least one value +-- exists in an @'NEIntSet' a@ at all times. +-- +-- Functions that /take/ an 'NEIntSet' can safely operate on it with the +-- assumption that it has at least one item. +-- +-- Functions that /return/ an 'NEIntSet' provide an assurance that the +-- result has at least one item. +-- +-- "Data.IntSet.NonEmpty" re-exports the API of "Data.IntSet", faithfully +-- reproducing asymptotics, typeclass constraints, and semantics. +-- Functions that ensure that input and output sets are both non-empty +-- (like 'Data.IntSet.NonEmpty.insert') return 'NEIntSet', but functions that +-- might potentially return an empty map (like 'Data.IntSet.NonEmpty.delete') +-- return a 'IntSet' instead. +-- +-- You can directly construct an 'NEIntSet' with the API from +-- "Data.IntSet.NonEmpty"; it's more or less the same as constructing a normal +-- 'IntSet', except you don't have access to 'Data.IntSet.empty'. There are also +-- a few ways to construct an 'NEIntSet' from a 'IntSet': +-- +-- 1. The 'nonEmptySet' smart constructor will convert a @'IntSet' a@ into +-- a @'Maybe' ('NEIntSet' a)@, returning 'Nothing' if the original 'IntSet' +-- was empty. +-- 2. You can use the 'Data.IntSet.NonEmpty.insertIntSet' family of functions to +-- insert a value into a 'IntSet' to create a guaranteed 'NEIntSet'. +-- 3. You can use the 'Data.IntSet.NonEmpty.IsNonEmpty' and +-- 'Data.IntSet.NonEmpty.IsEmpty' patterns to "pattern match" on a 'IntSet' +-- to reveal it as either containing a 'NEIntSet' or an empty map. +-- 4. 'withNonEmpty' offers a continuation-based interface +-- for deconstructing a 'IntSet' and treating it as if it were an 'NEIntSet'. +-- +-- You can convert an 'NEIntSet' into a 'IntSet' with 'toSet' or +-- 'Data.IntSet.NonEmpty.IsNonEmpty', essentially "obscuring" the non-empty +-- property from the type. +data NEIntSet = + NEIntSet { neisV0 :: !Key -- ^ invariant: must be smaller than smallest value in set + , neisIntSet :: !IntSet + } + deriving (Typeable) + +instance Eq NEIntSet where + t1 == t2 = S.size (neisIntSet t1) == S.size (neisIntSet t2) + && toList t1 == toList t2 + +instance Ord NEIntSet where + compare = compare `on` toList + (<) = (<) `on` toList + (>) = (>) `on` toList + (<=) = (<=) `on` toList + (>=) = (>=) `on` toList + +instance Show NEIntSet where + showsPrec p xs = showParen (p > 10) $ + showString "fromList (" . shows (toList xs) . showString ")" + +instance Read NEIntSet where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- parens . prec 10 $ readPrec + return (fromList xs) + + readListPrec = readListPrecDefault + +instance NFData NEIntSet where + rnf (NEIntSet x s) = rnf x `seq` rnf s + +-- Data instance code from Data.IntSet.Internal +-- +-- Copyright : (c) Daan Leijen 2002 +-- (c) Joachim Breitner 2011 +instance Data NEIntSet where + gfoldl f z is = z fromList `f` (toList is) + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = intSetDataType + +fromListConstr :: Constr +fromListConstr = mkConstr intSetDataType "fromList" [] Prefix + +intSetDataType :: DataType +intSetDataType = mkDataType "Data.IntSet.NonEmpty.Internal.NEIntSet" [fromListConstr] + + + + + +-- | /O(log n)/. Smart constructor for an 'NEIntSet' from a 'IntSet'. Returns +-- 'Nothing' if the 'IntSet' was originally actually empty, and @'Just' n@ +-- with an 'NEIntSet', if the 'IntSet' was not empty. +-- +-- 'nonEmptySet' and @'maybe' 'Data.IntSet.empty' 'toSet'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- See 'Data.IntSet.NonEmpty.IsNonEmpty' for a pattern synonym that lets you +-- "match on" the possiblity of a 'IntSet' being an 'NEIntSet'. +-- +-- > nonEmptySet (Data.IntSet.fromList [3,5]) == Just (fromList (3:|[5])) +nonEmptySet :: IntSet -> Maybe NEIntSet +nonEmptySet = (fmap . uncurry) NEIntSet . S.minView +{-# INLINE nonEmptySet #-} + +-- | /O(log n)/. A general continuation-based way to consume a 'IntSet' as if +-- it were an 'NEIntSet'. @'withNonEmpty' def f@ will take a 'IntSet'. If set is +-- empty, it will evaluate to @def@. Otherwise, a non-empty set 'NEIntSet' +-- will be fed to the function @f@ instead. +-- +-- @'nonEmptySet' == 'withNonEmpty' 'Nothing' 'Just'@ +withNonEmpty + :: r -- ^ value to return if set is empty + -> (NEIntSet -> r) -- ^ function to apply if set is not empty + -> IntSet + -> r +withNonEmpty def f = maybe def f . nonEmptySet +{-# INLINE withNonEmpty #-} + +-- | /O(log n)/. +-- Convert a non-empty set back into a normal possibly-empty map, for usage +-- with functions that expect 'IntSet'. +-- +-- Can be thought of as "obscuring" the non-emptiness of the set in its +-- type. See the 'Data.IntSet.NonEmpty.IsNotEmpty' pattern. +-- +-- 'nonEmptySet' and @'maybe' 'Data.IntSet.empty' 'toSet'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- > toSet (fromList ((3,"a") :| [(5,"b")])) == Data.IntSet.fromList [(3,"a"), (5,"b")] +toSet :: NEIntSet -> IntSet +toSet (NEIntSet x s) = insertMinSet x s +{-# INLINE toSet #-} + +-- | /O(1)/. Create a singleton set. +singleton :: Key -> NEIntSet +singleton x = NEIntSet x S.empty +{-# INLINE singleton #-} + +-- | /O(n*log n)/. Create a set from a list of elements. + +-- TODO: write manually and optimize to be equivalent to +-- 'fromDistinctAscList' if items are ordered, just like the actual +-- 'S.fromList'. +fromList :: NonEmpty Key -> NEIntSet +fromList (x :| s) = withNonEmpty (singleton x) (<> singleton x) + . S.fromList + $ s +{-# INLINE fromList #-} + +-- | /O(n)/. Convert the set to a non-empty list of elements. +toList :: NEIntSet -> NonEmpty Key +toList (NEIntSet x s) = x :| S.toList s +{-# INLINE toList #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. The union of two sets, preferring the first set when +-- equal elements are encountered. +union + :: NEIntSet + -> NEIntSet + -> NEIntSet +union n1@(NEIntSet x1 s1) n2@(NEIntSet x2 s2) = case compare x1 x2 of + LT -> NEIntSet x1 . S.union s1 . toSet $ n2 + EQ -> NEIntSet x1 . S.union s1 $ s2 + GT -> NEIntSet x2 . S.union (toSet n1) $ s2 +{-# INLINE union #-} + +-- | The union of a non-empty list of sets +unions + :: Foldable1 f + => f NEIntSet + -> NEIntSet +unions (F1.toNonEmpty->(s :| ss)) = F.foldl' union s ss +{-# INLINE unions #-} + +-- | Left-biased union +instance Semigroup NEIntSet where + (<>) = union + {-# INLINE (<>) #-} + sconcat = unions + {-# INLINE sconcat #-} + +-- | /O(n)/. Test if the internal set structure is valid. +valid :: NEIntSet -> Bool +valid (NEIntSet x s) = all ((x <) . fst) (S.minView s) + + + + + + + +-- | /O(log n)/. Insert new value into a set where values are +-- /strictly greater than/ the new values That is, the new value must be +-- /strictly less than/ all values present in the 'IntSet'. /The precondition +-- is not checked./ +-- +-- At the moment this is simply an alias for @Data.IntSet.insert@, but it's +-- left here as a placeholder in case this eventually gets implemented in +-- a more efficient way. + +-- TODO: implementation +insertMinSet :: Key -> IntSet -> IntSet +insertMinSet = S.insert +{-# INLINABLE insertMinSet #-} + +-- | /O(log n)/. Insert new value into a set where values are /strictly +-- less than/ the new value. That is, the new value must be /strictly +-- greater than/ all values present in the 'IntSet'. /The precondition is not +-- checked./ +-- +-- At the moment this is simply an alias for @Data.IntSet.insert@, but it's +-- left here as a placeholder in case this eventually gets implemented in +-- a more efficient way. + +-- TODO: implementation +insertMaxSet :: Key -> IntSet -> IntSet +insertMaxSet = S.insert +{-# INLINABLE insertMaxSet #-} + +-- --------------------------------------------- +-- | CPP for new functions not in old containers +-- --------------------------------------------- + +-- | Comptability layer for 'Data.IntSet.disjoint'. +disjointSet :: IntSet -> IntSet -> Bool +#if MIN_VERSION_containers(0,5,11) +disjointSet = S.disjoint +#else +disjointSet xs = S.null . S.intersection xs +#endif +{-# INLINE disjointSet #-} + diff --git a/src/Data/Map/NonEmpty.hs b/src/Data/Map/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/Map/NonEmpty.hs @@ -0,0 +1,2372 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.Map.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Finite Maps (lazy interface) +-- +-- The @'NEMap' k v@ type represents a non-empty finite map (sometimes +-- called a dictionary) from keys of type @k@ to values of type @v@. +-- An 'NEMap' is strict in its keys but lazy in its values. +-- +-- See documentation for 'NEMap' for information on how to convert and +-- manipulate such non-empty maps. +-- +-- This module essentially re-imports the API of "Data.Map.Lazy" and its +-- 'Map' type, along with semantics and asymptotics. In most situations, +-- asymptotics are different only by a constant factor. In some +-- situations, asmyptotics are even better (constant-time instead of +-- log-time). All typeclass constraints are identical to their "Data.Map" +-- counterparts. +-- +-- Because 'NEMap' is implemented using 'Map', all of the caveats of using +-- 'Map' apply (such as the limitation of the maximum size of maps). +-- +-- All functions take non-empty maps as inputs. In situations where their +-- results can be guarunteed to also be non-empty, they also return +-- non-empty maps. In situations where their results could potentially be +-- empty, 'Map' is returned instead. +-- +-- Some variants of functions (like 'alter'', 'alterF'', 'adjustAt', +-- 'adjustMin', 'adjustMax', 'adjustMinWithKey', 'adjustMaxWithKey') are +-- provided in a way restructured to preserve guaruntees of non-empty maps +-- being returned. +-- +-- Some functions (like 'mapEither', 'partition', 'spanAntitone', 'split') +-- have modified return types to account for possible configurations of +-- non-emptiness. +-- +-- This module is intended to be imported qualified, to avoid name clashes with +-- "Prelude" and "Data.Map" functions: +-- +-- > import qualified Data.Map.NonEmpty as NEM +-- +-- At the moment, this package does not provide a variant strict on values +-- for these functions, like /containers/ does. This is a planned future +-- implementation (PR's are appreciated). For now, you can simulate +-- a strict interface by manually forcing values before returning results. +module Data.Map.NonEmpty ( + -- * Non-Empty Map type + NEMap + -- ** Conversions between empty and non-empty maps + , pattern IsNonEmpty + , pattern IsEmpty + , nonEmptyMap + , toMap + , withNonEmpty + , insertMap + , insertMapWith + , insertMapWithKey + , insertMapMin + , insertMapMax + , unsafeFromMap + + -- * Construction + , singleton + , fromSet + + -- ** From Unordered Lists + , fromList + , fromListWith + , fromListWithKey + + -- ** From Ascending Lists + , fromAscList + , fromAscListWith + , fromAscListWithKey + , fromDistinctAscList + + -- ** From Descending Lists + , fromDescList + , fromDescListWith + , fromDescListWithKey + , fromDistinctDescList + + -- * Insertion + , insert + , insertWith + , insertWithKey + , insertLookupWithKey + + -- * Deletion\/Update + , delete + , adjust + , adjustWithKey + , update + , updateWithKey + , updateLookupWithKey + , alter + , alterF + , alter' + , alterF' + + -- * Query + -- ** Lookup + , lookup + , (!?) + , (!) + , findWithDefault + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + + -- ** Size + , size + + -- * Combine + + -- ** Union + , union + , unionWith + , unionWithKey + , unions + , unionsWith + + -- ** Difference + , difference + , (\\) + , differenceWith + , differenceWithKey + + -- ** Intersection + , intersection + , intersectionWith + , intersectionWithKey + + -- -- ** Unsafe general combining function + -- , mergeWithKey + + -- * Traversal + -- ** Map + , map + , mapWithKey + , traverseWithKey1 + , traverseWithKey + , traverseMaybeWithKey1 + , traverseMaybeWithKey + , mapAccum + , mapAccumWithKey + , mapAccumRWithKey + , mapKeys + , mapKeysWith + , mapKeysMonotonic + + -- * Folds + , foldr + , foldl + , foldr1 + , foldl1 + , foldrWithKey + , foldlWithKey + , foldMapWithKey + + -- ** Strict folds + , foldr' + , foldr1' + , foldl' + , foldl1' + , foldrWithKey' + , foldlWithKey' + + -- * Conversion + , elems + , keys + , assocs + , keysSet + + -- ** Lists + , toList + + -- ** Ordered lists + , toAscList + , toDescList + + -- * Filter + , filter + , filterWithKey + , restrictKeys + , withoutKeys + , partition + , partitionWithKey + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + + , mapMaybe + , mapMaybeWithKey + , mapEither + , mapEitherWithKey + + , split + , splitLookup + , splitRoot + + -- * Submap + , isSubmapOf, isSubmapOfBy + , isProperSubmapOf, isProperSubmapOfBy + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , updateAt + , adjustAt + , deleteAt + , take + , drop + , splitAt + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + , updateMin + , updateMax + , adjustMin + , adjustMax + , updateMinWithKey + , updateMaxWithKey + , adjustMinWithKey + , adjustMaxWithKey + , minView + , maxView + + -- * Debugging + , valid + ) where + +import Control.Applicative +import Data.Bifunctor +import Data.Function +import Data.Functor.Apply +import Data.Functor.Identity +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Map (Map) +import Data.Map.NonEmpty.Internal +import Data.Maybe hiding (mapMaybe) +import Data.Semigroup.Foldable (Foldable1) +import Data.Set (Set) +import Data.Set.NonEmpty.Internal (NESet(..)) +import Data.These +import Prelude hiding (lookup, foldr1, foldl1, foldr, foldl, filter, map, take, drop, splitAt) +import qualified Data.Foldable as F +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import qualified Data.Maybe as Maybe +import qualified Data.Semigroup.Foldable as F1 +import qualified Data.Set as S + +-- | /O(1)/ match, /O(log n)/ usage of contents. The 'IsNonEmpty' and +-- 'IsEmpty' patterns allow you to treat a 'Map' as if it were either +-- a @'IsNonEmpty' n@ (where @n@ is a 'NEMap') or an 'IsEmpty'. +-- +-- For example, you can pattern match on a 'Map': +-- +-- @ +-- myFunc :: 'Map' K X -> Y +-- myFunc ('IsNonEmpty' n) = -- here, the user provided a non-empty map, and @n@ is the 'NEMap' +-- myFunc 'IsEmpty' = -- here, the user provided an empty map. +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original 'Map' was /not/ +-- empty, and you have a verified-non-empty 'NEMap' @n@ to use. +-- +-- Note that patching on this pattern is /O(1)/. However, using the +-- contents requires a /O(log n)/ cost that is deferred until after the +-- pattern is matched on (and is not incurred at all if the contents are +-- never used). +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a 'NEMap' back into a 'Map', obscuring its non-emptiness (see 'toMap'). +pattern IsNonEmpty :: NEMap k a -> Map k a +pattern IsNonEmpty n <- (nonEmptyMap->Just n) + where + IsNonEmpty n = toMap n + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'Map' as if it were either a @'IsNonEmpty' n@ (where @n@ is +-- a 'NEMap') or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original 'Map' was empty. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'Data.Map.empty'. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: Map k a +pattern IsEmpty <- (M.null->True) + where + IsEmpty = M.empty + +{-# COMPLETE IsNonEmpty, IsEmpty #-} + +-- | /O(log n)/. Unsafe version of 'nonEmptyMap'. Coerces a 'Map' into an +-- 'NEMap', but is undefined (throws a runtime exception when evaluation is +-- attempted) for an empty 'Map'. +unsafeFromMap + :: Map k a + -> NEMap k a +unsafeFromMap = withNonEmpty e id + where + e = errorWithoutStackTrace "NEMap.unsafeFromMap: empty map" +{-# INLINE unsafeFromMap #-} + +-- | /O(n)/. Build a non-empty map from a non-empty set of keys and +-- a function which for each key computes its value. +-- +-- > fromSet (\k -> replicate k 'a') (Data.Set.NonEmpty.fromList (3 :| [5])) == fromList ((5,"aaaaa") :| [(3,"aaa")]) +fromSet + :: (k -> a) + -> NESet k + -> NEMap k a +fromSet f (NESet k ks) = NEMap k (f k) (M.fromSet f ks) +{-# INLINE fromSet #-} + +-- | /O(log n)/. Lookup the value at a key in the map. +-- +-- The function will return the corresponding value as @('Just' value)@, +-- or 'Nothing' if the key isn't in the map. +-- +-- An example of using @lookup@: +-- +-- > import Prelude hiding (lookup) +-- > import Data.Map.NonEmpty +-- > +-- > employeeDept = fromList (("John","Sales") :| [("Bob","IT")]) +-- > deptCountry = fromList (("IT","USA") :| [("Sales","France")]) +-- > countryCurrency = fromList (("USA", "Dollar") :| [("France", "Euro")]) +-- > +-- > employeeCurrency :: String -> Maybe String +-- > employeeCurrency name = do +-- > dept <- lookup name employeeDept +-- > country <- lookup dept deptCountry +-- > lookup country countryCurrency +-- > +-- > main = do +-- > putStrLn $ "John's currency: " ++ (show (employeeCurrency "John")) +-- > putStrLn $ "Pete's currency: " ++ (show (employeeCurrency "Pete")) +-- +-- The output of this program: +-- +-- > John's currency: Just "Euro" +-- > Pete's currency: Nothing +lookup + :: Ord k + => k + -> NEMap k a + -> Maybe a +lookup k (NEMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Just v + GT -> M.lookup k m +{-# INLINE lookup #-} + +-- | /O(log n)/. Find the value at a key. Returns 'Nothing' when the +-- element can not be found. +-- +-- prop> fromList ((5, 'a') :| [(3, 'b')]) !? 1 == Nothing +-- prop> fromList ((5, 'a') :| [(3, 'b')]) !? 5 == Just 'a' +(!?) :: Ord k => NEMap k a -> k -> Maybe a +(!?) = flip lookup +{-# INLINE (!?) #-} + +-- | /O(log n)/. Find the value at a key. Calls 'error' when the element +-- can not be found. +-- +-- > fromList ((5,'a') :| [(3,'b')]) ! 1 Error: element not in the map +-- > fromList ((5,'a') :| [(3,'b')]) ! 5 == 'a' +(!) :: Ord k => NEMap k a -> k -> a +(!) m k = fromMaybe e $ m !? k + where + e = error "NEMap.!: given key is not an element in the map" +{-# INLINE (!) #-} + +infixl 9 !? +infixl 9 ! + +-- | /O(log n)/. The expression @('findWithDefault' def k map)@ returns +-- the value at key @k@ or returns default value @def@ +-- when the key is not in the map. +-- +-- > findWithDefault 'x' 1 (fromList ((5,'a') :| [(3,'b')])) == 'x' +-- > findWithDefault 'x' 5 (fromList ((5,'a') :| [(3,'b')])) == 'a' +findWithDefault + :: Ord k + => a + -> k + -> NEMap k a + -> a +findWithDefault def k (NEMap k0 v m) = case compare k k0 of + LT -> def + EQ -> v + GT -> M.findWithDefault def k m +{-# INLINE findWithDefault #-} + +-- | /O(log n)/. Is the key a member of the map? See also 'notMember'. +-- +-- > member 5 (fromList ((5,'a') :| [(3,'b')])) == True +-- > member 1 (fromList ((5,'a') :| [(3,'b')])) == False +member :: Ord k => k -> NEMap k a -> Bool +member k (NEMap k0 _ m) = case compare k k0 of + LT -> False + EQ -> True + GT -> M.member k m +{-# INLINE member #-} + +-- | /O(log n)/. Is the key not a member of the map? See also 'member'. +-- +-- > notMember 5 (fromList ((5,'a') :| [(3,'b')])) == False +-- > notMember 1 (fromList ((5,'a') :| [(3,'b')])) == True +notMember :: Ord k => k -> NEMap k a -> Bool +notMember k (NEMap k0 _ m) = case compare k k0 of + LT -> True + EQ -> False + GT -> M.notMember k m +{-# INLINE notMember #-} + +-- | /O(log n)/. Find largest key smaller than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupLT 3 (fromList ((3,'a') :| [(5,'b')])) == Nothing +-- > lookupLT 4 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +lookupLT :: Ord k => k -> NEMap k a -> Maybe (k, a) +lookupLT k (NEMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Nothing + GT -> M.lookupLT k m <|> Just (k0, v) +{-# INLINE lookupLT #-} + +-- | /O(log n)/. Find smallest key greater than the given one and return the +-- corresponding (key, value) pair. +-- +-- > lookupGT 4 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +-- > lookupGT 5 (fromList ((3,'a') :| [(5,'b')])) == Nothing +lookupGT :: Ord k => k -> NEMap k a -> Maybe (k, a) +lookupGT k (NEMap k0 v m) = case compare k k0 of + LT -> Just (k0, v) + EQ -> M.lookupMin m + GT -> M.lookupGT k m +{-# INLINE lookupGT #-} + +-- | /O(log n)/. Find largest key smaller or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupLE 2 (fromList ((3,'a') :| [(5,'b')])) == Nothing +-- > lookupLE 4 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +-- > lookupLE 5 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +lookupLE :: Ord k => k -> NEMap k a -> Maybe (k, a) +lookupLE k (NEMap k0 v m) = case compare k k0 of + LT -> Nothing + EQ -> Just (k0, v) + GT -> M.lookupLE k m <|> Just (k0, v) +{-# INLINE lookupLE #-} + +-- | /O(log n)/. Find smallest key greater or equal to the given one and return +-- the corresponding (key, value) pair. +-- +-- > lookupGE 3 (fromList ((3,'a') :| [(5,'b')])) == Just (3, 'a') +-- > lookupGE 4 (fromList ((3,'a') :| [(5,'b')])) == Just (5, 'b') +-- > lookupGE 6 (fromList ((3,'a') :| [(5,'b')])) == Nothing +lookupGE :: Ord k => k -> NEMap k a -> Maybe (k, a) +lookupGE k (NEMap k0 v m) = case compare k k0 of + LT -> Just (k0, v) + EQ -> Just (k0, v) + GT -> M.lookupGE k m +{-# INLINE lookupGE #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Union with a combining function. +-- +-- > unionWith (++) (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "aA"), (7, "C")]) +unionWith + :: Ord k + => (a -> a -> a) + -> NEMap k a + -> NEMap k a + -> NEMap k a +unionWith f n1@(NEMap k1 v1 m1) n2@(NEMap k2 v2 m2) = case compare k1 k2 of + LT -> NEMap k1 v1 . M.unionWith f m1 . toMap $ n2 + EQ -> NEMap k1 (f v1 v2) . M.unionWith f m1 $ m2 + GT -> NEMap k2 v2 . M.unionWith f (toMap n1) $ m2 +{-# INLINE unionWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- Union with a combining function, given the matching key. +-- +-- > let f key left_value right_value = (show key) ++ ":" ++ left_value ++ "|" ++ right_value +-- > unionWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "5:a|A"), (7, "C")]) +unionWithKey + :: Ord k + => (k -> a -> a -> a) + -> NEMap k a + -> NEMap k a + -> NEMap k a +unionWithKey f n1@(NEMap k1 v1 m1) n2@(NEMap k2 v2 m2) = case compare k1 k2 of + LT -> NEMap k1 v1 . M.unionWithKey f m1 . toMap $ n2 + EQ -> NEMap k1 (f k1 v1 v2) . M.unionWithKey f m1 $ m2 + GT -> NEMap k2 v2 . M.unionWithKey f (toMap n1) $ m2 +{-# INLINE unionWithKey #-} + +-- | The union of a non-empty list of maps, with a combining operation: +-- (@'unionsWith' f == 'Data.Foldable.foldl1' ('unionWith' f)@). +-- +-- > unionsWith (++) (fromList ((5, "a") :| [(3, "b")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "A3") :| [(3, "B3")])]) +-- > == fromList ((3, "bB3") :| [(5, "aAA3"), (7, "C")]) +unionsWith + :: (Foldable1 f, Ord k) + => (a -> a -> a) + -> f (NEMap k a) + -> NEMap k a +unionsWith f (F1.toNonEmpty->(m :| ms)) = F.foldl' (unionWith f) m ms +{-# INLINE unionsWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two maps. +-- Return elements of the first map not existing in the second map. +-- +-- Returns a potentially empty map ('Map'), in case the first map is +-- a subset of the second map. +-- +-- > difference (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.Map.singleton 3 "b" +difference + :: Ord k + => NEMap k a + -> NEMap k b + -> Map k a +difference n1@(NEMap k1 v1 m1) n2@(NEMap k2 _ m2) = case compare k1 k2 of + -- k1 is not in n2, so cannot be deleted + LT -> insertMinMap k1 v1 $ m1 `M.difference` toMap n2 + -- k2 deletes k1, and only k1 + EQ -> m1 `M.difference` m2 + -- k2 is not in n1, so cannot delete anything, so we can just difference n1 // m2. + GT -> toMap n1 `M.difference` m2 +{-# INLINE difference #-} + +-- | Same as 'difference'. +(\\) + :: Ord k + => NEMap k a + -> NEMap k b + -> Map k a +(\\) = difference +{-# INLINE (\\) #-} + +-- | /O(n+m)/. Difference with a combining function. +-- When two equal keys are +-- encountered, the combining function is applied to the values of these keys. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- Returns a potentially empty map ('Map'), in case the first map is +-- a subset of the second map and the function returns 'Nothing' for every +-- pair. +-- +-- > let f al ar = if al == "b" then Just (al ++ ":" ++ ar) else Nothing +-- > differenceWith f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(3, "B"), (7, "C")])) +-- > == Data.Map.singleton 3 "b:B" +differenceWith + :: Ord k + => (a -> b -> Maybe a) + -> NEMap k a + -> NEMap k b + -> Map k a +differenceWith f = differenceWithKey (const f) +{-# INLINE differenceWith #-} + +-- | /O(n+m)/. Difference with a combining function. When two equal keys are +-- encountered, the combining function is applied to the key and both values. +-- If it returns 'Nothing', the element is discarded (proper set difference). If +-- it returns (@'Just' y@), the element is updated with a new value @y@. +-- +-- Returns a potentially empty map ('Map'), in case the first map is +-- a subset of the second map and the function returns 'Nothing' for every +-- pair. +-- +-- > let f k al ar = if al == "b" then Just ((show k) ++ ":" ++ al ++ "|" ++ ar) else Nothing +-- > differenceWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(3, "B"), (10, "C")])) +-- > == Data.Map.singleton 3 "3:b|B" +differenceWithKey + :: Ord k + => (k -> a -> b -> Maybe a) + -> NEMap k a + -> NEMap k b + -> Map k a +differenceWithKey f n1@(NEMap k1 v1 m1) n2@(NEMap k2 v2 m2) = case compare k1 k2 of + -- k1 is not in n2, so cannot be deleted + LT -> insertMinMap k1 v1 $ M.differenceWithKey f m1 (toMap n2) + -- k2 deletes k1, and only k1 + EQ -> ($ M.differenceWithKey f m1 m2) . maybe id (insertMinMap k1) $ f k1 v1 v2 + -- k2 is not in n1, so cannot delete anything, so we can just difference n1 // m2. + GT -> M.differenceWithKey f (toMap n1) m2 +{-# INLINE differenceWithKey #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection of two maps. +-- Return data in the first map for the keys existing in both maps. +-- (@'intersection' m1 m2 == 'intersectionWith' 'const' m1 m2@). +-- +-- Returns a potentially empty map ('Map'), in case the two maps share no +-- keys in common. +-- +-- > intersection (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.Map.singleton 5 "a" +intersection + :: Ord k + => NEMap k a + -> NEMap k b + -> Map k a +intersection n1@(NEMap k1 v1 m1) n2@(NEMap k2 _ m2) = case compare k1 k2 of + -- k1 is not in n2 + LT -> m1 `M.intersection` toMap n2 + -- k1 and k2 are a part of the result + EQ -> insertMinMap k1 v1 $ m1 `M.intersection` m2 + -- k2 is not in n1 + GT -> toMap n1 `M.intersection` m2 +{-# INLINE intersection #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- Returns a potentially empty map ('Map'), in case the two maps share no +-- keys in common. +-- +-- > intersectionWith (++) (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.Map.singleton 5 "aA" +intersectionWith + :: Ord k + => (a -> b -> c) + -> NEMap k a + -> NEMap k b + -> Map k c +intersectionWith f = intersectionWithKey (const f) +{-# INLINE intersectionWith #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Intersection with a combining function. +-- +-- Returns a potentially empty map ('Map'), in case the two maps share no +-- keys in common. +-- +-- > let f k al ar = (show k) ++ ":" ++ al ++ "|" ++ ar +-- > intersectionWithKey f (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == Data.Map.singleton 5 "5:a|A" +intersectionWithKey + :: Ord k + => (k -> a -> b -> c) + -> NEMap k a + -> NEMap k b + -> Map k c +intersectionWithKey f n1@(NEMap k1 v1 m1) n2@(NEMap k2 v2 m2) = case compare k1 k2 of + -- k1 is not in n2 + LT -> M.intersectionWithKey f m1 (toMap n2) + -- k1 and k2 are a part of the result + EQ -> insertMinMap k1 (f k1 v1 v2) $ M.intersectionWithKey f m1 m2 + -- k2 is not in n1 + GT -> M.intersectionWithKey f (toMap n1) m2 +{-# INLINE intersectionWithKey #-} + +-- | /O(n)/. A strict version of 'foldr1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr1' :: (a -> a -> a) -> NEMap k a -> a +foldr1' f (NEMap _ v m) = case M.maxView m of + Nothing -> v + Just (y, m') -> let !z = M.foldr' f y m' in v `f` z +{-# INLINE foldr1' #-} + +-- | /O(n)/. A strict version of 'foldl1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl1' :: (a -> a -> a) -> NEMap k a -> a +foldl1' f (NEMap _ v m) = M.foldl' f v m +{-# INLINE foldl1' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given right-associative +-- binary operator, such that +-- @'foldrWithKey' f z == 'Prelude.foldr' ('uncurry' f) z . 'toAscList'@. +-- +-- For example, +-- +-- > keysList map = foldrWithKey (\k x ks -> k:ks) [] map +foldrWithKey :: (k -> a -> b -> b) -> b -> NEMap k a -> b +foldrWithKey f z (NEMap k v m) = f k v . M.foldrWithKey f z $ m +{-# INLINE foldrWithKey #-} + +-- | /O(n)/. A strict version of 'foldrWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldrWithKey' :: (k -> a -> b -> b) -> b -> NEMap k a -> b +foldrWithKey' f z (NEMap k v m) = f k v y + where + !y = M.foldrWithKey f z m +{-# INLINE foldrWithKey' #-} + +-- | /O(n)/. Fold the keys and values in the map using the given left-associative +-- binary operator, such that +-- @'foldlWithKey' f z == 'Prelude.foldl' (\\z' (kx, x) -> f z' kx x) z . 'toAscList'@. +-- +-- For example, +-- +-- > keysList = reverse . foldlWithKey (\ks k x -> k:ks) [] +foldlWithKey :: (a -> k -> b -> a) -> a -> NEMap k b -> a +foldlWithKey f z (NEMap k v m) = M.foldlWithKey f (f z k v) m +{-# INLINE foldlWithKey #-} + +-- | /O(n)/. A strict version of 'foldlWithKey'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldlWithKey' :: (a -> k -> b -> a) -> a -> NEMap k b -> a +foldlWithKey' f z (NEMap k v m) = M.foldlWithKey' f x m + where + !x = f z k v +{-# INLINE foldlWithKey' #-} + +-- | /O(n)/. Return all keys of the map in ascending order. +-- +-- > keys (fromList ((5,"a") :| [(3,"b")])) == (3 :| [5]) +keys :: NEMap k a -> NonEmpty k +keys (NEMap k _ m) = k :| M.keys m +{-# INLINE keys #-} + +-- | /O(n)/. An alias for 'toAscList'. Return all key\/value pairs in the map +-- in ascending key order. +-- +-- > assocs (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +assocs :: NEMap k a -> NonEmpty (k, a) +assocs = toList +{-# INLINE assocs #-} + +-- | /O(n)/. The non-empty set of all keys of the map. +-- +-- > keysSet (fromList ((5,"a") :| [(3,"b")])) == Data.Set.NonEmpty.fromList (3 :| [5]) +keysSet :: NEMap k a -> NESet k +keysSet (NEMap k _ m) = NESet k (M.keysSet m) +{-# INLINE keysSet #-} + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > let f key x = (show key) ++ ":" ++ x +-- > mapWithKey f (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "3:b") :| [(5, "5:a")]) +mapWithKey :: (k -> a -> b) -> NEMap k a -> NEMap k b +mapWithKey f (NEMap k v m) = NEMap k (f k v) (M.mapWithKey f m) +{-# NOINLINE [1] mapWithKey #-} +{-# RULES +"mapWithKey/mapWithKey" forall f g xs . mapWithKey f (mapWithKey g xs) = + mapWithKey (\k a -> f k (g k a)) xs +"mapWithKey/map" forall f g xs . mapWithKey f (map g xs) = + mapWithKey (\k a -> f k (g a)) xs +"map/mapWithKey" forall f g xs . map f (mapWithKey g xs) = + mapWithKey (\k a -> f (g k a)) xs + #-} + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys are +-- in ascending order. +-- +-- > toAscList (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +toAscList :: NEMap k a -> NonEmpty (k, a) +toAscList = toList +{-# INLINE toAscList #-} + +-- | /O(n)/. Convert the map to a list of key\/value pairs where the keys +-- are in descending order. +-- +-- > toDescList (fromList ((5,"a") :| [(3,"b")])) == ((5,"a") :| [(3,"b")]) +toDescList :: NEMap k a -> NonEmpty (k, a) +toDescList (NEMap k0 v0 m) = M.foldlWithKey' go ((k0, v0) :| []) m + where + go xs k v = (k, v) NE.<| xs +{-# INLINE toDescList #-} + +-- | /O(log n)/. Convert a 'Map' into an 'NEMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. If key is already present, +-- will overwrite the original value. +-- +-- See 'insertMapMin' for a version that is constant-time if the new key is +-- /strictly smaller than/ all keys in the original map. +-- +-- > insertMap 4 "c" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(4,"c"), (5,"a")]) +-- > insertMap 4 "c" Data.Map.empty == singleton 4 "c" +insertMap :: Ord k => k -> a -> Map k a -> NEMap k a +insertMap k v = withNonEmpty (singleton k v) (insert k v) +{-# INLINE insertMap #-} + +-- | /O(log n)/. Convert a 'Map' into an 'NEMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. Uses a combining function +-- with the new value as the first argument if the key is already present. +-- +-- > insertMapWith (++) 4 "c" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(4,"c"), (5,"a")]) +-- > insertMapWith (++) 5 "c" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(5,"ca")]) +insertMapWith + :: Ord k + => (a -> a -> a) + -> k + -> a + -> Map k a + -> NEMap k a +insertMapWith f k v = withNonEmpty (singleton k v) (insertWith f k v) +{-# INLINE insertMapWith #-} + +-- | /O(log n)/. Convert a 'Map' into an 'NEMap' by adding a key-value +-- pair. Because of this, we know that the map must have at least one +-- element, and so therefore cannot be empty. Uses a combining function +-- with the key and new value as the first and second arguments if the key +-- is already present. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3, "b") :| [(5, "5:xxx|a")]) +-- > insertWithKey f 7 "xxx" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +-- > insertWithKey f 5 "xxx" Data.Map.empty == singleton 5 "xxx" +insertMapWithKey + :: Ord k + => (k -> a -> a -> a) + -> k + -> a + -> Map k a + -> NEMap k a +insertMapWithKey f k v = withNonEmpty (singleton k v) (insertWithKey f k v) +{-# INLINE insertMapWithKey #-} + +-- | /O(1)/ Convert a 'Map' into an 'NEMap' by adding a key-value pair +-- where the key is /strictly less than/ all keys in the input map. The +-- keys in the original map must all be /strictly greater than/ the new +-- key. /The precondition is not checked./ +-- +-- > insertMapMin 2 "c" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((2,"c") :| [(3,"b"), (5,"a")]) +-- > valid (insertMapMin 2 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == True +-- > valid (insertMapMin 7 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == False +-- > valid (insertMapMin 3 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == False +insertMapMin + :: k + -> a + -> Map k a + -> NEMap k a +insertMapMin = NEMap +{-# INLINE insertMapMin #-} + +-- | /O(log n)/ Convert a 'Map' into an 'NEMap' by adding a key-value pair +-- where the key is /strictly greater than/ all keys in the input map. The +-- keys in the original map must all be /strictly less than/ the new +-- key. /The precondition is not checked./ +-- +-- While this has the same asymptotics as 'insertMap', it saves a constant +-- factor for key comparison (so may be helpful if comparison is expensive) +-- and also does not require an 'Ord' instance for the key type. +-- +-- > insertMap 7 "c" (Data.Map.fromList [(5,"a"), (3,"b")]) == fromList ((3,"b") :| [(5,"a"), (7,"c")]) +-- > valid (insertMap 7 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == True +-- > valid (insertMap 2 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == False +-- > valid (insertMap 5 "c" (Data.Map.fromList [(5,"a"), (3,"b")])) == False +insertMapMax + :: k + -> a + -> Map k a + -> NEMap k a +insertMapMax k v = withNonEmpty (singleton k v) go + where + go (NEMap k0 v0 m0) = NEMap k0 v0 . insertMaxMap k v $ m0 +{-# INLINE insertMapMax #-} + + +-- | /O(log n)/. Insert a new key and value in the map. +-- If the key is already present in the map, the associated value is +-- replaced with the supplied value. 'insert' is equivalent to +-- @'insertWith' 'const'@. +-- +-- See 'insertMap' for a version where the first argument is a 'Map'. +-- +-- > insert 5 'x' (fromList ((5,'a') :| [(3,'b')])) == fromList ((3, 'b') :| [(5, 'x')]) +-- > insert 7 'x' (fromList ((5,'a') :| [(3,'b')])) == fromList ((3, 'b') :| [(5, 'a'), (7, 'x')]) +insert + :: Ord k + => k + -> a + -> NEMap k a + -> NEMap k a +insert k v n@(NEMap k0 v0 m) = case compare k k0 of + LT -> NEMap k v . toMap $ n + EQ -> NEMap k v m + GT -> NEMap k0 v0 . M.insert k v $ m +{-# INLINE insert #-} + +-- | /O(log n)/. Insert with a function, combining key, new value and old +-- value. @'insertWithKey' f key value mp@ will insert the pair (key, +-- value) into @mp@ if key does not exist in the map. If the key does +-- exist, the function will insert the pair @(key,f key new_value +-- old_value)@. Note that the key passed to f is the same key passed to +-- 'insertWithKey'. +-- +-- See 'insertMapWithKey' for a version where the first argument is a 'Map'. +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertWithKey f 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "5:xxx|a")]) +-- > insertWithKey f 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +insertWithKey + :: Ord k + => (k -> a -> a -> a) + -> k + -> a + -> NEMap k a + -> NEMap k a +insertWithKey f k v n@(NEMap k0 v0 m) = case compare k k0 of + LT -> NEMap k v . toMap $ n + EQ -> NEMap k (f k v v0) m + GT -> NEMap k0 v0 $ M.insertWithKey f k v m +{-# INLINE insertWithKey #-} + +-- | /O(log n)/. Combines insert operation with old value retrieval. The +-- expression (@'insertLookupWithKey' f k x map@) is a pair where the first +-- element is equal to (@'lookup' k map@) and the second element equal to +-- (@'insertWithKey' f k x map@). +-- +-- > let f key new_value old_value = (show key) ++ ":" ++ new_value ++ "|" ++ old_value +-- > insertLookupWithKey f 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == (Just "a", fromList ((3, "b") :| [(5, "5:xxx|a")])) +-- > insertLookupWithKey f 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == (Nothing, fromList ((3, "b") :| [(5, "a"), (7, "xxx")])) +-- +-- This is how to define @insertLookup@ using @insertLookupWithKey@: +-- +-- > let insertLookup kx x t = insertLookupWithKey (\_ a _ -> a) kx x t +-- > insertLookup 5 "x" (fromList ((5,"a") :| [(3,"b")])) == (Just "a", fromList ((3, "b") :| [(5, "x")])) +-- > insertLookup 7 "x" (fromList ((5,"a") :| [(3,"b")])) == (Nothing, fromList ((3, "b") :| [(5, "a"), (7, "x")])) +insertLookupWithKey + :: Ord k + => (k -> a -> a -> a) + -> k + -> a + -> NEMap k a + -> (Maybe a, NEMap k a) +insertLookupWithKey f k v n@(NEMap k0 v0 m) = case compare k k0 of + LT -> (Nothing, NEMap k v . toMap $ n ) + EQ -> (Just v , NEMap k (f k v v0) m ) + GT -> NEMap k0 v0 <$> M.insertLookupWithKey f k v m +{-# INLINE insertLookupWithKey #-} + +-- | /O(n*log n)/. Build a map from a non-empty list of key\/value pairs +-- with a combining function. See also 'fromAscListWith'. +-- +-- > fromListWith (++) ((5,"a") :| [(5,"b"), (3,"b"), (3,"a"), (5,"a")]) == fromList ((3, "ab") :| [(5, "aba")]) +fromListWith + :: Ord k + => (a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromListWith f = fromListWithKey (const f) +{-# INLINE fromListWith #-} + +-- | /O(n*log n)/. Build a map from a non-empty list of key\/value pairs +-- with a combining function. See also 'fromAscListWithKey'. +-- +-- > let f k a1 a2 = (show k) ++ a1 ++ a2 +-- > fromListWithKey f ((5,"a") :| [(5,"b"), (3,"b"), (3,"a"), (5,"a")]) == fromList ((3, "3ab") :| [(5, "5a5ba")]) +fromListWithKey + :: Ord k + => (k -> a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromListWithKey f ((k0, v0) :| xs) = F.foldl' go (singleton k0 v0) xs + where + go m (k, v) = insertWithKey f k v m + {-# INLINE go #-} +{-# INLINE fromListWithKey #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time. +-- /The precondition (input list is ascending) is not checked./ +-- +-- > fromAscList ((3,"b") :| [(5,"a")]) == fromList ((3, "b") :| [(5, "a")]) +-- > fromAscList ((3,"b") :| [(5,"a"), (5,"b")]) == fromList ((3, "b") :| [(5, "b")]) +-- > valid (fromAscList ((3,"b") :| [(5,"a"), (5,"b")])) == True +-- > valid (fromAscList ((5,"a") :| [(3,"b"), (5,"b")])) == False +fromAscList + :: Eq k + => NonEmpty (k, a) + -> NEMap k a +fromAscList = fromDistinctAscList . combineEq +{-# INLINE fromAscList #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is ascending) is not checked./ +-- +-- > fromAscListWith (++) ((3,"b") :| [(5,"a"), (5,"b")]) == fromList ((3, "b") :| [(5, "ba")]) +-- > valid (fromAscListWith (++) ((3,"b") :| [(5,"a"), (5,"b"))]) == True +-- > valid (fromAscListWith (++) ((5,"a") :| [(3,"b"), (5,"b"))]) == False +fromAscListWith + :: Eq k + => (a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromAscListWith f = fromAscListWithKey (const f) +{-# INLINE fromAscListWith #-} + +-- | /O(n)/. Build a map from an ascending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is ascending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromAscListWithKey f ((3,"b") :| [(5,"a"), (5,"b"), (5,"b")]) == fromList ((3, "b") :| [(5, "5:b5:ba")]) +-- > valid (fromAscListWithKey f ((3,"b") :| [(5,"a"), (5,"b"), (5,"b")])) == True +-- > valid (fromAscListWithKey f ((5,"a") :| [(3,"b"), (5,"b"), (5,"b")])) == False +fromAscListWithKey + :: Eq k + => (k -> a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromAscListWithKey f = fromDistinctAscList . combineEqWith f +{-# INLINE fromAscListWithKey #-} + +-- | /O(n)/. Build a map from an ascending non-empty list of distinct +-- elements in linear time. /The precondition is not checked./ +-- +-- > fromDistinctAscList ((3,"b") :| [(5,"a")]) == fromList ((3, "b") :| [(5, "a")]) +-- > valid (fromDistinctAscList ((3,"b") :| [(5,"a")])) == True +-- > valid (fromDistinctAscList ((3,"b") :| [(5,"a"), (5,"b")])) == False +fromDistinctAscList :: NonEmpty (k, a) -> NEMap k a +fromDistinctAscList ((k, v) :| xs) = insertMapMin k v + . M.fromDistinctAscList + $ xs +{-# INLINE fromDistinctAscList #-} + +-- | /O(n)/. Build a map from a descending non-empty list in linear time. +-- /The precondition (input list is descending) is not checked./ +-- +-- > fromDescList ((5,"a") :| [(3,"b")]) == fromList ((3, "b") :| [(5, "a")]) +-- > fromDescList ((5,"a") :| [(5,"b"), (3,"b")]) == fromList ((3, "b") :| [(5, "b")]) +-- > valid (fromDescList ((5,"a") :| [(5,"b"), (3,"b")])) == True +-- > valid (fromDescList ((5,"a") :| [(3,"b"), (5,"b")])) == False +fromDescList + :: Eq k + => NonEmpty (k, a) + -> NEMap k a +fromDescList = fromDistinctDescList . combineEq +{-# INLINE fromDescList #-} + +-- | /O(n)/. Build a map from a descending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is descending) is not checked./ +-- +-- > fromDescListWith (++) ((5,"a") :| [(5,"b"), (3,"b")]) == fromList ((3, "b") :| [(5, "ba")]) +-- > valid (fromDescListWith (++) ((5,"a") :| [(5,"b"), (3,"b")])) == True +-- > valid (fromDescListWith (++) ((5,"a") :| [(3,"b"), (5,"b")])) == False +fromDescListWith + :: Eq k + => (a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromDescListWith f = fromDescListWithKey (const f) +{-# INLINE fromDescListWith #-} + +-- | /O(n)/. Build a map from a descending non-empty list in linear time +-- with a combining function for equal keys. /The precondition (input list +-- is descending) is not checked./ +-- +-- > let f k a1 a2 = (show k) ++ ":" ++ a1 ++ a2 +-- > fromDescListWithKey f ((5,"a") :| [(5,"b"), (5,"b"), (3,"b")]) == fromList ((3, "b") :| [(5, "5:b5:ba")]) +-- > valid (fromDescListWithKey f ((5,"a") :| [(5,"b"), (5,"b"), (3,"b")])) == True +-- > valid (fromDescListWithKey f ((5,"a") :| [(3,"b"), (5,"b"), (5,"b")])) == False +fromDescListWithKey + :: Eq k + => (k -> a -> a -> a) + -> NonEmpty (k, a) + -> NEMap k a +fromDescListWithKey f = fromDistinctDescList . combineEqWith f +{-# INLINE fromDescListWithKey #-} + +-- | /O(n)/. Build a map from a descending list of distinct elements in linear time. +-- /The precondition is not checked./ +-- +-- > fromDistinctDescList ((5,"a") :| [(3,"b")]) == fromList ((3, "b") :| [(5, "a")]) +-- > valid (fromDistinctDescList ((5,"a") :| [(3,"b")])) == True +-- > valid (fromDistinctDescList ((5,"a") :| [(5,"b"), (3,"b")])) == False +-- +-- @since 0.5.8 +fromDistinctDescList :: NonEmpty (k, a) -> NEMap k a +fromDistinctDescList ((k, v) :| xs) = insertMapMax k v + . M.fromDistinctDescList + $ xs +{-# INLINE fromDistinctDescList #-} + +-- | /O(log n)/. Delete a key and its value from the non-empty map. +-- A potentially empty map ('Map') is returned, since this might delete the +-- last item in the 'NEMap'. When the key is not a member of the map, is +-- equivalent to 'toMap'. +-- +-- > delete 5 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +-- > delete 7 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.Singleton [(3, "b"), (5, "a")] +delete :: Ord k => k -> NEMap k a -> Map k a +delete k n@(NEMap k0 v m) = case compare k k0 of + LT -> toMap n + EQ -> m + GT -> insertMinMap k0 v . M.delete k $ m +{-# INLINE delete #-} + +-- | /O(log n)/. Update a value at a specific key with the result of the +-- provided function. When the key is not a member of the map, the original +-- map is returned. +-- +-- > adjust ("new " ++) 5 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "new a")]) +-- > adjust ("new " ++) 7 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a")]) +adjust + :: Ord k + => (a -> a) + -> k + -> NEMap k a + -> NEMap k a +adjust f = adjustWithKey (const f) +{-# INLINE adjust #-} + +-- | /O(log n)/. Adjust a value at a specific key. When the key is not +-- a member of the map, the original map is returned. +-- +-- > let f key x = (show key) ++ ":new " ++ x +-- > adjustWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "5:new a")]) +-- > adjustWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a")]) +adjustWithKey + :: Ord k + => (k -> a -> a) + -> k + -> NEMap k a + -> NEMap k a +adjustWithKey f k n@(NEMap k0 v m) = case compare k k0 of + LT -> n + EQ -> NEMap k0 (f k0 v) m + GT -> NEMap k0 v . M.adjustWithKey f k $ m +{-# INLINE adjustWithKey #-} + +-- | /O(log n)/. The expression (@'update' f k map@) updates the value @x@ +-- at @k@ (if it is in the map). If (@f x@) is 'Nothing', the element is +-- deleted. If it is (@'Just' y@), the key @k@ is bound to the new value @y@. +-- +-- Returns a potentially empty map ('Map'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEMap'. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > update f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "new a")] +-- > update f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "a")] +-- > update f 3 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +update + :: Ord k + => (a -> Maybe a) + -> k + -> NEMap k a + -> Map k a +update f = updateWithKey (const f) +{-# INLINE update #-} + +-- | /O(log n)/. The expression (@'updateWithKey' f k map@) updates the +-- value @x@ at @k@ (if it is in the map). If (@f k x@) is 'Nothing', +-- the element is deleted. If it is (@'Just' y@), the key @k@ is bound +-- to the new value @y@. +-- +-- Returns a potentially empty map ('Map'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEMap'. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "5:new a")] +-- > updateWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "a")] +-- > updateWithKey f 3 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +updateWithKey + :: Ord k + => (k -> a -> Maybe a) + -> k + -> NEMap k a + -> Map k a +updateWithKey f k n@(NEMap k0 v m) = case compare k k0 of + LT -> toMap n + EQ -> maybe m (flip (insertMinMap k0) m) . f k0 $ v + GT -> insertMinMap k0 v . M.updateWithKey f k $ m +{-# INLINE updateWithKey #-} + +-- | /O(log n)/. Lookup and update. See also 'updateWithKey'. +-- The function returns changed value, if it is updated. +-- Returns the original key value if the map entry is deleted. +-- +-- Returns a potentially empty map ('Map') in the case that we delete the +-- final key of a singleton map. +-- +-- > let f k x = if x == "a" then Just ((show k) ++ ":new a") else Nothing +-- > updateLookupWithKey f 5 (fromList ((5,"a") :| [(3,"b")])) == (Just "5:new a", Data.Map.fromList ((3, "b") :| [(5, "5:new a")])) +-- > updateLookupWithKey f 7 (fromList ((5,"a") :| [(3,"b")])) == (Nothing, Data.Map.fromList ((3, "b") :| [(5, "a")])) +-- > updateLookupWithKey f 3 (fromList ((5,"a") :| [(3,"b")])) == (Just "b", Data.Map.singleton 5 "a") +updateLookupWithKey + :: Ord k + => (k -> a -> Maybe a) + -> k + -> NEMap k a + -> (Maybe a, Map k a) +updateLookupWithKey f k n@(NEMap k0 v m) = case compare k k0 of + LT -> (Nothing, toMap n) + EQ -> let u = f k0 v + in (u <|> Just v, maybe m (flip (insertMinMap k0) m) u) + GT -> fmap (insertMinMap k0 v) . M.updateLookupWithKey f k $ m +{-# INLINE updateLookupWithKey #-} + +-- | /O(log n)/. The expression (@'alter' f k map@) alters the value @x@ at +-- @k@, or absence thereof. 'alter' can be used to insert, delete, or +-- update a value in a 'Map'. In short : @Data.Map.lookup k ('alter' +-- f k m) = f ('lookup' k m)@. +-- +-- Returns a potentially empty map ('Map'), because we can't know ahead of +-- time if the function returns 'Nothing' and deletes the final item in the +-- 'NEMap'. +-- +-- See 'alterF'' for a version that disallows deletion, and so therefore +-- can return 'NEMap'. +-- +-- > let f _ = Nothing +-- > alter f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "a")] +-- > alter f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +-- > +-- > let f _ = Just "c" +-- > alter f 7 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "a"), (7, "c")] +-- > alter f 5 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "c")] +alter + :: Ord k + => (Maybe a -> Maybe a) + -> k + -> NEMap k a + -> Map k a +alter f k n@(NEMap k0 v m) = case compare k k0 of + LT -> ($ toMap n) . maybe id (insertMinMap k ) $ f Nothing + EQ -> ($ m ) . maybe id (insertMinMap k0) $ f (Just v) + GT -> insertMinMap k0 v . M.alter f k $ m +{-# INLINE alter #-} + +-- | /O(log n)/. The expression (@'alterF' f k map@) alters the value @x@ +-- at @k@, or absence thereof. 'alterF' can be used to inspect, insert, +-- delete, or update a value in a 'Map'. In short: @Data.Map.lookup +-- k \<$\> 'alterF' f k m = f ('lookup' k m)@. +-- +-- Example: +-- +-- @ +-- interactiveAlter :: Int -> NEMap Int String -> IO (Map Int String) +-- interactiveAlter k m = alterF f k m where +-- f Nothing = do +-- putStrLn $ show k ++ +-- " was not found in the map. Would you like to add it?" +-- getUserResponse1 :: IO (Maybe String) +-- f (Just old) = do +-- putStrLn $ "The key is currently bound to " ++ show old ++ +-- ". Would you like to change or delete it?" +-- getUserResponse2 :: IO (Maybe String) +-- @ +-- +-- Like @Data.Map.alterF@ for 'Map', 'alterF' can be considered +-- to be a unifying generalization of 'lookup' and 'delete'; however, as +-- a constrast, it cannot be used to implement 'insert', because it must +-- return a 'Map' instead of an 'NEMap' (because the function might delete +-- the final item in the 'NEMap'). When used with trivial functors like +-- 'Identity' and 'Const', it is often slightly slower than +-- specialized 'lookup' and 'delete'. However, when the functor is +-- non-trivial and key comparison is not particularly cheap, it is the +-- fastest way. +-- +-- See 'alterF'' for a version that disallows deletion, and so therefore +-- can return 'NEMap' and be used to implement 'insert' +-- +-- Note on rewrite rules: +-- +-- This module includes GHC rewrite rules to optimize 'alterF' for +-- the 'Const' and 'Identity' functors. In general, these rules +-- improve performance. The sole exception is that when using +-- 'Identity', deleting a key that is already absent takes longer +-- than it would without the rules. If you expect this to occur +-- a very large fraction of the time, you might consider using a +-- private copy of the 'Identity' type. +-- +-- Note: Unlike @Data.Map.alterF@ for 'Map', 'alterF' is /not/ a flipped +-- version of the 'Control.Lens.At.at' combinator from "Control.Lens.At". +-- However, it match the shape expected from most functions expecting +-- lenses, getters, and setters, so can be thought of as a "psuedo-lens", +-- with virtually the same practical applications as a legitimate lens. +alterF + :: (Ord k, Functor f) + => (Maybe a -> f (Maybe a)) + -> k + -> NEMap k a + -> f (Map k a) +alterF f k n@(NEMap k0 v m) = case compare k k0 of + LT -> ($ toMap n) . maybe id (insertMinMap k ) <$> f Nothing + EQ -> ($ m ) . maybe id (insertMinMap k0) <$> f (Just v) + GT -> insertMinMap k0 v <$> M.alterF f k m +{-# INLINABLE [2] alterF #-} + +-- if f ~ Const b, it's a lookup +{-# RULES +"alterF/Const" forall k (f :: Maybe a -> Const b (Maybe a)) . alterF f k = \m -> Const . getConst . f $ lookup k m + #-} +-- if f ~ Identity, it's an 'alter' +{-# RULES +"alterF/Identity" forall k (f :: Maybe a -> Identity (Maybe a)) . alterF f k = Identity . alter (runIdentity . f) k + #-} + +-- | /O(log n)/. Variant of 'alter' that disallows deletion. Allows us to +-- guarantee that the result is also a non-empty Map. +alter' + :: Ord k + => (Maybe a -> a) + -> k + -> NEMap k a + -> NEMap k a +alter' f k n@(NEMap k0 v m) = case compare k k0 of + LT -> NEMap k (f Nothing) . toMap $ n + EQ -> NEMap k0 (f (Just v)) $ m + GT -> NEMap k0 v . M.alter (Just . f) k $ m +{-# INLINE alter' #-} + +-- | /O(log n)/. Variant of 'alterF' that disallows deletion. Allows us to +-- guarantee that the result is also a non-empty Map. +-- +-- Like @Data.Map.alterF@ for 'Map', can be used to generalize and unify +-- 'lookup' and 'insert'. However, because it disallows deletion, it +-- cannot be used to implement 'delete'. +-- +-- See 'alterF' for usage information and caveats. +-- +-- Note: Neither 'alterF' nor 'alterF'' can be considered flipped versions +-- of the 'Control.Lens.At.at' combinator from "Control.Lens.At". However, +-- this can match the shape expected from most functions expecting lenses, +-- getters, and setters, so can be thought of as a "psuedo-lens", with +-- virtually the same practical applications as a legitimate lens. +-- +-- __WARNING__: The rewrite rule for 'Identity' exposes an inconsistency in +-- undefined behavior for "Data.Map". @Data.Map.alterF@ will actually +-- /maintain/ the original key in the map when used with 'Identity'; +-- however, @Data.Map.insertWith@ will /replace/ the orginal key in the +-- map. The rewrite rule for 'alterF'' has chosen to be faithful to +-- @Data.Map.insertWith@, and /not/ @Data.Map.alterF@, for the sake of +-- a cleaner implementation. +alterF' + :: (Ord k, Functor f) + => (Maybe a -> f a) + -> k + -> NEMap k a + -> f (NEMap k a) +alterF' f k n@(NEMap k0 v m) = case compare k k0 of + LT -> flip (NEMap k ) (toMap n) <$> f Nothing + EQ -> flip (NEMap k0) m <$> f (Just v) + GT -> NEMap k0 v <$> M.alterF (fmap Just . f) k m +{-# INLINABLE [2] alterF' #-} + +-- if f ~ Const b, it's a lookup +{-# RULES +"alterF'/Const" forall k (f :: Maybe a -> Const b a) . alterF' f k = \m -> Const . getConst . f $ lookup k m + #-} +-- if f ~ Identity, it's an insertWith +{-# RULES +"alterF'/Identity" forall k (f :: Maybe a -> Identity a) . alterF' f k = Identity . insertWith (\_ -> runIdentity . f . Just) k (runIdentity (f Nothing)) + #-} + +-- | /O(n)/. Traverse keys\/values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('Map'), our function might return +-- 'Nothing' on every item in the 'NEMap'. +-- +-- /Use 'traverseMaybeWithKey1'/ whenever possible (if your 'Applicative' +-- also has 'Apply' instance). This version is provided only for types +-- that do not have 'Apply' instance, since 'Apply' is not at the moment +-- (and might not ever be) an official superclass of 'Applicative'. +traverseMaybeWithKey + :: Applicative t + => (k -> a -> t (Maybe b)) + -> NEMap k a + -> t (Map k b) +traverseMaybeWithKey f (NEMap k0 v m0) = + combine <$> f k0 v <*> M.traverseMaybeWithKey f m0 + where + combine Nothing = id + combine (Just v') = insertMinMap k0 v' +{-# INLINE traverseMaybeWithKey #-} + +-- | /O(n)/. Traverse keys\/values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('Map'), our function might return +-- 'Nothing' on every item in the 'NEMap'. +-- +-- Is more general than 'traverseWithKey', since works with all 'Apply', +-- and not just 'Applicative'. + +-- TODO: benchmark against M.maxView version +traverseMaybeWithKey1 + :: Apply t + => (k -> a -> t (Maybe b)) + -> NEMap k a + -> t (Map k b) +traverseMaybeWithKey1 f (NEMap k0 v m0) = case runMaybeApply m1 of + Left m2 -> combine <$> f k0 v <.> m2 + Right m2 -> (`combine` m2) <$> f k0 v + where + m1 = M.traverseMaybeWithKey (\k -> MaybeApply . Left . f k) m0 + combine Nothing = id + combine (Just v') = insertMinMap k0 v' +{-# INLINE traverseMaybeWithKey1 #-} + +-- | /O(n)/. The function 'mapAccum' threads an accumulating argument +-- through the map in ascending order of keys. +-- +-- > let f a b = (a ++ b, b ++ "X") +-- > mapAccum f "Everything: " (fromList ((5,"a") :| [(3,"b")])) == ("Everything: ba", fromList ((3, "bX") :| [(5, "aX")])) +mapAccum + :: (a -> b -> (a, c)) + -> a + -> NEMap k b + -> (a, NEMap k c) +mapAccum f = mapAccumWithKey (\x _ -> f x) +{-# INLINE mapAccum #-} + +-- | /O(n)/. The function 'mapAccumWithKey' threads an accumulating +-- argument through the map in ascending order of keys. +-- +-- > let f a k b = (a ++ " " ++ (show k) ++ "-" ++ b, b ++ "X") +-- > mapAccumWithKey f "Everything:" (fromList ((5,"a") :| [(3,"b")])) == ("Everything: 3-b 5-a", fromList ((3, "bX") :| [(5, "aX")])) +mapAccumWithKey + :: (a -> k -> b -> (a, c)) + -> a + -> NEMap k b + -> (a, NEMap k c) +mapAccumWithKey f z0 (NEMap k v m) = (z2, NEMap k v' m') + where + ~(z1, v') = f z0 k v + ~(z2, m') = M.mapAccumWithKey f z1 m +{-# INLINE mapAccumWithKey #-} + +-- | /O(n)/. The function 'mapAccumRWithKey' threads an accumulating +-- argument through the map in descending order of keys. +mapAccumRWithKey + :: (a -> k -> b -> (a, c)) + -> a + -> NEMap k b + -> (a, NEMap k c) +mapAccumRWithKey f z0 (NEMap k v m) = (z2, NEMap k v' m') + where + ~(z1, m') = M.mapAccumRWithKey f z0 m + ~(z2, v') = f z1 k v +{-# INLINE mapAccumRWithKey #-} +-- TODO: what other situations can we take advantage of lazy tuple pattern +-- matching? + +-- | /O(n*log n)/. +-- @'mapKeys' f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the value at the greatest of the +-- original keys is retained. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeys (+ 1) (fromList ((5,"a") :| [(3,"b")])) == fromList ((4, "b") :| [(6, "a")]) +-- > mapKeys (\ _ -> 1) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 1 "c" +-- > mapKeys (\ _ -> 3) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 3 "c" +mapKeys + :: Ord k2 + => (k1 -> k2) + -> NEMap k1 a + -> NEMap k2 a +mapKeys f (NEMap k0 v0 m) = fromListWith const + . ((f k0, v0) :|) + . M.foldrWithKey (\k v kvs -> (f k, v) : kvs) [] + $ m +{-# INLINABLE mapKeys #-} + +-- | /O(n*log n)/. +-- @'mapKeysWith' c f s@ is the map obtained by applying @f@ to each key of @s@. +-- +-- The size of the result may be smaller if @f@ maps two or more distinct +-- keys to the same new key. In this case the associated values will be +-- combined using @c@. The value at the greater of the two original keys +-- is used as the first argument to @c@. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeysWith (++) (\ _ -> 1) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 1 "cdab" +-- > mapKeysWith (++) (\ _ -> 3) (fromList ((1,"b") :| [(2,"a"), (3,"d"), (4,"c")])) == singleton 3 "cdab" +mapKeysWith + :: Ord k2 + => (a -> a -> a) + -> (k1 -> k2) + -> NEMap k1 a + -> NEMap k2 a +mapKeysWith c f (NEMap k0 v0 m) = fromListWith c + . ((f k0, v0) :|) + . M.foldrWithKey (\k v kvs -> (f k, v) : kvs) [] + $ m +{-# INLINABLE mapKeysWith #-} + +-- | /O(n)/. +-- @'mapKeysMonotonic' f s == 'mapKeys' f s@, but works only when @f@ +-- is strictly monotonic. +-- That is, for any values @x@ and @y@, if @x@ < @y@ then @f x@ < @f y@. +-- /The precondition is not checked./ +-- Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapKeysMonotonic f s == mapKeys f s +-- > where ls = keys s +-- +-- This means that @f@ maps distinct original keys to distinct resulting keys. +-- This function has better performance than 'mapKeys'. +-- +-- While the size of the result map may be smaller than the input map, the +-- output map is still guaranteed to be non-empty if the input map is +-- non-empty. +-- +-- > mapKeysMonotonic (\ k -> k * 2) (fromList ((5,"a") :| [(3,"b")])) == fromList ((6, "b") :| [(10, "a")]) +-- > valid (mapKeysMonotonic (\ k -> k * 2) (fromList ((5,"a") :| [(3,"b")]))) == True +-- > valid (mapKeysMonotonic (\ _ -> 1) (fromList ((5,"a") :| [(3,"b")]))) == False +mapKeysMonotonic + :: (k1 -> k2) + -> NEMap k1 a + -> NEMap k2 a +mapKeysMonotonic f (NEMap k v m) = NEMap (f k) v + . M.mapKeysMonotonic f + $ m +{-# INLINE mapKeysMonotonic #-} + +-- | /O(n)/. Filter all values that satisfy the predicate. +-- +-- Returns a potentially empty map ('Map'), because we could +-- potentailly filter out all items in the original 'NEMap'. +-- +-- > filter (> "a") (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +-- > filter (> "x") (fromList ((5,"a") :| [(3,"b")])) == Data.Map.empty +-- > filter (< "a") (fromList ((5,"a") :| [(3,"b")])) == Data.Map.empty +filter + :: (a -> Bool) + -> NEMap k a + -> Map k a +filter f (NEMap k v m) + | f v = insertMinMap k v . M.filter f $ m + | otherwise = M.filter f m +{-# INLINE filter #-} + +-- | /O(n)/. Filter all keys\/values that satisfy the predicate. +-- +-- Returns a potentially empty map ('Map'), because we could +-- potentailly filter out all items in the original 'NEMap'. +-- +-- > filterWithKey (\k _ -> k > 4) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +filterWithKey + :: (k -> a -> Bool) + -> NEMap k a + -> Map k a +filterWithKey f (NEMap k v m) + | f k v = insertMinMap k v . M.filterWithKey f $ m + | otherwise = M.filterWithKey f m +{-# INLINE filterWithKey #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Restrict an 'NEMap' to only those keys +-- found in a 'Data.Set.Set'. +-- +-- @ +-- m \`restrictKeys\` s = 'filterWithKey' (\k _ -> k ``Set.member`` s) m +-- m \`restrictKeys\` s = m ``intersection`` 'fromSet' (const ()) s +-- @ +restrictKeys + :: Ord k + => NEMap k a + -> Set k + -> Map k a +restrictKeys n@(NEMap k v m) xs = case S.minView xs of + Nothing -> M.empty + Just (y, ys) -> case compare k y of + -- k is not in xs + LT -> m `M.restrictKeys` xs + -- k and y are a part of the result + EQ -> insertMinMap k v $ m `M.restrictKeys` ys + -- y is not in m + GT -> toMap n `M.restrictKeys` ys +{-# INLINE restrictKeys #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Remove all keys in a 'Data.Set.Set' from +-- an 'NEMap'. +-- +-- @ +-- m \`withoutKeys\` s = 'filterWithKey' (\k _ -> k ``Set.notMember`` s) m +-- m \`withoutKeys\` s = m ``difference`` 'fromSet' (const ()) s +-- @ +withoutKeys + :: Ord k + => NEMap k a + -> Set k + -> Map k a +withoutKeys n@(NEMap k v m) xs = case S.minView xs of + Nothing -> toMap n + Just (y, ys) -> case compare k y of + -- k is not in xs, so cannot be deleted + LT -> insertMinMap k v $ m `M.withoutKeys` xs + -- y deletes k, and only k + EQ -> m `M.withoutKeys` ys + -- y is not in n, so cannot delete anything, so we can just difference n and ys + GT -> toMap n `M.withoutKeys` ys +{-# INLINE withoutKeys #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the predicate was true for all items. +-- * @'That' n2@ means that the predicate was false for all items. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partition (> "a") (fromList ((5,"a") :| [(3,"b")])) == These (singleton 3 "b") (singleton 5 "a") +-- > partition (< "x") (fromList ((5,"a") :| [(3,"b")])) == This (fromList ((3, "b") :| [(5, "a")])) +-- > partition (> "x") (fromList ((5,"a") :| [(3,"b")])) == That (fromList ((3, "b") :| [(5, "a")])) +partition + :: (a -> Bool) + -> NEMap k a + -> These (NEMap k a) (NEMap k a) +partition f = partitionWithKey (const f) +{-# INLINE partition #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the predicate was true for all items, +-- returning the original map. +-- * @'That' n2@ means that the predicate was false for all items, +-- returning the original map. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partitionWithKey (\ k _ -> k > 3) (fromList ((5,"a") :| [(3,"b")])) == These (singleton 5 "a") (singleton 3 "b") +-- > partitionWithKey (\ k _ -> k < 7) (fromList ((5,"a") :| [(3,"b")])) == This (fromList ((3, "b") :| [(5, "a")])) +-- > partitionWithKey (\ k _ -> k > 7) (fromList ((5,"a") :| [(3,"b")])) == That (fromList ((3, "b") :| [(5, "a")])) +partitionWithKey + :: (k -> a -> Bool) + -> NEMap k a + -> These (NEMap k a) (NEMap k a) +partitionWithKey f n@(NEMap k v m0) = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) + | f k v -> This n + | otherwise -> That n + (Just n1, Nothing) + | f k v -> This n + | otherwise -> These n1 (singleton k v) + (Nothing, Just n2) + | f k v -> These (singleton k v) n2 + | otherwise -> That n + (Just n1, Just n2) + | f k v -> These (insertMapMin k v m1) n2 + | otherwise -> These n1 (insertMapMin k v m2) + where + (m1, m2) = M.partitionWithKey f m0 +{-# INLINABLE partitionWithKey #-} + +-- | /O(log n)/. Take while a predicate on the keys holds. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- Returns a potentially empty map ('Map'), because the predicate might +-- fail on the first input. +-- +-- @ +-- takeWhileAntitone p = Data.Map.fromDistinctAscList . Data.List.takeWhile (p . fst) . Data.Foldable.toList +-- takeWhileAntitone p = 'filterWithKey' (\k _ -> p k) +-- @ +takeWhileAntitone + :: (k -> Bool) + -> NEMap k a + -> Map k a +takeWhileAntitone f (NEMap k v m) + | f k = insertMinMap k v . M.takeWhileAntitone f $ m + | otherwise = M.empty +{-# INLINE takeWhileAntitone #-} + +-- | /O(log n)/. Drop while a predicate on the keys holds. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- @ +-- dropWhileAntitone p = Data.Map.fromDistinctAscList . Data.List.dropWhile (p . fst) . Data.Foldable.toList +-- dropWhileAntitone p = 'filterWithKey' (\k -> not (p k)) +-- @ +dropWhileAntitone + :: (k -> Bool) + -> NEMap k a + -> Map k a +dropWhileAntitone f n@(NEMap k _ m) + | f k = M.dropWhileAntitone f m + | otherwise = toMap n +{-# INLINE dropWhileAntitone #-} + +-- | /O(log n)/. Divide a map at the point where a predicate on the keys stops holding. +-- The user is responsible for ensuring that for all keys @j@ and @k@ in the map, +-- @j \< k ==\> p j \>= p k@. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the predicate never failed for any item, +-- returning the original map. +-- * @'That' n2@ means that the predicate failed for the first item, +-- returning the original map. +-- * @'These' n1 n2@ gives @n1@ (the map up to the point where the +-- predicate on the keys stops holding) and @n2@ (the map starting from +-- the point where the predicate stops holding) +-- +-- @ +-- spanAntitone p xs = partitionWithKey (\k _ -> p k) xs +-- @ +-- +-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the map +-- at some /unspecified/ point where the predicate switches from holding to not +-- holding (where the predicate is seen to hold before the first key and to fail +-- after the last key). +spanAntitone + :: (k -> Bool) + -> NEMap k a + -> These (NEMap k a) (NEMap k a) +spanAntitone f n@(NEMap k v m0) + | f k = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> This n + (Just _ , Nothing) -> This n + (Nothing, Just n2) -> These (singleton k v) n2 + (Just _ , Just n2) -> These (insertMapMin k v m1) n2 + | otherwise = That n + where + (m1, m2) = M.spanAntitone f m0 +{-# INLINABLE spanAntitone #-} + +-- | /O(n)/. Map values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('Map'), because the function could +-- potentially return 'Nothing' on all items in the 'NEMap'. +-- +-- > let f x = if x == "a" then Just "new a" else Nothing +-- > mapMaybe f (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "new a" +mapMaybe + :: (a -> Maybe b) + -> NEMap k a + -> Map k b +mapMaybe f = mapMaybeWithKey (const f) +{-# INLINE mapMaybe #-} + +-- | /O(n)/. Map keys\/values and collect the 'Just' results. +-- +-- Returns a potentially empty map ('Map'), because the function could +-- potentially return 'Nothing' on all items in the 'NEMap'. +-- +-- > let f k _ = if k < 5 then Just ("key : " ++ (show k)) else Nothing +-- > mapMaybeWithKey f (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "key : 3" +mapMaybeWithKey + :: (k -> a -> Maybe b) + -> NEMap k a + -> Map k b +mapMaybeWithKey f (NEMap k v m) = ($ M.mapMaybeWithKey f m) + . maybe id (insertMinMap k) + $ f k v +{-# INLINE mapMaybeWithKey #-} + +-- | /O(n)/. Map values and separate the 'Left' and 'Right' results. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the results were all 'Left'. +-- * @'That' n2@ means that the results were all 'Right'. +-- * @'These' n1 n2@ gives @n1@ (the map where the results were 'Left') +-- and @n2@ (the map where the results were 'Right') +-- +-- > let f a = if a < "c" then Left a else Right a +-- > mapEither f (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == These (fromList ((3,"b") :| [(5,"a")])) (fromList ((1,"x") :| [(7,"z")])) +-- > +-- > mapEither (\ a -> Right a) (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == That (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +mapEither + :: (a -> Either b c) + -> NEMap k a + -> These (NEMap k b) (NEMap k c) +mapEither f = mapEitherWithKey (const f) +{-# INLINE mapEither #-} + +-- | /O(n)/. Map keys\/values and separate the 'Left' and 'Right' results. +-- +-- Returns a 'These' with potentially two non-empty maps: +-- +-- * @'This' n1@ means that the results were all 'Left'. +-- * @'That' n2@ means that the results were all 'Right'. +-- * @'These' n1 n2@ gives @n1@ (the map where the results were 'Left') +-- and @n2@ (the map where the results were 'Right') +-- +-- > let f k a = if k < 5 then Left (k * 2) else Right (a ++ a) +-- > mapEitherWithKey f (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == These (fromList ((1,2) :| [(3,6)])) (fromList ((5,"aa") :| [(7,"zz")])) +-- > +-- > mapEitherWithKey (\_ a -> Right a) (fromList ((5,"a") :| [(3,"b"), (1,"x"), (7,"z")])) +-- > == That (fromList ((1,"x") :| [(3,"b"), (5,"a"), (7,"z")])) +mapEitherWithKey + :: (k -> a -> Either b c) + -> NEMap k a + -> These (NEMap k b) (NEMap k c) +mapEitherWithKey f (NEMap k v m0) = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> case f k v of + Left v' -> This (singleton k v') + Right v' -> That (singleton k v') + (Just n1, Nothing) -> case f k v of + Left v' -> This (insertMapMin k v' m1) + Right v' -> These n1 (singleton k v') + (Nothing, Just n2) -> case f k v of + Left v' -> These (singleton k v') n2 + Right v' -> That (insertMapMin k v' m2) + (Just n1, Just n2) -> case f k v of + Left v' -> These (insertMapMin k v' m1) n2 + Right v' -> These n1 (insertMapMin k v' m2) + where + (m1, m2) = M.mapEitherWithKey f m0 +{-# INLINABLE mapEitherWithKey #-} + +-- | /O(log n)/. The expression (@'split' k map@) is potentially a 'These' +-- containing up to two 'NEMap's based on splitting the map into maps +-- containing items before and after the given key @k@. It will never +-- return a map that contains @k@ itself. +-- +-- * 'Nothing' means that @k@ was the only key in the the original map, +-- and so there are no items before or after it. +-- * @'Just' ('This' n1)@ means @k@ was larger than or equal to all items +-- in the map, and @n1@ is the entire original map (minus @k@, if it was +-- present) +-- * @'Just' ('That' n2)@ means @k@ was smaller than or equal to all +-- items in the map, and @n2@ is the entire original map (minus @k@, if +-- it was present) +-- * @'Just' ('These' n1 n2)@ gives @n1@ (the map of all keys from the +-- original map less than @k@) and @n2@ (the map of all keys from the +-- original map greater than @k@) +-- +-- > split 2 (fromList ((5,"a") :| [(3,"b")])) == Just (That (fromList ((3,"b") :| [(5,"a")])) ) +-- > split 3 (fromList ((5,"a") :| [(3,"b")])) == Just (That (singleton 5 "a") ) +-- > split 4 (fromList ((5,"a") :| [(3,"b")])) == Just (These (singleton 3 "b") (singleton 5 "a")) +-- > split 5 (fromList ((5,"a") :| [(3,"b")])) == Just (This (singleton 3 "b") ) +-- > split 6 (fromList ((5,"a") :| [(3,"b")])) == Just (This (fromList ((3,"b") :| [(5,"a")])) ) +-- > split 5 (singleton 5 "a") == Nothing +split + :: Ord k + => k + -> NEMap k a + -> Maybe (These (NEMap k a) (NEMap k a)) +split k n@(NEMap k0 v m0) = case compare k k0 of + LT -> Just $ That n + EQ -> That <$> nonEmptyMap m0 + GT -> case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> Just $ This (singleton k0 v) + (Just _ , Nothing) -> Just $ This (insertMapMin k0 v m1) + (Nothing, Just n2) -> Just $ These (singleton k0 v) n2 + (Just _ , Just n2) -> Just $ These (insertMapMin k0 v m1) n2 + where + (m1, m2) = M.split k m0 +{-# INLINABLE split #-} + +-- | /O(log n)/. The expression (@'splitLookup' k map@) splits a map just +-- like 'split' but also returns @'lookup' k map@, as a @'Maybe' a@. +-- +-- > splitLookup 2 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (That (fromList ((3,"b") :| [(5,"a")])))) +-- > splitLookup 3 (fromList ((5,"a") :| [(3,"b")])) == (Just "b", Just (That (singleton 5 "a"))) +-- > splitLookup 4 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (These (singleton 3 "b") (singleton 5 "a"))) +-- > splitLookup 5 (fromList ((5,"a") :| [(3,"b")])) == (Just "a", Just (This (singleton 3 "b")) +-- > splitLookup 6 (fromList ((5,"a") :| [(3,"b")])) == (Nothing , Just (This (fromList ((3,"b") :| [(5,"a")]))) +-- > splitLookup 5 (singleton 5 "a") == (Just "a", Nothing) +splitLookup + :: Ord k + => k + -> NEMap k a + -> (Maybe a, Maybe (These (NEMap k a) (NEMap k a))) +splitLookup k n@(NEMap k0 v0 m0) = case compare k k0 of + LT -> (Nothing, Just $ That n) + EQ -> (Just v0, That <$> nonEmptyMap m0) + GT -> (v ,) $ case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> Just $ This (singleton k0 v0) + (Just _ , Nothing) -> Just $ This (insertMapMin k0 v0 m1) + (Nothing, Just n2) -> Just $ These (singleton k0 v0) n2 + (Just _ , Just n2) -> Just $ These (insertMapMin k0 v0 m1) n2 + where + (m1, v, m2) = M.splitLookup k m0 +{-# INLINABLE splitLookup #-} + +-- | /O(1)/. Decompose a map into pieces based on the structure of the +-- underlying tree. This function is useful for consuming a map in +-- parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that +-- the pieces returned will be in ascending order (all elements in the +-- first submap less than all elements in the second, and so on). +-- +-- Note that the current implementation does not return more than four +-- submaps, but you should not depend on this behaviour because it can +-- change in the future without notice. +splitRoot + :: NEMap k a + -> NonEmpty (NEMap k a) +splitRoot (NEMap k v m) = singleton k v + :| Maybe.mapMaybe nonEmptyMap (M.splitRoot m) +{-# INLINE splitRoot #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- This function is defined as (@'isSubmapOf' = 'isSubmapOfBy' (==)@). +isSubmapOf :: (Ord k, Eq a) => NEMap k a -> NEMap k a -> Bool +isSubmapOf = isSubmapOfBy (==) +{-# INLINE isSubmapOf #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- The expression (@'isSubmapOfBy' f t1 t2@) returns 'True' if +-- all keys in @t1@ are in tree @t2@, and when @f@ returns 'True' when +-- applied to their respective values. For example, the following +-- expressions are all 'True': +-- +-- > isSubmapOfBy (==) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (<=) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (==) (fromList (('a',1) :| [('b',2)])) (fromList (('a',1) :| [('b',2)])) +-- +-- But the following are all 'False': +-- +-- > isSubmapOfBy (==) (singleton 'a' 2) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (<) (singleton 'a' 1) (fromList (('a',1) :| [('b',2)])) +-- > isSubmapOfBy (==) (fromList (('a',1) :| [('b',2)])) (singleton 'a' 1) +isSubmapOfBy + :: Ord k + => (a -> b -> Bool) + -> NEMap k a + -> NEMap k b + -> Bool +isSubmapOfBy f (NEMap k v m0) (toMap->m1) = kvSub + && M.isSubmapOfBy f m0 m1 + where + kvSub = case M.lookup k m1 of + Just v0 -> f v v0 + Nothing -> False +{-# INLINE isSubmapOfBy #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap +-- but not equal). Defined as (@'isProperSubmapOf' = 'isProperSubmapOfBy' +-- (==)@). +isProperSubmapOf :: (Ord k, Eq a) => NEMap k a -> NEMap k a -> Bool +isProperSubmapOf = isProperSubmapOfBy (==) +{-# INLINE isProperSubmapOf #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Is this a proper submap? (ie. a submap +-- but not equal). The expression (@'isProperSubmapOfBy' f m1 m2@) returns +-- 'True' when @m1@ and @m2@ are not equal, all keys in @m1@ are in @m2@, +-- and when @f@ returns 'True' when applied to their respective values. For +-- example, the following expressions are all 'True': +-- +-- > isProperSubmapOfBy (==) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +-- > isProperSubmapOfBy (<=) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +-- +-- But the following are all 'False': +-- +-- > isProperSubmapOfBy (==) (fromList ((1,1) :| [(2,2)])) (fromList ((1,1) :| [(2,2)])) +-- > isProperSubmapOfBy (==) (fromList ((1,1) :| [(2,2)])) (singleton 1 1)) +-- > isProperSubmapOfBy (<) (singleton 1 1) (fromList ((1,1) :| [(2,2)])) +isProperSubmapOfBy + :: Ord k + => (a -> b -> Bool) + -> NEMap k a + -> NEMap k b + -> Bool +isProperSubmapOfBy f m1 m2 = M.size (nemMap m1) < M.size (nemMap m2) + && isSubmapOfBy f m1 m2 +{-# INLINE isProperSubmapOfBy #-} + +-- | /O(log n)/. Lookup the /index/ of a key, which is its zero-based index +-- in the sequence sorted by keys. The index is a number from /0/ up to, +-- but not including, the 'size' of the map. +-- +-- > isJust (lookupIndex 2 (fromList ((5,"a") :| [(3,"b")]))) == False +-- > fromJust (lookupIndex 3 (fromList ((5,"a") :| [(3,"b")]))) == 0 +-- > fromJust (lookupIndex 5 (fromList ((5,"a") :| [(3,"b")]))) == 1 +-- > isJust (lookupIndex 6 (fromList ((5,"a") :| [(3,"b")]))) == False +lookupIndex + :: Ord k + => k + -> NEMap k a + -> Maybe Int +lookupIndex k (NEMap k0 _ m) = case compare k k0 of + LT -> Nothing + EQ -> Just 0 + GT -> (+ 1) <$> M.lookupIndex k m +{-# INLINE lookupIndex #-} + +-- | /O(log n)/. Return the /index/ of a key, which is its zero-based index +-- in the sequence sorted by keys. The index is a number from /0/ up to, +-- but not including, the 'size' of the map. Calls 'error' when the key is +-- not a 'member' of the map. +-- +-- > findIndex 2 (fromList ((5,"a") :| [(3,"b")])) Error: element is not in the map +-- > findIndex 3 (fromList ((5,"a") :| [(3,"b")])) == 0 +-- > findIndex 5 (fromList ((5,"a") :| [(3,"b")])) == 1 +-- > findIndex 6 (fromList ((5,"a") :| [(3,"b")])) Error: element is not in the map +findIndex + :: Ord k + => k + -> NEMap k a + -> Int +findIndex k = fromMaybe e . lookupIndex k + where + e = error "NEMap.findIndex: element is not in the map" +{-# INLINE findIndex #-} + +-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based +-- index in the sequence sorted by keys. If the /index/ is out of range +-- (less than zero, greater or equal to 'size' of the map), 'error' is +-- called. +-- +-- > elemAt 0 (fromList ((5,"a") :| [(3,"b")])) == (3,"b") +-- > elemAt 1 (fromList ((5,"a") :| [(3,"b")])) == (5, "a") +-- > elemAt 2 (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +elemAt + :: Int + -> NEMap k a + -> (k, a) +elemAt 0 (NEMap k v _) = (k, v) +elemAt i (NEMap _ _ m) = M.elemAt (i - 1) m +{-# INLINABLE elemAt #-} + +-- | /O(log n)/. Update the element at /index/, i.e. by its zero-based index in +-- the sequence sorted by keys. If the /index/ is out of range (less than zero, +-- greater or equal to 'size' of the map), 'error' is called. +-- +-- Returns a possibly empty map ('Map'), because the function might end up +-- deleting the last key in the map. See 'adjustAt' for a version that +-- disallows deletion, guaranteeing that the result is also a non-empty +-- Map. +-- +-- > updateAt (\ _ _ -> Just "x") 0 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "x"), (5, "a")] +-- > updateAt (\ _ _ -> Just "x") 1 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "x")] +-- > updateAt (\ _ _ -> Just "x") 2 (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +-- > updateAt (\ _ _ -> Just "x") (-1) (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +-- > updateAt (\_ _ -> Nothing) 0 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +-- > updateAt (\_ _ -> Nothing) 1 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +-- > updateAt (\_ _ -> Nothing) 2 (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +-- > updateAt (\_ _ -> Nothing) (-1) (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +updateAt + :: (k -> a -> Maybe a) + -> Int + -> NEMap k a + -> Map k a +updateAt f 0 (NEMap k v m) = maybe m (flip (insertMinMap k) m) $ f k v +updateAt f i (NEMap k v m) = insertMinMap k v . M.updateAt f (i - 1) $ m +{-# INLINABLE updateAt #-} + +-- | /O(log n)/. Variant of 'updateAt' that disallows deletion. Allows us +-- to guarantee that the result is also a non-empty Map. +adjustAt + :: (k -> a -> a) + -> Int + -> NEMap k a + -> NEMap k a +adjustAt f 0 (NEMap k0 v m) = NEMap k0 (f k0 v) m +adjustAt f i (NEMap k0 v m) = NEMap k0 v + . M.updateAt (\k -> Just . f k) (i - 1) + $ m +{-# INLINABLE adjustAt #-} + +-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based +-- index in the sequence sorted by keys. If the /index/ is out of range +-- (less than zero, greater or equal to 'size' of the map), 'error' is +-- called. +-- +-- Returns a potentially empty map ('Map') because of the possibility of +-- deleting the last item in a map. +-- +-- > deleteAt 0 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +-- > deleteAt 1 (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +-- > deleteAt 2 (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +-- > deleteAt (-1) (fromList ((5,"a") :| [(3,"b")])) Error: index out of range +deleteAt + :: Int + -> NEMap k a + -> Map k a +deleteAt 0 (NEMap _ _ m) = m +deleteAt i (NEMap k v m) = insertMinMap k v . M.deleteAt (i - 1) $ m +{-# INLINABLE deleteAt #-} + +-- | Take a given number of entries in key order, beginning with the +-- smallest keys. +-- +-- Returns a possibly empty map ('Map'), which can only happen if we call +-- @take 0@. +-- +-- @ +-- take n = Data.Map.fromDistinctAscList . Data.List.NonEmpty.take n . 'toList' +-- @ +take + :: Int + -> NEMap k a + -> Map k a +take 0 NEMap{} = M.empty +take i (NEMap k v m) = insertMinMap k v . M.take (i - 1) $ m +{-# INLINABLE take #-} + +-- | Drop a given number of entries in key order, beginning +-- with the smallest keys. +-- +-- Returns a possibly empty map ('Map'), in case we drop all of the +-- elements (which can happen if we drop a number greater than or equal to +-- the number of items in the map) +-- +-- @ +-- drop n = Data.Map.fromDistinctAscList . Data.List.NonEmpty.drop' n . 'toList' +-- @ +drop + :: Int + -> NEMap k a + -> Map k a +drop 0 n = toMap n +drop i (NEMap _ _ m) = M.drop (i - 1) m +{-# INLINABLE drop #-} + +-- | /O(log n)/. Split a map at a particular index @i@. +-- +-- * @'This' n1@ means that there are less than @i@ items in the map, and +-- @n1@ is the original map. +-- * @'That' n2@ means @i@ was 0; we dropped 0 items, so @n2@ is the +-- original map. +-- * @'These' n1 n2@ gives @n1@ (taking @i@ items from the original map) +-- and @n2@ (dropping @i@ items from the original map)) +splitAt + :: Int + -> NEMap k a + -> These (NEMap k a) (NEMap k a) +splitAt 0 n = That n +splitAt i n@(NEMap k v m0) = case (nonEmptyMap m1, nonEmptyMap m2) of + (Nothing, Nothing) -> This (singleton k v) + (Just _ , Nothing) -> This n + (Nothing, Just n2) -> These (singleton k v) n2 + (Just _ , Just n2) -> These (insertMapMin k v m1) n2 + where + (m1, m2) = M.splitAt (i - 1) m0 +{-# INLINABLE splitAt #-} + +-- | /O(1)/. The minimal key of the map. Note that this is total, making +-- 'Data.Map.lookupMin' obsolete. It is constant-time, so has better +-- asymptotics than @Data.Map.lookupMin@ and @Data.Map.findMin@, as well. +-- +-- > findMin (fromList ((5,"a") :| [(3,"b")])) == (3,"b") +findMin :: NEMap k a -> (k, a) +findMin (NEMap k v _) = (k, v) +{-# INLINE findMin #-} + +-- | /O(log n)/. The maximal key of the map. Note that this is total, making +-- 'Data.Map.lookupMin' obsolete. +-- +-- > findMax (fromList ((5,"a") :| [(3,"b")])) == (5,"a") +findMax :: NEMap k a -> (k, a) +findMax (NEMap k v m) = fromMaybe (k, v) . M.lookupMax $ m +{-# INLINE findMax #-} + +-- | /O(1)/. Delete the minimal key. Returns a potentially empty map +-- ('Map'), because we might end up deleting the final key in a singleton +-- map. It is constant-time, so has better asymptotics than +-- 'Data.Map.deleteMin'. +-- +-- > deleteMin (fromList ((5,"a") :| [(3,"b"), (7,"c")])) == Data.Map.fromList [(5,"a"), (7,"c")] +-- > deleteMin (singleton 5 "a") == Data.Map.empty +deleteMin :: NEMap k a -> Map k a +deleteMin (NEMap _ _ m) = m +{-# INLINE deleteMin #-} + +-- | /O(log n)/. Delete the maximal key. Returns a potentially empty map +-- ('Map'), because we might end up deleting the final key in a singleton +-- map. +-- +-- > deleteMax (fromList ((5,"a") :| [(3,"b"), (7,"c")])) == Data.Map.fromList [(3,"b"), (5,"a")] +-- > deleteMax (singleton 5 "a") == Data.Map.empty +deleteMax :: NEMap k a -> Map k a +deleteMax (NEMap k v m) = insertMinMap k v . M.deleteMax $ m +{-# INLINE deleteMax #-} + +-- | /O(1)/ if delete, /O(log n)/ otherwise. Update the value at the +-- minimal key. Returns a potentially empty map ('Map'), because we might +-- end up deleting the final key in the map if the function returns +-- 'Nothing'. See 'adjustMin' for a version that can guaruntee that we +-- return a non-empty map. +-- +-- > updateMin (\ a -> Just ("X" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "Xb"), (5, "a")] +-- > updateMin (\ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +updateMin :: (a -> Maybe a) -> NEMap k a -> Map k a +updateMin f = updateMinWithKey (const f) +{-# INLINE updateMin #-} + +-- | /O(1)/. A version of 'updateMin' that disallows deletion, allowing us +-- to guarantee that the result is also non-empty. +adjustMin :: (a -> a) -> NEMap k a -> NEMap k a +adjustMin f = adjustMinWithKey (const f) +{-# INLINE adjustMin #-} + +-- | /O(1)/ if delete, /O(log n)/ otherwise. Update the value at the +-- minimal key. Returns a potentially empty map ('Map'), because we might +-- end up deleting the final key in the map if the function returns +-- 'Nothing'. See 'adjustMinWithKey' for a version that guaruntees +-- a non-empty map. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +updateMinWithKey :: (k -> a -> Maybe a) -> NEMap k a -> Map k a +updateMinWithKey f (NEMap k v m) = ($ m) . maybe id (insertMinMap k) $ f k v +{-# INLINE updateMinWithKey #-} + +-- | /O(1)/. A version of 'adjustMaxWithKey' that disallows deletion, +-- allowing us to guarantee that the result is also non-empty. Note that +-- it also is able to have better asymptotics than 'updateMinWithKey' in +-- general. +adjustMinWithKey :: (k -> a -> a) -> NEMap k a -> NEMap k a +adjustMinWithKey f (NEMap k v m) = NEMap k (f k v) m +{-# INLINE adjustMinWithKey #-} + +-- | /O(log n)/. Update the value at the maximal key. Returns +-- a potentially empty map ('Map'), because we might end up deleting the +-- final key in the map if the function returns 'Nothing'. See 'adjustMax' +-- for a version that can guarantee that we return a non-empty map. +-- +-- > updateMax (\ a -> Just ("X" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3, "b"), (5, "Xa")] +-- > updateMax (\ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 3 "b" +updateMax :: (a -> Maybe a) -> NEMap k a -> Map k a +updateMax f = updateMaxWithKey (const f) +{-# INLINE updateMax #-} + +-- | /O(log n)/. A version of 'updateMax' that disallows deletion, allowing +-- us to guarantee that the result is also non-empty. +adjustMax :: (a -> a) -> NEMap k a -> NEMap k a +adjustMax f = adjustMaxWithKey (const f) +{-# INLINE adjustMax #-} + +-- | /O(log n)/. Update the value at the maximal key. Returns +-- a potentially empty map ('Map'), because we might end up deleting the +-- final key in the map if the function returns 'Nothing'. See +-- 'adjustMaxWithKey' for a version that guaruntees a non-empty map. +-- +-- > updateMinWithKey (\ k a -> Just ((show k) ++ ":" ++ a)) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.fromList [(3,"3:b"), (5,"a")] +-- > updateMinWithKey (\ _ _ -> Nothing) (fromList ((5,"a") :| [(3,"b")])) == Data.Map.singleton 5 "a" +updateMaxWithKey :: (k -> a -> Maybe a) -> NEMap k a -> Map k a +updateMaxWithKey f (NEMap k v m) + | M.null m = maybe m (M.singleton k) $ f k v + | otherwise = insertMinMap k v + . M.updateMaxWithKey f + $ m +{-# INLINE updateMaxWithKey #-} + +-- | /O(log n)/. A version of 'updateMaxWithKey' that disallows deletion, +-- allowing us to guarantee that the result is also non-empty. +adjustMaxWithKey :: (k -> a -> a) -> NEMap k a -> NEMap k a +adjustMaxWithKey f (NEMap k0 v m) + | M.null m = NEMap k0 (f k0 v) m + | otherwise = insertMapMin k0 v + . M.updateMaxWithKey (\k -> Just . f k) + $ m +{-# INLINE adjustMaxWithKey #-} + +-- | /O(1)/. Retrieves the value associated with minimal key of the +-- map, and the map stripped of that element. It is constant-time, so has +-- better asymptotics than @Data.Map.minView@ for 'Map'. +-- +-- Note that unlike @Data.Map.minView@ for 'Map', this cannot ever fail, +-- so doesn't need to return in a 'Maybe'. However, the result 'Map' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > minView (fromList ((5,"a") :| [(3,"b")])) == ("b", Data.Map.singleton 5 "a") +minView :: NEMap k a -> (a, Map k a) +minView = first snd . deleteFindMin +{-# INLINE minView #-} + +-- | /O(1)/. Delete and find the minimal key-value pair. It is +-- constant-time, so has better asymptotics that @Data.Map.minView@ for +-- 'Map'. +-- +-- Note that unlike @Data.Map.deleteFindMin@ for 'Map', this cannot ever +-- fail, and so is a total function. However, the result 'Map' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > deleteFindMin (fromList ((5,"a") :| [(3,"b"), (10,"c")])) == ((3,"b"), Data.Map.fromList [(5,"a"), (10,"c")]) +deleteFindMin :: NEMap k a -> ((k, a), Map k a) +deleteFindMin (NEMap k v m) = ((k, v), m) +{-# INLINE deleteFindMin #-} + +-- | /O(log n)/. Retrieves the value associated with maximal key of the +-- map, and the map stripped of that element. +-- +-- Note that unlike @Data.Map.maxView@ from 'Map', this cannot ever fail, +-- so doesn't need to return in a 'Maybe'. However, the result 'Map' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > maxView (fromList ((5,"a") :| [(3,"b")])) == ("a", Data.Map.singleton 3 "b") +maxView :: NEMap k a -> (a, Map k a) +maxView = first snd . deleteFindMax +{-# INLINE maxView #-} + +-- | /O(log n)/. Delete and find the minimal key-value pair. +-- +-- Note that unlike @Data.Map.deleteFindMax@ for 'Map', this cannot ever +-- fail, and so is a total function. However, the result 'Map' is +-- potentially empty, since the original map might have contained just +-- a single item. +-- +-- > deleteFindMax (fromList ((5,"a") :| [(3,"b"), (10,"c")])) == ((10,"c"), Data.Map.fromList [(3,"b"), (5,"a")]) +deleteFindMax :: NEMap k a -> ((k, a), Map k a) +deleteFindMax (NEMap k v m) = maybe ((k, v), M.empty) (second (insertMinMap k v)) + . M.maxViewWithKey + $ m +{-# INLINE deleteFindMax #-} + +-- --------------------------- +-- Combining functions +-- --------------------------- +-- +-- Code comes from "Data.Map.Internal" from containers, modified slightly +-- to work with NonEmpty +-- +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 + +combineEq :: Eq a => NonEmpty (a, b) -> NonEmpty (a, b) +combineEq = \case + x :| [] -> x :| [] + x :| xx@(_:_) -> go x xx + where + go z [] = z :| [] + go z@(kz,_) (x@(kx,xx):xs') + | kx==kz = go (kx,xx) xs' + | otherwise = z NE.<| go x xs' + +combineEqWith + :: Eq a + => (a -> b -> b -> b) + -> NonEmpty (a, b) + -> NonEmpty (a, b) +combineEqWith f = \case + x :| [] -> x :| [] + x :| xx@(_:_) -> go x xx + where + go z [] = z :| [] + go z@(kz,zz) (x@(kx,xx):xs') + | kx==kz = let yy = f kx xx zz in go (kx,yy) xs' + | otherwise = z NE.<| go x xs' diff --git a/src/Data/Map/NonEmpty/Internal.hs b/src/Data/Map/NonEmpty/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Map/NonEmpty/Internal.hs @@ -0,0 +1,593 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Data.Map.NonEmpty.Internal +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Unsafe internal-use functions used in the implementation of +-- "Data.Map.NonEmpty". These functions can potentially be used to break +-- the abstraction of 'NEMap' and produce unsound maps, so be wary! +module Data.Map.NonEmpty.Internal ( + -- * Non-Empty Map type + NEMap(..) + , singleton + , nonEmptyMap + , withNonEmpty + , fromList + , toList + , map + , insertWith + , union + , unions + , elems + , size + , toMap + -- * Folds + , foldr + , foldr' + , foldr1 + , foldl + , foldl' + , foldl1 + -- * Traversals + , traverseWithKey + , traverseWithKey1 + , foldMapWithKey + -- * Unsafe Map Functions + , insertMinMap + , insertMaxMap + -- * Debug + , valid + ) where + +import Control.Applicative +import Control.DeepSeq +import Data.Coerce +import Data.Data +import Data.Function +import Data.Functor.Apply +import Data.Functor.Classes +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Map.Internal (Map(..)) +import Data.Maybe +import Data.Semigroup +import Data.Semigroup.Foldable (Foldable1(fold1)) +import Data.Semigroup.Traversable (Traversable1(..)) +import Data.Typeable (Typeable) +import Prelude hiding (foldr1, foldl1, foldr, foldl, map) +import Text.Read +import qualified Data.Foldable as F +import qualified Data.Map as M +import qualified Data.Map.Internal as M +import qualified Data.Semigroup.Foldable as F1 + +-- | A non-empty (by construction) map from keys @k@ to values @a@. At +-- least one key-value pair exists in an @'NEMap' k v@ at all times. +-- +-- Functions that /take/ an 'NEMap' can safely operate on it with the +-- assumption that it has at least one key-value pair. +-- +-- Functions that /return/ an 'NEMap' provide an assurance that the result +-- has at least one key-value pair. +-- +-- "Data.Map.NonEmpty" re-exports the API of "Data.Map", faithfully +-- reproducing asymptotics, typeclass constraints, and semantics. +-- Functions that ensure that input and output maps are both non-empty +-- (like 'Data.Map.NonEmpty.insert') return 'NEMap', but functions that +-- might potentially return an empty map (like 'Data.Map.NonEmpty.delete') +-- return a 'Map' instead. +-- +-- You can directly construct an 'NEMap' with the API from +-- "Data.Map.NonEmpty"; it's more or less the same as constructing a normal +-- 'Map', except you don't have access to 'Data.Map.empty'. There are also +-- a few ways to construct an 'NEMap' from a 'Map': +-- +-- 1. The 'nonEmptyMap' smart constructor will convert a @'Map' k a@ into +-- a @'Maybe' ('NEMap' k a)@, returning 'Nothing' if the original 'Map' +-- was empty. +-- 2. You can use the 'Data.Map.NonEmpty.insertMap' family of functions to +-- insert a value into a 'Map' to create a guaranteed 'NEMap'. +-- 3. You can use the 'Data.Map.NonEmpty.IsNonEmpty' and +-- 'Data.Map.NonEmpty.IsEmpty' patterns to "pattern match" on a 'Map' +-- to reveal it as either containing a 'NEMap' or an empty map. +-- 4. 'withNonEmpty' offers a continuation-based interface for +-- deconstructing a 'Map' and treating it as if it were an 'NEMap'. +-- +-- You can convert an 'NEMap' into a 'Map' with 'toMap' or +-- 'Data.Map.NonEmpty.IsNonEmpty', essentially "obscuring" the non-empty +-- property from the type. +data NEMap k a = + NEMap { nemK0 :: !k -- ^ invariant: must be smaller than smallest key in map + , nemV0 :: a + , nemMap :: !(Map k a) + } + deriving (Typeable) + +instance (Eq k, Eq a) => Eq (NEMap k a) where + t1 == t2 = M.size (nemMap t1) == M.size (nemMap t2) + && toList t1 == toList t2 + +instance (Ord k, Ord a) => Ord (NEMap k a) where + compare = compare `on` toList + (<) = (<) `on` toList + (>) = (>) `on` toList + (<=) = (<=) `on` toList + (>=) = (>=) `on` toList + +instance Eq2 NEMap where + liftEq2 eqk eqv m n = + size m == size n && liftEq (liftEq2 eqk eqv) (toList m) (toList n) + +instance Eq k => Eq1 (NEMap k) where + liftEq = liftEq2 (==) + +instance Ord2 NEMap where + liftCompare2 cmpk cmpv m n = + liftCompare (liftCompare2 cmpk cmpv) (toList m) (toList n) + +instance Ord k => Ord1 (NEMap k) where + liftCompare = liftCompare2 compare + +instance Show2 NEMap where + liftShowsPrec2 spk slk spv slv d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) + where + sp = liftShowsPrec2 spk slk spv slv + sl = liftShowList2 spk slk spv slv + +instance Show k => Show1 (NEMap k) where + liftShowsPrec = liftShowsPrec2 showsPrec showList + +instance (Ord k, Read k) => Read1 (NEMap k) where + liftReadsPrec rp rl = readsData $ + readsUnaryWith (liftReadsPrec rp' rl') "fromList" fromList + where + rp' = liftReadsPrec rp rl + rl' = liftReadList rp rl + +instance (Ord k, Read k, Read e) => Read (NEMap k e) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- parens . prec 10 $ readPrec + return (fromList xs) + readListPrec = readListPrecDefault + +instance (Show k, Show a) => Show (NEMap k a) where + showsPrec d m = showParen (d > 10) $ + showString "fromList (" . shows (toList m) . showString ")" + +instance (NFData k, NFData a) => NFData (NEMap k a) where + rnf (NEMap k v a) = rnf k `seq` rnf v `seq` rnf a + +-- Data instance code from Data.Map.Internal +-- +-- Copyright : (c) Daan Leijen 2002 +-- (c) Andriy Palamarchuk 2008 +instance (Data k, Data a, Ord k) => Data (NEMap k a) where + gfoldl f z m = z fromList `f` toList m + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = mapDataType + dataCast2 = gcast2 + +fromListConstr :: Constr +fromListConstr = mkConstr mapDataType "fromList" [] Prefix + +mapDataType :: DataType +mapDataType = mkDataType "Data.Map.NonEmpty.NonEmpty.Internal.NEMap" [fromListConstr] + +-- | /O(n)/. Fold the values in the map using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'elems'@. +-- +-- > elemsList map = foldr (:) [] map +-- +-- > let f a len = len + (length a) +-- > foldr f 0 (fromList ((5,"a") :| [(3,"bbb")])) == 4 +foldr :: (a -> b -> b) -> b -> NEMap k a -> b +foldr f z (NEMap _ v m) = v `f` M.foldr f z m +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (a -> b -> b) -> b -> NEMap k a -> b +foldr' f z (NEMap _ v m) = v `f` y + where + !y = M.foldr' f z m +{-# INLINE foldr' #-} + +-- | /O(n)/. A version of 'foldr' that uses the value at the maximal key in +-- the map as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldr1' for 'Map', this function is +-- total if the input function is total. +foldr1 :: (a -> a -> a) -> NEMap k a -> a +foldr1 f (NEMap _ v m) = maybe v (f v . uncurry (M.foldr f)) + . M.maxView + $ m +{-# INLINE foldr1 #-} + +-- | /O(n)/. Fold the values in the map using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'elems'@. +-- +-- > elemsList = reverse . foldl (flip (:)) [] +-- +-- > let f len a = len + (length a) +-- > foldl f 0 (fromList ((5,"a") :| [(3,"bbb")])) == 4 +foldl :: (a -> b -> a) -> a -> NEMap k b -> a +foldl f z (NEMap _ v m) = M.foldl f (f z v) m +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> NEMap k b -> a +foldl' f z (NEMap _ v m) = M.foldl' f x m + where + !x = f z v +{-# INLINE foldl' #-} + +-- | /O(n)/. A version of 'foldl' that uses the value at the minimal key in +-- the map as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldl1' for 'Map', this function is +-- total if the input function is total. +foldl1 :: (a -> a -> a) -> NEMap k a -> a +foldl1 f (NEMap _ v m) = M.foldl f v m +{-# INLINE foldl1 #-} + +-- | /O(n)/. Fold the keys and values in the map using the given semigroup, +-- such that +-- +-- @'foldMapWithKey' f = 'Data.Semigroup.Foldable.fold1' . 'Data.Map.NonEmpty.mapWithKey' f@ +-- +-- This can be an asymptotically faster than +-- 'Data.Map.NonEmpty.foldrWithKey' or 'Data.Map.NonEmpty.foldlWithKey' for +-- some monoids. + +-- TODO: benchmark against maxView method +foldMapWithKey + :: Semigroup m + => (k -> a -> m) + -> NEMap k a + -> m +foldMapWithKey f (NEMap k0 v m) = maybe (f k0 v) (f k0 v <>) + . getOption + . M.foldMapWithKey (\k -> Option . Just . f k) + $ m +{-# INLINE foldMapWithKey #-} + +-- | /O(n)/. Map a function over all values in the map. +-- +-- > map (++ "x") (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "bx") :| [(5, "ax")]) +map :: (a -> b) -> NEMap k a -> NEMap k b +map f (NEMap k0 v m) = NEMap k0 (f v) (M.map f m) +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs + #-} +{-# RULES +"map/coerce" map coerce = coerce + #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. +-- The expression (@'union' t1 t2@) takes the left-biased union of @t1@ and +-- @t2@. It prefers @t1@ when duplicate keys are encountered, i.e. +-- (@'union' == 'Data.Map.NonEmpty.unionWith' 'const'@). +-- +-- > union (fromList ((5, "a") :| [(3, "b")])) (fromList ((5, "A") :| [(7, "C")])) == fromList ((3, "b") :| [(5, "a"), (7, "C")]) +union + :: Ord k + => NEMap k a + -> NEMap k a + -> NEMap k a +union n1@(NEMap k1 v1 m1) n2@(NEMap k2 v2 m2) = case compare k1 k2 of + LT -> NEMap k1 v1 . M.union m1 . toMap $ n2 + EQ -> NEMap k1 v1 . M.union m1 $ m2 + GT -> NEMap k2 v2 . M.union (toMap n1) $ m2 +{-# INLINE union #-} + +-- | The left-biased union of a non-empty list of maps. +-- +-- > unions (fromList ((5, "a") :| [(3, "b")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "A3") :| [(3, "B3")])]) +-- > == fromList [(3, "b"), (5, "a"), (7, "C")] +-- > unions (fromList ((5, "A3") :| [(3, "B3")]) :| [fromList ((5, "A") :| [(7, "C")]), fromList ((5, "a") :| [(3, "b")])]) +-- > == fromList ((3, "B3") :| [(5, "A3"), (7, "C")]) +unions + :: (Foldable1 f, Ord k) + => f (NEMap k a) + -> NEMap k a +unions (F1.toNonEmpty->(m :| ms)) = F.foldl' union m ms +{-# INLINE unions #-} + +-- | /O(n)/. +-- Return all elements of the map in the ascending order of their keys. +-- +-- > elems (fromList ((5,"a") :| [(3,"b")])) == ("b" :| ["a"]) +elems :: NEMap k a -> NonEmpty a +elems (NEMap _ v m) = v :| M.elems m +{-# INLINE elems #-} + +-- | /O(1)/. The number of elements in the map. Guaranteed to be greater +-- than zero. +-- +-- > size (singleton 1 'a') == 1 +-- > size (fromList ((1,'a') :| [(2,'c'), (3,'b')])) == 3 +size :: NEMap k a -> Int +size (NEMap _ _ m) = 1 + M.size m +{-# INLINE size #-} + +-- | /O(log n)/. +-- Convert a non-empty map back into a normal possibly-empty map, for usage +-- with functions that expect 'Map'. +-- +-- Can be thought of as "obscuring" the non-emptiness of the map in its +-- type. See the 'Data.Map.NonEmpty.IsNotEmpty' pattern. +-- +-- 'nonEmptyMap' and @'maybe' 'Data.Map.empty' 'toMap'@ form an isomorphism: they +-- are perfect structure-preserving inverses of eachother. +-- +-- > toMap (fromList ((3,"a") :| [(5,"b")])) == Data.Map.fromList [(3,"a"), (5,"b")] +toMap :: NEMap k a -> Map k a +toMap (NEMap k v m) = insertMinMap k v m +{-# INLINE toMap #-} + +-- | /O(n)/. +-- @'traverseWithKey' f m == 'fromList' <$> 'traverse' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- That is, behaves exactly like a regular 'traverse' except that the traversing +-- function also has access to the key associated with a value. +-- +-- /Use 'traverseWithKey1'/ whenever possible (if your 'Applicative' +-- also has 'Apply' instance). This version is provided only for types +-- that do not have 'Apply' instance, since 'Apply' is not at the moment +-- (and might not ever be) an official superclass of 'Applicative'. +-- +-- @ +-- 'traverseWithKey' f = 'unwrapApplicative' . 'traverseWithKey1' (\\k -> WrapApplicative . f k) +-- @ +traverseWithKey + :: Applicative t + => (k -> a -> t b) + -> NEMap k a + -> t (NEMap k b) +traverseWithKey f (NEMap k v m0) = NEMap k <$> f k v <*> M.traverseWithKey f m0 +{-# INLINE traverseWithKey #-} + +-- | /O(n)/. +-- @'traverseWithKey1' f m == 'fromList' <$> 'traverse1' (\(k, v) -> (,) k <$> f k v) ('toList' m)@ +-- +-- That is, behaves exactly like a regular 'traverse1' except that the traversing +-- function also has access to the key associated with a value. +-- +-- Is more general than 'traverseWithKey', since works with all 'Apply', +-- and not just 'Applicative'. + +-- TODO: benchmark against maxView-based methods +traverseWithKey1 + :: Apply t + => (k -> a -> t b) + -> NEMap k a + -> t (NEMap k b) +traverseWithKey1 f (NEMap k0 v m0) = case runMaybeApply m1 of + Left m2 -> NEMap k0 <$> f k0 v <.> m2 + Right m2 -> flip (NEMap k0) m2 <$> f k0 v + where + m1 = M.traverseWithKey (\k -> MaybeApply . Left . f k) m0 +{-# INLINABLE traverseWithKey1 #-} + +-- | /O(n)/. Convert the map to a non-empty list of key\/value pairs. +-- +-- > toList (fromList ((5,"a") :| [(3,"b")])) == ((3,"b") :| [(5,"a")]) +toList :: NEMap k a -> NonEmpty (k, a) +toList (NEMap k v m) = (k,v) :| M.toList m +{-# INLINE toList #-} + +-- | /O(log n)/. Smart constructor for an 'NEMap' from a 'Map'. Returns +-- 'Nothing' if the 'Map' was originally actually empty, and @'Just' n@ +-- with an 'NEMap', if the 'Map' was not empty. +-- +-- 'nonEmptyMap' and @'maybe' 'Data.Map.empty' 'toMap'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- See 'Data.Map.NonEmpty.IsNonEmpty' for a pattern synonym that lets you +-- "match on" the possiblity of a 'Map' being an 'NEMap'. +-- +-- > nonEmptyMap (Data.Map.fromList [(3,"a"), (5,"b")]) == Just (fromList ((3,"a") :| [(5,"b")])) +nonEmptyMap :: Map k a -> Maybe (NEMap k a) +nonEmptyMap = (fmap . uncurry . uncurry) NEMap . M.minViewWithKey +{-# INLINE nonEmptyMap #-} + +-- | /O(log n)/. A general continuation-based way to consume a 'Map' as if +-- it were an 'NEMap'. @'withNonEmpty' def f@ will take a 'Map'. If map is +-- empty, it will evaluate to @def@. Otherwise, a non-empty map 'NEMap' +-- will be fed to the function @f@ instead. +-- +-- @'nonEmptyMap' == 'withNonEmpty' 'Nothing' 'Just'@ +withNonEmpty + :: r -- ^ value to return if map is empty + -> (NEMap k a -> r) -- ^ function to apply if map is not empty + -> Map k a + -> r +withNonEmpty def f = maybe def f . nonEmptyMap +{-# INLINE withNonEmpty #-} + +-- | /O(n*log n)/. Build a non-empty map from a non-empty list of +-- key\/value pairs. See also 'Data.Map.NonEmpty.fromAscList'. If the list +-- contains more than one value for the same key, the last value for the +-- key is retained. +-- +-- > fromList ((5,"a") :| [(3,"b"), (5, "c")]) == fromList ((5,"c") :| [(3,"b")]) +-- > fromList ((5,"c") :| [(3,"b"), (5, "a")]) == fromList ((5,"a") :| [(3,"b")]) + +-- TODO: write manually and optimize to be equivalent to +-- 'fromDistinctAscList' if items are ordered, just like the actual +-- 'M.fromList'. +fromList :: Ord k => NonEmpty (k, a) -> NEMap k a +fromList ((k, v) :| xs) = withNonEmpty (singleton k v) (insertWith (const id) k v) + . M.fromList + $ xs +{-# INLINE fromList #-} + +-- | /O(1)/. A map with a single element. +-- +-- > singleton 1 'a' == fromList ((1, 'a') :| []) +-- > size (singleton 1 'a') == 1 +singleton :: k -> a -> NEMap k a +singleton k v = NEMap k v M.empty +{-# INLINE singleton #-} + +-- | /O(log n)/. Insert with a function, combining new value and old value. +-- @'insertWith' f key value mp@ will insert the pair (key, value) into +-- @mp@ if key does not exist in the map. If the key does exist, the +-- function will insert the pair @(key, f new_value old_value)@. +-- +-- See 'Data.Map.NonEmpty.insertMapWith' for a version where the first +-- argument is a 'Map'. +-- +-- > insertWith (++) 5 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "xxxa")]) +-- > insertWith (++) 7 "xxx" (fromList ((5,"a") :| [(3,"b")])) == fromList ((3, "b") :| [(5, "a"), (7, "xxx")]) +insertWith + :: Ord k + => (a -> a -> a) + -> k + -> a + -> NEMap k a + -> NEMap k a +insertWith f k v n@(NEMap k0 v0 m) = case compare k k0 of + LT -> NEMap k v . toMap $ n + EQ -> NEMap k (f v v0) m + GT -> NEMap k0 v0 $ M.insertWith f k v m +{-# INLINE insertWith #-} + + +-- | Left-biased union +instance Ord k => Semigroup (NEMap k a) where + (<>) = union + {-# INLINE (<>) #-} + sconcat = unions + {-# INLINE sconcat #-} + +instance Functor (NEMap k) where + fmap = map + {-# INLINE fmap #-} + x <$ NEMap k _ m = NEMap k x (x <$ m) + {-# INLINE (<$) #-} + +-- | Traverses elements in order of ascending keys +-- +-- 'Data.Foldable.foldr1', 'Data.Foldable.foldl1', 'Data.Foldable.minimum', +-- 'Data.Foldable.maximum' are all total. +instance Foldable (NEMap k) where +#if MIN_VERSION_base(4,11,0) + fold (NEMap _ v m) = v <> F.fold m + {-# INLINE fold #-} + foldMap f (NEMap _ v m) = f v <> foldMap f m + {-# INLINE foldMap #-} +#else + fold (NEMap _ v m) = v `mappend` F.fold m + {-# INLINE fold #-} + foldMap f (NEMap _ v m) = f v `mappend` foldMap f m + {-# INLINE foldMap #-} +#endif + foldr = foldr + {-# INLINE foldr #-} + foldr' = foldr' + {-# INLINE foldr' #-} + foldr1 = foldr1 + {-# INLINE foldr1 #-} + foldl = foldl + {-# INLINE foldl #-} + foldl' = foldl' + {-# INLINE foldl' #-} + foldl1 = foldl1 + {-# INLINE foldl1 #-} + null _ = False + {-# INLINE null #-} + length = size + {-# INLINE length #-} + elem x (NEMap _ v m) = F.elem x m + || x == v + {-# INLINE elem #-} + -- TODO: use build + toList = F.toList . elems + {-# INLINE toList #-} + +-- | Traverses elements in order of ascending keys +instance Traversable (NEMap k) where + traverse f (NEMap k v m) = NEMap k <$> f v <*> traverse f m + {-# INLINE traverse #-} + sequenceA (NEMap k v m) = NEMap k <$> v <*> sequenceA m + {-# INLINE sequenceA #-} + +-- | Traverses elements in order of ascending keys +instance Foldable1 (NEMap k) where + fold1 (NEMap _ v m) = maybe v (v <>) + . getOption + . F.foldMap (Option . Just) + $ m + {-# INLINE fold1 #-} + foldMap1 f = foldMapWithKey (const f) + {-# INLINE foldMap1 #-} + toNonEmpty = elems + {-# INLINE toNonEmpty #-} + +-- | Traverses elements in order of ascending keys +instance Traversable1 (NEMap k) where + traverse1 f = traverseWithKey1 (const f) + {-# INLINE traverse1 #-} + sequence1 (NEMap k v m0) = case runMaybeApply m1 of + Left m2 -> NEMap k <$> v <.> m2 + Right m2 -> flip (NEMap k) m2 <$> v + where + m1 = traverse (MaybeApply . Left) m0 + {-# INLINABLE sequence1 #-} + +-- | /O(n)/. Test if the internal map structure is valid. +valid :: Ord k => NEMap k a -> Bool +valid (NEMap k _ m) = M.valid m + && all ((k <) . fst . fst) (M.minViewWithKey m) + + + + + +-- | /O(log n)/. Insert new key and value into a map where keys are +-- /strictly greater than/ the new key. That is, the new key must be +-- /strictly less than/ all keys present in the 'Map'. /The precondition +-- is not checked./ +-- +-- While this has the same asymptotics as @Data.Map.insert@, it saves +-- a constant factor for key comparison (so may be helpful if comparison is +-- expensive) and also does not require an 'Ord' instance for the key type. +insertMinMap :: k -> a -> Map k a -> Map k a +insertMinMap kx x = \case + Tip -> M.singleton kx x + Bin _ ky y l r -> M.balanceL ky y (insertMinMap kx x l) r +{-# INLINABLE insertMinMap #-} + +-- | /O(log n)/. Insert new key and value into a map where keys are +-- /strictly less than/ the new key. That is, the new key must be +-- /strictly greater than/ all keys present in the 'Map'. /The +-- precondition is not checked./ +-- +-- While this has the same asymptotics as @Data.Map.insert@, it saves +-- a constant factor for key comparison (so may be helpful if comparison is +-- expensive) and also does not require an 'Ord' instance for the key type. +insertMaxMap :: k -> a -> Map k a -> Map k a +insertMaxMap kx x = \case + Tip -> M.singleton kx x + Bin _ ky y l r -> M.balanceR ky y l (insertMaxMap kx x r) +{-# INLINABLE insertMaxMap #-} diff --git a/src/Data/Sequence/NonEmpty.hs b/src/Data/Sequence/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/Sequence/NonEmpty.hs @@ -0,0 +1,1012 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.Sequence.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Finite Sequences +-- +-- | An @'NESeq' a@ is a non-empty (but finite) sequence of values of type +-- @a@. Generally has the same interface as 'Data.List.NonEmpty.NonEmpty'. +-- This is a non-empty version of 'Data.Sequence.Seq' from "Data.Sequence". +-- +-- The main differences between this type and 'Data.List.NonEmpty.NonEmpty' +-- are: +-- +-- * You cannot have infinite 'NESeq's +-- * You have constant-time consing from either end, and constant-time +-- unconsing as well (through '<|', '|>', ':<||', and ':||>') +-- * Concatenation ('><', '|><', '><|') is logarithmic-time. +-- * You have logarithmic-time indexing and updating at a given index. +-- +-- While asymptotics are often better than for 'Data.List.NonEmpty.NonEmpty', there is +-- a decent constant factor involved in most operations. +-- +-- See documentation for 'NESeq' for information on how to convert and +-- manipulate such non-empty sequences +-- +-- This module essentially re-imports the API of "Data.Sequence.Lazy" and its +-- 'Seq' type, along with semantics and asymptotics. +-- +-- Because 'NESeq' is implemented using 'Seq', all of the caveats of using +-- 'Seq' apply. +-- +-- All functions take non-empty sequences as inputs. In situations where +-- their results can be guarunteed to also be non-empty, they also return +-- non-empty maps. In situations where their results could potentially be +-- empty, 'Seq' is returned instead. +-- +-- Some functions (like 'spanl', 'spanr', 'breakl', 'breakr', 'partition', +-- 'splitAt') have modified return types to account for possible +-- configurations of non-emptiness. +-- +-- Some functions ('head', 'last', 'tail', 'init') are provided because +-- they are total for non-empty sequences. +-- +-- This module is intended to be imported qualified, to avoid name clashes with +-- "Prelude" and "Data.Sequence" functions: +-- +-- > import qualified Data.Sequence.NonEmpty as NESeq +module Data.Sequence.NonEmpty ( + -- * Finite sequences + NESeq ((:<||), (:||>)) + -- ** Conversions between empty and non-empty sequences + , pattern IsNonEmpty + , pattern IsEmpty + , nonEmptySeq + , toSeq + , withNonEmpty + , unsafeFromSeq + , insertSeqAt + -- * Construction + , singleton + , (<|) + , (|>) + , (><) + , (|><) + , (><|) + , fromList + , fromFunction + -- ** Repetition + , replicate + , replicateA + , replicateA1 + , replicateM + , cycleTaking + -- ** Iterative construction + , iterateN + , unfoldr + , unfoldl + -- * Deconstruction + -- | Additional functions for deconstructing sequences are available + -- via the 'Foldable' instance of 'NESeq'. + , head + , tail + , last + , init + -- ** Queries + , length + + -- * Scans + , scanl + , scanl1 + , scanr + , scanr1 + -- * Sublists + , tails + , inits + , chunksOf + -- ** Sequential searches + , takeWhileL + , takeWhileR + , dropWhileL + , dropWhileR + , spanl + , spanr + , breakl + , breakr + , partition + , filter + -- * Sorting + , sort + , sortBy + , sortOn + , unstableSort + , unstableSortBy + , unstableSortOn + -- * Indexing + , lookup + , (!?) + , index + , adjust + , adjust' + , update + , take + , drop + , insertAt + , deleteAt + , splitAt + -- ** Indexing with predicates + -- | These functions perform sequential searches from the left + -- or right ends of the sequence returning indices of matching + -- elements. + , elemIndexL + , elemIndicesL + , elemIndexR + , elemIndicesR + , findIndexL + , findIndicesL + , findIndexR + , findIndicesR + -- * Folds + -- | General folds are available via the 'Foldable' instance of 'Seq'. + , foldMapWithIndex + , foldlWithIndex + , foldrWithIndex + -- * Transformations + , mapWithIndex + , traverseWithIndex + , traverseWithIndex1 + , reverse + , intersperse + -- ** Zips and unzip + , zip + , zipWith + , zip3 + , zipWith3 + , zip4 + , zipWith4 + , unzip + , unzipWith + ) where + +import Control.Applicative +import Control.Monad hiding (replicateM) +import Data.Bifunctor +import Data.Functor.Apply +import Data.Sequence (Seq(..)) +import Data.Sequence.NonEmpty.Internal +import Data.These +import Prelude hiding (length, scanl, scanl1, scanr, scanr1, splitAt, zip, zipWith, zip3, zipWith3, unzip, replicate, filter, reverse, lookup, take, drop, head, tail, init, last, map) +import qualified Data.Sequence as Seq + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'Seq' as if it were either a @'IsNonEmpty' n@ (where @n@ is a 'NESeq') +-- or an 'IsEmpty'. +-- +-- For example, you can pattern match on a 'Seq': +-- +-- @ +-- safeHead :: 'Seq' Int -> Int +-- safeHead ('IsNonEmpty' (x :<|| _)) = x -- here, user provided a non-empty sequence, and @n@ is the 'NESeq' +-- safeHead 'IsEmpty' = 0 -- here the user provided an empty sequence +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original 'Seq' was /not/ +-- empty, and you have a verified-non-empty 'NESeq' @n@ to use. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a 'NESeq' back into a 'Seq', obscuring its non-emptiness (see 'toSeq'). +pattern IsNonEmpty :: NESeq a -> Seq a +pattern IsNonEmpty n <- (nonEmptySeq->Just n) + where + IsNonEmpty n = toSeq n + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'Seq' as if it were either a @'IsNonEmpty' n@ (where @n@ is +-- a 'NESeq') or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original 'Seq' was empty. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'Data.Seq.empty'. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: Seq a +pattern IsEmpty <- (Seq.null->True) + where + IsEmpty = Seq.empty + +{-# COMPLETE IsNonEmpty, IsEmpty #-} + +-- | /O(1)/. Smart constructor for an 'NESeq' from a 'Seq'. Returns +-- 'Nothing' if the 'Seq' was originally actually empty, and @'Just' n@ +-- with an 'NESeq', if the 'Seq' was not empty. +-- +-- 'nonEmptySeq' and @'maybe' 'Data.Sequence.empty' 'toSeq'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- See 'Data.Sequence.NonEmpty.IsNonEmpty' for a pattern synonym that lets +-- you "match on" the possiblity of a 'Seq' being an 'NESeq'. +-- +-- > nonEmptySeq (Data.Sequence.fromList [1,2,3]) == Just (fromList (1) :| [2,3]) +nonEmptySeq :: Seq a -> Maybe (NESeq a) +nonEmptySeq (x :<| xs) = Just $ x :<|| xs +nonEmptySeq Empty = Nothing +{-# INLINE nonEmptySeq #-} + +-- | /O(1)/. Unsafe version of 'nonEmptySeq'. Coerces a 'Seq' into an +-- 'NESeq', but is undefined (throws a runtime exception when evaluation is +-- attempted) for an empty 'Seq'. +unsafeFromSeq :: Seq a -> NESeq a +unsafeFromSeq (x :<| xs) = x :<|| xs +unsafeFromSeq Empty = errorWithoutStackTrace "NESeq.unsafeFromSeq: empty seq" +{-# INLINE unsafeFromSeq #-} + +-- | Turn a 'Seq' into a guarantted non-empty 'NESeq' by adding an element +-- at a given index. +-- +-- > insertSeqAt 1 0 (Data.Sequence.fromList [1,2,3]) == fromList (1 :| [0,2,3]) +insertSeqAt :: Int -> a -> Seq a -> NESeq a +insertSeqAt i y + | i <= 0 = (y :<||) + | otherwise = \case + x :<| xs -> x :<|| Seq.insertAt (i - 1) y xs + Empty -> y :<|| Seq.empty +{-# INLINE insertSeqAt #-} + +-- | \( O(1) \). Add an element to the right end of a non-empty sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(|>) :: NESeq a -> a -> NESeq a +(x :<|| xs) |> y = x :<|| (xs Seq.|> y) +{-# INLINE (|>) #-} + +-- | \( O(\log(\min(n_1,n_2))) \). Concatenate a non-empty sequence with +-- a potentially empty sequence ('Seq'), to produce a guaranteed non-empty +-- sequence. Mnemonic: like '><', but a pipe for the guarunteed non-empty +-- side. +(><|) :: Seq a -> NESeq a -> NESeq a +xs ><| ys = withNonEmpty ys (>< ys) xs +{-# INLINE (><|) #-} + +infixl 5 |> +infixr 5 ><| + +-- | 'replicateA' is an 'Applicative' version of 'replicate', and makes \( +-- O(\log n) \) calls to 'liftA2' and 'pure'. Is only defined when @n@ is +-- positive. +-- +-- > replicateA n x = sequenceA (replicate n x) +-- +-- Is a more restrictive version of 'replicateA1'. 'replicateA1' should be +-- preferred whenever possible. +replicateA :: Applicative f => Int -> f a -> f (NESeq a) +replicateA n x + | n < 1 = error "NESeq.replicateA: must take a positive integer argument" + | otherwise = liftA2 (:<||) x (Seq.replicateA (n - 1) x) +{-# INLINE replicateA #-} + +-- | 'replicateA' is an 'Apply' version of 'replicate', and makes \( O(\log +-- n) \) calls to '<.>'. Is only defined when @n@ is positive. +-- +-- > replicateA1 n x = sequence1 (replicate n x) +replicateA1 :: Apply f => Int -> f a -> f (NESeq a) +replicateA1 n x + | n < 1 = error "NESeq.replicateA1: must take a positive integer argument" + | otherwise = case runMaybeApply (Seq.replicateA (n - 1) (MaybeApply (Left x))) of + Left xs -> (:<||) <$> x <.> xs + Right xs -> (:<|| xs) <$> x +{-# INLINE replicateA1 #-} + +-- | An alias of 'replicateA'. +replicateM :: Applicative m => Int -> m a -> m (NESeq a) +replicateM = replicateA +{-# INLINE replicateM #-} + +-- | /O(/log/ k)/. @'cycleTaking' k xs@ forms a sequence of length @k@ by +-- repeatedly concatenating @xs@ with itself. Is only defined when @k@ is +-- positive. +-- +-- prop> cycleTaking k = fromList . fromJust . nonEmpty . take k . cycle . toList + +-- If you wish to concatenate a non-empty sequence @xs@ with itself precisely +-- @k@ times, you can use @cycleTaking (k * length xs)@ or just +-- @replicate k () *> xs@. +cycleTaking :: Int -> NESeq a -> NESeq a +cycleTaking n xs0@(x :<|| xs) + | n < 1 = error "NESeq.cycleTaking: must take a positive integer argument" + | n < Seq.length xs = x :<|| Seq.take (n - 1) xs + | otherwise = xs0 |>< Seq.cycleTaking (n - length xs0) (toSeq xs0) +{-# INLINE cycleTaking #-} + +-- | \( O(n) \). Constructs a sequence by repeated application of +-- a function to a seed value. Is only defined if given a positive value. +-- +-- > iterateN n f x = fromList (fromJust (nonEmpty ((Prelude.take n (Prelude.iterate f x))))) +iterateN :: Int -> (a -> a) -> a -> NESeq a +iterateN n f x + | n < 1 = error "NESeq.iterateN: must take a positive integer argument" + | otherwise = x :<|| Seq.iterateN (n - 1) f (f x) +{-# INLINE iterateN #-} + +-- | Builds a sequence from a seed value. Takes time linear in the +-- number of generated elements. /WARNING:/ If the number of generated +-- elements is infinite, this method will not terminate. +unfoldr :: (b -> (a, Maybe b)) -> b -> NESeq a +unfoldr f = go + where + go x0 = y :<|| maybe Seq.empty (toSeq . go) x1 + where + (y, x1) = f x0 +{-# INLINE unfoldr #-} + +-- | @'unfoldl' f x@ is equivalent to @'reverse' ('unfoldr' ('fmap' swap . f) x)@. +unfoldl :: (b -> (Maybe b, a)) -> b -> NESeq a +unfoldl f = go + where + go x0 = maybe Seq.empty (toSeq . go) x1 :||> y + where + (x1, y) = f x0 +{-# INLINE unfoldl #-} + +-- | /O(1)/. Retrieve the left-most item in a non-empty sequence. Note +-- that this function is total. +head :: NESeq a -> a +head (x :<|| _) = x +{-# INLINE head #-} + +-- | /O(1)/. Delete the left-most item in a non-empty sequence. Returns +-- a potentially empty sequence ('Seq') in the case that the original +-- 'NESeq' contained only a single element. Note that this function is +-- total. +tail :: NESeq a -> Seq a +tail (_ :<|| xs) = xs +{-# INLINE tail #-} + +-- | /O(1)/. Retrieve the right-most item in a non-empty sequence. Note +-- that this function is total. +last :: NESeq a -> a +last (_ :||> x) = x +{-# INLINE last #-} + +-- | /O(1)/. Delete the right-most item in a non-empty sequence. Returns +-- a potentially empty sequence ('Seq') in the case that the original +-- 'NESeq' contained only a single element. Note that this function is +-- total. +init :: NESeq a -> Seq a +init (xs :||> _) = xs +{-# INLINE init #-} + + +-- | 'scanl' is similar to 'foldl', but returns a sequence of reduced +-- values from the left: +-- +-- > scanl f z (fromList [x1, x2, ...]) = fromList [z, z `f` x1, (z `f` x1) `f` x2, ...] +scanl :: (a -> b -> a) -> a -> NESeq b -> NESeq a +scanl f y0 (x :<|| xs) = y0 :<|| Seq.scanl f (f y0 x) xs +{-# INLINE scanl #-} + +-- | 'scanl1' is a variant of 'scanl' that has no starting value argument: +-- +-- > scanl1 f (fromList [x1, x2, ...]) = fromList [x1, x1 `f` x2, ...] +scanl1 :: (a -> a -> a) -> NESeq a -> NESeq a +scanl1 f (x :<|| xs) = withNonEmpty (singleton x) (scanl f x) xs +{-# INLINE scanl1 #-} + +-- | 'scanr' is the right-to-left dual of 'scanl'. +scanr :: (a -> b -> b) -> b -> NESeq a -> NESeq b +scanr f y0 (xs :||> x) = Seq.scanr f (f x y0) xs :||> y0 +{-# INLINE scanr #-} + +-- | 'scanr1' is a variant of 'scanr' that has no starting value argument. +scanr1 :: (a -> a -> a) -> NESeq a -> NESeq a +scanr1 f (xs :||> x) = withNonEmpty (singleton x) (scanr f x) xs +{-# INLINE scanr1 #-} + +-- | \( O(n) \). Returns a sequence of all non-empty prefixes of this +-- sequence, shortest first. For example, +-- +-- > tails (fromList (1:|[2,3])) = fromList (fromList (1:|[]) :| [fromList (1:|[2]), fromList (1:|[2,3])) +-- +-- Evaluating the \( i \)th prefix takes \( O(\log(\min(i, n-i))) \), but evaluating +-- every prefix in the sequence takes \( O(n) \) due to sharing. + +-- TODO: is this true? +inits :: NESeq a -> NESeq (NESeq a) +inits xs@(ys :||> _) = withNonEmpty (singleton xs) ((|> xs) . inits) ys +{-# INLINABLE inits #-} + +-- | \(O \Bigl(\bigl(\frac{n}{c}\bigr) \log c\Bigr)\). @chunksOf c xs@ splits @xs@ into chunks of size @c>0@. +-- If @c@ does not divide the length of @xs@ evenly, then the last element +-- of the result will be short. Is only defined if @c@ is a positive +-- number. +-- +-- Side note: the given performance bound is missing some messy terms that only +-- really affect edge cases. Performance degrades smoothly from \( O(1) \) (for +-- \( c = n \)) to \( O(n) \) (for \( c = 1 \)). The true bound is more like +-- \( O \Bigl( \bigl(\frac{n}{c} - 1\bigr) (\log (c + 1)) + 1 \Bigr) \) + +-- TODO: is this true? +chunksOf :: Int -> NESeq a -> NESeq (NESeq a) +chunksOf n = go + where + go xs = case splitAt n xs of + This ys -> singleton ys + That _ -> e + These ys zs -> ys <| go zs + e = error "chunksOf: A non-empty sequence can only be broken up into positively-sized chunks." +{-# INLINABLE chunksOf #-} + +-- | \( O(i) \) where \( i \) is the prefix length. 'takeWhileL', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest prefix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +-- +-- Returns a possibly empty sequence ('Seq') in the case that the predicate +-- fails on the first item. +takeWhileL :: (a -> Bool) -> NESeq a -> Seq a +takeWhileL p (x :<|| xs) + | p x = x Seq.<| Seq.takeWhileL p xs + | otherwise = Seq.empty +{-# INLINE takeWhileL #-} + +-- | \( O(i) \) where \( i \) is the suffix length. 'takeWhileR', applied +-- to a predicate @p@ and a sequence @xs@, returns the longest suffix +-- (possibly empty) of @xs@ of elements that satisfy @p@. +-- +-- Returns a possibly empty sequence ('Seq') in the case that the predicate +-- fails on the first item. +-- +-- @'takeWhileR' p xs@ is equivalent to @'reverse' ('takeWhileL' p ('reverse' xs))@. +takeWhileR :: (a -> Bool) -> NESeq a -> Seq a +takeWhileR p (xs :||> x) + | p x = Seq.takeWhileR p xs Seq.|> x + | otherwise = Seq.empty +{-# INLINE takeWhileR #-} + +-- | \( O(i) \) where \( i \) is the prefix length. @'dropWhileL' p xs@ returns +-- the suffix remaining after @'takeWhileL' p xs@. +-- +-- Returns a possibly empty sequence ('Seq') in the case that the predicate +-- passes for all items. +dropWhileL :: (a -> Bool) -> NESeq a -> Seq a +dropWhileL p xs0@(x :<|| xs) + | p x = Seq.dropWhileL p xs + | otherwise = toSeq xs0 +{-# INLINE dropWhileL #-} + +-- | \( O(i) \) where \( i \) is the suffix length. @'dropWhileR' p xs@ returns +-- the prefix remaining after @'takeWhileR' p xs@. +-- +-- Returns a possibly empty sequence ('Seq') in the case that the predicate +-- passes for all items. +-- +-- @'dropWhileR' p xs@ is equivalent to @'reverse' ('dropWhileL' p ('reverse' xs))@. +dropWhileR :: (a -> Bool) -> NESeq a -> Seq a +dropWhileR p xs0@(xs :||> x) + | p x = Seq.dropWhileR p xs + | otherwise = toSeq xs0 +{-# INLINE dropWhileR #-} + +-- | \( O(i) \) where \( i \) is the prefix length. 'spanl', applied to +-- a predicate @p@ and a sequence @xs@, returns a 'These' based on the +-- point where the predicate fails: +-- +-- * @'This' ys@ means that the predicate was true for all items, and +-- @ys@ is the entire original sequence. +-- * @'That' zs@ means that the predicate failed on the first item, and +-- @zs@ is the entire original sequence. +-- * @'These' ys zs@ gives @ys@ (the prefix of elements that satisfy the +-- predicae) and @zs@ (the remainder of the sequence) +spanl :: (a -> Bool) -> NESeq a -> These (NESeq a) (NESeq a) +spanl p xs0@(x :<|| xs) + | p x = case (nonEmptySeq ys, nonEmptySeq zs) of + (Nothing , Nothing ) -> This (singleton x) + (Just _ , Nothing ) -> This xs0 + (Nothing , Just zs') -> These (singleton x) zs' + (Just ys', Just zs') -> These (x <| ys') zs' + | otherwise = That xs0 + where + (ys, zs) = Seq.spanl p xs +{-# INLINABLE spanl #-} + +-- | \( O(i) \) where \( i \) is the suffix length. 'spanr', applied to +-- a predicate @p@ and a sequence @xs@, returns a 'These' based on the +-- point where the predicate fails: +-- +-- * @'This' ys@ means that the predicate was true for all items, and +-- @ys@ is the entire original sequence. +-- * @'That' zs@ means that the predicate failed on the first item, and +-- @zs@ is the entire original sequence. +-- * @'These' ys zs@ gives @ys@ (the suffix of elements that satisfy the +-- predicae) and @zs@ (the remainder of the sequence, before the suffix) +spanr :: (a -> Bool) -> NESeq a -> These (NESeq a) (NESeq a) +spanr p xs0@(xs :||> x) + | p x = case (nonEmptySeq ys, nonEmptySeq zs) of + (Nothing , Nothing ) -> This (singleton x) + (Just _ , Nothing ) -> This xs0 + (Nothing , Just zs') -> These (singleton x) zs' + (Just ys', Just zs') -> These (ys' |> x ) zs' + | otherwise = That xs0 + where + (ys, zs) = Seq.spanr p xs +{-# INLINABLE spanr #-} + +-- | \( O(i) \) where \( i \) is the breakpoint index. +-- +-- @'breakl' p@ is @'spanl' (not . p)@. +breakl :: (a -> Bool) -> NESeq a -> These (NESeq a) (NESeq a) +breakl p = spanl (not . p) +{-# INLINE breakl #-} + +-- | \( O(i) \) where \( i \) is the breakpoint index. +-- +-- @'breakr' p@ is @'spanr' (not . p)@. +breakr :: (a -> Bool) -> NESeq a -> These (NESeq a) (NESeq a) +breakr p = spanr (not . p) +{-# INLINE breakr #-} + +-- | \( O(n) \). The 'partition' function takes a predicate @p@ and a +-- sequence @xs@ and returns sequences of those elements which do and +-- do not satisfy the predicate, as a 'These': +-- +-- * @'This' ys@ means that the predicate was true for all items, and +-- @ys@ is the entire original sequence. +-- * @'That' zs@ means that the predicate failed on the first item, and +-- @zs@ is the entire original sequence. +-- * @'These' ys zs@ gives @ys@ (the sequence of elements for which the +-- predicate was true) and @zs@ (the sequence of elements for which the +-- predicate was false). +partition :: (a -> Bool) -> NESeq a -> These (NESeq a) (NESeq a) +partition p xs0@(x :<|| xs) = case (nonEmptySeq ys, nonEmptySeq zs) of + (Nothing , Nothing ) + | p x -> This (singleton x) + | otherwise -> That (singleton x) + (Just ys', Nothing ) + | p x -> This xs0 + | otherwise -> These ys' (singleton x) + (Nothing, Just zs' ) + | p x -> These (singleton x) zs' + | otherwise -> That xs0 + (Just ys', Just zs') + | p x -> These (x <| ys') zs' + | otherwise -> These ys' (x <| zs') + where + (ys, zs) = Seq.partition p xs +{-# INLINABLE partition #-} + +-- | \( O(n) \). The 'filter' function takes a predicate @p@ and a sequence +-- @xs@ and returns a sequence of those elements which satisfy the +-- predicate. +-- +-- Returns a potentially empty sequence ('Seq') in the case that the +-- predicate fails for all items in the sequence. +filter :: (a -> Bool) -> NESeq a -> Seq a +filter p (x :<|| xs) + | p x = x Seq.<| Seq.filter p xs + | otherwise = Seq.filter p xs +{-# INLINE filter #-} + +-- | \( O(n \log n) \). 'sort' sorts the specified 'NESeq' by the natural +-- ordering of its elements. The sort is stable. If stability is not +-- required, 'unstableSort' can be slightly faster. +sort :: Ord a => NESeq a -> NESeq a +sort = sortBy compare +{-# INLINE sort #-} + +-- | \( O(n \log n) \). 'sortBy' sorts the specified 'NESeq' according to +-- the specified comparator. The sort is stable. If stability is not +-- required, 'unstableSortBy' can be slightly faster. + +-- TODO: benchmark against just unsafe unwrapping and wrapping +sortBy :: (a -> a -> Ordering) -> NESeq a -> NESeq a +sortBy c (x :<|| xs) = withNonEmpty (singleton x) (insertBy c x) + . Seq.sortBy c + $ xs +{-# INLINE sortBy #-} + +-- | \( O(n \log n) \). 'sortOn' sorts the specified 'NESeq' by comparing +-- the results of a key function applied to each element. @'sortOn' f@ is +-- equivalent to @'sortBy' ('compare' ``Data.Function.on`` f)@, but has the +-- performance advantage of only evaluating @f@ once for each element in +-- the input list. This is called the decorate-sort-undecorate paradigm, or +-- Schwartzian transform. +-- +-- An example of using 'sortOn' might be to sort a 'NESeq' of strings +-- according to their length: +-- +-- > sortOn length (fromList ("alligator" :| ["monkey", "zebra"])) == fromList ("zebra" :| ["monkey", "alligator"]) +-- +-- If, instead, 'sortBy' had been used, 'length' would be evaluated on +-- every comparison, giving \( O(n \log n) \) evaluations, rather than +-- \( O(n) \). +-- +-- If @f@ is very cheap (for example a record selector, or 'fst'), +-- @'sortBy' ('compare' ``Data.Function.on`` f)@ will be faster than +-- @'sortOn' f@. + +-- TODO: benchmark against just unsafe unwrapping and wrapping +sortOn :: Ord b => (a -> b) -> NESeq a -> NESeq a +sortOn f (x :<|| xs) = withNonEmpty (singleton x) (insertOn f x) + . sortOnSeq f + $ xs +{-# INLINE sortOn #-} + +-- | \( O(n \log n) \). 'unstableSort' sorts the specified 'NESeq' by the +-- natural ordering of its elements, but the sort is not stable. This +-- algorithm is frequently faster and uses less memory than 'sort'. +unstableSort :: Ord a => NESeq a -> NESeq a +unstableSort = unstableSortBy compare +{-# INLINE unstableSort #-} + +-- | \( O(n \log n) \). A generalization of 'unstableSort', +-- 'unstableSortBy' takes an arbitrary comparator and sorts the specified +-- sequence. The sort is not stable. This algorithm is frequently faster +-- and uses less memory than 'sortBy'. + +-- TODO: figure out how to make it match 'Data.Sequence.unstableSortBy' +-- without unsafe wrapping/unwrapping +unstableSortBy :: (a -> a -> Ordering) -> NESeq a -> NESeq a +unstableSortBy c = unsafeFromSeq . Seq.unstableSortBy c . toSeq +-- unstableSortBy c (x :<|| xs) = withNonEmpty (singleton x) (insertBy c x) +-- . Seq.unstableSortBy c +-- $ xs +{-# INLINE unstableSortBy #-} + +-- | \( O(n \log n) \). 'unstableSortOn' sorts the specified 'NESeq' by +-- comparing the results of a key function applied to each element. +-- @'unstableSortOn' f@ is equivalent to @'unstableSortBy' ('compare' ``Data.Function.on`` f)@, +-- but has the performance advantage of only evaluating @f@ once for each +-- element in the input list. This is called the +-- decorate-sort-undecorate paradigm, or Schwartzian transform. +-- +-- An example of using 'unstableSortOn' might be to sort a 'NESeq' of strings +-- according to their length. +-- +-- > unstableSortOn length (fromList ("alligator" :| ["monkey", "zebra"])) == fromList ("zebra" :| ["monkey", "alligator]") +-- +-- If, instead, 'unstableSortBy' had been used, 'length' would be evaluated on +-- every comparison, giving \( O(n \log n) \) evaluations, rather than +-- \( O(n) \). +-- +-- If @f@ is very cheap (for example a record selector, or 'fst'), +-- @'unstableSortBy' ('compare' ``Data.Function.on`` f)@ will be faster than +-- @'unstableSortOn' f@. + +-- TODO: figure out how to make it match 'Data.Sequence.unstableSortBy' +-- without unsafe wrapping/unwrapping +unstableSortOn :: Ord b => (a -> b) -> NESeq a -> NESeq a +unstableSortOn f = unsafeFromSeq . unstableSortOnSeq f . toSeq +-- unstableSortOn f (x :<|| xs) = withNonEmpty (singleton x) (insertOn f x) +-- . Seq.unstableSortOn f +-- $ xs +{-# INLINE unstableSortOn #-} + +insertBy :: (a -> a -> Ordering) -> a -> NESeq a -> NESeq a +insertBy c x xs = case spanl ltx xs of + This ys -> ys |> x + That zs -> x <| zs + These ys zs -> ys >< (x <| zs) + where + ltx y = c x y == GT +{-# INLINABLE insertBy #-} + +insertOn :: Ord b => (a -> b) -> a -> NESeq a -> NESeq a +insertOn f x xs = case spanl ltx xs of + This ys -> ys |> x + That zs -> x <| zs + These ys zs -> ys >< (x <| zs) + where + fx = f x + ltx y = fx > f y +{-# INLINABLE insertOn #-} + +-- | \( O(\log(\min(i,n-i))) \). The element at the specified position, +-- counting from 0. If the specified position is negative or at +-- least the length of the sequence, 'lookup' returns 'Nothing'. +-- +-- Unlike 'index', this can be used to retrieve an element without +-- forcing it. +lookup :: Int -> NESeq a -> Maybe a +lookup 0 (x :<|| _ ) = Just x +lookup i (_ :<|| xs) = Seq.lookup (i - 1) xs +{-# INLINE lookup #-} + +-- | \( O(\log(\min(i,n-i))) \). A flipped, infix version of `lookup`. +(!?) :: NESeq a -> Int -> Maybe a +(!?) = flip lookup +{-# INLINE (!?) #-} + +-- | \( O(\log(\min(i,n-i))) \). Update the element at the specified position. If +-- the position is out of range, the original sequence is returned. 'adjust' +-- can lead to poor performance and even memory leaks, because it does not +-- force the new value before installing it in the sequence. 'adjust'' should +-- usually be preferred. +adjust :: (a -> a) -> Int -> NESeq a -> NESeq a +adjust f 0 (x :<|| xs) = f x :<|| xs +adjust f i (x :<|| xs) = x :<|| Seq.adjust f (i - 1) xs +{-# INLINE adjust #-} + +-- | \( O(\log(\min(i,n-i))) \). Update the element at the specified position. +-- If the position is out of range, the original sequence is returned. +-- The new value is forced before it is installed in the sequence. +-- +-- @ +-- adjust' f i xs = +-- case xs !? i of +-- Nothing -> xs +-- Just x -> let !x' = f x +-- in update i x' xs +-- @ +adjust' :: (a -> a) -> Int -> NESeq a -> NESeq a +adjust' f 0 (x :<|| xs) = let !y = f x in y :<|| xs +adjust' f i (x :<|| xs) = x :<|| Seq.adjust f (i - 1) xs +{-# INLINE adjust' #-} + +-- | \( O(\log(\min(i,n-i))) \). Replace the element at the specified position. +-- If the position is out of range, the original sequence is returned. +update :: Int -> a -> NESeq a -> NESeq a +update 0 y (_ :<|| xs) = y :<|| xs +update i y (x :<|| xs) = x :<|| Seq.update (i - 1) y xs +{-# INLINE update #-} + +-- | \( O(\log(\min(i,n-i))) \). The first @i@ elements of a sequence. +-- If @i@ is negative, @'take' i s@ yields the empty sequence. +-- If the sequence contains fewer than @i@ elements, the whole sequence +-- is returned. +take :: Int -> NESeq a -> Seq a +take i (x :<|| xs) + | i <= 0 = Seq.empty + | otherwise = x Seq.<| Seq.take (i - 1) xs +{-# INLINE take #-} + +-- | \( O(\log(\min(i,n-i))) \). Elements of a sequence after the first @i@. +-- If @i@ is negative, @'drop' i s@ yields the whole sequence. +-- If the sequence contains fewer than @i@ elements, the empty sequence +-- is returned. +drop :: Int -> NESeq a -> Seq a +drop i xs0@(_ :<|| xs) + | i <= 0 = toSeq xs0 + | otherwise = Seq.drop (i - 1) xs +{-# INLINE drop #-} + +-- | \( O(\log(\min(i,n-i))) \). @'insertAt' i x xs@ inserts @x@ into @xs@ +-- at the index @i@, shifting the rest of the sequence over. +-- +-- @ +-- insertAt 2 x (fromList (a:|[b,c,d])) = fromList (a:|[b,x,c,d]) +-- insertAt 4 x (fromList (a:|[b,c,d])) = insertAt 10 x (fromList (a:|[b,c,d])) +-- = fromList (a:|[b,c,d,x]) +-- @ +-- +-- prop> insertAt i x xs = take i xs >< singleton x >< drop i xs +insertAt :: Int -> a -> NESeq a -> NESeq a +insertAt i y xs0@(x :<|| xs) + | i <= 0 = y <| xs0 + | otherwise = x :<|| Seq.insertAt (i - 1) y xs +{-# INLINE insertAt #-} + +-- | \( O(\log(\min(i,n-i))) \). Delete the element of a sequence at a given +-- index. Return the original sequence if the index is out of range. +-- +-- @ +-- deleteAt 2 (a:|[b,c,d]) = a:|[b,d] +-- deleteAt 4 (a:|[b,c,d]) = deleteAt (-1) (a:|[b,c,d]) = a:|[b,c,d] +-- @ +deleteAt :: Int -> NESeq a -> Seq a +deleteAt i xs0@(x :<|| xs) = case compare i 0 of + LT -> toSeq xs0 + EQ -> xs + GT -> x Seq.<| Seq.deleteAt (i - 1) xs +{-# INLINE deleteAt #-} + +-- | \( O(\log(\min(i,n-i))) \). Split a sequence at a given position. +-- +-- * @'This' ys@ means that the given position was longer than the length +-- of the list, and @ys@ is the entire original system. +-- * @'That' zs@ means that the given position was zero or smaller, and +-- so @zs@ is the entire original sequence. +-- * @'These' ys zs@ gives @ys@ (the sequence of elements before the +-- given position, @take n xs@) and @zs@ (the sequence of elements +-- after the given position, @drop n xs@). +splitAt :: Int -> NESeq a -> These (NESeq a) (NESeq a) +splitAt n xs0@(x :<|| xs) + | n <= 0 = That xs0 + | otherwise = case (nonEmptySeq ys, nonEmptySeq zs) of + (Nothing , Nothing ) -> This (singleton x) + (Just _ , Nothing ) -> This xs0 + (Nothing , Just zs') -> These (singleton x) zs' + (Just ys', Just zs') -> These (x <| ys') zs' + where + (ys, zs) = Seq.splitAt (n - 1) xs +{-# INLINABLE splitAt #-} + +-- | 'elemIndexL' finds the leftmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexL :: Eq a => a -> NESeq a -> Maybe Int +elemIndexL x = findIndexL (== x) +{-# INLINE elemIndexL #-} + +-- | 'elemIndexR' finds the rightmost index of the specified element, +-- if it is present, and otherwise 'Nothing'. +elemIndexR :: Eq a => a -> NESeq a -> Maybe Int +elemIndexR x = findIndexR (== x) +{-# INLINE elemIndexR #-} + +-- | 'elemIndicesL' finds the indices of the specified element, from +-- left to right (i.e. in ascending order). +elemIndicesL :: Eq a => a -> NESeq a -> [Int] +elemIndicesL x = findIndicesL (== x) +{-# INLINE elemIndicesL #-} + +-- | 'elemIndicesR' finds the indices of the specified element, from +-- right to left (i.e. in descending order). +elemIndicesR :: Eq a => a -> NESeq a -> [Int] +elemIndicesR x = findIndicesR (== x) +{-# INLINE elemIndicesR #-} + +-- | @'findIndexL' p xs@ finds the index of the leftmost element that +-- satisfies @p@, if any exist. +findIndexL :: (a -> Bool) -> NESeq a -> Maybe Int +findIndexL p (x :<|| xs) = here_ <|> there_ + where + here_ = 0 <$ guard (p x) + there_ = (+ 1) <$> Seq.findIndexL p xs +{-# INLINE findIndexL #-} + +-- | @'findIndexR' p xs@ finds the index of the rightmost element that +-- satisfies @p@, if any exist. +findIndexR :: (a -> Bool) -> NESeq a -> Maybe Int +findIndexR p (xs :||> x) = here_ <|> there_ + where + here_ = Seq.length xs <$ guard (p x) + there_ = Seq.findIndexR p xs +{-# INLINE findIndexR #-} + +-- | @'findIndicesL' p@ finds all indices of elements that satisfy @p@, +-- in ascending order. + +-- TODO: use build +findIndicesL :: (a -> Bool) -> NESeq a -> [Int] +findIndicesL p (x :<|| xs) + | p x = 0 : ixs + | otherwise = ixs + where + ixs = (+ 1) <$> Seq.findIndicesL p xs +{-# INLINE findIndicesL #-} + +-- | @'findIndicesR' p@ finds all indices of elements that satisfy @p@, +-- in descending order. + +-- TODO: use build +findIndicesR :: (a -> Bool) -> NESeq a -> [Int] +findIndicesR p (xs :||> x) + | p x = Seq.length xs : ixs + | otherwise = ixs + where + ixs = Seq.findIndicesR p xs +{-# INLINE findIndicesR #-} + +-- | 'foldlWithIndex' is a version of 'foldl' that also provides access +-- to the index of each element. +foldlWithIndex :: (b -> Int -> a -> b) -> b -> NESeq a -> b +foldlWithIndex f z (xs :||> x) = (\z' -> f z' (Seq.length xs) x) . Seq.foldlWithIndex f z $ xs +{-# INLINE foldlWithIndex #-} + +-- | 'foldrWithIndex' is a version of 'foldr' that also provides access +-- to the index of each element. +foldrWithIndex :: (Int -> a -> b -> b) -> b -> NESeq a -> b +foldrWithIndex f z (x :<|| xs) = f 0 x . Seq.foldrWithIndex (f . (+ 1)) z $ xs +{-# INLINE foldrWithIndex #-} + +-- | A generalization of 'fmap', 'mapWithIndex' takes a mapping +-- function that also depends on the element's index, and applies it to every +-- element in the sequence. +mapWithIndex :: (Int -> a -> b) -> NESeq a -> NESeq b +mapWithIndex f (x :<|| xs) = f 0 x :<|| Seq.mapWithIndex (f . (+ 1)) xs +{-# NOINLINE [1] mapWithIndex #-} +{-# RULES +"mapWithIndex/mapWithIndex" forall f g xs . mapWithIndex f (mapWithIndex g xs) = + mapWithIndex (\k a -> f k (g k a)) xs +"mapWithIndex/map" forall f g xs . mapWithIndex f (map g xs) = + mapWithIndex (\k a -> f k (g a)) xs +"map/mapWithIndex" forall f g xs . map f (mapWithIndex g xs) = + mapWithIndex (\k a -> f (g k a)) xs + #-} + +-- | 'traverseWithIndex' is a version of 'traverse' that also offers +-- access to the index of each element. +-- +-- Is a more restrictive version of 'traverseWithIndex1'; +-- 'traverseWithIndex1' should be used whenever possible. +traverseWithIndex :: Applicative f => (Int -> a -> f b) -> NESeq a -> f (NESeq b) +traverseWithIndex f (x :<|| xs) = (:<||) <$> f 0 x <*> Seq.traverseWithIndex (f . (+ 1)) xs +{-# NOINLINE [1] traverseWithIndex #-} +{-# RULES +"travWithIndex/mapWithIndex" forall f g xs . traverseWithIndex f (mapWithIndex g xs) = + traverseWithIndex (\k a -> f k (g k a)) xs +"travWithIndex/map" forall f g xs . traverseWithIndex f (map g xs) = + traverseWithIndex (\k a -> f k (g a)) xs + #-} + +-- | \( O(n) \). The reverse of a sequence. +reverse :: NESeq a -> NESeq a +reverse (x :<|| xs) = Seq.reverse xs :||> x +{-# NOINLINE [1] reverse #-} + +-- | \( O(n) \). Reverse a sequence while mapping over it. This is not +-- currently exported, but is used in rewrite rules. +mapReverse :: (a -> b) -> NESeq a -> NESeq b +mapReverse f (x :<|| xs) = fmap f (Seq.reverse xs) :||> f x + +{-# RULES +"map/reverse" forall f xs . map f (reverse xs) = mapReverse f xs +"reverse/map" forall f xs . reverse (map f xs) = mapReverse f xs + #-} + +-- | \( O(n) \). Intersperse an element between the elements of a sequence. +-- +-- @ +-- intersperse a empty = empty +-- intersperse a (singleton x) = singleton x +-- intersperse a (fromList [x,y]) = fromList [x,a,y] +-- intersperse a (fromList [x,y,z]) = fromList [x,a,y,a,z] +-- @ +intersperse :: a -> NESeq a -> NESeq a +intersperse z (x :<|| xs) = x :<|| (z Seq.<| Seq.intersperse z xs) +{-# INLINE intersperse #-} + +-- | \( O(\min(n_1,n_2,n_3)) \). 'zip3' takes three sequences and returns a +-- sequence of triples, analogous to 'zip'. +zip3 :: NESeq a -> NESeq b -> NESeq c -> NESeq (a, b, c) +zip3 (x :<|| xs) (y :<|| ys) (z :<|| zs) = (x, y, z) :<|| Seq.zip3 xs ys zs +{-# INLINE zip3 #-} + +-- | \( O(\min(n_1,n_2,n_3)) \). 'zipWith3' takes a function which combines +-- three elements, as well as three sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith3 :: (a -> b -> c -> d) -> NESeq a -> NESeq b -> NESeq c -> NESeq d +zipWith3 f (x :<|| xs) (y :<|| ys) (z :<|| zs) = f x y z :<|| Seq.zipWith3 f xs ys zs +{-# INLINE zipWith3 #-} + +-- | \( O(\min(n_1,n_2,n_3,n_4)) \). 'zip4' takes four sequences and returns a +-- sequence of quadruples, analogous to 'zip'. +zip4 :: NESeq a -> NESeq b -> NESeq c -> NESeq d -> NESeq (a, b, c, d) +zip4 (x :<|| xs) (y :<|| ys) (z :<|| zs) (r :<|| rs) = (x, y, z, r) :<|| Seq.zip4 xs ys zs rs +{-# INLINE zip4 #-} + +-- | \( O(\min(n_1,n_2,n_3,n_4)) \). 'zipWith4' takes a function which combines +-- four elements, as well as four sequences and returns a sequence of +-- their point-wise combinations, analogous to 'zipWith'. +zipWith4 :: (a -> b -> c -> d -> e) -> NESeq a -> NESeq b -> NESeq c -> NESeq d -> NESeq e +zipWith4 f (x :<|| xs) (y :<|| ys) (z :<|| zs) (r :<|| rs) = f x y z r :<|| Seq.zipWith4 f xs ys zs rs +{-# INLINE zipWith4 #-} + +-- | \( O(n) \). Unzip a sequence using a function to divide elements. +-- +-- @ unzipWith f xs == 'unzip' ('fmap' f xs) @ +-- +-- Efficiency note: +-- +-- @unzipWith@ produces its two results in lockstep. If you calculate +-- @ unzipWith f xs @ and fully force /either/ of the results, then the +-- entire structure of the /other/ one will be built as well. This +-- behavior allows the garbage collector to collect each calculated +-- pair component as soon as it dies, without having to wait for its mate +-- to die. If you do not need this behavior, you may be better off simply +-- calculating the sequence of pairs and using 'fmap' to extract each +-- component sequence. +unzipWith :: (a -> (b, c)) -> NESeq a -> (NESeq b, NESeq c) +unzipWith f (x :<|| xs) = bimap (y :<||) (z :<||) . unzipWithSeq f $ xs + where + ~(y, z) = f x +{-# NOINLINE [1] unzipWith #-} + +{-# RULES +"unzipWith/map" forall f g xs. unzipWith f (map g xs) = + unzipWith (f . g) xs + #-} diff --git a/src/Data/Sequence/NonEmpty/Internal.hs b/src/Data/Sequence/NonEmpty/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Sequence/NonEmpty/Internal.hs @@ -0,0 +1,541 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE DeriveTraversable #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Data.Sequence.NonEmpty.Internal +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Unsafe internal-use functions used in the implementation of +-- "Data.Sequence.NonEmpty". These functions can potentially be used to +-- break the abstraction of 'NESeq' and produce unsound sequences, so be +-- wary! +module Data.Sequence.NonEmpty.Internal ( + NESeq(..) + , pattern (:<||) + , pattern (:||>) + , withNonEmpty + , toSeq + , singleton + , length + , fromList + , fromFunction + , replicate + , index + , (<|), (><), (|><) + , map + , foldMapWithIndex + , traverseWithIndex1 + , tails + , zip + , zipWith + , unzip + , sortOnSeq + , unstableSortOnSeq + , unzipSeq + , unzipWithSeq + ) where + +import Control.Comonad +import Control.DeepSeq +import Control.Monad.Fix +import Control.Monad.Zip +import Data.Bifunctor +import Data.Coerce +import Data.Data +import Data.Foldable (Foldable) +import Data.Functor.Alt +import Data.Functor.Bind +import Data.Functor.Classes +import Data.Functor.Extend +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup +import Data.Semigroup.Foldable +import Data.Semigroup.Traversable +import Data.Sequence (Seq(..)) +import Prelude hiding (length, zipWith, unzip, zip, map, replicate) +import Text.Read +import qualified Data.Foldable as F +import qualified Data.Sequence as Seq + +-- | A general-purpose non-empty (by construction) finite sequence type. +-- +-- Non-emptiness means that: +-- +-- * Functions that /take/ an 'NESeq' can safely operate on it with the +-- assumption that it has at least value. +-- * Functions that /return/ an 'NESeq' provide an assurance that the +-- result has at least one value. +-- +-- "Data.Sequence.NonEmpty" re-exports the API of "Data.Sequence", +-- faithfully reproducing asymptotics, typeclass constraints, and +-- semantics. Functions that ensure that input and output maps are both +-- non-empty (like 'Data.Sequence.NonEmpty.<|') return 'NESeq', but +-- functions that might potentially return an empty map (like +-- 'Data.Sequence.NonEmpty.tail') return a 'Seq' instead. +-- +-- You can directly construct an 'NESeq' with the API from +-- "Data.Sequence.NonEmpty"; it's more or less the same as constructing +-- a normal 'Seq', except you don't have access to 'Data.Seq.empty'. There +-- are also a few ways to construct an 'NESeq' from a 'Seq': +-- +-- 1. The 'Data.Sequence.NonEmpty.nonEmptySeq' smart constructor will +-- convert a @'Seq' a@ into a @'Maybe' ('NESeq' a)@, returning 'Nothing' if +-- the original 'Seq' was empty. +-- 2. You can use 'Data.Sequence.NonEmpty.:<||', +-- 'Data.Sequence.NonEmpty.:||>', and +-- 'Data.Sequence.NonEmpty.insertSeqAt' to insert a value into a 'Seq' +-- to create a guaranteed 'NESeq'. +-- 3. You can use the 'Data.Sequence.NonEmpty.IsNonEmpty' and +-- 'Data.Sequence.NonEmpty.IsEmpty' patterns to "pattern match" on +-- a 'Seq' to reveal it as either containing a 'NESeq' or an empty +-- sequence. +-- 4. 'Data.Sequence.NonEmpty.withNonEmpty' offers a continuation-based +-- interface for deconstructing a 'Seq' and treating it as if it were an +-- 'NESeq'. +-- +-- You can convert an 'NESeq' into a 'Seq' with 'toSeq' or +-- 'Data.Sequence.NonEmpty.IsNonEmpty', essentially "obscuring" the +-- non-empty property from the type. +data NESeq a = NESeq { nesHead :: a + , nesTail :: !(Seq a) + } + deriving (Traversable, Typeable) + +-- | /O(1)/. An abstract constructor for an 'NESeq' that consists of +-- a "head" @a@ and a "tail" @'Seq' a@. Similar to ':|' for 'NonEmpty'. +-- +-- Can be used to match on the head and tail of an 'NESeq', and also used +-- to /construct/ an 'NESeq' by consing an item to the beginnong of +-- a 'Seq', ensuring that the result is non-empty. +pattern (:<||) :: a -> Seq a -> NESeq a +pattern x :<|| xs = NESeq x xs +{-# COMPLETE (:<||) #-} + +unsnoc :: NESeq a -> (Seq a, a) +unsnoc (x :<|| (xs :|> y)) = (x :<| xs, y) +unsnoc (x :<|| Empty ) = (Empty , x) +{-# INLINE unsnoc #-} + +-- | /O(1)/. An abstract constructor for an 'NESeq' that consists of +-- a "init" @'Seq' a@ and a "last" @a@. Similar to ':|' for 'NonEmpty', +-- but at the end of the list instead of at the beginning. +-- +-- Can be used to match on the init and last of an 'NESeq', and also used +-- to /construct/ an 'NESeq' by snocing an item to the end of a 'Seq', +-- ensuring that the result is non-empty. +pattern (:||>) :: Seq a -> a -> NESeq a +pattern xs :||> x <- (unsnoc->(!xs, x)) + where + (x :<| xs) :||> y = x :<|| (xs :|> y) + Empty :||> y = y :<|| Empty +{-# COMPLETE (:||>) #-} + +infixr 5 `NESeq` +infixr 5 :<|| +infixl 5 :||> + +instance Show a => Show (NESeq a) where + showsPrec p xs = showParen (p > 10) $ + showString "fromList (" . shows (toNonEmpty xs) . showString ")" + +instance Read a => Read (NESeq a) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- parens . prec 10 $ readPrec + return (fromList xs) + readListPrec = readListPrecDefault + +instance Eq a => Eq (NESeq a) where + xs == ys = length xs == length ys + && toNonEmpty xs == toNonEmpty ys + +instance Show1 NESeq where + liftShowsPrec sp sl d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toNonEmpty m) + +instance Read1 NESeq where + liftReadsPrec _rp readLst p = readParen (p > 10) $ \r -> do + ("fromList",s) <- lex r + (xs, t) <- liftReadsPrec _rp readLst 10 s + pure (fromList xs, t) + +instance Eq1 NESeq where + liftEq eq xs ys = length xs == length ys && liftEq eq (toNonEmpty xs) (toNonEmpty ys) + +instance Ord1 NESeq where + liftCompare cmp xs ys = liftCompare cmp (toNonEmpty xs) (toNonEmpty ys) + +instance Data a => Data (NESeq a) where + gfoldl f z (x :<|| xs) = z (:<||) `f` x `f` xs + gunfold k z _ = k (k (z (:<||))) + toConstr _ = consConstr + dataTypeOf _ = seqDataType + dataCast1 = gcast1 + +consConstr :: Constr +consConstr = mkConstr seqDataType ":<||" [] Infix + +seqDataType :: DataType +seqDataType = mkDataType "Data.Sequence.NonEmpty.Internal.NESeq" [consConstr] + +-- | /O(log n)/. A general continuation-based way to consume a 'Seq' as if +-- it were an 'NESeq'. @'withNonEmpty' def f@ will take a 'Seq'. If map is +-- empty, it will evaluate to @def@. Otherwise, a non-empty map 'NESeq' +-- will be fed to the function @f@ instead. +-- +-- @'Data.Sequence.NonEmpty.nonEmptySeq' == 'withNonEmpty' 'Nothing' 'Just'@ +withNonEmpty :: r -> (NESeq a -> r) -> Seq a -> r +withNonEmpty def f = \case + x :<| xs -> f (x :<|| xs) + Empty -> def +{-# INLINE withNonEmpty #-} + +-- | /O(1)/. +-- Convert a non-empty sequence back into a normal possibly-empty sequence, +-- for usage with functions that expect 'Seq'. +-- +-- Can be thought of as "obscuring" the non-emptiness of the map in its +-- type. See the 'Data.Sequence.NonEmpty.IsNotEmpty' pattern. +-- +-- 'Data.Sequence.NonEmpty.nonEmptySeq' and @'maybe' 'Data.Seq.empty' +-- 'toSeq'@ form an isomorphism: they are perfect structure-preserving +-- inverses of eachother. +toSeq :: NESeq a -> Seq a +toSeq (x :<|| xs) = x :<| xs +{-# INLINE toSeq #-} + +-- | \( O(1) \). A singleton sequence. +singleton :: a -> NESeq a +singleton = (:<|| Seq.empty) +{-# INLINE singleton #-} + +-- | \( O(1) \). The number of elements in the sequence. +length :: NESeq a -> Int +length (_ :<|| xs) = 1 + Seq.length xs +{-# INLINE length #-} + +-- | \( O(n) \). Create a sequence from a finite list of elements. There +-- is a function 'toNonEmpty' in the opposite direction for all instances +-- of the 'Foldable1' class, including 'NESeq'. +fromList :: NonEmpty a -> NESeq a +fromList (x :| xs) = x :<|| Seq.fromList xs +{-# INLINE fromList #-} + +-- | \( O(n) \). Convert a given sequence length and a function representing that +-- sequence into a sequence. +fromFunction :: Int -> (Int -> a) -> NESeq a +fromFunction n f + | n < 1 = error "NESeq.fromFunction: must take a positive integer argument" + | otherwise = f 0 :<|| Seq.fromFunction (n - 1) (f . (+ 1)) + +-- | \( O(\log n) \). @replicate n x@ is a sequence consisting of @n@ +-- copies of @x@. Is only defined when @n@ is positive. +replicate :: Int -> a -> NESeq a +replicate n x + | n < 1 = error "NESeq.replicate: must take a positive integer argument" + | otherwise = x :<|| Seq.replicate (n - 1) x +{-# INLINE replicate #-} + +-- | \( O(\log(\min(i,n-i))) \). The element at the specified position, +-- counting from 0. The argument should thus be a non-negative +-- integer less than the size of the sequence. +-- If the position is out of range, 'index' fails with an error. +-- +-- prop> xs `index` i = toList xs !! i +-- +-- Caution: 'index' necessarily delays retrieving the requested +-- element until the result is forced. It can therefore lead to a space +-- leak if the result is stored, unforced, in another structure. To retrieve +-- an element immediately without forcing it, use 'lookup' or '(!?)'. +index :: NESeq a -> Int -> a +index (x :<|| _ ) 0 = x +index (_ :<|| xs) i = xs `Seq.index` (i - 1) +{-# INLINE index #-} + +-- | \( O(1) \). Add an element to the left end of a non-empty sequence. +-- Mnemonic: a triangle with the single element at the pointy end. +(<|) :: a -> NESeq a -> NESeq a +x <| xs = x :<|| toSeq xs +{-# INLINE (<|) #-} + +-- | \( O(\log(\min(n_1,n_2))) \). Concatenate two non-empty sequences. +(><) :: NESeq a -> NESeq a -> NESeq a +(x :<|| xs) >< ys = x :<|| (xs Seq.>< toSeq ys) +{-# INLINE (><) #-} + +-- | \( O(\log(\min(n_1,n_2))) \). Concatenate a non-empty sequence with +-- a potentially empty sequence ('Seq'), to produce a guaranteed non-empty +-- sequence. Mnemonic: like '><', but a pipe for the guarunteed non-empty +-- side. +(|><) :: NESeq a -> Seq a -> NESeq a +(x :<|| xs) |>< ys = x :<|| (xs Seq.>< ys) +{-# INLINE (|><) #-} + +infixr 5 <| +infixr 5 >< +infixr 5 |>< + +-- | Defined here but hidden; intended for use with RULES pragma. +map :: (a -> b) -> NESeq a -> NESeq b +map f (x :<|| xs) = f x :<|| fmap f xs +{-# NOINLINE [1] map #-} +{-# RULES +"map/map" forall f g xs . map f (map g xs) = map (f . g) xs + #-} +{-# RULES +"map/coerce" map coerce = coerce + #-} + +-- | /O(n)/. A generalization of 'foldMap1', 'foldMapWithIndex' takes +-- a folding function that also depends on the element's index, and applies +-- it to every element in the sequence. +foldMapWithIndex :: Semigroup m => (Int -> a -> m) -> NESeq a -> m +foldMapWithIndex f (x :<|| xs) = maybe (f 0 x) (f 0 x <>) + . getOption + . Seq.foldMapWithIndex (\i -> Option . Just . f (i + 1)) + $ xs +{-# INLINE foldMapWithIndex #-} + +-- | /O(n)/. 'traverseWithIndex1' is a version of 'traverse1' that also +-- offers access to the index of each element. +traverseWithIndex1 :: Apply f => (Int -> a -> f b) -> NESeq a -> f (NESeq b) +traverseWithIndex1 f (x :<|| xs) = case runMaybeApply xs' of + Left ys -> (:<||) <$> f 0 x <.> ys + Right ys -> (:<|| ys) <$> f 0 x + where + xs' = Seq.traverseWithIndex (\i -> MaybeApply . Left . f (i+1)) xs +{-# INLINABLE traverseWithIndex1 #-} + +-- | \( O(n) \). Returns a sequence of all non-empty suffixes of this +-- sequence, longest first. For example, +-- +-- > tails (fromList (1:|[2,3])) = fromList (fromList (1:|[2,3]) :| [fromList (2:|[3]), fromList (3:|[])]) +-- +-- Evaluating the \( i \)th suffix takes \( O(\log(\min(i, n-i))) \), but evaluating +-- every suffix in the sequence takes \( O(n) \) due to sharing. + +-- TODO: is this true? +tails :: NESeq a -> NESeq (NESeq a) +tails xs@(_ :<|| ys) = withNonEmpty (singleton xs) ((xs <|) . tails) ys +{-# INLINABLE tails #-} + +-- | \( O(\min(n_1,n_2)) \). 'zip' takes two sequences and returns +-- a sequence of corresponding pairs. If one input is short, excess +-- elements are discarded from the right end of the longer sequence. +zip :: NESeq a -> NESeq b -> NESeq (a, b) +zip (x :<|| xs) (y :<|| ys) = (x, y) :<|| Seq.zip xs ys +{-# INLINE zip #-} + +-- | \( O(\min(n_1,n_2)) \). 'zipWith' generalizes 'zip' by zipping with the +-- function given as the first argument, instead of a tupling function. +-- For example, @zipWith (+)@ is applied to two sequences to take the +-- sequence of corresponding sums. +zipWith :: (a -> b -> c) -> NESeq a -> NESeq b -> NESeq c +zipWith f (x :<|| xs) (y :<|| ys) = f x y :<|| Seq.zipWith f xs ys +{-# INLINE zipWith #-} + +-- | Unzip a sequence of pairs. +-- +-- @ +-- unzip ps = ps ``seq`` ('fmap' 'fst' ps) ('fmap' 'snd' ps) +-- @ +-- +-- Example: +-- +-- @ +-- unzip $ fromList ((1,"a") :| [(2,"b"), (3,"c")]) = +-- (fromList (1:|[2,3]), fromList ("a":|["b","c"])) +-- @ +-- +-- See the note about efficiency at 'Data.Sequence.NonEmpty.unzipWith'. +unzip :: NESeq (a, b) -> (NESeq a, NESeq b) +unzip ((x, y) :<|| xys) = bimap (x :<||) (y :<||) . unzipSeq $ xys +{-# INLINE unzip #-} + +instance Semigroup (NESeq a) where + (<>) = (><) + {-# INLINE (<>) #-} + +instance Functor NESeq where + fmap = map + {-# INLINE fmap #-} + x <$ xs = replicate (length xs) x + {-# INLINE (<$) #-} + +instance Apply NESeq where + (f :<|| fs) <.> xs = fxs |>< fsxs + where + fxs = f <$> xs + fsxs = fs <.> toSeq xs + {-# INLINABLE (<.>) #-} + +instance Applicative NESeq where + pure = singleton + {-# INLINE pure #-} + (<*>) = (<.>) + {-# INLINABLE (<*>) #-} + +instance Alt NESeq where + () = (><) + {-# INLINE () #-} + +instance Bind NESeq where + NESeq x xs >>- f = withNonEmpty (f x) ((f x ><) . (>>- f)) xs + {-# INLINABLE (>>-) #-} + +instance Monad NESeq where + return = pure + {-# INLINE return #-} + (>>=) = (>>-) + {-# INLINABLE (>>=) #-} + +instance Extend NESeq where + duplicated = tails + {-# INLINE duplicated #-} + extended f xs0@(_ :<|| xs) = withNonEmpty (singleton (f xs0)) + ((f xs0 <|) . extend f) + xs + {-# INLINE extended #-} + +instance Comonad NESeq where + extract (x :<|| _) = x + {-# INLINE extract #-} + duplicate = duplicated + {-# INLINE duplicate #-} + extend = extended + {-# INLINE extend #-} + +-- | 'foldr1', 'foldl', 'maximum', and 'minimum' are all total, unlike for +-- 'Seq'. +instance Foldable NESeq where +#if MIN_VERSION_base(4,11,0) + fold (x :<|| xs) = x <> F.fold xs + {-# INLINE fold #-} + foldMap f (x :<|| xs) = f x <> F.foldMap f xs + {-# INLINE foldMap #-} +#else + fold (x :<|| xs) = x `mappend` F.fold xs + {-# INLINE fold #-} + foldMap f (x :<|| xs) = f x `mappend` F.foldMap f xs + {-# INLINE foldMap #-} +#endif + foldr f z (x :<|| xs) = x `f` foldr f z xs + {-# INLINE foldr #-} + foldr' f z (xs :||> x) = F.foldr' f y xs + where + !y = f x z + {-# INLINE foldr' #-} + foldl f z (xs :||> x) = foldl f z xs `f` x + {-# INLINE foldl #-} + foldl' f z (x :<|| xs) = F.foldl' f y xs + where + !y = f z x + {-# INLINE foldl' #-} + foldr1 f (xs :||> x) = foldr f x xs + {-# INLINE foldr1 #-} + foldl1 f (x :<|| xs) = foldl f x xs + {-# INLINE foldl1 #-} + null _ = False + {-# INLINE null #-} + length = length + {-# INLINE length #-} + +instance Foldable1 NESeq where + fold1 (x :<|| xs) = maybe x (x <>) + . getOption + . F.foldMap (Option . Just) + $ xs + {-# INLINE fold1 #-} + foldMap1 f = foldMapWithIndex (const f) + {-# INLINE foldMap1 #-} + -- TODO: use build + toNonEmpty (x :<|| xs) = x :| F.toList xs + {-# INLINE toNonEmpty #-} + +instance Traversable1 NESeq where + traverse1 f = traverseWithIndex1 (const f) + {-# INLINE traverse1 #-} + sequence1 (x :<|| xs) = case runMaybeApply xs' of + Left ys -> (:<||) <$> x <.> ys + Right ys -> (:<|| ys) <$> x + where + xs' = traverse (MaybeApply . Left) xs + {-# INLINABLE sequence1 #-} + +-- | @mzipWith = zipWith@ +-- +-- @munzip = unzip@ +instance MonadZip NESeq where + mzipWith = zipWith + munzip = unzip + +instance MonadFix NESeq where + mfix = mfixSeq + +mfixSeq :: (a -> NESeq a) -> NESeq a +mfixSeq f = fromFunction (length (f err)) (\k -> fix (\xk -> f xk `index` k)) + where + err = error "mfix for Data.Sequence.NonEmpty.NESeq applied to strict function" + +instance NFData a => NFData (NESeq a) where + rnf (x :<|| xs) = rnf x `seq` rnf xs `seq` () + +-- --------------------------------------------- +-- | CPP for new functions not in old containers +-- --------------------------------------------- + +-- | Compatibility layer for 'Data.Sequence.sortOn'. +sortOnSeq :: Ord b => (a -> b) -> Seq a -> Seq a +#if MIN_VERSION_containers(0,5,11) +sortOnSeq = Seq.sortOn +#else +sortOnSeq f = Seq.sortBy (\x y -> f x `compare` f y) +#endif +{-# INLINE sortOnSeq #-} + +-- | Compatibility layer for 'Data.Sequence.unstableSortOn'. +unstableSortOnSeq :: Ord b => (a -> b) -> Seq a -> Seq a +#if MIN_VERSION_containers(0,5,11) +unstableSortOnSeq = Seq.unstableSortOn +#else +unstableSortOnSeq f = Seq.unstableSortBy (\x y -> f x `compare` f y) +#endif +{-# INLINE unstableSortOnSeq #-} + +-- | Compatibility layer for 'Data.Sequence.unzip'. +unzipSeq :: Seq (a, b) -> (Seq a, Seq b) +#if MIN_VERSION_containers(0,5,11) +unzipSeq = Seq.unzip +{-# INLINE unzipSeq #-} +#else +unzipSeq = \case + (x, y) :<| xys -> bimap (x :<|) (y :<|) . unzipSeq $ xys + Empty -> (Empty, Empty) +{-# INLINABLE unzipSeq #-} +#endif + +-- | Compatibility layer for 'Data.Sequence.unzipWith'. +unzipWithSeq :: (a -> (b, c)) -> Seq a -> (Seq b, Seq c) +#if MIN_VERSION_containers(0,5,11) +unzipWithSeq = Seq.unzipWith +{-# INLINE unzipWithSeq #-} +#else +unzipWithSeq f = go + where + go = \case + x :<| xs -> let ~(y, z) = f x + in bimap (y :<|) (z :<|) . go $ xs + Empty -> (Empty, Empty) +{-# INLINABLE unzipWithSeq #-} +#endif diff --git a/src/Data/Set/NonEmpty.hs b/src/Data/Set/NonEmpty.hs new file mode 100644 --- /dev/null +++ b/src/Data/Set/NonEmpty.hs @@ -0,0 +1,1038 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE PatternSynonyms #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE ViewPatterns #-} + +-- | +-- Module : Data.Set.NonEmpty +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- = Non-Empty Finite Sets +-- +-- The @'NESet' e@ type represents a non-empty set of elements of type @e@. +-- Most operations require that @e@ be an instance of the 'Ord' class. +-- A 'NESet' is strict in its elements. +-- +-- See documentation for 'NESet' for information on how to convert and +-- manipulate such non-empty set. +-- +-- This module essentially re-imports the API of "Data.Set" and its 'Set' +-- type, along with semantics and asymptotics. In most situations, +-- asymptotics are different only by a constant factor. In some +-- situations, asmyptotics are even better (constant-time instead of +-- log-time). All typeclass constraints are identical to their "Data.Set" +-- counterparts. +-- +-- Because 'NESet' is implemented using 'Set', all of the caveats of using +-- 'Set' apply (such as the limitation of the maximum size of sets). +-- +-- All functions take non-empty sets as inputs. In situations where their +-- results can be guarunteed to also be non-empty, they also return +-- non-empty sets. In situations where their results could potentially be +-- empty, 'Set' is returned instead. +-- +-- Some functions ('partition', 'spanAntitone', 'split') have modified +-- return types to account for possible configurations of non-emptiness. +-- +-- This module is intended to be imported qualified, to avoid name clashes +-- with "Prelude" and "Data.Set" functions: +-- +-- > import qualified Data.Set.NonEmpty as NES +module Data.Set.NonEmpty ( + -- * Non-Empty Set Type + NESet + -- ** Conversions between empty and non-empty sets + , pattern IsNonEmpty + , pattern IsEmpty + , nonEmptySet + , toSet + , withNonEmpty + , insertSet + , insertSetMin + , insertSetMax + , unsafeFromSet + + -- * Construction + , singleton + , fromList + , fromAscList + , fromDescList + , fromDistinctAscList + , fromDistinctDescList + , powerSet + + -- * Insertion + , insert + + -- * Deletion + , delete + + -- * Query + , member + , notMember + , lookupLT + , lookupGT + , lookupLE + , lookupGE + , size + , isSubsetOf + , isProperSubsetOf + , disjoint + + -- * Combine + , union + , unions + , difference + , (\\) + , intersection + , cartesianProduct + , disjointUnion + + -- * Filter + , filter + , takeWhileAntitone + , dropWhileAntitone + , spanAntitone + , partition + , split + , splitMember + , splitRoot + + -- * Indexed + , lookupIndex + , findIndex + , elemAt + , deleteAt + , take + , drop + , splitAt + + -- * Map + , map + , mapMonotonic + + -- * Folds + , foldr + , foldl + , foldr1 + , foldl1 + -- ** Strict folds + , foldr' + , foldl' + , foldr1' + , foldl1' + + -- * Min\/Max + , findMin + , findMax + , deleteMin + , deleteMax + , deleteFindMin + , deleteFindMax + + -- * Conversion + + -- ** List + , elems + , toList + , toAscList + , toDescList + + -- * Debugging + , valid + ) where + +import Control.Applicative +import Data.Bifunctor +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Maybe +import Data.Set (Set) +import Data.Set.NonEmpty.Internal +import Data.These +import Prelude hiding (foldr, foldl, filter, map, take, drop, splitAt) +import qualified Data.List.NonEmpty as NE +import qualified Data.Semigroup.Foldable as F1 +import qualified Data.Set as S + +-- | /O(1)/ match, /O(log n)/ usage of contents. The 'IsNonEmpty' and +-- 'IsEmpty' patterns allow you to treat a 'Set' as if it were either +-- a @'IsNonEmpty' n@ (where @n@ is a 'NESet') or an 'IsEmpty'. +-- +-- For example, you can pattern match on a 'Set': +-- +-- @ +-- myFunc :: 'Set' X -> Y +-- myFunc ('IsNonEmpty' n) = -- here, the user provided a non-empty set, and @n@ is the 'NESet' +-- myFunc 'IsEmpty' = -- here, the user provided an empty set +-- @ +-- +-- Matching on @'IsNonEmpty' n@ means that the original 'Set' was /not/ +-- empty, and you have a verified-non-empty 'NESet' @n@ to use. +-- +-- Note that patching on this pattern is /O(1)/. However, using the +-- contents requires a /O(log n)/ cost that is deferred until after the +-- pattern is matched on (and is not incurred at all if the contents are +-- never used). +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsNonEmpty' to convert +-- a 'NESet' back into a 'Set', obscuring its non-emptiness (see 'toSet'). +pattern IsNonEmpty :: NESet a -> Set a +pattern IsNonEmpty n <- (nonEmptySet->Just n) + where + IsNonEmpty n = toSet n + +-- | /O(1)/. The 'IsNonEmpty' and 'IsEmpty' patterns allow you to treat +-- a 'Set' as if it were either a @'IsNonEmpty' n@ (where @n@ is +-- a 'NESet') or an 'IsEmpty'. +-- +-- Matching on 'IsEmpty' means that the original 'Set' was empty. +-- +-- A case statement handling both 'IsNonEmpty' and 'IsEmpty' provides +-- complete coverage. +-- +-- This is a bidirectional pattern, so you can use 'IsEmpty' as an +-- expression, and it will be interpreted as 'Data.Set.empty'. +-- +-- See 'IsNonEmpty' for more information. +pattern IsEmpty :: Set a +pattern IsEmpty <- (S.null->True) + where + IsEmpty = S.empty + +{-# COMPLETE IsNonEmpty, IsEmpty #-} + +-- | /O(log n)/. Unsafe version of 'nonEmptySet'. Coerces a 'Set' into an +-- 'NESet', but is undefined (throws a runtime exception when evaluation is +-- attempted) for an empty 'Set'. +unsafeFromSet + :: Set a + -> NESet a +unsafeFromSet = withNonEmpty e id + where + e = errorWithoutStackTrace "NESet.unsafeFromSet: empty set" +{-# INLINE unsafeFromSet #-} + +-- | /O(log n)/. Convert a 'Set' into an 'NESet' by adding a value. +-- Because of this, we know that the set must have at least one +-- element, and so therefore cannot be empty. +-- +-- See 'insertSetMin' for a version that is constant-time if the new value is +-- /strictly smaller than/ all values in the original set +-- +-- > insertSet 4 (Data.Set.fromList [5, 3]) == fromList (3 :| [4, 5]) +-- > insertSet 4 Data.Set.empty == singleton 4 "c" +insertSet :: Ord a => a -> Set a -> NESet a +insertSet x = withNonEmpty (singleton x) (insert x) +{-# INLINE insertSet #-} + +-- | /O(1)/ Convert a 'Set' into an 'NESet' by adding a value where the +-- value is /strictly less than/ all values in the input set The values in +-- the original map must all be /strictly greater than/ the new value. +-- /The precondition is not checked./ +-- +-- > insertSetMin 2 (Data.Set.fromList [5, 3]) == fromList (2 :| [3, 5]) +-- > valid (insertSetMin 2 (Data.Set.fromList [5, 3])) == True +-- > valid (insertSetMin 7 (Data.Set.fromList [5, 3])) == False +-- > valid (insertSetMin 3 (Data.Set.fromList [5, 3])) == False +insertSetMin :: a -> Set a -> NESet a +insertSetMin = NESet +{-# INLINE insertSetMin #-} + +-- | /O(log n)/ Convert a 'Set' into an 'NESet' by adding a value where the +-- value is /strictly less than/ all values in the input set The values in +-- the original map must all be /strictly greater than/ the new value. +-- /The precondition is not checked./ +-- +-- While this has the same asymptotics as 'insertSet', it saves a constant +-- factor for key comparison (so may be helpful if comparison is expensive) +-- and also does not require an 'Ord' instance for the key type. +-- +-- > insertSetMin 7 (Data.Set.fromList [5, 3]) == fromList (3 :| [5, 7]) +-- > valid (insertSetMin 7 (Data.Set.fromList [5, 3])) == True +-- > valid (insertSetMin 2 (Data.Set.fromList [5, 3])) == False +-- > valid (insertSetMin 5 (Data.Set.fromList [5, 3])) == False +insertSetMax :: a -> Set a -> NESet a +insertSetMax x = withNonEmpty (singleton x) go + where + go (NESet x0 s0) = NESet x0 . insertMaxSet x $ s0 +{-# INLINE insertSetMax #-} + +-- | /O(n)/. Build a set from an ascending list in linear time. /The +-- precondition (input list is ascending) is not checked./ +fromAscList :: Eq a => NonEmpty a -> NESet a +fromAscList = fromDistinctAscList . combineEq +{-# INLINE fromAscList #-} + +-- | /O(n)/. Build a set from an ascending list of distinct elements in linear time. +-- /The precondition (input list is strictly ascending) is not checked./ +fromDistinctAscList :: NonEmpty a -> NESet a +fromDistinctAscList (x :| xs) = insertSetMin x + . S.fromDistinctAscList + $ xs +{-# INLINE fromDistinctAscList #-} + +-- | /O(n)/. Build a set from a descending list in linear time. +-- /The precondition (input list is descending) is not checked./ +fromDescList :: Eq a => NonEmpty a -> NESet a +fromDescList = fromDistinctDescList . combineEq +{-# INLINE fromDescList #-} + +-- | /O(n)/. Build a set from a descending list of distinct elements in linear time. +-- /The precondition (input list is strictly descending) is not checked./ +fromDistinctDescList :: NonEmpty a -> NESet a +fromDistinctDescList (x :| xs) = insertSetMax x + . S.fromDistinctDescList + $ xs +{-# INLINE fromDistinctDescList #-} + +-- | Calculate the power set of a non-empty: the set of all its (non-empty) +-- subsets. +-- +-- @ +-- t ``member`` powerSet s == t ``isSubsetOf`` s +-- @ +-- +-- Example: +-- +-- @ +-- powerSet (fromList (1 :| [2,3])) = +-- fromList (singleton 1 :| [ singleton 2 +-- , singleton 3 +-- , fromList (1 :| [2]) +-- , fromList (1 :| [3]) +-- , fromList (2 :| [3]) +-- , fromList (1 :| [2,3]) +-- ] +-- ) +-- @ +-- +-- We know that the result is non-empty because the result will always at +-- least contain the original set. +powerSet + :: forall a. () + => NESet a + -> NESet (NESet a) +powerSet (NESet x s0) = case nonEmptySet p1 of + -- s0 was empty originally + Nothing -> singleton (singleton x) + -- s1 was not empty originally + Just p2 -> mapMonotonic (insertSetMin x) p0 + `merge` p2 + where + -- powerset should never be empty + p0 :: NESet (Set a) + p0@(NESet _ p0s) = forSure $ powerSetSet s0 + p1 :: Set (NESet a) + p1 = S.mapMonotonic forSure p0s -- only minimal element is empty, so the rest aren't + forSure = withNonEmpty (errorWithoutStackTrace "NESet.powerSet: internal error") + id +{-# INLINABLE powerSet #-} + +-- | /O(log n)/. Insert an element in a set. +-- If the set already contains an element equal to the given value, +-- it is replaced with the new value. +insert :: Ord a => a -> NESet a -> NESet a +insert x n@(NESet x0 s) = case compare x x0 of + LT -> NESet x $ toSet n + EQ -> NESet x s + GT -> NESet x0 $ S.insert x s +{-# INLINE insert #-} + +-- | /O(log n)/. Delete an element from a set. +delete :: Ord a => a -> NESet a -> Set a +delete x n@(NESet x0 s) = case compare x x0 of + LT -> toSet n + EQ -> s + GT -> insertMinSet x0 . S.delete x $ s +{-# INLINE delete #-} + +-- | /O(log n)/. Is the element in the set? +member :: Ord a => a -> NESet a -> Bool +member x (NESet x0 s) = case compare x x0 of + LT -> False + EQ -> True + GT -> S.member x s +{-# INLINE member #-} + +-- | /O(log n)/. Is the element not in the set? +notMember :: Ord a => a -> NESet a -> Bool +notMember x (NESet x0 s) = case compare x x0 of + LT -> True + EQ -> False + GT -> S.notMember x s +{-# INLINE notMember #-} + +-- | /O(log n)/. Find largest element smaller than the given one. +-- +-- > lookupLT 3 (fromList (3 :| [5])) == Nothing +-- > lookupLT 5 (fromList (3 :| [5])) == Just 3 +lookupLT :: Ord a => a -> NESet a -> Maybe a +lookupLT x (NESet x0 s) = case compare x x0 of + LT -> Nothing + EQ -> Nothing + GT -> S.lookupLT x s <|> Just x0 +{-# INLINE lookupLT #-} + +-- | /O(log n)/. Find smallest element greater than the given one. +-- +-- > lookupLT 4 (fromList (3 :| [5])) == Just 5 +-- > lookupLT 5 (fromList (3 :| [5])) == Nothing +lookupGT :: Ord a => a -> NESet a -> Maybe a +lookupGT x (NESet x0 s) = case compare x x0 of + LT -> Just x0 + EQ -> S.lookupMin s + GT -> S.lookupGT x s +{-# INLINE lookupGT #-} + +-- | /O(log n)/. Find largest element smaller or equal to the given one. +-- +-- > lookupLT 2 (fromList (3 :| [5])) == Nothing +-- > lookupLT 4 (fromList (3 :| [5])) == Just 3 +-- > lookupLT 5 (fromList (3 :| [5])) == Just 5 +lookupLE :: Ord a => a -> NESet a -> Maybe a +lookupLE x (NESet x0 s) = case compare x x0 of + LT -> Nothing + EQ -> Just x0 + GT -> S.lookupLE x s <|> Just x0 +{-# INLINE lookupLE #-} + +-- | /O(log n)/. Find smallest element greater or equal to the given one. +-- +-- > lookupLT 3 (fromList (3 :| [5])) == Just 3 +-- > lookupLT 4 (fromList (3 :| [5])) == Just 5 +-- > lookupLT 6 (fromList (3 :| [5])) == Nothing +lookupGE :: Ord a => a -> NESet a -> Maybe a +lookupGE x (NESet x0 s) = case compare x x0 of + LT -> Just x0 + EQ -> Just x0 + GT -> S.lookupGE x s +{-# INLINE lookupGE #-} + +-- | /O(n+m)/. Is this a subset? +-- @(s1 \`isSubsetOf\` s2)@ tells whether @s1@ is a subset of @s2@. +isSubsetOf + :: Ord a + => NESet a + -> NESet a + -> Bool +isSubsetOf (NESet x s0) (toSet->s1) = x `S.member` s1 + && s0 `S.isSubsetOf` s1 +{-# INLINE isSubsetOf #-} + +-- | /O(n+m)/. Is this a proper subset? (ie. a subset but not equal). +isProperSubsetOf + :: Ord a + => NESet a + -> NESet a + -> Bool +isProperSubsetOf s0 s1 = S.size (nesSet s0) < S.size (nesSet s1) + && s0 `isSubsetOf` s1 +{-# INLINE isProperSubsetOf #-} + +-- | /O(n+m)/. Check whether two sets are disjoint (i.e. their intersection +-- is empty). +-- +-- > disjoint (fromList (2:|[4,6])) (fromList (1:|[3])) == True +-- > disjoint (fromList (2:|[4,6,8])) (fromList (2:|[3,5,7])) == False +-- > disjoint (fromList (1:|[2])) (fromList (1:|[2,3,4])) == False +disjoint + :: Ord a + => NESet a + -> NESet a + -> Bool +disjoint n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of + -- x1 is not in n2 + LT -> s1 `disjointSet` toSet n2 + -- k1 and k2 are a part of the result + EQ -> False + -- k2 is not in n1 + GT -> toSet n1 `disjointSet` s2 +{-# INLINE disjoint #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. Difference of two sets. +-- +-- Returns a potentially empty set ('Set') because the first set might be +-- a subset of the second set, and therefore have all of its elements +-- removed. +difference + :: Ord a + => NESet a + -> NESet a + -> Set a +difference n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of + -- x1 is not in n2, so cannot be deleted + LT -> insertMinSet x1 $ s1 `S.difference` toSet n2 + -- x2 deletes x1, and only x1 + EQ -> s1 `S.difference` s2 + -- x2 is not in n1, so cannot delete anything, so we can just difference n1 // s2. + GT -> toSet n1 `S.difference` s2 +{-# INLINE difference #-} + +-- | Same as 'difference'. +(\\) + :: Ord a + => NESet a + -> NESet a + -> Set a +(\\) = difference +{-# INLINE (\\) #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. The intersection of two sets. +-- +-- Returns a potentially empty set ('Set'), because the two sets might have +-- an empty intersection. +-- +-- Elements of the result come from the first set, so for example +-- +-- > import qualified Data.Set.NonEmpty as NES +-- > data AB = A | B deriving Show +-- > instance Ord AB where compare _ _ = EQ +-- > instance Eq AB where _ == _ = True +-- > main = print (NES.singleton A `NES.intersection` NES.singleton B, +-- > NES.singleton B `NES.intersection` NES.singleton A) +-- +-- prints @(fromList (A:|[]),fromList (B:|[]))@. +intersection + :: Ord a + => NESet a + -> NESet a + -> Set a +intersection n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of + -- x1 is not in n2 + LT -> s1 `S.intersection` toSet n2 + -- x1 and x2 are a part of the result + EQ -> insertMinSet x1 $ s1 `S.intersection` s2 + -- x2 is not in n1 + GT -> toSet n1 `S.intersection` s2 +{-# INLINE intersection #-} + +-- | Calculate the Cartesian product of two sets. +-- +-- @ +-- cartesianProduct xs ys = fromList $ liftA2 (,) (toList xs) (toList ys) +-- @ +-- +-- Example: +-- +-- @ +-- cartesianProduct (fromList (1:|[2])) (fromList (\'a\':|[\'b\'])) = +-- fromList ((1,\'a\') :| [(1,\'b\'), (2,\'a\'), (2,\'b\')]) +-- @ +cartesianProduct + :: NESet a + -> NESet b + -> NESet (a, b) +cartesianProduct n1 n2 = getMergeNESet + . F1.foldMap1 (\x -> MergeNESet $ mapMonotonic (x,) n2) + $ n1 +{-# INLINE cartesianProduct #-} + +-- | Calculate the disjoint union of two sets. +-- +-- @ disjointUnion xs ys = map Left xs ``union`` map Right ys @ +-- +-- Example: +-- +-- @ +-- disjointUnion (fromList (1:|[2])) (fromList ("hi":|["bye"])) = +-- fromList (Left 1 :| [Left 2, Right "hi", Right "bye"]) +-- @ +disjointUnion + :: NESet a + -> NESet b + -> NESet (Either a b) +disjointUnion (NESet x1 s1) n2 = NESet (Left x1) + (s1 `disjointUnionSet` toSet n2) +{-# INLINE disjointUnion #-} + +-- | /O(n)/. Filter all elements that satisfy the predicate. +-- +-- Returns a potentially empty set ('Set') because the predicate might +-- filter out all items in the original non-empty set. +filter + :: (a -> Bool) + -> NESet a + -> Set a +filter f (NESet x s1) + | f x = insertMinSet x . S.filter f $ s1 + | otherwise = S.filter f s1 +{-# INLINE filter #-} + +-- | /O(log n)/. Take while a predicate on the elements holds. The user is +-- responsible for ensuring that for all elements @j@ and @k@ in the set, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- Returns a potentially empty set ('Set') because the predicate might fail +-- on the first input. +-- +-- @ +-- takeWhileAntitone p = Data.Set.fromDistinctAscList . Data.List.NonEmpty.takeWhile p . 'toList' +-- takeWhileAntitone p = 'filter' p +-- @ +takeWhileAntitone + :: (a -> Bool) + -> NESet a + -> Set a +takeWhileAntitone f (NESet x s) + | f x = insertMinSet x . S.takeWhileAntitone f $ s + | otherwise = S.empty +{-# INLINE takeWhileAntitone #-} + +-- | /O(log n)/. Drop while a predicate on the elements holds. The user is +-- responsible for ensuring that for all elements @j@ and @k@ in the set, +-- @j \< k ==\> p j \>= p k@. See note at 'spanAntitone'. +-- +-- Returns a potentially empty set ('Set') because the predicate might be +-- true for all items. +-- +-- @ +-- dropWhileAntitone p = Data.Set.fromDistinctAscList . Data.List.NonEmpty.dropWhile p . 'toList' +-- dropWhileAntitone p = 'filter' (not . p) +-- @ +dropWhileAntitone + :: (a -> Bool) + -> NESet a + -> Set a +dropWhileAntitone f n@(NESet x s) + | f x = S.dropWhileAntitone f s + | otherwise = toSet n +{-# INLINE dropWhileAntitone #-} + +-- | /O(log n)/. Divide a set at the point where a predicate on the +-- elements stops holding. The user is responsible for ensuring that for +-- all elements @j@ and @k@ in the set, @j \< k ==\> p j \>= p k@. +-- +-- Returns a 'These' with potentially two non-empty sets: +-- +-- * @'This' n1@ means that the predicate never failed for any item, +-- returning the original set +-- * @'That' n2@ means that the predicate failed for the first item, +-- returning the original set +-- * @'These' n1 n2@ gives @n1@ (the set up to the point where the +-- predicate stops holding) and @n2@ (the set starting from +-- the point where the predicate stops holding) +-- +-- @ +-- spanAntitone p xs = partition p xs +-- @ +-- +-- Note: if @p@ is not actually antitone, then @spanAntitone@ will split the set +-- at some /unspecified/ point where the predicate switches from holding to not +-- holding (where the predicate is seen to hold before the first element and to fail +-- after the last element). +spanAntitone + :: (a -> Bool) + -> NESet a + -> These (NESet a) (NESet a) +spanAntitone f n@(NESet x s0) + | f x = case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> This n + (Just _ , Nothing) -> This n + (Nothing, Just n2) -> These (singleton x) n2 + (Just _ , Just n2) -> These (insertSetMin x s1) n2 + | otherwise = That n + where + (s1, s2) = S.spanAntitone f s0 +{-# INLINABLE spanAntitone #-} + +-- | /O(n)/. Partition the map according to a predicate. +-- +-- Returns a 'These' with potentially two non-empty sets: +-- +-- * @'This' n1@ means that the predicate was true for all items. +-- * @'That' n2@ means that the predicate was false for all items. +-- * @'These' n1 n2@ gives @n1@ (all of the items that were true for the +-- predicate) and @n2@ (all of the items that were false for the +-- predicate). +-- +-- See also 'split'. +-- +-- > partition (> 3) (fromList (5 :| [3])) == These (singleton 5) (singleton 3) +-- > partition (< 7) (fromList (5 :| [3])) == This (fromList (3 :| [5])) +-- > partition (> 7) (fromList (5 :| [3])) == That (fromList (3 :| [5])) +partition + :: (a -> Bool) + -> NESet a + -> These (NESet a) (NESet a) +partition f n@(NESet x s0) = case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) + | f x -> This n + | otherwise -> That n + (Just n1, Nothing) + | f x -> This n + | otherwise -> These n1 (singleton x) + (Nothing, Just n2) + | f x -> These (singleton x) n2 + | otherwise -> That n + (Just n1, Just n2) + | f x -> These (insertSetMin x s1) n2 + | otherwise -> These n1 (insertSetMin x s2) + where + (s1, s2) = S.partition f s0 +{-# INLINABLE partition #-} + +-- | /O(log n)/. The expression (@'split' x set@) is potentially a 'These' +-- containing up to two 'NESet's based on splitting the set into sets +-- containing items before and after the value @x@. It will never return +-- a set that contains @x@ itself. +-- +-- * 'Nothing' means that @x@ was the only value in the the original set, +-- and so there are no items before or after it. +-- * @'Just' ('This' n1)@ means @x@ was larger than or equal to all items +-- in the set, and @n1@ is the entire original set (minus @x@, if it +-- was present) +-- * @'Just' ('That' n2)@ means @x@ was smaller than or equal to all +-- items in the set, and @n2@ is the entire original set (minus @x@, if +-- it was present) +-- * @'Just' ('These' n1 n2)@ gives @n1@ (the set of all values from the +-- original set less than @x@) and @n2@ (the set of all values from the +-- original set greater than @x@). +-- +-- > split 2 (fromList (5 :| [3])) == Just (That (fromList (3 :| [5])) ) +-- > split 3 (fromList (5 :| [3])) == Just (That (singleton 5) ) +-- > split 4 (fromList (5 :| [3])) == Just (These (singleton 3) (singleton 5)) +-- > split 5 (fromList (5 :| [3])) == Just (This (singleton 3) ) +-- > split 6 (fromList (5 :| [3])) == Just (This (fromList (3 :| [5])) ) +-- > split 5 (singleton 5) == Nothing +split + :: Ord a + => a + -> NESet a + -> Maybe (These (NESet a) (NESet a)) +split x n@(NESet x0 s0) = case compare x x0 of + LT -> Just $ That n + EQ -> That <$> nonEmptySet s0 + GT -> case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> Just $ This (singleton x0) + (Just _ , Nothing) -> Just $ This (insertSetMin x0 s1) + (Nothing, Just n2) -> Just $ These (singleton x0) n2 + (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2 + where + (s1, s2) = S.split x s0 +{-# INLINABLE split #-} + +-- | /O(log n)/. The expression (@'splitMember' x set@) splits a set just +-- like 'split' but also returns @'member' x set@ (whether or not @x@ was +-- in @set@) +-- +-- > splitMember 2 (fromList (5 :| [3])) == (False, Just (That (fromList (3 :| [5)])))) +-- > splitMember 3 (fromList (5 :| [3])) == (True , Just (That (singleton 5))) +-- > splitMember 4 (fromList (5 :| [3])) == (False, Just (These (singleton 3) (singleton 5))) +-- > splitMember 5 (fromList (5 :| [3])) == (True , Just (This (singleton 3)) +-- > splitMember 6 (fromList (5 :| [3])) == (False, Just (This (fromList (3 :| [5]))) +-- > splitMember 5 (singleton 5) == (True , Nothing) +splitMember + :: Ord a + => a + -> NESet a + -> (Bool, Maybe (These (NESet a) (NESet a))) +splitMember x n@(NESet x0 s0) = case compare x x0 of + LT -> (False, Just $ That n) + EQ -> (True , That <$> nonEmptySet s0) + GT -> (mem ,) $ case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> Just $ This (singleton x0) + (Just _ , Nothing) -> Just $ This (insertSetMin x0 s1) + (Nothing, Just n2) -> Just $ These (singleton x0) n2 + (Just _ , Just n2) -> Just $ These (insertSetMin x0 s1) n2 + where + (s1, mem, s2) = S.splitMember x s0 +{-# INLINABLE splitMember #-} + +-- | /O(1)/. Decompose a set into pieces based on the structure of the underlying +-- tree. This function is useful for consuming a set in parallel. +-- +-- No guarantee is made as to the sizes of the pieces; an internal, but +-- deterministic process determines this. However, it is guaranteed that +-- the pieces returned will be in ascending order (all elements in the +-- first subset less than all elements in the second, and so on). +-- +-- Note that the current implementation does not return more than four +-- subsets, but you should not depend on this behaviour because it can +-- change in the future without notice. +splitRoot + :: NESet a + -> NonEmpty (NESet a) +splitRoot (NESet x s) = singleton x + :| mapMaybe nonEmptySet (S.splitRoot s) +{-# INLINE splitRoot #-} + +-- | /O(log n)/. Lookup the /index/ of an element, which is its zero-based +-- index in the sorted sequence of elements. The index is a number from /0/ +-- up to, but not including, the 'size' of the set. +-- +-- > isJust (lookupIndex 2 (fromList (5:|[3]))) == False +-- > fromJust (lookupIndex 3 (fromList (5:|[3]))) == 0 +-- > fromJust (lookupIndex 5 (fromList (5:|[3]))) == 1 +-- > isJust (lookupIndex 6 (fromList (5:|[3]))) == False +lookupIndex + :: Ord a + => a + -> NESet a + -> Maybe Int +lookupIndex x (NESet x0 s) = case compare x x0 of + LT -> Nothing + EQ -> Just 0 + GT -> (+ 1) <$> S.lookupIndex x s +{-# INLINE lookupIndex #-} + +-- | /O(log n)/. Return the /index/ of an element, which is its zero-based +-- index in the sorted sequence of elements. The index is a number from /0/ +-- up to, but not including, the 'size' of the set. Calls 'error' when the +-- element is not a 'member' of the set. +-- +-- > findIndex 2 (fromList (5:|[3])) Error: element is not in the set +-- > findIndex 3 (fromList (5:|[3])) == 0 +-- > findIndex 5 (fromList (5:|[3])) == 1 +-- > findIndex 6 (fromList (5:|[3])) Error: element is not in the set +findIndex + :: Ord a + => a + -> NESet a + -> Int +findIndex k = fromMaybe e . lookupIndex k + where + e = error "NESet.findIndex: element is not in the set" +{-# INLINE findIndex #-} + +-- | /O(log n)/. Retrieve an element by its /index/, i.e. by its zero-based +-- index in the sorted sequence of elements. If the /index/ is out of range +-- (less than zero, greater or equal to 'size' of the set), 'error' is +-- called. +-- +-- > elemAt 0 (fromList (5:|[3])) == 3 +-- > elemAt 1 (fromList (5:|[3])) == 5 +-- > elemAt 2 (fromList (5:|[3])) Error: index out of range +elemAt + :: Int + -> NESet a + -> a +elemAt 0 (NESet x _) = x +elemAt i (NESet _ s) = S.elemAt (i - 1) s +{-# INLINE elemAt #-} + +-- | /O(log n)/. Delete the element at /index/, i.e. by its zero-based +-- index in the sorted sequence of elements. If the /index/ is out of range +-- (less than zero, greater or equal to 'size' of the set), 'error' is +-- called. +-- +-- Returns a potentially empty set ('Set'), because this could potentailly +-- delete the final element in a singleton set. +-- +-- > deleteAt 0 (fromList (5:|[3])) == singleton 5 +-- > deleteAt 1 (fromList (5:|[3])) == singleton 3 +-- > deleteAt 2 (fromList (5:|[3])) Error: index out of range +-- > deleteAt (-1) (fromList (5:|[3])) Error: index out of range +deleteAt + :: Int + -> NESet a + -> Set a +deleteAt 0 (NESet _ s) = s +deleteAt i (NESet x s) = insertMinSet x . S.deleteAt (i - 1) $ s +{-# INLINABLE deleteAt #-} + +-- | Take a given number of elements in order, beginning +-- with the smallest ones. +-- +-- Returns a potentailly empty set ('Set'), which can only happen when +-- calling @take 0@. +-- +-- @ +-- take n = Data.Set.fromDistinctAscList . Data.List.NonEmpty.take n . 'toAscList' +-- @ +take + :: Int + -> NESet a + -> Set a +take 0 (NESet _ _) = S.empty +take i (NESet x s) = insertMinSet x . S.take (i - 1) $ s +{-# INLINABLE take #-} + +-- | Drop a given number of elements in order, beginning +-- with the smallest ones. +-- +-- Returns a potentailly empty set ('Set'), in the case that 'drop' is +-- called with a number equal to or greater the number of items in the set, +-- and we drop every item. +-- +-- @ +-- drop n = Data.Set.fromDistinctAscList . Data.List.NonEmpty.drop n . 'toAscList' +-- @ +drop + :: Int + -> NESet a + -> Set a +drop 0 n = toSet n +drop n (NESet _ s) = S.drop (n - 1) s +{-# INLINABLE drop #-} + +-- | /O(log n)/. Split a set at a particular index @i@. +-- +-- * @'This' n1@ means that there are less than @i@ items in the set, and +-- @n1@ is the original set. +-- * @'That' n2@ means @i@ was 0; we dropped 0 items, so @n2@ is the +-- original set. +-- * @'These' n1 n2@ gives @n1@ (taking @i@ items from the original set) +-- and @n2@ (dropping @i@ items from the original set)) +splitAt + :: Int + -> NESet a + -> These (NESet a) (NESet a) +splitAt 0 n = That n +splitAt i n@(NESet x s0) = case (nonEmptySet s1, nonEmptySet s2) of + (Nothing, Nothing) -> This (singleton x) + (Just _ , Nothing) -> This n + (Nothing, Just n2) -> These (singleton x) n2 + (Just _ , Just n2) -> These (insertSetMin x s1) n2 + where + (s1, s2) = S.splitAt (i - 1) s0 +{-# INLINABLE splitAt #-} + +-- | /O(n*log n)/. +-- @'map' f s@ is the set obtained by applying @f@ to each element of @s@. +-- +-- It's worth noting that the size of the result may be smaller if, +-- for some @(x,y)@, @x \/= y && f x == f y@ +map :: Ord b + => (a -> b) + -> NESet a + -> NESet b +map f (NESet x0 s) = fromList + . (f x0 :|) + . S.foldr (\x xs -> f x : xs) [] + $ s +{-# INLINE map #-} + +-- | /O(n)/. +-- @'mapMonotonic' f s == 'map' f s@, but works only when @f@ is strictly +-- increasing. /The precondition is not checked./ Semi-formally, we have: +-- +-- > and [x < y ==> f x < f y | x <- ls, y <- ls] +-- > ==> mapMonotonic f s == map f s +-- > where ls = Data.Foldable.toList s +mapMonotonic + :: (a -> b) + -> NESet a + -> NESet b +mapMonotonic f (NESet x s) = NESet (f x) (S.mapMonotonic f s) +{-# INLINE mapMonotonic #-} + +-- | /O(n)/. A strict version of 'foldr1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr1' :: (a -> a -> a) -> NESet a -> a +foldr1' f (NESet x s) = case S.maxView s of + Nothing -> x + Just (y, s') -> let !z = S.foldr' f y s' in x `f` z +{-# INLINE foldr1' #-} + +-- | /O(n)/. A strict version of 'foldl1'. Each application of the operator +-- is evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl1' :: (a -> a -> a) -> NESet a -> a +foldl1' f (NESet x s) = S.foldl' f x s +{-# INLINE foldl1' #-} + +-- | /O(1)/. The minimal element of a set. Note that this is total, making +-- 'Data.Set.lookupMin' obsolete. It is constant-time, so has better +-- asymptotics than @Data.Set.lookupMin@ and @Data.Map.findMin@ as well. +-- +-- > findMin (fromList (5 :| [3])) == 3 +findMin :: NESet a -> a +findMin (NESet x _) = x +{-# INLINE findMin #-} + +-- | /O(log n)/. The maximal key of a set Note that this is total, +-- making 'Data.Set.lookupMin' obsolete. +-- +-- > findMax (fromList (5 :| [3])) == 5 +findMax :: NESet a -> a +findMax (NESet x s) = fromMaybe x . S.lookupMax $ s +{-# INLINE findMax #-} + +-- | /O(1)/. Delete the minimal element. Returns a potentially empty set +-- ('Set'), because we might delete the final item in a singleton set. It +-- is constant-time, so has better asymptotics than @Data.Set.deleteMin@. +-- +-- > deleteMin (fromList (5 :| [3, 7])) == Data.Set.fromList [5, 7] +-- > deleteMin (singleton 5) == Data.Set.empty +deleteMin :: NESet a -> Set a +deleteMin (NESet _ s) = s +{-# INLINE deleteMin #-} + +-- | /O(log n)/. Delete the maximal element. Returns a potentially empty +-- set ('Set'), because we might delete the final item in a singleton set. +-- +-- > deleteMax (fromList (5 :| [3, 7])) == Data.Set.fromList [3, 5] +-- > deleteMax (singleton 5) == Data.Set.empty +deleteMax :: NESet a -> Set a +deleteMax (NESet x s) = insertMinSet x . S.deleteMax $ s +{-# INLINE deleteMax #-} + +-- | /O(1)/. Delete and find the minimal element. It is constant-time, so +-- has better asymptotics that @Data.Set.minView@ for 'Set'. +-- +-- Note that unlike @Data.Set.deleteFindMin@ for 'Set', this cannot ever +-- fail, and so is a total function. However, the result 'Set' is +-- potentially empty, since the original set might have contained just +-- a single item. +-- +-- > deleteFindMin (fromList (5 :| [3, 10])) == (3, Data.Set.fromList [5, 10]) +deleteFindMin :: NESet a -> (a, Set a) +deleteFindMin (NESet x s) = (x, s) +{-# INLINE deleteFindMin #-} + +-- | /O(log n)/. Delete and find the minimal element. +-- +-- Note that unlike @Data.Set.deleteFindMax@ for 'Set', this cannot ever +-- fail, and so is a total function. However, the result 'Set' is +-- potentially empty, since the original set might have contained just +-- a single item. +-- +-- > deleteFindMax (fromList (5 :| [3, 10])) == (10, Data.Set.fromList [3, 5]) +deleteFindMax :: NESet a -> (a, Set a) +deleteFindMax (NESet x s) = maybe (x, S.empty) (second (insertMinSet x)) + . S.maxView + $ s +{-# INLINE deleteFindMax #-} + +-- | /O(n)/. An alias of 'toAscList'. The elements of a set in ascending +-- order. +elems :: NESet a -> NonEmpty a +elems = toList +{-# INLINE elems #-} + +-- | /O(n)/. Convert the set to an ascending non-empty list of elements. +toAscList :: NESet a -> NonEmpty a +toAscList = toList +{-# INLINE toAscList #-} + +-- | /O(n)/. Convert the set to a descending non-empty list of elements. +toDescList :: NESet a -> NonEmpty a +toDescList (NESet x s) = S.foldl' (flip (NE.<|)) (x :| []) s +{-# INLINE toDescList #-} + +-- --------------------------- +-- Combining functions +-- --------------------------- +-- +-- Code comes from "Data.Set.Internal" from containers, modified slightly +-- to work with NonEmpty +-- +-- Copyright : (c) Daan Leijen 2002 + +combineEq :: Eq a => NonEmpty a -> NonEmpty a +combineEq (x :| xs) = go x xs + where + go z [] = z :| [] + go z (y:ys) + | z == y = go z ys + | otherwise = z NE.<| go y ys diff --git a/src/Data/Set/NonEmpty/Internal.hs b/src/Data/Set/NonEmpty/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Set/NonEmpty/Internal.hs @@ -0,0 +1,561 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveDataTypeable #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE ViewPatterns #-} +{-# OPTIONS_HADDOCK not-home #-} + +-- | +-- Module : Data.Set.NonEmpty.Internal +-- Copyright : (c) Justin Le 2018 +-- License : BSD3 +-- +-- Maintainer : justin@jle.im +-- Stability : experimental +-- Portability : non-portable +-- +-- Unsafe internal-use functions used in the implementation of +-- "Data.Set.NonEmpty". These functions can potentially be used to break +-- the abstraction of 'NESet' and produce unsound sets, so be wary! +module Data.Set.NonEmpty.Internal ( + NESet(..) + , nonEmptySet + , withNonEmpty + , toSet + , singleton + , fromList + , toList + , size + , union + , unions + , foldr + , foldl + , foldr' + , foldl' + , MergeNESet(..) + , merge + , valid + , insertMinSet + , insertMaxSet + , disjointSet + , powerSetSet + , disjointUnionSet + , cartesianProductSet + ) where + +import Control.DeepSeq +import Data.Data +import Data.Function +import Data.Functor.Classes +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup +import Data.Semigroup.Foldable (Foldable1) +import Data.Set.Internal (Set(..)) +import Data.Typeable (Typeable) +import Prelude hiding (foldr, foldr1, foldl, foldl1) +import Text.Read +import qualified Data.Foldable as F +import qualified Data.Semigroup.Foldable as F1 +import qualified Data.Set as S +import qualified Data.Set.Internal as S + +#if !MIN_VERSION_containers(0,5,11) +import Utils.Containers.Internal.StrictPair +#endif + +-- | A non-empty (by construction) set of values @a@. At least one value +-- exists in an @'NESet' a@ at all times. +-- +-- Functions that /take/ an 'NESet' can safely operate on it with the +-- assumption that it has at least one item. +-- +-- Functions that /return/ an 'NESet' provide an assurance that the result +-- has at least one item. +-- +-- "Data.Set.NonEmpty" re-exports the API of "Data.Set", faithfully +-- reproducing asymptotics, typeclass constraints, and semantics. +-- Functions that ensure that input and output sets are both non-empty +-- (like 'Data.Set.NonEmpty.insert') return 'NESet', but functions that +-- might potentially return an empty map (like 'Data.Set.NonEmpty.delete') +-- return a 'Set' instead. +-- +-- You can directly construct an 'NESet' with the API from +-- "Data.Set.NonEmpty"; it's more or less the same as constructing a normal +-- 'Set', except you don't have access to 'Data.Set.empty'. There are also +-- a few ways to construct an 'NESet' from a 'Set': +-- +-- 1. The 'nonEmptySet' smart constructor will convert a @'Set' a@ into +-- a @'Maybe' ('NESet' a)@, returning 'Nothing' if the original 'Set' +-- was empty. +-- 2. You can use the 'Data.Set.NonEmpty.insertSet' family of functions to +-- insert a value into a 'Set' to create a guaranteed 'NESet'. +-- 3. You can use the 'Data.Set.NonEmpty.IsNonEmpty' and +-- 'Data.Set.NonEmpty.IsEmpty' patterns to "pattern match" on a 'Set' +-- to reveal it as either containing a 'NESet' or an empty map. +-- 4. 'withNonEmpty' offers a continuation-based interface for +-- deconstructing a 'Set' and treating it as if it were an 'NESet'. +-- +-- You can convert an 'NESet' into a 'Set' with 'toSet' or +-- 'Data.Set.NonEmpty.IsNonEmpty', essentially "obscuring" the non-empty +-- property from the type. +data NESet a = + NESet { nesV0 :: !a -- ^ invariant: must be smaller than smallest value in set + , nesSet :: !(Set a) + } + deriving (Typeable) + +instance Eq a => Eq (NESet a) where + t1 == t2 = S.size (nesSet t1) == S.size (nesSet t2) + && toList t1 == toList t2 + +instance Ord a => Ord (NESet a) where + compare = compare `on` toList + (<) = (<) `on` toList + (>) = (>) `on` toList + (<=) = (<=) `on` toList + (>=) = (>=) `on` toList + +instance Show a => Show (NESet a) where + showsPrec p xs = showParen (p > 10) $ + showString "fromList (" . shows (toList xs) . showString ")" + +instance (Read a, Ord a) => Read (NESet a) where + readPrec = parens $ prec 10 $ do + Ident "fromList" <- lexP + xs <- parens . prec 10 $ readPrec + return (fromList xs) + + readListPrec = readListPrecDefault + +instance Eq1 NESet where + liftEq eq m n = + size m == size n && liftEq eq (toList m) (toList n) + +instance Ord1 NESet where + liftCompare cmp m n = + liftCompare cmp (toList m) (toList n) + +instance Show1 NESet where + liftShowsPrec sp sl d m = + showsUnaryWith (liftShowsPrec sp sl) "fromList" d (toList m) + +instance NFData a => NFData (NESet a) where + rnf (NESet x s) = rnf x `seq` rnf s + +-- Data instance code from Data.Set.Internal +-- +-- Copyright : (c) Daan Leijen 2002 +instance (Data a, Ord a) => Data (NESet a) where + gfoldl f z set = z fromList `f` toList set + toConstr _ = fromListConstr + gunfold k z c = case constrIndex c of + 1 -> k (z fromList) + _ -> error "gunfold" + dataTypeOf _ = setDataType + dataCast1 = gcast1 + +fromListConstr :: Constr +fromListConstr = mkConstr setDataType "fromList" [] Prefix + +setDataType :: DataType +setDataType = mkDataType "Data.Set.NonEmpty.Internal.NESet" [fromListConstr] + + + + + +-- | /O(log n)/. Smart constructor for an 'NESet' from a 'Set'. Returns +-- 'Nothing' if the 'Set' was originally actually empty, and @'Just' n@ +-- with an 'NESet', if the 'Set' was not empty. +-- +-- 'nonEmptySet' and @'maybe' 'Data.Set.empty' 'toSet'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- See 'Data.Set.NonEmpty.IsNonEmpty' for a pattern synonym that lets you +-- "match on" the possiblity of a 'Set' being an 'NESet'. +-- +-- > nonEmptySet (Data.Set.fromList [3,5]) == Just (fromList (3:|[5])) +nonEmptySet :: Set a -> Maybe (NESet a) +nonEmptySet = (fmap . uncurry) NESet . S.minView +{-# INLINE nonEmptySet #-} + +-- | /O(log n)/. A general continuation-based way to consume a 'Set' as if +-- it were an 'NESet'. @'withNonEmpty' def f@ will take a 'Set'. If set is +-- empty, it will evaluate to @def@. Otherwise, a non-empty set 'NESet' +-- will be fed to the function @f@ instead. +-- +-- @'nonEmptySet' == 'withNonEmpty' 'Nothing' 'Just'@ +withNonEmpty + :: r -- ^ value to return if set is empty + -> (NESet a -> r) -- ^ function to apply if set is not empty + -> Set a + -> r +withNonEmpty def f = maybe def f . nonEmptySet +{-# INLINE withNonEmpty #-} + +-- | /O(log n)/. +-- Convert a non-empty set back into a normal possibly-empty map, for usage +-- with functions that expect 'Set'. +-- +-- Can be thought of as "obscuring" the non-emptiness of the set in its +-- type. See the 'Data.Set.NonEmpty.IsNotEmpty' pattern. +-- +-- 'nonEmptySet' and @'maybe' 'Data.Set.empty' 'toSet'@ form an +-- isomorphism: they are perfect structure-preserving inverses of +-- eachother. +-- +-- > toSet (fromList ((3,"a") :| [(5,"b")])) == Data.Set.fromList [(3,"a"), (5,"b")] +toSet :: NESet a -> Set a +toSet (NESet x s) = insertMinSet x s +{-# INLINE toSet #-} + +-- | /O(1)/. Create a singleton set. +singleton :: a -> NESet a +singleton x = NESet x S.empty +{-# INLINE singleton #-} + +-- | /O(n*log n)/. Create a set from a list of elements. + +-- TODO: write manually and optimize to be equivalent to +-- 'fromDistinctAscList' if items are ordered, just like the actual +-- 'S.fromList'. +fromList :: Ord a => NonEmpty a -> NESet a +fromList (x :| s) = withNonEmpty (singleton x) (<> singleton x) + . S.fromList + $ s +{-# INLINE fromList #-} + +-- | /O(n)/. Convert the set to a non-empty list of elements. +toList :: NESet a -> NonEmpty a +toList (NESet x s) = x :| S.toList s +{-# INLINE toList #-} + +-- | /O(1)/. The number of elements in the set. Guaranteed to be greater +-- than zero. +size :: NESet a -> Int +size (NESet _ s) = 1 + S.size s +{-# INLINE size #-} + +-- | /O(n)/. Fold the elements in the set using the given right-associative +-- binary operator, such that @'foldr' f z == 'Prelude.foldr' f z . 'Data.Set.NonEmpty.toAscList'@. +-- +-- For example, +-- +-- > elemsList set = foldr (:) [] set +foldr :: (a -> b -> b) -> b -> NESet a -> b +foldr f z (NESet x s) = x `f` S.foldr f z s +{-# INLINE foldr #-} + +-- | /O(n)/. A strict version of 'foldr'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldr' :: (a -> b -> b) -> b -> NESet a -> b +foldr' f z (NESet x s) = x `f` y + where + !y = S.foldr' f z s +{-# INLINE foldr' #-} + +-- | /O(n)/. A version of 'foldr' that uses the value at the maximal value +-- in the set as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldr1' for 'Set', this function is +-- total if the input function is total. +foldr1 :: (a -> a -> a) -> NESet a -> a +foldr1 f (NESet x s) = maybe x (f x . uncurry (S.foldr f)) + . S.maxView + $ s +{-# INLINE foldr1 #-} + +-- | /O(n)/. Fold the elements in the set using the given left-associative +-- binary operator, such that @'foldl' f z == 'Prelude.foldl' f z . 'Data.Set.NonEmpty.toAscList'@. +-- +-- For example, +-- +-- > descElemsList set = foldl (flip (:)) [] set +foldl :: (a -> b -> a) -> a -> NESet b -> a +foldl f z (NESet x s) = S.foldl f (f z x) s +{-# INLINE foldl #-} + +-- | /O(n)/. A strict version of 'foldl'. Each application of the operator is +-- evaluated before using the result in the next application. This +-- function is strict in the starting value. +foldl' :: (a -> b -> a) -> a -> NESet b -> a +foldl' f z (NESet x s) = S.foldl' f y s + where + !y = f z x +{-# INLINE foldl' #-} + +-- | /O(n)/. A version of 'foldl' that uses the value at the minimal value +-- in the set as the starting value. +-- +-- Note that, unlike 'Data.Foldable.foldl1' for 'Set', this function is +-- total if the input function is total. +foldl1 :: (a -> a -> a) -> NESet a -> a +foldl1 f (NESet x s) = S.foldl f x s +{-# INLINE foldl1 #-} + +-- | /O(m*log(n\/m + 1)), m <= n/. The union of two sets, preferring the first set when +-- equal elements are encountered. +union + :: Ord a + => NESet a + -> NESet a + -> NESet a +union n1@(NESet x1 s1) n2@(NESet x2 s2) = case compare x1 x2 of + LT -> NESet x1 . S.union s1 . toSet $ n2 + EQ -> NESet x1 . S.union s1 $ s2 + GT -> NESet x2 . S.union (toSet n1) $ s2 +{-# INLINE union #-} + +-- | The union of a non-empty list of sets +unions + :: (Foldable1 f, Ord a) + => f (NESet a) + -> NESet a +unions (F1.toNonEmpty->(s :| ss)) = F.foldl' union s ss +{-# INLINE unions #-} + +-- | Left-biased union +instance Ord a => Semigroup (NESet a) where + (<>) = union + {-# INLINE (<>) #-} + sconcat = unions + {-# INLINE sconcat #-} + +-- | Traverses elements in ascending order +-- +-- 'Data.Foldable.foldr1', 'Data.Foldable.foldl1', 'Data.Foldable.minimum', +-- 'Data.Foldable.maximum' are all total. +instance Foldable NESet where +#if MIN_VERSION_base(4,11,0) + fold (NESet x s) = x <> F.fold s + {-# INLINE fold #-} + foldMap f (NESet x s) = f x <> foldMap f s + {-# INLINE foldMap #-} +#else + fold (NESet x s) = x `mappend` F.fold s + {-# INLINE fold #-} + foldMap f (NESet x s) = f x `mappend` foldMap f s + {-# INLINE foldMap #-} +#endif + foldr = foldr + {-# INLINE foldr #-} + foldr' = foldr' + {-# INLINE foldr' #-} + foldr1 = foldr1 + {-# INLINE foldr1 #-} + foldl = foldl + {-# INLINE foldl #-} + foldl' = foldl' + {-# INLINE foldl' #-} + foldl1 = foldl1 + {-# INLINE foldl1 #-} + null _ = False + {-# INLINE null #-} + length = size + {-# INLINE length #-} + elem x (NESet x0 s) = F.elem x s + || x == x0 + {-# INLINE elem #-} + minimum (NESet x _) = x + {-# INLINE minimum #-} + maximum (NESet x s) = maybe x fst . S.maxView $ s + {-# INLINE maximum #-} + -- TODO: use build + toList = F.toList . toList + {-# INLINE toList #-} + +-- | Traverses elements in ascending order +instance Foldable1 NESet where + fold1 (NESet x s) = maybe x (x <>) + . getOption + . F.foldMap (Option . Just) + $ s + {-# INLINE fold1 #-} + -- TODO: benchmark against maxView-based method + foldMap1 f (NESet x s) = maybe (f x) (f x <>) + . getOption + . F.foldMap (Option . Just . f) + $ s + {-# INLINE foldMap1 #-} + toNonEmpty = toList + {-# INLINE toNonEmpty #-} + + +-- | Used for 'Data.Set.NonEmpty.cartesianProduct' +newtype MergeNESet a = MergeNESet { getMergeNESet :: NESet a } + +instance Semigroup (MergeNESet a) where + MergeNESet n1 <> MergeNESet n2 = MergeNESet (merge n1 n2) + {-# INLINE (<>) #-} + +-- | Unsafely merge two disjoint sets. Only legal if all items in the +-- first set are less than all items in the second set +merge :: NESet a -> NESet a -> NESet a +merge (NESet x1 s1) n2 = NESet x1 $ s1 `S.merge` toSet n2 + +-- | /O(n)/. Test if the internal set structure is valid. +valid :: Ord a => NESet a -> Bool +valid (NESet x s) = S.valid s + && all ((x <) . fst) (S.minView s) + + + + +-- | /O(log n)/. Insert new value into a set where values are +-- /strictly greater than/ the new values That is, the new value must be +-- /strictly less than/ all values present in the 'Set'. /The precondition +-- is not checked./ +-- +-- While this has the same asymptotics as @Data.Set.insert@, it saves +-- a constant factor for value comparison (so may be helpful if comparison +-- is expensive) and also does not require an 'Ord' instance for the value +-- type. +insertMinSet :: a -> Set a -> Set a +insertMinSet x = \case + Tip -> S.singleton x + Bin _ y l r -> balanceL y (insertMinSet x l) r +{-# INLINABLE insertMinSet #-} + +-- | /O(log n)/. Insert new value into a set where values are /strictly +-- less than/ the new value. That is, the new value must be /strictly +-- greater than/ all values present in the 'Set'. /The precondition is not +-- checked./ +-- +-- While this has the same asymptotics as @Data.Set.insert@, it saves +-- a constant factor for value comparison (so may be helpful if comparison +-- is expensive) and also does not require an 'Ord' instance for the value +-- type. +insertMaxSet :: a -> Set a -> Set a +insertMaxSet x = \case + Tip -> S.singleton x + Bin _ y l r -> balanceR y l (insertMaxSet x r) +{-# INLINABLE insertMaxSet #-} + +-- --------------------------------------------- +-- | CPP for new functions not in old containers +-- --------------------------------------------- + +-- | Comptability layer for 'Data.Set.disjoint'. +disjointSet :: Ord a => Set a -> Set a -> Bool +#if MIN_VERSION_containers(0,5,11) +disjointSet = S.disjoint +#else +disjointSet xs = S.null . S.intersection xs +#endif +{-# INLINE disjointSet #-} + +-- | Comptability layer for 'Data.Set.powerSet'. +powerSetSet :: Set a -> Set (Set a) +#if MIN_VERSION_containers(0,5,11) +powerSetSet = S.powerSet +{-# INLINE powerSetSet #-} +#else +powerSetSet xs0 = insertMinSet S.empty (S.foldr' step' Tip xs0) where + step' x pxs = insertMinSet (S.singleton x) (insertMinSet x `S.mapMonotonic` pxs) `glue` pxs +{-# INLINABLE powerSetSet #-} + +minViewSure :: a -> Set a -> Set a -> StrictPair a (Set a) +minViewSure = go + where + go x Tip r = x :*: r + go x (Bin _ xl ll lr) r = + case go xl ll lr of + xm :*: l' -> xm :*: balanceR x l' r + +maxViewSure :: a -> Set a -> Set a -> StrictPair a (Set a) +maxViewSure = go + where + go x l Tip = x :*: l + go x l (Bin _ xr rl rr) = + case go xr rl rr of + xm :*: r' -> xm :*: balanceL x l r' + +glue :: Set a -> Set a -> Set a +glue Tip r = r +glue l Tip = l +glue l@(Bin sl xl ll lr) r@(Bin sr xr rl rr) + | sl > sr = let !(m :*: l') = maxViewSure xl ll lr in balanceR m l' r + | otherwise = let !(m :*: r') = minViewSure xr rl rr in balanceL m l r' +#endif + +-- | Comptability layer for 'Data.Set.disjointUnion'. +disjointUnionSet :: Set a -> Set b -> Set (Either a b) +#if MIN_VERSION_containers(0,5,11) +disjointUnionSet = S.disjointUnion +#else +disjointUnionSet as bs = S.merge (S.mapMonotonic Left as) (S.mapMonotonic Right bs) +#endif +{-# INLINE disjointUnionSet #-} + +-- | Comptability layer for 'Data.Set.cartesianProduct'. +cartesianProductSet :: Set a -> Set b -> Set (a, b) +#if MIN_VERSION_containers(0,5,11) +cartesianProductSet = S.cartesianProduct +#else +cartesianProductSet as bs = + getMergeSet $ foldMap (\a -> MergeSet $ S.mapMonotonic ((,) a) bs) as + +newtype MergeSet a = MergeSet { getMergeSet :: Set a } + +instance Semigroup (MergeSet a) where + MergeSet xs <> MergeSet ys = MergeSet (S.merge xs ys) + +instance Monoid (MergeSet a) where + mempty = MergeSet S.empty + mappend = (<>) +#endif +{-# INLINE cartesianProductSet #-} + + + +-- ------------------------------------------ +-- | Unexported code from "Data.Set.Internal" +-- ------------------------------------------ + +balanceR :: a -> Set a -> Set a -> Set a +balanceR x l r = case l of + Tip -> case r of + Tip -> Bin 1 x Tip Tip + Bin _ _ Tip Tip -> Bin 2 x Tip r + Bin _ rx Tip rr@Bin{} -> Bin 3 rx (Bin 1 x Tip Tip) rr + Bin _ rx (Bin _ rlx _ _) Tip -> Bin 3 rlx (Bin 1 x Tip Tip) (Bin 1 rx Tip Tip) + Bin rs rx rl@(Bin rls rlx rll rlr) rr@(Bin rrs _ _ _) + | rls < ratio*rrs -> Bin (1+rs) rx (Bin (1+rls) x Tip rl) rr + | otherwise -> Bin (1+rs) rlx (Bin (1+S.size rll) x Tip rll) (Bin (1+rrs+S.size rlr) rx rlr rr) + Bin ls _ _ _ -> case r of + Tip -> Bin (1+ls) x l Tip + Bin rs rx rl rr + | rs > delta*ls -> case (rl, rr) of + (Bin rls rlx rll rlr, Bin rrs _ _ _) + | rls < ratio*rrs -> Bin (1+ls+rs) rx (Bin (1+ls+rls) x l rl) rr + | otherwise -> Bin (1+ls+rs) rlx (Bin (1+ls+S.size rll) x l rll) (Bin (1+rrs+S.size rlr) rx rlr rr) + (_, _) -> error "Failure in Data.Map.balanceR" + | otherwise -> Bin (1+ls+rs) x l r +{-# NOINLINE balanceR #-} + +balanceL :: a -> Set a -> Set a -> Set a +balanceL x l r = case r of + Tip -> case l of + Tip -> Bin 1 x Tip Tip + Bin _ _ Tip Tip -> Bin 2 x l Tip + Bin _ lx Tip (Bin _ lrx _ _) -> Bin 3 lrx (Bin 1 lx Tip Tip) (Bin 1 x Tip Tip) + Bin _ lx ll@Bin{} Tip -> Bin 3 lx ll (Bin 1 x Tip Tip) + Bin ls lx ll@(Bin lls _ _ _) lr@(Bin lrs lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls) lx ll (Bin (1+lrs) x lr Tip) + | otherwise -> Bin (1+ls) lrx (Bin (1+lls+S.size lrl) lx ll lrl) (Bin (1+S.size lrr) x lrr Tip) + Bin rs _ _ _ -> case l of + Tip -> Bin (1+rs) x Tip r + Bin ls lx ll lr + | ls > delta*rs -> case (ll, lr) of + (Bin lls _ _ _, Bin lrs lrx lrl lrr) + | lrs < ratio*lls -> Bin (1+ls+rs) lx ll (Bin (1+rs+lrs) x lr r) + | otherwise -> Bin (1+ls+rs) lrx (Bin (1+lls+S.size lrl) lx ll lrl) (Bin (1+rs+S.size lrr) x lrr r) + (_, _) -> error "Failure in Data.Set.NonEmpty.Internal.balanceL" + | otherwise -> Bin (1+ls+rs) x l r +{-# NOINLINE balanceL #-} + +delta,ratio :: Int +delta = 3 +ratio = 2 diff --git a/test/Spec.hs b/test/Spec.hs new file mode 100644 --- /dev/null +++ b/test/Spec.hs @@ -0,0 +1,25 @@ + +import Test.Tasty +import Test.Tasty.Hedgehog +import Test.Tasty.Ingredients.ConsoleReporter +import Tests.IntMap +import Tests.IntSet +import Tests.Map +import Tests.Sequence +import Tests.Set + +setOpts :: TestTree -> TestTree +setOpts = id +-- setOpts = localOption (HedgehogTestLimit (Just 500)) +-- . localOption (HedgehogDiscardLimit (Just 500)) +-- . localOption (HideSuccesses True ) + +main :: IO () +main = defaultMain . setOpts $ + testGroup "Tests" [ mapTests + , setTests + , intMapTests + , intSetTests + , sequenceTests + ] + diff --git a/test/Tests/IntMap.hs b/test/Tests/IntMap.hs new file mode 100644 --- /dev/null +++ b/test/Tests/IntMap.hs @@ -0,0 +1,711 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} + +module Tests.IntMap (intMapTests) where + +import Control.Applicative +import Data.Coerce +import Data.Foldable +import Data.Functor.Identity +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup.Foldable +import Data.Text (Text) +import Hedgehog +import Test.Tasty +import Tests.Util +import qualified Data.IntMap as M +import qualified Data.IntMap.NonEmpty as NEM +import qualified Data.IntMap.NonEmpty.Internal as NEM +import qualified Data.List.NonEmpty as NE +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range + +intMapTests :: TestTree +intMapTests = groupTree $$(discover) + + + + + +prop_valid :: Property +prop_valid = property $ + assert . NEM.valid =<< forAll neIntMapGen + +-- | We cannot implement these because there is no 'valid' for IntSet +-- prop_valid_toMap :: Property +-- prop_valid_toMap = property $ +-- assert . M.valid . NEM.toMap =<< forAll neIntMapGen + +-- prop_valid_insertMinIntMap :: Property +-- prop_valid_insertMinIntMap = property $ do +-- n <- forAll $ do +-- m <- intMapGen +-- let k = maybe 0 (subtract 1 . fst) $ M.lookupMin m +-- v <- valGen +-- pure $ NEM.insertMinIntMap k v m +-- assert $ M.valid n + +-- prop_valid_insertMaxIntMap :: Property +-- prop_valid_insertMaxIntMap = property $ do +-- n <- forAll $ do +-- m <- intMapGen +-- let k = maybe 0 ((+ 1) . fst) $ M.lookupMax m +-- v <- valGen +-- pure $ NEM.insertMaxIntMap k v m +-- assert $ M.valid n + +prop_valid_insertMapMin :: Property +prop_valid_insertMapMin = property $ do + n <- forAll $ do + m <- intMapGen + let k = maybe 0 (subtract 1 . fst) $ NEM.lookupMinMap m + v <- valGen + pure $ NEM.insertMapMin k v m + assert $ NEM.valid n + +prop_valid_insertMapMax :: Property +prop_valid_insertMapMax = property $ do + n <- forAll $ do + m <- intMapGen + let k = maybe 0 ((+ 1) . fst) $ NEM.lookupMaxMap m + v <- valGen + pure $ NEM.insertMapMax k v m + assert $ NEM.valid n + +prop_toMapIso1 :: Property +prop_toMapIso1 = property $ do + m0 <- forAll intMapGen + tripping m0 (NEM.nonEmptyMap) + (Identity . maybe M.empty NEM.toMap) + +prop_toMapIso2 :: Property +prop_toMapIso2 = property $ do + m0 <- forAll $ Gen.maybe neIntMapGen + tripping m0 (maybe M.empty NEM.toMap) + (Identity . NEM.nonEmptyMap) + +prop_read_show :: Property +prop_read_show = readShow neIntMapGen + +prop_read1_show1 :: Property +prop_read1_show1 = readShow1 neIntMapGen + +prop_show_show1 :: Property +prop_show_show1 = showShow1 neIntMapGen + +prop_splitRoot :: Property +prop_splitRoot = property $ do + n <- forAll neIntMapGen + let rs = NEM.splitRoot n + allItems = foldMap1 NEM.keys rs + n' = NEM.unions rs + assert $ ascending allItems + mapM_ (assert . (`NEM.isSubmapOf` n)) rs + length allItems === length n' + n === n' + where + ascending (x :| xs) = case NE.nonEmpty xs of + Nothing -> True + Just ys@(y :| _) -> x < y && ascending ys + + + + + + + +prop_insertMapWithKey :: Property +prop_insertMapWithKey = ttProp (gf3 valGen :?> GTIntKey :-> GTVal :-> GTIntMap :-> TTNEIntMap) + M.insertWithKey + NEM.insertMapWithKey + +prop_singleton :: Property +prop_singleton = ttProp (GTIntKey :-> GTVal :-> TTNEIntMap) + M.singleton + NEM.singleton + +prop_fromSet :: Property +prop_fromSet = ttProp (gf1 valGen :?> GTNEIntSet :-> TTNEIntMap) + M.fromSet + NEM.fromSet + +prop_fromAscList :: Property +prop_fromAscList = ttProp (GTSorted STAsc (GTNEList Nothing (GTIntKey :&: GTVal)) :-> TTNEIntMap) + M.fromAscList + NEM.fromAscList + +prop_fromAscListWithKey :: Property +prop_fromAscListWithKey = ttProp (gf3 valGen :?> GTSorted STAsc (GTNEList Nothing (GTIntKey :&: GTVal)) :-> TTNEIntMap) + M.fromAscListWithKey + NEM.fromAscListWithKey + +prop_fromDistinctAscList :: Property +prop_fromDistinctAscList = ttProp (GTSorted STDistinctAsc (GTNEList Nothing (GTIntKey :&: GTVal)) :-> TTNEIntMap) + M.fromDistinctAscList + NEM.fromDistinctAscList + +prop_fromListWithKey :: Property +prop_fromListWithKey = ttProp (gf3 valGen :?> GTNEList Nothing (GTIntKey :&: GTVal) :-> TTNEIntMap) + M.fromListWithKey + NEM.fromListWithKey + +prop_insert :: Property +prop_insert = ttProp (GTIntKey :-> GTVal :-> GTNEIntMap :-> TTNEIntMap) + M.insert + NEM.insert + +prop_insertWithKey :: Property +prop_insertWithKey = ttProp (gf3 valGen :?> GTIntKey :-> GTVal :-> GTNEIntMap :-> TTNEIntMap) + M.insertWithKey + NEM.insertWithKey + +prop_delete :: Property +prop_delete = ttProp (GTIntKey :-> GTNEIntMap :-> TTOther) + M.delete + NEM.delete + +prop_adjustWithKey :: Property +prop_adjustWithKey = ttProp (gf2 valGen :?> GTIntKey :-> GTNEIntMap :-> TTNEIntMap) + M.adjustWithKey + NEM.adjustWithKey + +prop_updateWithKey :: Property +prop_updateWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTIntKey :-> GTNEIntMap :-> TTOther) + M.updateWithKey + NEM.updateWithKey + +prop_updateLookupWithKey :: Property +prop_updateLookupWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTIntKey :-> GTNEIntMap :-> TTMaybe TTVal :*: TTOther) + M.updateLookupWithKey + NEM.updateLookupWithKey + +prop_alter :: Property +prop_alter = ttProp (gf1 (Gen.maybe valGen) :?> GTIntKey :-> GTNEIntMap :-> TTOther) + M.alter + NEM.alter + +prop_alter' :: Property +prop_alter' = ttProp (gf1 valGen :?> GTIntKey :-> GTNEIntMap :-> TTNEIntMap) + (M.alter . fmap Just) + NEM.alter' + +prop_alterF :: Property +prop_alterF = ttProp ( gf1 (Gen.maybe valGen) + :?> GTIntKey + :-> GTNEIntMap + :-> TTCtx (GTMaybe GTVal :-> TTOther) (TTMaybe TTVal) + ) + (M.alterF . Context) + (NEM.alterF . Context) + +prop_alterF_rules_Const :: Property +prop_alterF_rules_Const = ttProp ( gf1 (Const <$> valGen) + :?> GTIntKey + :-> GTNEIntMap + :-> TTOther + ) + (\f k m -> getConst (M.alterF f k m)) + (\f k m -> getConst (NEM.alterF f k m)) + +prop_alterF_rules_Identity :: Property +prop_alterF_rules_Identity = ttProp ( gf1 (Identity <$> Gen.maybe valGen) + :?> GTIntKey + :-> GTNEIntMap + :-> TTOther + ) + (\f k m -> runIdentity (M.alterF f k m)) + (\f k m -> runIdentity (NEM.alterF f k m)) + +prop_alterF' :: Property +prop_alterF' = ttProp (gf1 valGen :?> GTIntKey :-> GTNEIntMap :-> TTCtx (GTVal :-> TTNEIntMap) (TTMaybe TTVal)) + (M.alterF . Context . fmap Just) + (NEM.alterF' . Context) + +prop_alterF'_rules_Const :: Property +prop_alterF'_rules_Const = ttProp ( gf1 (Const <$> valGen) + :?> GTIntKey + :-> GTNEIntMap + :-> TTOther + ) + (\f k m -> let f' = fmap Just . f in getConst (M.alterF f' k m)) + (\f k m -> getConst (NEM.alterF' f k m)) + +-- -- | This fails, but isn't possible to fix without copying-and-pasting more +-- -- in code from containers. +-- prop_alterF'_rules_Identity :: Property +-- prop_alterF'_rules_Identity = ttProp ( gf1 (Identity <$> valGen) +-- :?> GTIntKey +-- :-> GTNEIntMap +-- :-> TTNEIntMap +-- ) +-- (\f k m -> let f' = fmap Just . f in runIdentity (M.alterF f' k m)) +-- (\f k m -> runIdentity (NEM.alterF' f k m)) + +prop_lookup :: Property +prop_lookup = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe TTVal) + M.lookup + NEM.lookup + +prop_findWithDefault :: Property +prop_findWithDefault = ttProp (GTVal :-> GTIntKey :-> GTNEIntMap :-> TTVal) + M.findWithDefault + NEM.findWithDefault + +prop_member :: Property +prop_member = ttProp (GTIntKey :-> GTNEIntMap :-> TTOther) + M.member + NEM.member + +prop_notMember :: Property +prop_notMember = ttProp (GTIntKey :-> GTNEIntMap :-> TTOther) + M.notMember + NEM.notMember + +prop_lookupLT :: Property +prop_lookupLT = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe (TTOther :*: TTVal)) + M.lookupLT + NEM.lookupLT + +prop_lookupGT :: Property +prop_lookupGT = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe (TTOther :*: TTVal)) + M.lookupGT + NEM.lookupGT + +prop_lookupLE :: Property +prop_lookupLE = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe (TTOther :*: TTVal)) + M.lookupLE + NEM.lookupLE + +prop_lookupGE :: Property +prop_lookupGE = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe (TTOther :*: TTVal)) + M.lookupGE + NEM.lookupGE + +prop_size :: Property +prop_size = ttProp (GTNEIntMap :-> TTOther) + M.size + NEM.size + +prop_union :: Property +prop_union = ttProp (GTNEIntMap :-> GTNEIntMap :-> TTNEIntMap) + M.union + NEM.union + +prop_unionWith :: Property +prop_unionWith = ttProp (gf2 valGen :?> GTNEIntMap :-> GTNEIntMap :-> TTNEIntMap) + M.unionWith + NEM.unionWith + +prop_unionWithKey :: Property +prop_unionWithKey = ttProp (gf3 valGen :?> GTNEIntMap :-> GTNEIntMap :-> TTNEIntMap) + M.unionWithKey + NEM.unionWithKey + +prop_unions :: Property +prop_unions = ttProp (GTNEList (Just (Range.linear 2 5)) GTNEIntMap :-> TTNEIntMap) + M.unions + NEM.unions + +prop_unionsWith :: Property +prop_unionsWith = ttProp (gf2 valGen :?> GTNEList (Just (Range.linear 2 5)) GTNEIntMap :-> TTNEIntMap) + M.unionsWith + NEM.unionsWith + +prop_difference :: Property +prop_difference = ttProp (GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.difference + NEM.difference + +prop_differenceWithKey :: Property +prop_differenceWithKey = ttProp (gf3 (Gen.maybe valGen) :?> GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.differenceWithKey + NEM.differenceWithKey + +prop_intersection :: Property +prop_intersection = ttProp (GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.intersection + NEM.intersection + +prop_intersectionWithKey :: Property +prop_intersectionWithKey = ttProp (gf3 valGen :?> GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.intersectionWithKey + NEM.intersectionWithKey + +prop_map :: Property +prop_map = ttProp (gf1 valGen :?> GTNEIntMap :-> TTNEIntMap) + M.map + NEM.map + +prop_map_rules_map :: Property +prop_map_rules_map = ttProp (gf1 valGen :?> gf1 valGen :?> GTNEIntMap :-> TTNEIntMap) + (\f g xs -> M.map f (M.map g xs)) + (\f g xs -> NEM.map f (NEM.map g xs)) + +prop_map_rules_coerce :: Property +prop_map_rules_coerce = ttProp (GTNEIntMap :-> TTNEIntMap) + (M.map @Text @Text coerce) + (NEM.map @Text @Text coerce) + +prop_map_rules_mapWithKey :: Property +prop_map_rules_mapWithKey = ttProp (gf1 valGen :?> gf2 valGen :?> GTNEIntMap :-> TTNEIntMap) + (\f g xs -> M.map f (M.mapWithKey g xs)) + (\f g xs -> NEM.map f (NEM.mapWithKey g xs)) + +prop_mapWithKey :: Property +prop_mapWithKey = ttProp (gf2 valGen :?> GTNEIntMap :-> TTNEIntMap) + M.mapWithKey + NEM.mapWithKey + +prop_mapWithKey_rules_mapWithKey :: Property +prop_mapWithKey_rules_mapWithKey = ttProp (gf2 valGen :?> gf2 valGen :?> GTNEIntMap :-> TTNEIntMap) + (\f g xs -> M.mapWithKey f (M.mapWithKey g xs)) + (\f g xs -> NEM.mapWithKey f (NEM.mapWithKey g xs)) + +prop_mapWithKey_rules_map :: Property +prop_mapWithKey_rules_map = ttProp (gf2 valGen :?> gf1 valGen :?> GTNEIntMap :-> TTNEIntMap) + (\f g xs -> M.mapWithKey f (M.map g xs)) + (\f g xs -> NEM.mapWithKey f (NEM.map g xs)) + +-- | These intentionally do not match, because Foldable for IntMap is +-- inconsistent +-- prop_traverseWithKey1 :: Property +-- prop_traverseWithKey1 = ttProp (gf1 valGen :?> GTNEIntMap :-> TTBazaar GTVal TTNEIntMap TTVal) +-- (\f -> M.traverseWithKey (\k -> (`More` Done (f . (k,))))) +-- (\f -> NEM.traverseWithKey1 (\k -> (`More` Done (f . (k,))))) + +-- prop_traverseWithKey :: Property +-- prop_traverseWithKey = ttProp (gf1 valGen :?> GTNEIntMap :-> TTBazaar GTVal TTNEIntMap TTVal) +-- (\f -> M.traverseWithKey (\k -> (`More` Done (f . (k,))))) +-- (\f -> NEM.traverseWithKey (\k -> (`More` Done (f . (k,))))) + +-- prop_sequence1 :: Property +-- prop_sequence1 = ttProp (GTNEIntMap :-> TTBazaar GTVal TTNEIntMap TTVal) +-- (sequenceA . fmap (`More` Done id)) +-- (sequence1 . fmap (`More` Done id)) + +-- prop_sequenceA :: Property +-- prop_sequenceA = ttProp (GTNEIntMap :-> TTBazaar GTVal TTNEIntMap TTVal) +-- (sequenceA . fmap (`More` Done id)) +-- (sequenceA . fmap (`More` Done id)) + +-- prop_mapAccumWithKey :: Property +-- prop_mapAccumWithKey = ttProp ( gf3 ((,) <$> valGen <*> valGen) +-- :?> GTOther valGen +-- :-> GTNEIntMap +-- :-> TTOther :*: TTNEIntMap +-- ) +-- M.mapAccumWithKey +-- NEM.mapAccumWithKey + +-- prop_mapAccumRWithKey :: Property +-- prop_mapAccumRWithKey = ttProp ( gf3 ((,) <$> valGen <*> valGen) +-- :?> GTOther valGen +-- :-> GTNEIntMap +-- :-> TTOther :*: TTNEIntMap +-- ) +-- M.mapAccumRWithKey +-- NEM.mapAccumRWithKey + +prop_mapKeys :: Property +prop_mapKeys = ttProp (gf1 intKeyGen :?> GTNEIntMap :-> TTNEIntMap) + M.mapKeys + NEM.mapKeys + +prop_mapKeysWith :: Property +prop_mapKeysWith = ttProp ( gf2 valGen + :?> gf1 intKeyGen + :?> GTNEIntMap + :-> TTNEIntMap + ) + M.mapKeysWith + NEM.mapKeysWith + +prop_mapKeysMonotonic :: Property +prop_mapKeysMonotonic = ttProp (GTNEIntMap :-> TTNEIntMap) + (M.mapKeysMonotonic (*2)) + (NEM.mapKeysMonotonic (*2)) + +prop_foldr :: Property +prop_foldr = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldr + NEM.foldr + +prop_foldl :: Property +prop_foldl = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldl + NEM.foldl + +prop_foldr1 :: Property +prop_foldr1 = ttProp ( gf2 valGen + :?> GTNEIntMap + :-> TTOther + ) + foldr1 + NEM.foldr1 + +prop_foldl1 :: Property +prop_foldl1 = ttProp ( gf2 valGen + :?> GTNEIntMap + :-> TTOther + ) + foldl1 + NEM.foldl1 + +prop_foldrWithKey :: Property +prop_foldrWithKey = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldrWithKey + NEM.foldrWithKey + +prop_foldlWithKey :: Property +prop_foldlWithKey = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldlWithKey + NEM.foldlWithKey + +prop_foldMapWithKey :: Property +prop_foldMapWithKey = ttProp (gf2 valGen :?> GTNEIntMap :-> TTOther) + (\f -> foldMap (uncurry f) . M.toList) + NEM.foldMapWithKey + +prop_foldr' :: Property +prop_foldr' = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldr' + NEM.foldr' + +prop_foldl' :: Property +prop_foldl' = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldl' + NEM.foldl' + +prop_foldr1' :: Property +prop_foldr1' = ttProp ( gf2 valGen + :?> GTNEIntMap + :-> TTOther + ) + foldr1 + NEM.foldr1' + +prop_foldl1' :: Property +prop_foldl1' = ttProp ( gf2 valGen + :?> GTNEIntMap + :-> TTOther + ) + foldl1 + NEM.foldl1' + +prop_foldrWithKey' :: Property +prop_foldrWithKey' = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldrWithKey' + NEM.foldrWithKey' + +prop_foldlWithKey' :: Property +prop_foldlWithKey' = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEIntMap + :-> TTOther + ) + M.foldlWithKey' + NEM.foldlWithKey' + +prop_elems :: Property +prop_elems = ttProp (GTNEIntMap :-> TTNEList TTVal) + M.elems + NEM.elems + +prop_keys :: Property +prop_keys = ttProp (GTNEIntMap :-> TTNEList TTOther) + M.keys + NEM.keys + +prop_assocs :: Property +prop_assocs = ttProp (GTNEIntMap :-> TTNEList (TTOther :*: TTVal)) + M.assocs + NEM.assocs + +prop_keysSet :: Property +prop_keysSet = ttProp (GTNEIntMap :-> TTNEIntSet) + M.keysSet + NEM.keysSet + +prop_toList :: Property +prop_toList = ttProp (GTNEIntMap :-> TTNEList (TTOther :*: TTVal)) + M.toList + NEM.toList + +prop_toDescList :: Property +prop_toDescList = ttProp (GTNEIntMap :-> TTNEList (TTOther :*: TTVal)) + M.toDescList + NEM.toDescList + +prop_filter :: Property +prop_filter = ttProp (gf1 Gen.bool :?> GTNEIntMap :-> TTOther) + M.filter + NEM.filter + +prop_filterWithKey :: Property +prop_filterWithKey = ttProp (gf2 Gen.bool :?> GTNEIntMap :-> TTOther) + M.filterWithKey + NEM.filterWithKey + +prop_restrictKeys :: Property +prop_restrictKeys = ttProp (GTNEIntMap :-> GTIntSet :-> TTOther) + M.restrictKeys + NEM.restrictKeys + +prop_withoutKeys :: Property +prop_withoutKeys = ttProp (GTNEIntMap :-> GTIntSet :-> TTOther) + M.withoutKeys + NEM.withoutKeys + +prop_partitionWithKey :: Property +prop_partitionWithKey = ttProp (gf2 Gen.bool :?> GTNEIntMap :-> TTThese TTNEIntMap TTNEIntMap) + M.partitionWithKey + NEM.partitionWithKey + +prop_mapMaybeWithKey :: Property +prop_mapMaybeWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEIntMap :-> TTOther) + M.mapMaybeWithKey + NEM.mapMaybeWithKey + +prop_mapEitherWithKey :: Property +prop_mapEitherWithKey = ttProp ( gf2 (Gen.choice [Left <$> valGen, Right <$> valGen]) + :?> GTNEIntMap + :-> TTThese TTNEIntMap TTNEIntMap + ) + M.mapEitherWithKey + NEM.mapEitherWithKey + +prop_split :: Property +prop_split = ttProp (GTIntKey :-> GTNEIntMap :-> TTMThese TTNEIntMap TTNEIntMap) + M.split + NEM.split + +prop_splitLookup :: Property +prop_splitLookup = ttProp (GTIntKey :-> GTNEIntMap :-> TTMaybe TTVal :*: TTMThese TTNEIntMap TTNEIntMap) + (\k -> (\(x,y,z) -> (y,(x,z))) . M.splitLookup k) + NEM.splitLookup + +prop_isSubmapOfBy :: Property +prop_isSubmapOfBy = ttProp (gf2 Gen.bool :?> GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.isSubmapOfBy + NEM.isSubmapOfBy + +prop_isProperSubmapOfBy :: Property +prop_isProperSubmapOfBy = ttProp (gf2 Gen.bool :?> GTNEIntMap :-> GTNEIntMap :-> TTOther) + M.isProperSubmapOfBy + NEM.isProperSubmapOfBy + +prop_findMin :: Property +prop_findMin = ttProp (GTNEIntMap :-> TTOther :*: TTVal) + M.findMin + NEM.findMin + +prop_findMax :: Property +prop_findMax = ttProp (GTNEIntMap :-> TTOther :*: TTVal) + M.findMax + NEM.findMax + +prop_deleteMin :: Property +prop_deleteMin = ttProp (GTNEIntMap :-> TTOther) + M.deleteMin + NEM.deleteMin + +prop_deleteMax :: Property +prop_deleteMax = ttProp (GTNEIntMap :-> TTOther) + M.deleteMax + NEM.deleteMax + +prop_deleteFindMin :: Property +prop_deleteFindMin = ttProp (GTNEIntMap :-> (TTOther :*: TTVal) :*: TTOther) + M.deleteFindMin + NEM.deleteFindMin + +prop_deleteFindMax :: Property +prop_deleteFindMax = ttProp (GTNEIntMap :-> (TTOther :*: TTVal) :*: TTOther) + M.deleteFindMax + NEM.deleteFindMax + +prop_updateMinWithKey :: Property +prop_updateMinWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEIntMap :-> TTOther) + M.updateMinWithKey + NEM.updateMinWithKey + +prop_updateMaxWithKey :: Property +prop_updateMaxWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEIntMap :-> TTOther) + M.updateMaxWithKey + NEM.updateMaxWithKey + +prop_adjustMinWithKey :: Property +prop_adjustMinWithKey = ttProp (gf2 valGen :?> GTNEIntMap :-> TTNEIntMap) + (M.updateMinWithKey . (fmap . fmap) Just) + NEM.adjustMinWithKey + +prop_adjustMaxWithKey :: Property +prop_adjustMaxWithKey = ttProp (gf2 valGen :?> GTNEIntMap :-> TTNEIntMap) + (M.updateMaxWithKey . (fmap . fmap) Just) + NEM.adjustMaxWithKey + +prop_minView :: Property +prop_minView = ttProp (GTNEIntMap :-> TTMaybe (TTVal :*: TTOther)) + M.minView + (Just . NEM.minView) + +prop_maxView :: Property +prop_maxView = ttProp (GTNEIntMap :-> TTMaybe (TTVal :*: TTOther)) + M.maxView + (Just . NEM.maxView) + +prop_elem :: Property +prop_elem = ttProp (GTVal :-> GTNEIntMap :-> TTOther) + elem + elem + +prop_fold1 :: Property +prop_fold1 = ttProp (GTNEIntMap :-> TTVal) + (fold . toList) + fold1 + +prop_fold :: Property +prop_fold = ttProp (GTNEIntMap :-> TTVal) + (fold . toList) + fold + +prop_foldMap1 :: Property +prop_foldMap1 = ttProp (gf1 valGen :?> GTNEIntMap :-> TTOther) + (\f -> foldMap ((:[]) . f) . toList) + (\f -> foldMap1 ((:[]) . f)) + +prop_foldMap :: Property +prop_foldMap = ttProp (gf1 valGen :?> GTNEIntMap :-> TTOther) + (\f -> foldMap ((:[]) . f) . toList) + (\f -> foldMap ((:[]) . f)) + + diff --git a/test/Tests/IntSet.hs b/test/Tests/IntSet.hs new file mode 100644 --- /dev/null +++ b/test/Tests/IntSet.hs @@ -0,0 +1,343 @@ +{-# LANGUAGE TemplateHaskell #-} + +module Tests.IntSet (intSetTests) where + +import Data.Functor.Identity +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup.Foldable +import Hedgehog +import Test.Tasty +import Tests.Util +import qualified Data.IntSet as S +import qualified Data.IntSet.NonEmpty as NES +import qualified Data.IntSet.NonEmpty.Internal as NES +import qualified Data.List.NonEmpty as NE +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range + +intSetTests :: TestTree +intSetTests = groupTree $$(discover) + + + + + +prop_valid :: Property +prop_valid = property $ + assert . NES.valid =<< forAll neIntSetGen + + +-- | We cannot implement these because there is no 'valid' for IntSet +-- prop_valid_toSet :: Property +-- prop_valid_toSet = property $ do +-- assert . S.valid . NES.toSet =<< forAll neIntSetGen + +-- prop_valid_insertMinIntSet :: Property +-- prop_valid_insertMinIntSet = property $ do +-- n <- forAll $ do +-- m <- setGen +-- let k = maybe dummyKey (subtract 1 . fst) $ S.maxView m +-- pure $ NES.insertMinIntSet k m +-- assert $ S.valid n + +-- prop_valid_insertMaxIntSet :: Property +-- prop_valid_insertMaxIntSet = property $ do +-- n <- forAll $ do +-- m <- setGen +-- let k = maybe dummyKey ((+ 1) . fst) $ S.maxView m +-- pure $ NES.insertMaxIntSet k m +-- assert $ S.valid n + +prop_valid_insertSetMin :: Property +prop_valid_insertSetMin = property $ do + n <- forAll $ do + m <- intSetGen + let k = maybe 0 (subtract 1 . fst) $ S.minView m + pure $ NES.insertSetMin k m + assert $ NES.valid n + +prop_valid_insertSetMax :: Property +prop_valid_insertSetMax = property $ do + n <- forAll $ do + m <- intSetGen + let k = maybe 0 ((+ 1) . fst) $ S.maxView m + pure $ NES.insertSetMax k m + assert $ NES.valid n + +prop_toSetIso1 :: Property +prop_toSetIso1 = property $ do + m0 <- forAll intSetGen + tripping m0 NES.nonEmptySet + (Identity . maybe S.empty NES.toSet) + +prop_toSetIso2 :: Property +prop_toSetIso2 = property $ do + m0 <- forAll $ Gen.maybe neIntSetGen + tripping m0 (maybe S.empty NES.toSet) + (Identity . NES.nonEmptySet) + +prop_read_show :: Property +prop_read_show = readShow neIntSetGen + +prop_splitRoot :: Property +prop_splitRoot = property $ do + n <- forAll neIntSetGen + let rs = NES.splitRoot n + allItems = foldMap1 NES.toList rs + n' = NES.unions rs + assert $ ascending allItems + mapM_ (assert . (`NES.isSubsetOf` n)) rs + length allItems === NES.size n' + n === n' + where + ascending (x :| xs) = case NE.nonEmpty xs of + Nothing -> True + Just ys@(y :| _) -> x < y && ascending ys + + + + + + + + + + +prop_insertSet :: Property +prop_insertSet = ttProp (GTIntKey :-> GTIntSet :-> TTNEIntSet) + S.insert + NES.insertSet + +prop_singleton :: Property +prop_singleton = ttProp (GTIntKey :-> TTNEIntSet) + S.singleton + NES.singleton + +prop_fromAscList :: Property +prop_fromAscList = ttProp (GTSorted STAsc (GTNEList Nothing (GTIntKey :&: GTVal)) :-> TTNEIntSet) + (S.fromAscList . fmap fst) + (NES.fromAscList . fmap fst) + +prop_fromDistinctAscList :: Property +prop_fromDistinctAscList = ttProp (GTSorted STAsc (GTNEList Nothing GTIntKey) :-> TTNEIntSet) + S.fromDistinctAscList + NES.fromDistinctAscList + +prop_fromList :: Property +prop_fromList = ttProp (GTNEList Nothing GTIntKey :-> TTNEIntSet) + S.fromList + NES.fromList + +prop_insert :: Property +prop_insert = ttProp (GTIntKey :-> GTNEIntSet :-> TTNEIntSet) + S.insert + NES.insert + +prop_delete :: Property +prop_delete = ttProp (GTIntKey :-> GTNEIntSet :-> TTOther) + S.delete + NES.delete + +prop_member :: Property +prop_member = ttProp (GTIntKey :-> GTNEIntSet :-> TTOther) + S.member + NES.member + +prop_notMember :: Property +prop_notMember = ttProp (GTIntKey :-> GTNEIntSet :-> TTOther) + S.notMember + NES.notMember + +prop_lookupLT :: Property +prop_lookupLT = ttProp (GTIntKey :-> GTNEIntSet :-> TTMaybe TTOther) + S.lookupLT + NES.lookupLT + +prop_lookupGT :: Property +prop_lookupGT = ttProp (GTIntKey :-> GTNEIntSet :-> TTMaybe TTOther) + S.lookupGT + NES.lookupGT + +prop_lookupLE :: Property +prop_lookupLE = ttProp (GTIntKey :-> GTNEIntSet :-> TTMaybe TTOther) + S.lookupLE + NES.lookupLE + +prop_lookupGE :: Property +prop_lookupGE = ttProp (GTIntKey :-> GTNEIntSet :-> TTMaybe TTOther) + S.lookupGE + NES.lookupGE + +prop_size :: Property +prop_size = ttProp (GTNEIntSet :-> TTOther) + S.size + NES.size + +prop_isSubsetOf :: Property +prop_isSubsetOf = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTOther) + S.isSubsetOf + NES.isSubsetOf + +prop_isProperSubsetOf :: Property +prop_isProperSubsetOf = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTOther) + S.isProperSubsetOf + NES.isProperSubsetOf + +prop_disjoint :: Property +prop_disjoint = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTOther) + NES.disjointSet + NES.disjoint + +prop_union :: Property +prop_union = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTNEIntSet) + S.union + NES.union + +prop_unions :: Property +prop_unions = ttProp (GTNEList (Just (Range.linear 2 5)) GTNEIntSet :-> TTNEIntSet) + S.unions + NES.unions + +prop_difference :: Property +prop_difference = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTOther) + S.difference + NES.difference + +prop_intersection :: Property +prop_intersection = ttProp (GTNEIntSet :-> GTNEIntSet :-> TTOther) + S.intersection + NES.intersection + +prop_filter :: Property +prop_filter = ttProp (gf1 Gen.bool :?> GTNEIntSet :-> TTOther) + S.filter + NES.filter + +prop_partition :: Property +prop_partition = ttProp (gf1 Gen.bool :?> GTNEIntSet :-> TTThese TTNEIntSet TTNEIntSet) + S.partition + NES.partition + +prop_split :: Property +prop_split = ttProp (GTIntKey :-> GTNEIntSet :-> TTMThese TTNEIntSet TTNEIntSet) + S.split + NES.split + +prop_splitMember :: Property +prop_splitMember = ttProp (GTIntKey :-> GTNEIntSet :-> TTOther :*: TTMThese TTNEIntSet TTNEIntSet) + (\k -> (\(x,y,z) -> (y,(x,z))) . S.splitMember k) + NES.splitMember + +prop_map :: Property +prop_map = ttProp (gf1 intKeyGen :?> GTNEIntSet :-> TTNEIntSet) + S.map + NES.map + +prop_foldr :: Property +prop_foldr = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntSet + :-> TTOther + ) + S.foldr + NES.foldr + +prop_foldl :: Property +prop_foldl = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEIntSet + :-> TTOther + ) + S.foldl + NES.foldl + +prop_foldr1 :: Property +prop_foldr1 = ttProp ( gf2 intKeyGen + :?> GTNEIntSet + :-> TTOther + ) + (\f -> foldr1 f . S.toList) + NES.foldr1 + +prop_foldl1 :: Property +prop_foldl1 = ttProp ( gf2 intKeyGen + :?> GTNEIntSet + :-> TTOther + ) + (\f -> foldl1 f . S.toList) + NES.foldl1 + +prop_foldr' :: Property +prop_foldr' = ttProp ( gf2 intKeyGen + :?> GTOther intKeyGen + :-> GTNEIntSet + :-> TTOther + ) + S.foldr' + NES.foldr' + +prop_foldl' :: Property +prop_foldl' = ttProp ( gf2 intKeyGen + :?> GTOther intKeyGen + :-> GTNEIntSet + :-> TTOther + ) + S.foldl' + NES.foldl' + +prop_foldr1' :: Property +prop_foldr1' = ttProp ( gf2 intKeyGen + :?> GTNEIntSet + :-> TTOther + ) + (\f -> foldr1 f . S.toList) + NES.foldr1' + +prop_foldl1' :: Property +prop_foldl1' = ttProp ( gf2 intKeyGen + :?> GTNEIntSet + :-> TTOther + ) + (\f -> foldl1 f . S.toList) + NES.foldl1' + +prop_findMin :: Property +prop_findMin = ttProp (GTNEIntSet :-> TTOther) + S.findMin + NES.findMin + +prop_findMax :: Property +prop_findMax = ttProp (GTNEIntSet :-> TTOther) + S.findMax + NES.findMax + +prop_deleteMin :: Property +prop_deleteMin = ttProp (GTNEIntSet :-> TTOther) + S.deleteMin + NES.deleteMin + +prop_deleteMax :: Property +prop_deleteMax = ttProp (GTNEIntSet :-> TTOther) + S.deleteMax + NES.deleteMax + +prop_deleteFindMin :: Property +prop_deleteFindMin = ttProp (GTNEIntSet :-> TTOther :*: TTOther) + S.deleteFindMin + NES.deleteFindMin + +prop_deleteFindMax :: Property +prop_deleteFindMax = ttProp (GTNEIntSet :-> TTOther :*: TTOther) + S.deleteFindMax + NES.deleteFindMax + +prop_toList :: Property +prop_toList = ttProp (GTNEIntSet :-> TTNEList TTOther) + S.toList + NES.toList + +prop_toDescList :: Property +prop_toDescList = ttProp (GTNEIntSet :-> TTNEList TTOther) + S.toDescList + NES.toDescList + diff --git a/test/Tests/Map.hs b/test/Tests/Map.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Map.hs @@ -0,0 +1,792 @@ +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} + +module Tests.Map (mapTests) where + +import Control.Applicative +import Data.Coerce +import Data.Foldable +import Data.Functor.Identity +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Semigroup.Foldable +import Data.Semigroup.Traversable +import Data.Text (Text) +import Hedgehog +import Test.Tasty +import Tests.Util +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import qualified Data.Map.NonEmpty as NEM +import qualified Data.Map.NonEmpty.Internal as NEM +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range + +mapTests :: TestTree +mapTests = groupTree $$(discover) + + + + + +prop_valid :: Property +prop_valid = property $ + assert . NEM.valid =<< forAll neMapGen + +prop_valid_toMap :: Property +prop_valid_toMap = property $ + assert . M.valid . NEM.toMap =<< forAll neMapGen + +prop_valid_insertMinMap :: Property +prop_valid_insertMinMap = property $ do + n <- forAll $ do + m <- mapGen + let k = maybe dummyKey (subtract 1 . fst) $ M.lookupMin m + v <- valGen + pure $ NEM.insertMinMap k v m + assert $ M.valid n + +prop_valid_insertMaxMap :: Property +prop_valid_insertMaxMap = property $ do + n <- forAll $ do + m <- mapGen + let k = maybe dummyKey ((+ 1) . fst) $ M.lookupMax m + v <- valGen + pure $ NEM.insertMaxMap k v m + assert $ M.valid n + +prop_valid_insertMapMin :: Property +prop_valid_insertMapMin = property $ do + n <- forAll $ do + m <- mapGen + let k = maybe dummyKey (subtract 1 . fst) $ M.lookupMin m + v <- valGen + pure $ NEM.insertMapMin k v m + assert $ NEM.valid n + +prop_valid_insertMapMax :: Property +prop_valid_insertMapMax = property $ do + n <- forAll $ do + m <- mapGen + let k = maybe dummyKey ((+ 1) . fst) $ M.lookupMax m + v <- valGen + pure $ NEM.insertMapMax k v m + assert $ NEM.valid n + +prop_toMapIso1 :: Property +prop_toMapIso1 = property $ do + m0 <- forAll mapGen + tripping m0 NEM.nonEmptyMap + (Identity . maybe M.empty NEM.toMap) + +prop_toMapIso2 :: Property +prop_toMapIso2 = property $ do + m0 <- forAll $ Gen.maybe neMapGen + tripping m0 (maybe M.empty NEM.toMap) + (Identity . NEM.nonEmptyMap) + +prop_read_show :: Property +prop_read_show = readShow neMapGen + +prop_read1_show1 :: Property +prop_read1_show1 = readShow1 neMapGen + +prop_show_show1 :: Property +prop_show_show1 = showShow1 neMapGen + +prop_show_show2 :: Property +prop_show_show2 = showShow2 neMapGen + +prop_splitRoot :: Property +prop_splitRoot = property $ do + n <- forAll neMapGen + let rs = NEM.splitRoot n + allItems = foldMap1 NEM.keys rs + n' = NEM.unions rs + assert $ ascending allItems + mapM_ (assert . (`NEM.isSubmapOf` n)) rs + length allItems === length n' + n === n' + where + ascending (x :| xs) = case NE.nonEmpty xs of + Nothing -> True + Just ys@(y :| _) -> x < y && ascending ys + + + + + + + +prop_insertMapWithKey :: Property +prop_insertMapWithKey = ttProp (gf3 valGen :?> GTKey :-> GTVal :-> GTMap :-> TTNEMap) + M.insertWithKey + NEM.insertMapWithKey + +prop_singleton :: Property +prop_singleton = ttProp (GTKey :-> GTVal :-> TTNEMap) + M.singleton + NEM.singleton + +prop_fromSet :: Property +prop_fromSet = ttProp (gf1 valGen :?> GTNESet :-> TTNEMap) + M.fromSet + NEM.fromSet + +prop_fromAscList :: Property +prop_fromAscList = ttProp (GTSorted STAsc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromAscList + NEM.fromAscList + +prop_fromDescList :: Property +prop_fromDescList = ttProp (GTSorted STDesc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromDescList + NEM.fromDescList + +prop_fromAscListWithKey :: Property +prop_fromAscListWithKey = ttProp (gf3 valGen :?> GTSorted STAsc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromAscListWithKey + NEM.fromAscListWithKey + +prop_fromDescListWithKey :: Property +prop_fromDescListWithKey = ttProp (gf3 valGen :?> GTSorted STDesc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromDescListWithKey + NEM.fromDescListWithKey + +prop_fromDistinctAscList :: Property +prop_fromDistinctAscList = ttProp (GTSorted STDistinctAsc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromDistinctAscList + NEM.fromDistinctAscList + +prop_fromDistinctDescList :: Property +prop_fromDistinctDescList = ttProp (GTSorted STDistinctDesc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNEMap) + M.fromDistinctDescList + NEM.fromDistinctDescList + +prop_fromListWithKey :: Property +prop_fromListWithKey = ttProp (gf3 valGen :?> GTNEList Nothing (GTKey :&: GTVal) :-> TTNEMap) + M.fromListWithKey + NEM.fromListWithKey + +prop_insert :: Property +prop_insert = ttProp (GTKey :-> GTVal :-> GTNEMap :-> TTNEMap) + M.insert + NEM.insert + +prop_insertWithKey :: Property +prop_insertWithKey = ttProp (gf3 valGen :?> GTKey :-> GTVal :-> GTNEMap :-> TTNEMap) + M.insertWithKey + NEM.insertWithKey + +prop_delete :: Property +prop_delete = ttProp (GTKey :-> GTNEMap :-> TTMap) + M.delete + NEM.delete + +prop_adjustWithKey :: Property +prop_adjustWithKey = ttProp (gf2 valGen :?> GTKey :-> GTNEMap :-> TTNEMap) + M.adjustWithKey + NEM.adjustWithKey + +prop_updateWithKey :: Property +prop_updateWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTKey :-> GTNEMap :-> TTMap) + M.updateWithKey + NEM.updateWithKey + +prop_updateLookupWithKey :: Property +prop_updateLookupWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTKey :-> GTNEMap :-> TTMaybe TTVal :*: TTMap) + M.updateLookupWithKey + NEM.updateLookupWithKey + +prop_alter :: Property +prop_alter = ttProp (gf1 (Gen.maybe valGen) :?> GTKey :-> GTNEMap :-> TTMap) + M.alter + NEM.alter + +prop_alter' :: Property +prop_alter' = ttProp (gf1 valGen :?> GTKey :-> GTNEMap :-> TTNEMap) + (M.alter . fmap Just) + NEM.alter' + +prop_alterF :: Property +prop_alterF = ttProp ( gf1 (Gen.maybe valGen) + :?> GTKey + :-> GTNEMap + :-> TTCtx (GTMaybe GTVal :-> TTMap) (TTMaybe TTVal) + ) + (M.alterF . Context) + (NEM.alterF . Context) + +prop_alterF_rules_Const :: Property +prop_alterF_rules_Const = ttProp ( gf1 (Const <$> valGen) + :?> GTKey + :-> GTNEMap + :-> TTOther + ) + (\f k m -> getConst (M.alterF f k m)) + (\f k m -> getConst (NEM.alterF f k m)) + +prop_alterF_rules_Identity :: Property +prop_alterF_rules_Identity = ttProp ( gf1 (Identity <$> Gen.maybe valGen) + :?> GTKey + :-> GTNEMap + :-> TTMap + ) + (\f k m -> runIdentity (M.alterF f k m)) + (\f k m -> runIdentity (NEM.alterF f k m)) + +prop_alterF' :: Property +prop_alterF' = ttProp (gf1 valGen :?> GTKey :-> GTNEMap :-> TTCtx (GTVal :-> TTNEMap) (TTMaybe TTVal)) + (M.alterF . Context . fmap Just) + (NEM.alterF' . Context) + +prop_alterF'_rules_Const :: Property +prop_alterF'_rules_Const = ttProp ( gf1 (Const <$> valGen) + :?> GTKey + :-> GTNEMap + :-> TTOther + ) + (\f k m -> let f' = fmap Just . f in getConst (M.alterF f' k m)) + (\f k m -> getConst (NEM.alterF' f k m)) + +-- -- | This fails, but isn't possible to fix without copying-and-pasting more +-- -- in code from containers. +-- prop_alterF'_rules_Identity :: Property +-- prop_alterF'_rules_Identity = ttProp ( gf1 (Identity <$> valGen) +-- :?> GTKey +-- :-> GTNEMap +-- :-> TTNEMap +-- ) +-- (\f k m -> let f' = fmap Just . f in runIdentity (M.alterF f' k m)) +-- (\f k m -> runIdentity (NEM.alterF' f k m)) + +prop_lookup :: Property +prop_lookup = ttProp (GTKey :-> GTNEMap :-> TTMaybe TTVal) + M.lookup + NEM.lookup + +prop_findWithDefault :: Property +prop_findWithDefault = ttProp (GTVal :-> GTKey :-> GTNEMap :-> TTVal) + M.findWithDefault + NEM.findWithDefault + +prop_member :: Property +prop_member = ttProp (GTKey :-> GTNEMap :-> TTOther) + M.member + NEM.member + +prop_notMember :: Property +prop_notMember = ttProp (GTKey :-> GTNEMap :-> TTOther) + M.notMember + NEM.notMember + +prop_lookupLT :: Property +prop_lookupLT = ttProp (GTKey :-> GTNEMap :-> TTMaybe (TTKey :*: TTVal)) + M.lookupLT + NEM.lookupLT + +prop_lookupGT :: Property +prop_lookupGT = ttProp (GTKey :-> GTNEMap :-> TTMaybe (TTKey :*: TTVal)) + M.lookupGT + NEM.lookupGT + +prop_lookupLE :: Property +prop_lookupLE = ttProp (GTKey :-> GTNEMap :-> TTMaybe (TTKey :*: TTVal)) + M.lookupLE + NEM.lookupLE + +prop_lookupGE :: Property +prop_lookupGE = ttProp (GTKey :-> GTNEMap :-> TTMaybe (TTKey :*: TTVal)) + M.lookupGE + NEM.lookupGE + +prop_size :: Property +prop_size = ttProp (GTNEMap :-> TTOther) + M.size + NEM.size + +prop_union :: Property +prop_union = ttProp (GTNEMap :-> GTNEMap :-> TTNEMap) + M.union + NEM.union + +prop_unionWith :: Property +prop_unionWith = ttProp (gf2 valGen :?> GTNEMap :-> GTNEMap :-> TTNEMap) + M.unionWith + NEM.unionWith + +prop_unionWithKey :: Property +prop_unionWithKey = ttProp (gf3 valGen :?> GTNEMap :-> GTNEMap :-> TTNEMap) + M.unionWithKey + NEM.unionWithKey + +prop_unions :: Property +prop_unions = ttProp (GTNEList (Just (Range.linear 2 5)) GTNEMap :-> TTNEMap) + M.unions + NEM.unions + +prop_unionsWith :: Property +prop_unionsWith = ttProp (gf2 valGen :?> GTNEList (Just (Range.linear 2 5)) GTNEMap :-> TTNEMap) + M.unionsWith + NEM.unionsWith + +prop_difference :: Property +prop_difference = ttProp (GTNEMap :-> GTNEMap :-> TTMap) + M.difference + NEM.difference + +prop_differenceWithKey :: Property +prop_differenceWithKey = ttProp (gf3 (Gen.maybe valGen) :?> GTNEMap :-> GTNEMap :-> TTMap) + M.differenceWithKey + NEM.differenceWithKey + +prop_intersection :: Property +prop_intersection = ttProp (GTNEMap :-> GTNEMap :-> TTMap) + M.intersection + NEM.intersection + +prop_intersectionWithKey :: Property +prop_intersectionWithKey = ttProp (gf3 valGen :?> GTNEMap :-> GTNEMap :-> TTMap) + M.intersectionWithKey + NEM.intersectionWithKey + +prop_map :: Property +prop_map = ttProp (gf1 valGen :?> GTNEMap :-> TTNEMap) + M.map + NEM.map + +prop_map_rules_map :: Property +prop_map_rules_map = ttProp (gf1 valGen :?> gf1 valGen :?> GTNEMap :-> TTNEMap) + (\f g xs -> M.map f (M.map g xs)) + (\f g xs -> NEM.map f (NEM.map g xs)) + +prop_map_rules_coerce :: Property +prop_map_rules_coerce = ttProp (GTNEMap :-> TTNEMap) + (M.map @Text @Text coerce) + (NEM.map @Text @Text coerce) + +prop_map_rules_mapWithKey :: Property +prop_map_rules_mapWithKey = ttProp (gf1 valGen :?> gf2 valGen :?> GTNEMap :-> TTNEMap) + (\f g xs -> M.map f (M.mapWithKey g xs)) + (\f g xs -> NEM.map f (NEM.mapWithKey g xs)) + +prop_mapWithKey :: Property +prop_mapWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTNEMap) + M.mapWithKey + NEM.mapWithKey + +prop_mapWithKey_rules_mapWithKey :: Property +prop_mapWithKey_rules_mapWithKey = ttProp (gf2 valGen :?> gf2 valGen :?> GTNEMap :-> TTNEMap) + (\f g xs -> M.mapWithKey f (M.mapWithKey g xs)) + (\f g xs -> NEM.mapWithKey f (NEM.mapWithKey g xs)) + +prop_mapWithKey_rules_map :: Property +prop_mapWithKey_rules_map = ttProp (gf2 valGen :?> gf1 valGen :?> GTNEMap :-> TTNEMap) + (\f g xs -> M.mapWithKey f (M.map g xs)) + (\f g xs -> NEM.mapWithKey f (NEM.map g xs)) + +prop_traverseWithKey1 :: Property +prop_traverseWithKey1 = ttProp (gf2 valGen :?> GTNEMap :-> TTBazaar GTVal TTNEMap TTVal) + (\f -> M.traverseWithKey (\k -> (`More` Done (f k)))) + (\f -> NEM.traverseWithKey1 (\k -> (`More` Done (f k)))) + +prop_traverseWithKey :: Property +prop_traverseWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTBazaar GTVal TTNEMap TTVal) + (\f -> M.traverseWithKey (\k -> (`More` Done (f k)))) + (\f -> NEM.traverseWithKey (\k -> (`More` Done (f k)))) + +prop_traverseMaybeWithKey1 :: Property +prop_traverseMaybeWithKey1 = ttProp (gf2 valGen :?> GTNEMap :-> TTBazaar (GTMaybe GTVal) TTMap TTVal) + (\f -> M.traverseMaybeWithKey (\k -> (`More` Done (fmap (f k))))) + (\f -> NEM.traverseMaybeWithKey1 (\k -> (`More` Done (fmap (f k))))) + +prop_traverseMaybeWithKey :: Property +prop_traverseMaybeWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTBazaar (GTMaybe GTVal) TTMap TTVal) + (\f -> M.traverseMaybeWithKey (\k -> (`More` Done (fmap (f k))))) + (\f -> NEM.traverseMaybeWithKey (\k -> (`More` Done (fmap (f k))))) + +prop_sequence1 :: Property +prop_sequence1 = ttProp (GTNEMap :-> TTBazaar GTVal TTNEMap TTVal) + (sequenceA . fmap (`More` Done id)) + (sequence1 . fmap (`More` Done id)) + +prop_sequenceA :: Property +prop_sequenceA = ttProp (GTNEMap :-> TTBazaar GTVal TTNEMap TTVal) + (sequenceA . fmap (`More` Done id)) + (sequenceA . fmap (`More` Done id)) + +prop_mapAccumWithKey :: Property +prop_mapAccumWithKey = ttProp ( gf3 ((,) <$> valGen <*> valGen) + :?> GTOther valGen + :-> GTNEMap + :-> TTOther :*: TTNEMap + ) + M.mapAccumWithKey + NEM.mapAccumWithKey + +prop_mapAccumRWithKey :: Property +prop_mapAccumRWithKey = ttProp ( gf3 ((,) <$> valGen <*> valGen) + :?> GTOther valGen + :-> GTNEMap + :-> TTOther :*: TTNEMap + ) + M.mapAccumRWithKey + NEM.mapAccumRWithKey + +prop_mapKeys :: Property +prop_mapKeys = ttProp (gf1 keyGen :?> GTNEMap :-> TTNEMap) + M.mapKeys + NEM.mapKeys + +prop_mapKeysWith :: Property +prop_mapKeysWith = ttProp ( gf2 valGen + :?> gf1 keyGen + :?> GTNEMap + :-> TTNEMap + ) + M.mapKeysWith + NEM.mapKeysWith + +prop_mapKeysMonotonic :: Property +prop_mapKeysMonotonic = ttProp (GF valGen go :?> GTNEMap :-> TTNEMap) + M.mapKeysMonotonic + NEM.mapKeysMonotonic + where + go f (K i t) = K (i * 2) (f t) + +prop_foldr :: Property +prop_foldr = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldr + NEM.foldr + +prop_foldl :: Property +prop_foldl = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldl + NEM.foldl + +prop_foldr1 :: Property +prop_foldr1 = ttProp ( gf2 valGen + :?> GTNEMap + :-> TTOther + ) + foldr1 + NEM.foldr1 + +prop_foldl1 :: Property +prop_foldl1 = ttProp ( gf2 valGen + :?> GTNEMap + :-> TTOther + ) + foldl1 + NEM.foldl1 + +prop_foldrWithKey :: Property +prop_foldrWithKey = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldrWithKey + NEM.foldrWithKey + +prop_foldlWithKey :: Property +prop_foldlWithKey = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldlWithKey + NEM.foldlWithKey + +prop_foldMapWithKey :: Property +prop_foldMapWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTOther) + M.foldMapWithKey + NEM.foldMapWithKey + +prop_foldr' :: Property +prop_foldr' = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldr' + NEM.foldr' + +prop_foldl' :: Property +prop_foldl' = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldl' + NEM.foldl' + +prop_foldr1' :: Property +prop_foldr1' = ttProp ( gf2 valGen + :?> GTNEMap + :-> TTOther + ) + foldr1 + NEM.foldr1' + +prop_foldl1' :: Property +prop_foldl1' = ttProp ( gf2 valGen + :?> GTNEMap + :-> TTOther + ) + foldl1 + NEM.foldl1' + +prop_foldrWithKey' :: Property +prop_foldrWithKey' = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldrWithKey' + NEM.foldrWithKey' + +prop_foldlWithKey' :: Property +prop_foldlWithKey' = ttProp ( gf3 valGen + :?> GTOther valGen + :-> GTNEMap + :-> TTOther + ) + M.foldlWithKey' + NEM.foldlWithKey' + +prop_elems :: Property +prop_elems = ttProp (GTNEMap :-> TTNEList TTVal) + M.elems + NEM.elems + +prop_keys :: Property +prop_keys = ttProp (GTNEMap :-> TTNEList TTKey) + M.keys + NEM.keys + +prop_assocs :: Property +prop_assocs = ttProp (GTNEMap :-> TTNEList (TTKey :*: TTVal)) + M.assocs + NEM.assocs + +prop_keysSet :: Property +prop_keysSet = ttProp (GTNEMap :-> TTNESet) + M.keysSet + NEM.keysSet + +prop_toList :: Property +prop_toList = ttProp (GTNEMap :-> TTNEList (TTKey :*: TTVal)) + M.toList + NEM.toList + +prop_toDescList :: Property +prop_toDescList = ttProp (GTNEMap :-> TTNEList (TTKey :*: TTVal)) + M.toDescList + NEM.toDescList + +prop_filter :: Property +prop_filter = ttProp (gf1 Gen.bool :?> GTNEMap :-> TTMap) + M.filter + NEM.filter + +prop_filterWithKey :: Property +prop_filterWithKey = ttProp (gf2 Gen.bool :?> GTNEMap :-> TTMap) + M.filterWithKey + NEM.filterWithKey + +prop_restrictKeys :: Property +prop_restrictKeys = ttProp (GTNEMap :-> GTSet :-> TTMap) + M.restrictKeys + NEM.restrictKeys + +prop_withoutKeys :: Property +prop_withoutKeys = ttProp (GTNEMap :-> GTSet :-> TTMap) + M.withoutKeys + NEM.withoutKeys + +prop_partitionWithKey :: Property +prop_partitionWithKey = ttProp (gf2 Gen.bool :?> GTNEMap :-> TTThese TTNEMap TTNEMap) + M.partitionWithKey + NEM.partitionWithKey + +prop_takeWhileAntitone :: Property +prop_takeWhileAntitone = ttProp (GTNEMap :-> TTMap) + (M.takeWhileAntitone ((< 0) . getKX)) + (NEM.takeWhileAntitone ((< 0) . getKX)) + +prop_dropWhileAntitone :: Property +prop_dropWhileAntitone = ttProp (GTNEMap :-> TTMap) + (M.dropWhileAntitone ((< 0) . getKX)) + (NEM.dropWhileAntitone ((< 0) . getKX)) + +prop_spanAntitone :: Property +prop_spanAntitone = ttProp (GTNEMap :-> TTThese TTNEMap TTNEMap) + (M.spanAntitone ((< 0) . getKX)) + (NEM.spanAntitone ((< 0) . getKX)) + +prop_mapMaybeWithKey :: Property +prop_mapMaybeWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEMap :-> TTMap) + M.mapMaybeWithKey + NEM.mapMaybeWithKey + +prop_mapEitherWithKey :: Property +prop_mapEitherWithKey = ttProp ( gf2 (Gen.choice [Left <$> valGen, Right <$> valGen]) + :?> GTNEMap + :-> TTThese TTNEMap TTNEMap + ) + M.mapEitherWithKey + NEM.mapEitherWithKey + +prop_split :: Property +prop_split = ttProp (GTKey :-> GTNEMap :-> TTMThese TTNEMap TTNEMap) + M.split + NEM.split + +prop_splitLookup :: Property +prop_splitLookup = ttProp (GTKey :-> GTNEMap :-> TTMaybe TTVal :*: TTMThese TTNEMap TTNEMap) + (\k -> (\(x,y,z) -> (y,(x,z))) . M.splitLookup k) + NEM.splitLookup + +prop_isSubmapOfBy :: Property +prop_isSubmapOfBy = ttProp (gf2 Gen.bool :?> GTNEMap :-> GTNEMap :-> TTOther) + M.isSubmapOfBy + NEM.isSubmapOfBy + +prop_isProperSubmapOfBy :: Property +prop_isProperSubmapOfBy = ttProp (gf2 Gen.bool :?> GTNEMap :-> GTNEMap :-> TTOther) + M.isProperSubmapOfBy + NEM.isProperSubmapOfBy + +prop_lookupIndex :: Property +prop_lookupIndex = ttProp (GTKey :-> GTNEMap :-> TTMaybe TTOther) + M.lookupIndex + NEM.lookupIndex + +prop_elemAt :: Property +prop_elemAt = ttProp (GTSize :-> GTNEMap :-> TTKey :*: TTVal) + (\i m -> M.elemAt (i `mod` M.size m) m) + (\i m -> NEM.elemAt (i `mod` NEM.size m) m) + +prop_adjustAt :: Property +prop_adjustAt = ttProp (gf2 valGen :?> GTSize :-> GTNEMap :-> TTNEMap) + (\f i m -> M.updateAt (\k -> Just . f k) (i `mod` M.size m) m) + (\f i m -> NEM.adjustAt f (i `mod` NEM.size m) m) + +prop_updateAt :: Property +prop_updateAt = ttProp (gf2 (Gen.maybe valGen) :?> GTSize :-> GTNEMap :-> TTMap) + (\f i m -> M.updateAt f (i `mod` M.size m) m) + (\f i m -> NEM.updateAt f (i `mod` NEM.size m) m) + +prop_deleteAt :: Property +prop_deleteAt = ttProp (GTSize :-> GTNEMap :-> TTMap) + (\i m -> M.deleteAt (i `mod` M.size m) m) + (\i m -> NEM.deleteAt (i `mod` NEM.size m) m) + +prop_take :: Property +prop_take = ttProp (GTSize :-> GTNEMap :-> TTMap) + M.take + NEM.take + +prop_drop :: Property +prop_drop = ttProp (GTSize :-> GTNEMap :-> TTMap) + M.drop + NEM.drop + +prop_splitAt :: Property +prop_splitAt = ttProp (GTSize :-> GTNEMap :-> TTThese TTNEMap TTNEMap) + M.splitAt + NEM.splitAt + +prop_findMin :: Property +prop_findMin = ttProp (GTNEMap :-> TTKey :*: TTVal) + M.findMin + NEM.findMin + +prop_findMax :: Property +prop_findMax = ttProp (GTNEMap :-> TTKey :*: TTVal) + M.findMax + NEM.findMax + +prop_deleteMin :: Property +prop_deleteMin = ttProp (GTNEMap :-> TTMap) + M.deleteMin + NEM.deleteMin + +prop_deleteMax :: Property +prop_deleteMax = ttProp (GTNEMap :-> TTMap) + M.deleteMax + NEM.deleteMax + +prop_deleteFindMin :: Property +prop_deleteFindMin = ttProp (GTNEMap :-> (TTKey :*: TTVal) :*: TTMap) + M.deleteFindMin + NEM.deleteFindMin + +prop_deleteFindMax :: Property +prop_deleteFindMax = ttProp (GTNEMap :-> (TTKey :*: TTVal) :*: TTMap) + M.deleteFindMax + NEM.deleteFindMax + +prop_updateMinWithKey :: Property +prop_updateMinWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEMap :-> TTMap) + M.updateMinWithKey + NEM.updateMinWithKey + +prop_updateMaxWithKey :: Property +prop_updateMaxWithKey = ttProp (gf2 (Gen.maybe valGen) :?> GTNEMap :-> TTMap) + M.updateMaxWithKey + NEM.updateMaxWithKey + +prop_adjustMinWithKey :: Property +prop_adjustMinWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTNEMap) + (M.updateMinWithKey . (fmap . fmap) Just) + NEM.adjustMinWithKey + +prop_adjustMaxWithKey :: Property +prop_adjustMaxWithKey = ttProp (gf2 valGen :?> GTNEMap :-> TTNEMap) + (M.updateMaxWithKey . (fmap . fmap) Just) + NEM.adjustMaxWithKey + +prop_minView :: Property +prop_minView = ttProp (GTNEMap :-> TTMaybe (TTVal :*: TTMap)) + M.minView + (Just . NEM.minView) + +prop_maxView :: Property +prop_maxView = ttProp (GTNEMap :-> TTMaybe (TTVal :*: TTMap)) + M.maxView + (Just . NEM.maxView) + +prop_elem :: Property +prop_elem = ttProp (GTVal :-> GTNEMap :-> TTOther) + elem + elem + +prop_fold1 :: Property +prop_fold1 = ttProp (GTNEMap :-> TTVal) + fold + fold1 + +prop_fold :: Property +prop_fold = ttProp (GTNEMap :-> TTVal) + fold + fold + +prop_foldMap1 :: Property +prop_foldMap1 = ttProp (gf1 valGen :?> GTNEMap :-> TTOther) + (\f -> foldMap ((:[]) . f)) + (\f -> foldMap1 ((:[]) . f)) + +prop_foldMap :: Property +prop_foldMap = ttProp (gf1 valGen :?> GTNEMap :-> TTOther) + (\f -> foldMap ((:[]) . f)) + (\f -> foldMap ((:[]) . f)) + diff --git a/test/Tests/Sequence.hs b/test/Tests/Sequence.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Sequence.hs @@ -0,0 +1,555 @@ +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TupleSections #-} + +module Tests.Sequence (sequenceTests) where + +import Control.Applicative +import Control.Comonad +import Control.Monad +import Data.Bifunctor +import Data.Functor.Identity +import Data.Ord +import Data.Sequence (Seq(..)) +import Data.Sequence.NonEmpty (NESeq(..)) +import Data.Tuple +import Hedgehog +import Test.Tasty +import Tests.Util +import qualified Data.Foldable as F +import qualified Data.List.NonEmpty as NE +import qualified Data.Semigroup.Foldable as F1 +import qualified Data.Semigroup.Traversable as T1 +import qualified Data.Sequence as Seq +import qualified Data.Sequence.NonEmpty as NESeq +import qualified Data.Sequence.NonEmpty.Internal as NESeq +import qualified Hedgehog.Gen as Gen + +sequenceTests :: TestTree +sequenceTests = groupTree $$(discover) + +prop_toSeqIso1 :: Property +prop_toSeqIso1 = property $ do + m0 <- forAll seqGen + tripping m0 NESeq.nonEmptySeq + (Identity . maybe Seq.empty NESeq.toSeq) + +prop_toSeqIso2 :: Property +prop_toSeqIso2 = property $ do + m0 <- forAll $ Gen.maybe neSeqGen + tripping m0 (maybe Seq.empty NESeq.toSeq) + (Identity . NESeq.nonEmptySeq) + +prop_read_show :: Property +prop_read_show = readShow neSeqGen + +prop_read1_show1 :: Property +prop_read1_show1 = readShow1 neSeqGen + +prop_show_show1 :: Property +prop_show_show1 = showShow1 neSeqGen + + + + + +prop_cons :: Property +prop_cons = ttProp (GTVal :-> GTSeq :-> TTNESeq) + (:<|) + (:<||) + +prop_snoc :: Property +prop_snoc = ttProp (GTSeq :-> GTVal :-> TTNESeq) + (:|>) + (:||>) + +prop_insertSeqAt :: Property +prop_insertSeqAt = ttProp (GTIntKey :-> GTVal :-> GTSeq :-> TTNESeq) + Seq.insertAt + NESeq.insertSeqAt + +prop_singleton :: Property +prop_singleton = ttProp (GTVal :-> TTNESeq) + Seq.singleton + NESeq.singleton + +prop_consNE :: Property +prop_consNE = ttProp (GTVal :-> GTNESeq :-> TTNESeq) + (Seq.<|) + (NESeq.<|) + +prop_snocNE :: Property +prop_snocNE = ttProp (GTNESeq :-> GTVal :-> TTNESeq) + (Seq.|>) + (NESeq.|>) + +prop_append :: Property +prop_append = ttProp (GTNESeq :-> GTNESeq :-> TTNESeq) + (Seq.><) + (NESeq.><) + +prop_appendL :: Property +prop_appendL = ttProp (GTNESeq :-> GTSeq :-> TTNESeq) + (Seq.><) + (NESeq.|><) + +prop_appendR :: Property +prop_appendR = ttProp (GTSeq :-> GTNESeq :-> TTNESeq) + (Seq.><) + (NESeq.><|) + +prop_fromList :: Property +prop_fromList = ttProp (GTNEList Nothing GTVal :-> TTNESeq) + Seq.fromList + NESeq.fromList + +prop_fromFunction :: Property +prop_fromFunction = ttProp (GTSize :-> gf1 valGen :?> TTNESeq) + (Seq.fromFunction . (+ 1)) + (NESeq.fromFunction . (+ 1)) + +prop_replicate :: Property +prop_replicate = ttProp (GTSize :-> GTVal :-> TTNESeq) + (Seq.replicate . (+ 1)) + (NESeq.replicate . (+ 1)) + +prop_replicateA :: Property +prop_replicateA = ttProp (GTSize :-> GTVal :-> TTBazaar GTVal TTNESeq TTVal) + (\i x -> Seq.replicateA (i + 1) (x `More` Done id)) + (\i x -> NESeq.replicateA (i + 1) (x `More` Done id)) + +prop_replicateA1 :: Property +prop_replicateA1 = ttProp (GTSize :-> GTVal :-> TTBazaar GTVal TTNESeq TTVal) + (\i x -> Seq.replicateA (i + 1) (x `More` Done id)) + (\i x -> NESeq.replicateA1 (i + 1) (x `More` Done id)) + +prop_cycleTaking :: Property +prop_cycleTaking = ttProp (GTSize :-> GTNESeq :-> TTNESeq) + (Seq.cycleTaking . (* 5) . (+ 1)) + (NESeq.cycleTaking . (* 5) . (+ 1)) + +prop_iterateN :: Property +prop_iterateN = ttProp (GTSize :-> gf1 valGen :?> GTVal :-> TTNESeq) + (Seq.iterateN . (+ 1)) + (NESeq.iterateN . (+ 1)) + +prop_unfoldr :: Property +prop_unfoldr = ttProp ( GTSize + :-> gf1 ((,) <$> valGen <*> Gen.maybe intKeyGen) + :?> GTIntKey + :-> TTNESeqList + ) + (\i f -> NE.unfoldr (limiter f) . (i,)) + (\i f -> NESeq.unfoldr (limiter f) . (i,)) + +prop_unfoldl :: Property +prop_unfoldl = ttProp ( GTSize + :-> gf1 ((,) <$> valGen <*> Gen.maybe intKeyGen) + :?> GTIntKey + :-> TTNESeqList + ) + (\i f -> NE.reverse . NE.unfoldr ( limiter f) . (i,)) + (\i f -> NESeq.unfoldl (swap . limiter f) . (i,)) + +limiter + :: (a -> (b, Maybe a)) + -> (Int, a) + -> (b, Maybe (Int, a)) +limiter f (n, x) = second (go =<<) $ f x + where + go y + | n <= 0 = Nothing + | otherwise = Just (n - 1, y) + +prop_head :: Property +prop_head = ttProp (GTNESeq :-> TTMaybe TTVal) + (\case x :<| _ -> Just x; Empty -> Nothing) + (Just . NESeq.head) + +prop_tail :: Property +prop_tail = ttProp (GTNESeq :-> TTMaybe TTOther) + (\case _ :<| xs -> Just xs; Empty -> Nothing) + (Just . NESeq.tail) + +prop_last :: Property +prop_last = ttProp (GTNESeq :-> TTMaybe TTVal) + (\case _ :|> x -> Just x; Empty -> Nothing) + (Just . NESeq.last) + +prop_init :: Property +prop_init = ttProp (GTNESeq :-> TTMaybe TTOther) + (\case xs :|> _ -> Just xs; Empty -> Nothing) + (Just . NESeq.init) + +prop_length :: Property +prop_length = ttProp (GTNESeq :-> TTOther) + Seq.length + NESeq.length + +prop_scanl :: Property +prop_scanl = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTNESeq) + Seq.scanl + NESeq.scanl + +prop_scanl1 :: Property +prop_scanl1 = ttProp (gf2 valGen :?> GTNESeq :-> TTNESeq) + Seq.scanl1 + NESeq.scanl1 + +prop_scanr :: Property +prop_scanr = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTNESeq) + Seq.scanr + NESeq.scanr + +prop_scanr1 :: Property +prop_scanr1 = ttProp (gf2 valGen :?> GTNESeq :-> TTNESeq) + Seq.scanl1 + NESeq.scanl1 + +prop_tails :: Property +prop_tails = ttProp (GTNESeq :-> TTNESeq) + (Seq.filter (not . null) . Seq.tails) + (fmap NESeq.toSeq . NESeq.tails) + +prop_inits :: Property +prop_inits = ttProp (GTNESeq :-> TTNESeq) + (Seq.filter (not . null) . Seq.inits) + (fmap NESeq.toSeq . NESeq.inits) + +prop_chunksOf :: Property +prop_chunksOf = ttProp (GTSize :-> GTNESeq :-> TTNESeq) + (\i -> Seq.filter (not . null) . Seq.chunksOf (i + 1)) + (\i -> fmap NESeq.toSeq . NESeq.chunksOf (i + 1)) + +prop_takeWhileL :: Property +prop_takeWhileL = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.takeWhileL + NESeq.takeWhileL + +prop_takeWhileR :: Property +prop_takeWhileR = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.takeWhileR + NESeq.takeWhileR + +prop_dropWhileL :: Property +prop_dropWhileL = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.dropWhileL + NESeq.dropWhileL + +prop_dropWhileR :: Property +prop_dropWhileR = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.dropWhileR + NESeq.dropWhileR + +prop_spanl :: Property +prop_spanl = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.spanl + NESeq.spanl + +prop_spanr :: Property +prop_spanr = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.spanr + NESeq.spanr + +prop_breakl :: Property +prop_breakl = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.breakl + NESeq.breakl + +prop_breakr :: Property +prop_breakr = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.breakr + NESeq.breakr + +prop_partition :: Property +prop_partition = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.partition + NESeq.partition + +prop_filter :: Property +prop_filter = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.filter + NESeq.filter + +prop_sort :: Property +prop_sort = ttProp (GTNESeq :-> TTNESeq) + Seq.sort + NESeq.sort + +prop_sortBy :: Property +prop_sortBy = ttProp (gf1 valGen :?> GTNESeq :-> TTNESeq) + (Seq.sortBy . comparing) + (NESeq.sortBy . comparing) + +prop_sortOn :: Property +prop_sortOn = ttProp (gf1 valGen :?> GTNESeq :-> TTNESeq) + NESeq.sortOnSeq + NESeq.sortOn + +prop_unstableSort :: Property +prop_unstableSort = ttProp (GTNESeq :-> TTNESeq) + Seq.unstableSort + NESeq.unstableSort + +prop_unstableSortBy :: Property +prop_unstableSortBy = ttProp (gf1 valGen :?> GTNESeq :-> TTNESeq) + (Seq.unstableSortBy . comparing) + (NESeq.unstableSortBy . comparing) + +prop_unstableSortOn :: Property +prop_unstableSortOn = ttProp (gf1 valGen :?> GTNESeq :-> TTNESeq) + NESeq.unstableSortOnSeq + NESeq.unstableSortOn + +prop_lookup :: Property +prop_lookup = ttProp (GTIntKey :-> GTNESeq :-> TTMaybe TTVal) + Seq.lookup + NESeq.lookup + +prop_index :: Property +prop_index = ttProp (GTNESeq :-> GTIntKey :-> TTVal) + (\xs i -> xs `Seq.index` (i `mod` Seq.length xs )) + (\xs i -> xs `NESeq.index` (i `mod` NESeq.length xs)) + +prop_adjust :: Property +prop_adjust = ttProp (gf1 valGen :?> GTIntKey :-> GTNESeq :-> TTNESeq) + Seq.adjust + NESeq.adjust + +prop_adjust' :: Property +prop_adjust' = ttProp (gf1 valGen :?> GTIntKey :-> GTNESeq :-> TTNESeq) + Seq.adjust' + NESeq.adjust' + +prop_update :: Property +prop_update = ttProp (GTIntKey :-> GTVal :-> GTNESeq :-> TTNESeq) + Seq.update + NESeq.update + +prop_take :: Property +prop_take = ttProp (GTIntKey :-> GTNESeq :-> TTOther) + Seq.take + NESeq.take + +prop_drop :: Property +prop_drop = ttProp (GTIntKey :-> GTNESeq :-> TTOther) + Seq.drop + NESeq.drop + +prop_insertAt :: Property +prop_insertAt = ttProp (GTIntKey :-> GTVal :-> GTNESeq :-> TTNESeq) + Seq.insertAt + NESeq.insertAt + +prop_deleteAt :: Property +prop_deleteAt = ttProp (GTIntKey :-> GTNESeq :-> TTOther) + Seq.deleteAt + NESeq.deleteAt + +prop_splitAt :: Property +prop_splitAt = ttProp (GTIntKey :-> GTNESeq :-> TTThese TTNESeq TTNESeq) + Seq.splitAt + NESeq.splitAt + +prop_elemIndexL :: Property +prop_elemIndexL = ttProp (GTVal :-> GTNESeq :-> TTOther) + Seq.elemIndexL + NESeq.elemIndexL + +prop_elemIndicesL :: Property +prop_elemIndicesL = ttProp (GTVal :-> GTNESeq :-> TTOther) + Seq.elemIndicesL + NESeq.elemIndicesL + +prop_elemIndexR :: Property +prop_elemIndexR = ttProp (GTVal :-> GTNESeq :-> TTOther) + Seq.elemIndexR + NESeq.elemIndexR + +prop_elemIndicesR :: Property +prop_elemIndicesR = ttProp (GTVal :-> GTNESeq :-> TTOther) + Seq.elemIndicesR + NESeq.elemIndicesR + +prop_findIndexL :: Property +prop_findIndexL = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.findIndexL + NESeq.findIndexL + +prop_findIndicesL :: Property +prop_findIndicesL = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.findIndicesL + NESeq.findIndicesL + +prop_findIndexR :: Property +prop_findIndexR = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.findIndexR + NESeq.findIndexR + +prop_findIndicesR :: Property +prop_findIndicesR = ttProp (gf1 Gen.bool :?> GTNESeq :-> TTOther) + Seq.findIndicesR + NESeq.findIndicesR + +prop_foldMapWithIndex :: Property +prop_foldMapWithIndex = ttProp (gf2 valGen :?> GTNESeq :-> TTOther) + (\f -> Seq.foldMapWithIndex (\i -> (:[]) . f i)) + (\f -> NESeq.foldMapWithIndex (\i -> (:[]) . f i)) + +prop_foldlWithIndex :: Property +prop_foldlWithIndex = ttProp (gf3 valGen :?> GTVal :-> GTNESeq :-> TTVal) + Seq.foldlWithIndex + NESeq.foldlWithIndex + +prop_foldrWithIndex :: Property +prop_foldrWithIndex = ttProp (gf3 valGen :?> GTVal :-> GTNESeq :-> TTVal) + Seq.foldrWithIndex + NESeq.foldrWithIndex + +prop_mapWithIndex :: Property +prop_mapWithIndex = ttProp (gf2 valGen :?> GTNESeq :-> TTNESeq) + Seq.mapWithIndex + NESeq.mapWithIndex + +prop_traverseWithIndex :: Property +prop_traverseWithIndex = ttProp (gf2 valGen :?> GTNESeq :-> TTBazaar GTVal TTNESeq TTVal) + (\f -> Seq.traverseWithIndex (\k -> (`More` Done (f k)))) + (\f -> NESeq.traverseWithIndex (\k -> (`More` Done (f k)))) + +prop_traverseWithIndex1 :: Property +prop_traverseWithIndex1 = ttProp (gf2 valGen :?> GTNESeq :-> TTBazaar GTVal TTNESeq TTVal) + (\f -> Seq.traverseWithIndex (\k -> (`More` Done (f k)))) + (\f -> NESeq.traverseWithIndex1 (\k -> (`More` Done (f k)))) + +prop_reverse :: Property +prop_reverse = ttProp (GTNESeq :-> TTNESeq) + Seq.reverse + NESeq.reverse + +prop_intersperse :: Property +prop_intersperse = ttProp (GTVal :-> GTNESeq :-> TTNESeq) + Seq.intersperse + NESeq.intersperse + +prop_zip :: Property +prop_zip = ttProp (GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zip + NESeq.zip + +prop_zipWith :: Property +prop_zipWith = ttProp (gf2 valGen :?> GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zipWith + NESeq.zipWith + +prop_zip3 :: Property +prop_zip3 = ttProp (GTNESeq :-> GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zip3 + NESeq.zip3 + +prop_zipWith3 :: Property +prop_zipWith3 = ttProp (gf3 valGen :?> GTNESeq :-> GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zipWith3 + NESeq.zipWith3 + +prop_zip4 :: Property +prop_zip4 = ttProp (GTNESeq :-> GTNESeq :-> GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zip4 + NESeq.zip4 + +prop_zipWith4 :: Property +prop_zipWith4 = ttProp (gf4 valGen :?> GTNESeq :-> GTNESeq :-> GTNESeq :-> GTNESeq :-> TTNESeq) + Seq.zipWith4 + NESeq.zipWith4 + +prop_unzip :: Property +prop_unzip = ttProp (GTNESeq :-> GTNESeq :-> TTNESeq :*: TTNESeq) + (\xs -> NESeq.unzipSeq . Seq.zip xs) + (\xs -> NESeq.unzip . NESeq.zip xs) + +prop_unzipWith :: Property +prop_unzipWith = ttProp ( gf1 ((,) <$> valGen <*> valGen) + :?> GTNESeq + :-> TTNESeq :*: TTNESeq + ) + NESeq.unzipWithSeq + NESeq.unzipWith + +prop_liftA2 :: Property +prop_liftA2 = ttProp (gf2 valGen :?> GTNESeq :-> GTNESeq :-> TTNESeq) + liftA2 + liftA2 + +prop_liftM2 :: Property +prop_liftM2 = ttProp (gf2 valGen :?> GTNESeq :-> GTNESeq :-> TTNESeq) + liftM2 + liftM2 + +prop_duplicate :: Property +prop_duplicate = ttProp (GTNESeqList :-> TTNESeqList) + duplicate + (fmap F1.toNonEmpty . duplicate) + +prop_foldMap :: Property +prop_foldMap = ttProp (gf1 valGen :?> GTNESeq :-> TTOther) + (foldMap . fmap (:[])) + (foldMap . fmap (:[])) + +prop_foldl :: Property +prop_foldl = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTVal) + foldl + foldl + +prop_foldr :: Property +prop_foldr = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTVal) + foldr + foldr + +prop_foldl' :: Property +prop_foldl' = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTVal) + F.foldl' + F.foldl' + +prop_foldr' :: Property +prop_foldr' = ttProp (gf2 valGen :?> GTVal :-> GTNESeq :-> TTVal) + F.foldr' + F.foldr' + +prop_foldl1 :: Property +prop_foldl1 = ttProp (gf2 valGen :?> GTNESeq :-> TTVal) + foldl1 + foldl1 + +prop_foldr1 :: Property +prop_foldr1 = ttProp (gf2 valGen :?> GTNESeq :-> TTVal) + foldr1 + foldr1 + +prop_fold :: Property +prop_fold = ttProp (GTNESeq :-> TTVal) + F.fold + F.fold + +prop_fold1 :: Property +prop_fold1 = ttProp (GTNESeq :-> TTVal) + F.fold + F1.fold1 + +prop_toList :: Property +prop_toList = ttProp (GTNESeq :-> TTOther) + F.toList + F.toList + +prop_toNonEmpty :: Property +prop_toNonEmpty = ttProp (GTNESeq :-> TTNEList TTVal) + F.toList + F1.toNonEmpty + +prop_sequenceA :: Property +prop_sequenceA = ttProp (GTNESeq :-> TTBazaar GTVal TTNESeq TTVal) + (sequenceA . fmap (`More` Done id)) + (sequenceA . fmap (`More` Done id)) + +prop_sequence1 :: Property +prop_sequence1 = ttProp (GTNESeq :-> TTBazaar GTVal TTNESeq TTVal) + (sequenceA . fmap (`More` Done id)) + (T1.sequence1 . fmap (`More` Done id)) diff --git a/test/Tests/Set.hs b/test/Tests/Set.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Set.hs @@ -0,0 +1,432 @@ +{-# LANGUAGE TemplateHaskell #-} + +module Tests.Set (setTests) where + +import Data.Foldable +import Data.Functor.Identity +import Data.Semigroup.Foldable +import Hedgehog +import Test.Tasty +import Tests.Util +import qualified Data.Set as S +import qualified Data.Set.NonEmpty as NES +import qualified Data.Set.NonEmpty.Internal as NES +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range + +setTests :: TestTree +setTests = groupTree $$(discover) + + + + + +prop_valid :: Property +prop_valid = property $ + assert . NES.valid =<< forAll neSetGen + +prop_valid_toSet :: Property +prop_valid_toSet = property $ do + assert . S.valid . NES.toSet =<< forAll neSetGen + +prop_valid_insertMinSet :: Property +prop_valid_insertMinSet = property $ do + n <- forAll $ do + m <- setGen + let k = maybe dummyKey (subtract 1) $ S.lookupMin m + pure $ NES.insertMinSet k m + assert $ S.valid n + +prop_valid_insertMaxSet :: Property +prop_valid_insertMaxSet = property $ do + n <- forAll $ do + m <- setGen + let k = maybe dummyKey (+ 1) $ S.lookupMax m + pure $ NES.insertMaxSet k m + assert $ S.valid n + +prop_valid_insertSetMin :: Property +prop_valid_insertSetMin = property $ do + n <- forAll $ do + m <- setGen + let k = maybe dummyKey (subtract 1) $ S.lookupMin m + pure $ NES.insertSetMin k m + assert $ NES.valid n + +prop_valid_insertSetMax :: Property +prop_valid_insertSetMax = property $ do + n <- forAll $ do + m <- setGen + let k = maybe dummyKey (+ 1) $ S.lookupMax m + pure $ NES.insertSetMax k m + assert $ NES.valid n + +prop_toSetIso1 :: Property +prop_toSetIso1 = property $ do + m0 <- forAll setGen + tripping m0 NES.nonEmptySet + (Identity . maybe S.empty NES.toSet) + +prop_toSetIso2 :: Property +prop_toSetIso2 = property $ do + m0 <- forAll $ Gen.maybe neSetGen + tripping m0 (maybe S.empty NES.toSet) + (Identity . NES.nonEmptySet) + +prop_read_show :: Property +prop_read_show = readShow neSetGen + +prop_show_show1 :: Property +prop_show_show1 = showShow1 neSetGen + +prop_splitRoot :: Property +prop_splitRoot = property $ do + n <- forAll neSetGen + let rs = NES.splitRoot n + n' = foldl1 NES.merge rs + assert $ NES.valid n' + mapM_ (assert . (`NES.isSubsetOf` n)) rs + n === n' + + + + +prop_insertSet :: Property +prop_insertSet = ttProp (GTKey :-> GTSet :-> TTNESet) + S.insert + NES.insertSet + +prop_singleton :: Property +prop_singleton = ttProp (GTKey :-> TTNESet) + S.singleton + NES.singleton + +prop_fromAscList :: Property +prop_fromAscList = ttProp (GTSorted STAsc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNESet) + (S.fromAscList . fmap fst) + (NES.fromAscList . fmap fst) + +prop_fromDescList :: Property +prop_fromDescList = ttProp (GTSorted STDesc (GTNEList Nothing (GTKey :&: GTVal)) :-> TTNESet) + (S.fromDescList . fmap fst) + (NES.fromDescList . fmap fst) + +prop_fromDistinctAscList :: Property +prop_fromDistinctAscList = ttProp (GTSorted STAsc (GTNEList Nothing GTKey) :-> TTNESet) + S.fromDistinctAscList + NES.fromDistinctAscList + +prop_fromDistinctDescList :: Property +prop_fromDistinctDescList = ttProp (GTSorted STDesc (GTNEList Nothing GTKey) :-> TTNESet) + S.fromDistinctDescList + NES.fromDistinctDescList + +prop_fromList :: Property +prop_fromList = ttProp (GTNEList Nothing GTKey :-> TTNESet) + S.fromList + NES.fromList + +prop_powerSet :: Property +prop_powerSet = ttProp (GTNESet :-> TTNEList TTNESet) + (S.toList . S.drop 1 . NES.powerSetSet) + (NES.toList . NES.powerSet ) + +prop_insert :: Property +prop_insert = ttProp (GTKey :-> GTNESet :-> TTNESet) + S.insert + NES.insert + +prop_delete :: Property +prop_delete = ttProp (GTKey :-> GTNESet :-> TTSet) + S.delete + NES.delete + +prop_member :: Property +prop_member = ttProp (GTKey :-> GTNESet :-> TTOther) + S.member + NES.member + +prop_notMember :: Property +prop_notMember = ttProp (GTKey :-> GTNESet :-> TTOther) + S.notMember + NES.notMember + +prop_lookupLT :: Property +prop_lookupLT = ttProp (GTKey :-> GTNESet :-> TTMaybe TTKey) + S.lookupLT + NES.lookupLT + +prop_lookupGT :: Property +prop_lookupGT = ttProp (GTKey :-> GTNESet :-> TTMaybe TTKey) + S.lookupGT + NES.lookupGT + +prop_lookupLE :: Property +prop_lookupLE = ttProp (GTKey :-> GTNESet :-> TTMaybe TTKey) + S.lookupLE + NES.lookupLE + +prop_lookupGE :: Property +prop_lookupGE = ttProp (GTKey :-> GTNESet :-> TTMaybe TTKey) + S.lookupGE + NES.lookupGE + +prop_size :: Property +prop_size = ttProp (GTNESet :-> TTOther) + S.size + NES.size + +prop_isSubsetOf :: Property +prop_isSubsetOf = ttProp (GTNESet :-> GTNESet :-> TTOther) + S.isSubsetOf + NES.isSubsetOf + +prop_isProperSubsetOf :: Property +prop_isProperSubsetOf = ttProp (GTNESet :-> GTNESet :-> TTOther) + S.isProperSubsetOf + NES.isProperSubsetOf + +prop_disjoint :: Property +prop_disjoint = ttProp (GTNESet :-> GTNESet :-> TTOther) + NES.disjointSet + NES.disjoint + +prop_union :: Property +prop_union = ttProp (GTNESet :-> GTNESet :-> TTNESet) + S.union + NES.union + +prop_unions :: Property +prop_unions = ttProp (GTNEList (Just (Range.linear 2 5)) GTNESet :-> TTNESet) + S.unions + NES.unions + +prop_difference :: Property +prop_difference = ttProp (GTNESet :-> GTNESet :-> TTSet) + S.difference + NES.difference + +prop_intersection :: Property +prop_intersection = ttProp (GTNESet :-> GTNESet :-> TTSet) + S.intersection + NES.intersection + +prop_cartesianProduct :: Property +prop_cartesianProduct = ttProp (GTNESet :-> GTNESet :-> TTNEList (TTKey :*: TTKey)) + (\xs -> S.toList . NES.cartesianProductSet xs) + (\xs -> NES.toList . NES.cartesianProduct xs) + +prop_disjointUnion :: Property +prop_disjointUnion = ttProp (GTNESet :-> GTNESet :-> TTNEList (TTEither TTKey TTKey)) + (\xs -> S.toList . NES.disjointUnionSet xs) + (\xs -> NES.toList . NES.disjointUnion xs) + +prop_filter :: Property +prop_filter = ttProp (gf1 Gen.bool :?> GTNESet :-> TTSet) + S.filter + NES.filter + +prop_takeWhileAntitone :: Property +prop_takeWhileAntitone = ttProp (GTNESet :-> TTSet) + (S.takeWhileAntitone ((< 0) . getKX)) + (NES.takeWhileAntitone ((< 0) . getKX)) + +prop_dropWhileAntitone :: Property +prop_dropWhileAntitone = ttProp (GTNESet :-> TTSet) + (S.dropWhileAntitone ((< 0) . getKX)) + (NES.dropWhileAntitone ((< 0) . getKX)) + +prop_spanAntitone :: Property +prop_spanAntitone = ttProp (GTNESet :-> TTThese TTNESet TTNESet) + (S.spanAntitone ((< 0) . getKX)) + (NES.spanAntitone ((< 0) . getKX)) + +prop_partition :: Property +prop_partition = ttProp (gf1 Gen.bool :?> GTNESet :-> TTThese TTNESet TTNESet) + S.partition + NES.partition + +prop_split :: Property +prop_split = ttProp (GTKey :-> GTNESet :-> TTMThese TTNESet TTNESet) + S.split + NES.split + +prop_splitMember :: Property +prop_splitMember = ttProp (GTKey :-> GTNESet :-> TTOther :*: TTMThese TTNESet TTNESet) + (\k -> (\(x,y,z) -> (y,(x,z))) . S.splitMember k) + NES.splitMember + +prop_lookupIndex :: Property +prop_lookupIndex = ttProp (GTKey :-> GTNESet :-> TTMaybe TTOther) + S.lookupIndex + NES.lookupIndex + +prop_elemAt :: Property +prop_elemAt = ttProp (GTSize :-> GTNESet :-> TTKey) + (\i m -> S.elemAt (i `mod` S.size m) m) + (\i m -> NES.elemAt (i `mod` NES.size m) m) + +prop_deleteAt :: Property +prop_deleteAt = ttProp (GTSize :-> GTNESet :-> TTSet) + (\i m -> S.deleteAt (i `mod` S.size m) m) + (\i m -> NES.deleteAt (i `mod` NES.size m) m) + +prop_take :: Property +prop_take = ttProp (GTSize :-> GTNESet :-> TTSet) + S.take + NES.take + +prop_drop :: Property +prop_drop = ttProp (GTSize :-> GTNESet :-> TTSet) + S.drop + NES.drop + +prop_splitAt :: Property +prop_splitAt = ttProp (GTSize :-> GTNESet :-> TTThese TTNESet TTNESet) + S.splitAt + NES.splitAt + +prop_map :: Property +prop_map = ttProp (gf1 keyGen :?> GTNESet :-> TTNESet) + S.map + NES.map + +prop_mapMonotonic :: Property +prop_mapMonotonic = ttProp (GF valGen go :?> GTNESet :-> TTNESet) + S.mapMonotonic + NES.mapMonotonic + where + go f (K i t) = K (i * 2) (f t) + +prop_foldr :: Property +prop_foldr = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNESet + :-> TTOther + ) + S.foldr + NES.foldr + +prop_foldl :: Property +prop_foldl = ttProp ( gf2 valGen + :?> GTOther valGen + :-> GTNESet + :-> TTOther + ) + S.foldl + NES.foldl + +prop_foldr1 :: Property +prop_foldr1 = ttProp ( gf2 keyGen + :?> GTNESet + :-> TTOther + ) + foldr1 + NES.foldr1 + +prop_foldl1 :: Property +prop_foldl1 = ttProp ( gf2 keyGen + :?> GTNESet + :-> TTOther + ) + foldl1 + NES.foldl1 + +prop_foldr' :: Property +prop_foldr' = ttProp ( gf2 keyGen + :?> GTOther keyGen + :-> GTNESet + :-> TTOther + ) + S.foldr' + NES.foldr' + +prop_foldl' :: Property +prop_foldl' = ttProp ( gf2 keyGen + :?> GTOther keyGen + :-> GTNESet + :-> TTOther + ) + S.foldl' + NES.foldl' + +prop_foldr1' :: Property +prop_foldr1' = ttProp ( gf2 keyGen + :?> GTNESet + :-> TTOther + ) + foldr1 + NES.foldr1' + +prop_foldl1' :: Property +prop_foldl1' = ttProp ( gf2 keyGen + :?> GTNESet + :-> TTOther + ) + foldl1 + NES.foldl1' + +prop_findMin :: Property +prop_findMin = ttProp (GTNESet :-> TTKey) + S.findMin + NES.findMin + +prop_findMax :: Property +prop_findMax = ttProp (GTNESet :-> TTKey) + S.findMax + NES.findMax + +prop_deleteMin :: Property +prop_deleteMin = ttProp (GTNESet :-> TTSet) + S.deleteMin + NES.deleteMin + +prop_deleteMax :: Property +prop_deleteMax = ttProp (GTNESet :-> TTSet) + S.deleteMax + NES.deleteMax + +prop_deleteFindMin :: Property +prop_deleteFindMin = ttProp (GTNESet :-> TTKey :*: TTSet) + S.deleteFindMin + NES.deleteFindMin + +prop_deleteFindMax :: Property +prop_deleteFindMax = ttProp (GTNESet :-> TTKey :*: TTSet) + S.deleteFindMax + NES.deleteFindMax + +prop_toList :: Property +prop_toList = ttProp (GTNESet :-> TTNEList TTKey) + S.toList + NES.toList + +prop_toDescList :: Property +prop_toDescList = ttProp (GTNESet :-> TTNEList TTKey) + S.toDescList + NES.toDescList + +prop_elem :: Property +prop_elem = ttProp (GTKey :-> GTNESet :-> TTOther) + elem + elem + +prop_fold1 :: Property +prop_fold1 = ttProp (GTNESet :-> TTKey) + fold + fold1 + +prop_fold :: Property +prop_fold = ttProp (GTNESet :-> TTKey) + fold + fold + +prop_foldMap1 :: Property +prop_foldMap1 = ttProp (gf1 keyGen :?> GTNESet :-> TTOther) + (\f -> foldMap ((:[]) . f)) + (\f -> foldMap1 ((:[]) . f)) + +prop_foldMap :: Property +prop_foldMap = ttProp (gf1 keyGen :?> GTNESet :-> TTOther) + (\f -> foldMap ((:[]) . f)) + (\f -> foldMap ((:[]) . f)) diff --git a/test/Tests/Util.hs b/test/Tests/Util.hs new file mode 100644 --- /dev/null +++ b/test/Tests/Util.hs @@ -0,0 +1,554 @@ +{-# LANGUAGE CPP #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE DeriveGeneric #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GADTs #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE LambdaCase #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeSynonymInstances #-} +{-# OPTIONS_GHC -Wno-orphans #-} + +module Tests.Util ( + K(..), KeyType, overKX, dummyKey + , SortType(..) + , GenFunc(..), gf1, gf2, gf3, gf4 + , GenType(..) + , TestType(..) + , ttProp + , groupTree + , readShow, readShow1, showShow1, showShow2 + , Context(..) + , Bazaar(..) + , keyGen, valGen, mapSize, mapGen, neMapGen, setGen, neSetGen + , intKeyGen, intMapGen, neIntMapGen, intSetGen, neIntSetGen + , seqGen, neSeqGen + ) where + +import Control.Applicative +import Control.Monad +import Data.Bifunctor +import Data.Char +import Data.Foldable +import Data.Function +import Data.Functor.Apply +import Data.Functor.Classes +import Data.IntMap (IntMap) +import Data.IntMap.NonEmpty (NEIntMap) +import Data.IntSet (IntSet, Key) +import Data.IntSet.NonEmpty (NEIntSet) +import Data.Kind +import Data.List.NonEmpty (NonEmpty(..)) +import Data.Map (Map) +import Data.Map.NonEmpty (NEMap) +import Data.Maybe +import Data.Semigroup.Foldable +import Data.Sequence (Seq(..)) +import Data.Sequence.NonEmpty (NESeq(..)) +import Data.Set (Set) +import Data.Set.NonEmpty (NESet) +import Data.Text (Text) +import Data.These +import Hedgehog +import Hedgehog.Function hiding ((:*:)) +import Hedgehog.Internal.Property +import Test.Tasty +import Test.Tasty.Hedgehog +import Text.Read +import qualified Data.IntMap as IM +import qualified Data.IntMap.NonEmpty as NEIM +import qualified Data.IntSet as IS +import qualified Data.IntSet.NonEmpty as NEIS +import qualified Data.List.NonEmpty as NE +import qualified Data.Map as M +import qualified Data.Map.NonEmpty as NEM +import qualified Data.Sequence.NonEmpty as NESeq +import qualified Data.Set as S +import qualified Data.Set.NonEmpty as NES +import qualified Data.Text as T +import qualified Hedgehog.Gen as Gen +import qualified Hedgehog.Range as Range + +#if !MIN_VERSION_base(4,11,0) +import Data.Semigroup (Semigroup(..)) +#endif + +groupTree :: Group -> TestTree +groupTree Group{..} = testGroup (unGroupName groupName) + (map (uncurry go) groupProperties) + where + go :: PropertyName -> Property -> TestTree + go n = testProperty (mkName (unPropertyName n)) + mkName = map deUnderscore . drop (length @[] @Char "prop_") + deUnderscore '_' = ' ' + deUnderscore c = c + +-- | test for stability +data K a b = K { getKX :: !a, getKY :: !b } + deriving (Show, Read, Generic) + +withK :: (a -> b -> c) -> K a b -> c +withK f (K x y) = f x y + +overKX :: (a -> c) -> K a b -> K c b +overKX f (K x y) = K (f x) y + +instance Eq a => Eq (K a b) where + (==) = (==) `on` getKX + +instance Ord a => Ord (K a b) where + compare = compare `on` getKX + +instance (Vary a, Vary b) => Vary (K a b) +instance (Arg a, Arg b) => Arg (K a b) + +type KeyType = K Int Text + +instance Semigroup KeyType where + K x1 y1 <> K x2 y2 = K (x1 + x2) (y1 <> y2) + +instance Monoid KeyType where + mempty = K 0 "" + mappend = (<>) + +dummyKey :: KeyType +dummyKey = K 0 "hello" + + +#if MIN_VERSION_base(4,11,0) +instance (Num a, Monoid b) => Num (K a b) where +#else +instance (Num a, Semigroup b, Monoid b) => Num (K a b) where +#endif + K x1 y1 + K x2 y2 = K (x1 + x2) (y1 <> y2) + K x1 y1 - K x2 y2 = K (x1 - x2) (y1 <> y2) + K x1 y1 * K x2 y2 = K (x1 * x2) (y1 <> y2) + negate (K x y) = K (negate x) y + abs (K x y) = K (abs x) y + signum (K x y) = K (signum x) y + fromInteger n = K (fromInteger n) mempty + +data Context a b t = Context (b -> t) a + deriving Functor + +data Bazaar a b t = Done t + | More a (Bazaar a b (b -> t)) + deriving Functor + +instance Apply (Bazaar a b) where +#if MIN_VERSION_semigroupoids(5,2,2) + liftF2 f = \case + Done x -> fmap (f x) + More x b -> More x . liftA2 (\g r y -> f (g y) r) b +#else + (<.>) = \case + Done x -> fmap x + More x b -> More x . liftA2 (\g r y -> g y r) b +#endif + +instance Applicative (Bazaar a b) where + pure = Done + liftA2 = liftF2 + +data SortType :: Type -> Type where + STAsc :: Ord a => SortType a + STDesc :: Ord a => SortType a + STDistinctAsc :: Ord a => SortType (a, b) + STDistinctDesc :: Ord a => SortType (a, b) + +data GenType :: Type -> Type -> Type where + GTNEMap :: GenType (Map KeyType Text) (NEMap KeyType Text) + GTMap :: GenType (Map KeyType Text) (Map KeyType Text ) + GTNESet :: GenType (Set KeyType ) (NESet KeyType ) + GTNEIntMap :: GenType (IntMap Text ) (NEIntMap Text ) + GTNEIntSet :: GenType IntSet NEIntSet + GTIntMap :: GenType (IntMap Text ) (IntMap Text ) + GTNESeq :: GenType (Seq Text ) (NESeq Text ) + GTNESeqList :: GenType (NonEmpty Text ) (NESeq Text ) + GTSeq :: GenType (Seq Text ) (Seq Text ) + GTKey :: GenType KeyType KeyType + GTIntKey :: GenType Int Int + GTVal :: GenType Text Text + GTSize :: GenType Int Int + GTOther :: Gen a + -> GenType a a + GTMaybe :: GenType a b + -> GenType (Maybe a) (Maybe b) + (:&:) :: GenType a b + -> GenType c d + -> GenType (a, c) (b, d) + GTNEList :: Maybe (Range Int) + -> GenType a b + -> GenType [a] (NonEmpty b) + GTSet :: GenType (Set KeyType) (Set KeyType) + GTIntSet :: GenType IntSet IntSet + GTSorted :: SortType a + -> GenType [a] (NonEmpty a) + -> GenType [a] (NonEmpty a) + +data GenFunc :: Type -> Type -> Type -> Type where + GF :: (Show a, Arg a, Vary a, Show b) + => Gen b + -> ((a -> b) -> f) + -> GenFunc f c d + +gf1 :: (Show a, Arg a, Vary a, Show b) + => Gen b + -> GenFunc (a -> b) c d +gf1 = (`GF` id) + +gf2 :: (Show a, Show b, Arg a, Vary a, Arg b, Vary b, Show c) + => Gen c + -> GenFunc (a -> b -> c) d e +gf2 = (`GF` curry) + +gf3 :: (Show a, Show b, Show c, Arg a, Vary a, Arg b, Vary b, Arg c, Vary c, Show d) + => Gen d + -> GenFunc (a -> b -> c -> d) e f +gf3 = (`GF` (curry . curry)) + +gf4 :: (Show a, Show b, Show c, Arg a, Vary a, Arg b, Vary b, Arg c, Vary c, Show d, Show e, Arg d, Vary d) + => Gen e + -> GenFunc (a -> b -> c -> d -> e) f g +gf4 = (`GF` (curry . curry . curry)) + + + + +data TestType :: Type -> Type -> Type where + TTNEMap :: (Eq a, Show a) + => TestType (Map KeyType a) (NEMap KeyType a ) + TTNEIntMap :: (Eq a, Show a) + => TestType (IntMap a ) (NEIntMap a ) + TTNESet :: TestType (Set KeyType ) (NESet KeyType ) + TTNEIntSet :: TestType IntSet NEIntSet + TTMap :: (Eq a, Show a) + => TestType (Map KeyType a) (Map KeyType a ) + TTSet :: TestType (Set KeyType ) (Set KeyType ) + TTNESeq :: (Eq a, Show a) + => TestType (Seq a ) (NESeq a ) + TTNESeqList :: (Eq a, Show a) + => TestType (NonEmpty a ) (NESeq a ) + TTKey :: TestType KeyType KeyType + TTVal :: TestType Text Text + TTOther :: (Eq a, Show a) + => TestType a a + TTThese :: (Eq a, Show a, Monoid a, Eq c, Show c, Monoid c) + => TestType a b + -> TestType c d + -> TestType (a, c) (These b d) + TTMThese :: (Eq a, Show a, Monoid a, Eq c, Show c, Monoid c) + => TestType a b + -> TestType c d + -> TestType (a, c) (Maybe (These b d)) + TTMaybe :: TestType a b + -> TestType (Maybe a) (Maybe b) + TTEither :: TestType a b + -> TestType c d + -> TestType (Either a c) (Either b d) + TTNEList :: TestType a b + -> TestType [a] (NonEmpty b) + TTCtx :: TestType (c -> t) (d -> u) + -> TestType a b + -> TestType (Context a c t) (Context b d u) + TTBazaar :: (Show a, Show b, Show c, Show d) + => GenType c d + -> TestType t u + -> TestType a b + -> TestType (Bazaar a c t) (Bazaar b d u) + (:*:) :: (Eq a, Eq b, Eq c, Eq d, Show a, Show b, Show c, Show d) + => TestType a b + -> TestType c d + -> TestType (a, c) (b, d) + (:?>) :: GenFunc f c d + -> TestType c d + -> TestType (f -> c) (f -> d) + (:->) :: (Show a, Show b) + => GenType a b + -> TestType c d + -> TestType (a -> c) (b -> d) + +infixr 2 :&: +infixr 1 :-> +infixr 1 :?> +infixr 2 :*: + +runSorter + :: SortType a + -> [a] + -> [a] +runSorter = \case + STAsc -> S.toAscList . S.fromList + STDesc -> S.toDescList . S.fromList + STDistinctAsc -> M.toAscList . M.fromList + STDistinctDesc -> M.toDescList . M.fromList + +runGT :: GenType a b -> Gen (a, b) +runGT = \case + GTNEMap -> (\n -> (NEM.IsNonEmpty n, n)) <$> neMapGen + GTMap -> join (,) <$> mapGen + GTNESet -> (\n -> (NES.IsNonEmpty n, n)) <$> neSetGen + GTNEIntMap -> (\n -> (NEIM.IsNonEmpty n, n)) <$> neIntMapGen + GTNEIntSet -> (\n -> (NEIS.IsNonEmpty n, n)) <$> neIntSetGen + GTIntMap -> join (,) <$> intMapGen + GTSet -> join (,) <$> setGen + GTIntSet -> join (,) <$> intSetGen + GTNESeq -> (\n -> (NESeq.IsNonEmpty n, n)) <$> neSeqGen + GTNESeqList -> (\n -> (toNonEmpty n, n)) <$> neSeqGen + GTSeq -> join (,) <$> seqGen + GTKey -> join (,) <$> keyGen + GTIntKey -> join (,) <$> intKeyGen + GTVal -> join (,) <$> valGen + GTSize -> join (,) <$> Gen.int mapSize + GTOther g -> join (,) <$> g + GTMaybe g -> maybe (Nothing, Nothing) (bimap Just Just) <$> + Gen.maybe (runGT g) + g1 :&: g2 -> do + (x1, y1) <- runGT g1 + (x2, y2) <- runGT g2 + pure ((x1,x2), (y1,y2)) + GTNEList r g -> first toList . NE.unzip <$> + Gen.nonEmpty (fromMaybe mapSize r) (runGT g) + GTSorted s g -> bimap (runSorter s) (fromJust . NE.nonEmpty . runSorter s . toList) <$> + runGT g + +runTT :: Monad m => TestType a b -> a -> b -> PropertyT m () +runTT = \case + TTNEMap -> \x y -> do + assert $ NEM.valid y + unKMap x === unKMap (NEM.IsNonEmpty y) + TTNEIntMap -> \x y -> do + assert $ NEIM.valid y + x === NEIM.IsNonEmpty y + TTNESet -> \x y -> do + assert $ NES.valid y + unKSet x === unKSet (NES.IsNonEmpty y) + TTNEIntSet -> \x y -> do + assert $ NEIS.valid y + x === NEIS.IsNonEmpty y + TTMap -> \x y -> + unKMap x === unKMap y + TTSet -> \x y -> + unKSet x === unKSet y + TTNESeq -> \x y -> + x === NESeq.IsNonEmpty y + TTNESeqList -> \x y -> + x === toNonEmpty y + TTKey -> \(K x1 y1) (K x2 y2) -> do + x1 === x2 + y1 === y2 + TTVal -> (===) + TTOther -> (===) + TTThese t1 t2 -> \(x1, x2) -> \case + This y1 -> do + runTT t1 x1 y1 + x2 === mempty + That y2 -> do + x1 === mempty + runTT t2 x2 y2 + These y1 y2 -> do + runTT t1 x1 y1 + runTT t2 x2 y2 + TTMThese t1 t2 -> \(x1, x2) -> \case + Nothing -> do + x1 === mempty + x2 === mempty + Just (This y1) -> do + runTT t1 x1 y1 + x2 === mempty + Just (That y2) -> do + x1 === mempty + runTT t2 x2 y2 + Just (These y1 y2) -> do + runTT t1 x1 y1 + runTT t2 x2 y2 + TTMaybe tt -> \x y -> do + isJust y === isJust y + traverse_ (uncurry (runTT tt)) $ liftA2 (,) x y + TTEither tl tr -> \case + Left x -> \case + Left y -> runTT tl x y + Right _ -> annotate "Left -> Right" *> failure + Right x -> \case + Left _ -> annotate "Right -> Left" *> failure + Right y -> runTT tr x y + TTNEList tt -> \xs ys -> do + length xs === length ys + zipWithM_ (runTT tt) xs (toList ys) + TTCtx tSet tView -> \(Context xS xV) (Context yS yV) -> do + runTT tSet xS yS + runTT tView xV yV + TTBazaar gNew tRes tView -> testBazaar gNew tRes tView + t1 :*: t2 -> \(x1, x2) (y1, y2) -> do + runTT t1 x1 y1 + runTT t2 x2 y2 + GF gt c :?> tt -> \gx gy -> do + f <- c <$> forAllFn (fn gt) + runTT tt (gx f) (gy f) + gt :-> tt -> \f g -> do + (x, y) <- forAll $ runGT gt + runTT tt (f x) (g y) + where + unKMap :: (Ord k, Ord j) => Map (K k j) c -> Map (k, j) c + unKMap = M.mapKeys (withK (,)) + unKSet :: (Ord k, Ord j) => Set (K k j) -> Set (k, j) + unKSet = S.map (withK (,)) + +testBazaar + :: forall a b c d t u m. (Show a, Show b, Show c, Show d, Monad m) + => GenType c d + -> TestType t u + -> TestType a b + -> Bazaar a c t + -> Bazaar b d u + -> PropertyT m () +testBazaar gNew tRes0 tView = go [] [] tRes0 + where + go :: [a] -> [b] -> TestType t' u' -> Bazaar a c t' -> Bazaar b d u' -> PropertyT m () + go xs ys tRes = \case + Done xRes -> \case + Done yRes -> do + annotate "The final result matches" + runTT tRes xRes yRes + More yView _ -> do + annotate "ys had more elements than xs" + annotate $ show xs + annotate $ show ys + annotate $ show yView + failure + More xView xNext -> \case + Done _ -> do + annotate "xs had more elements than ys" + annotate $ show xs + annotate $ show ys + annotate $ show xView + failure + More yView yNext -> do + annotate "Each individual piece matches pair-wise" + runTT tView xView yView + annotate "The remainders also match" + go (xView:xs) (yView:ys) (gNew :-> tRes) xNext yNext + + +-- --------------------- +-- Properties +-- --------------------- + +ttProp :: TestType a b -> a -> b -> Property +ttProp tt x = property . runTT tt x + +readShow + :: (Show a, Read a, Eq a) + => Gen a + -> Property +readShow g = property $ do + m0 <- forAll g + tripping m0 show readMaybe + +readShow1 + :: (Eq (f a), Show1 f, Show a, Show (f a), Read1 f, Read a) + => Gen (f a) + -> Property +readShow1 g = property $ do + m0 <- forAll g + tripping m0 (($ "") . showsPrec1 0) (fmap fst . listToMaybe . readsPrec1 0) + +showShow1 + :: (Show1 f, Show a, Show (f a)) + => Gen (f a) + -> Property +showShow1 g = property $ do + m0 <- forAll g + let s0 = show m0 + s1 = showsPrec1 0 m0 "" + s0 === s1 + +showShow2 + :: (Show2 f, Show a, Show b, Show (f a b)) + => Gen (f a b) + -> Property +showShow2 g = property $ do + m0 <- forAll g + let s0 = show m0 + s2 = showsPrec2 0 m0 "" + s0 === s2 + +-- readShow2 +-- :: (Eq (f a b), Show2 f, Show a, Show b, Show (f a b), Read2 f, Read a, Read b) +-- => Gen (f a b) +-- -> Property +-- readShow2 g = property $ do +-- m0 <- forAll g +-- tripping m0 (($ "") . showsPrec2 0) (fmap fst . listToMaybe . readsPrec2 0) + +-- --------------------- +-- Generators +-- --------------------- + +keyGen :: MonadGen m => m KeyType +keyGen = K <$> intKeyGen + <*> Gen.text (Range.linear 0 5) Gen.alphaNum + +valGen :: MonadGen m => m Text +valGen = Gen.text (Range.linear 0 5) Gen.alphaNum + +mapSize :: Range Int +mapSize = Range.exponential 4 8 + +mapGen :: MonadGen m => m (Map KeyType Text) +mapGen = Gen.map mapSize $ (,) <$> keyGen <*> valGen + +neMapGen :: MonadGen m => m (NEMap KeyType Text) +neMapGen = Gen.just $ NEM.nonEmptyMap <$> mapGen + +setGen :: MonadGen m => m (Set KeyType) +setGen = Gen.set mapSize keyGen + +neSetGen :: MonadGen m => m (NESet KeyType) +neSetGen = Gen.just $ NES.nonEmptySet <$> setGen + +intKeyGen :: MonadGen m => m Key +intKeyGen = Gen.int (Range.linear (-100) 100) + +intMapGen :: MonadGen m => m (IntMap Text) +intMapGen = IM.fromDistinctAscList . M.toList <$> Gen.map mapSize ((,) <$> intKeyGen <*> valGen) + +neIntMapGen :: MonadGen m => m (NEIntMap Text) +neIntMapGen = Gen.just $ NEIM.nonEmptyMap <$> intMapGen + +intSetGen :: MonadGen m => m IntSet +intSetGen = IS.fromDistinctAscList . S.toList <$> Gen.set mapSize intKeyGen + +neIntSetGen :: MonadGen m => m NEIntSet +neIntSetGen = Gen.just $ NEIS.nonEmptySet <$> intSetGen + +seqGen :: MonadGen m => m (Seq Text) +seqGen = Gen.seq mapSize valGen + +neSeqGen :: MonadGen m => m (NESeq Text) +neSeqGen = Gen.just $ NESeq.nonEmptySeq <$> seqGen + + + + + +-- --------------------- +-- Orphans +-- --------------------- + +instance Arg Char where + build = via ord chr + +instance Arg Text where + build = via T.unpack T.pack + +instance Vary Char where + vary = contramap ord vary + +instance Vary Text where + vary = contramap T.unpack vary +