diff --git a/AUTHORS b/AUTHORS new file mode 100644 --- /dev/null +++ b/AUTHORS @@ -0,0 +1,16 @@ +=== Haskell bytestring-trie package AUTHORS/THANKS file === + +The bytestring-trie package was written by wren gayle romano and is +released under the terms in the LICENSE file. I would also like to +give thanks to the following contributers: + +Maxime Henrion --- for the Binary (Trie a) instance and extensive + debugging work including almost all of the QuickCheck properties. + +Don Stewart --- for fostering the idea and offering feedback on + when the unsafe is safe in the internals of ByteStrings. + +Mark Wotton --- for benchmarking work comparing Trie to Map ByteString, + and for debugging. + +Gregory Crosswhite --- finding the critical bug in mergeBy diff --git a/CHANGELOG b/CHANGELOG new file mode 100644 --- /dev/null +++ b/CHANGELOG @@ -0,0 +1,104 @@ +0.2.7.6 (2025-02-11): + * Technical Debt: + - Relaxed upper-bounds for GHC 9.12 + - Silencing some warnings on GHC 9.10 +0.2.7.5 (2024-08-29): + * Technical Debt: + - Relaxed upper-bounds for the test suite +0.2.7.4 (2024-08-29): + * Technical Debt: + - Relaxed upper-bounds for GHC 9.10 + - Tightened lower-bounds to match only what CI still verifies. +0.2.7.3 (2024-04-05): + * Technical Debt: + - Updated for GHC 9.8 +0.2.7.2 (2023-03-19): + * Technical Debt: + - Updated for GHC 9.6 +0.2.7.1 (2022-08-28): + * Technical Debt: + - Updated for GHC 9.4 +0.2.7 (2022-03-14): + * Added: + - Instances: Eq1, Ord, Ord1, IsList, Show1, Read, Read1, NFData. + - Methods: Semigroup(stimes), Foldable(fold,foldMap') + - Functions: Data.Trie.Internal.{filter,filterA,wither} + * Performance Improvements: + - Significantly sped up `size` + - Misc other changes to the Foldable instance. + - Reordered data constructors, for better branch-prediction. + - Fixed quadratic slowdown for priority-queue functions. + - Reduced quadratic slowdown for functions that reconstruct keys; + however, it's still asymptotically quadratic: + + * Breakage: + - Changed the type of Data.Trie.Internal.lookupBy_ + * Technical Debt: + - Updated the *.cabal file to use Cabal-2.2 format/features. + - Greatly increased code-coverage of the test suite + - Added a bunch of benchmarks, including for regression testing. + - Added Safe/Trustworthy annotations. + - Cleaning up OPTIONS_GHC and imports. + - Cleaning up Haddock and adding more @since annotations. +0.2.6 (2021-11-22): + * Added: + - Data.Trie.minMatch + - Data.Trie.Internal.{cata,cata_} + - intersection functions (HT: Kevin Brubeck Unhammer) + - Data.Trie.deleteSubmap (HT: YongJoon Joe) + * Fixed: + - Adjusting the Applicative, Monad, and Monoid instances to + avoid warnings from GHC 9.2.1 regarding: + + + - Fixed a bug in mergeBy + * Documentation: + - Documented the order of results for `matches`/`matches_` + * Technical Debt: + - Numerous minor tweaks with smart-constructors + - Greatly increased code-coverage of the test suite +0.2.5.3 (2021-11-02): + * Fixed: + - Increasing upper bounds for GHC 9.2.1 +0.2.5.2 (2021-10-16): + * Fixed: + - Increasing upper bounds to Hackage latest, instead of just GHC 9.0.1 +0.2.5.1 (2021-10-16): + * Fixed: + - Fixing things to compile under GHC 9.0.1 +0.2.5.0 (2019-02-25): + * Added: + - Instance: Semigroup + * Fixed: + - Fixing things to compile under GHC 8.4 and 8.6. + * Technical Debt: + - Removed the obsolete/unused "useCinternal" Cabal flag +0.2.4.3 (2019-02-24): + * Technical Debt: + - Moved VERSION to CHANGELOG + * Fixed: + - Fixing things to compile under GHC 8.0 and 8.2. N.B., still + doesn't compile under 8.4 or 8.6, due to the version limit on `base`. +0.2.4.1 (2015-04-04): + * Fixed: + - Data.Trie.Internal: adjusted imports to compile under GHC 7.10.1 +0.2.4 (2014-10-09): + * Added: + - Data.Trie.Internal.{match_,matches_}, Data.Trie.{match,matches} +0.2.3 (2010-02-12): + * Added: + - Data.Trie.Internal.alterBy_ + - Data.Trie.Internal.{contextualMap, contextualMap', contextualFilterMap, contextualMapBy} + - Data.Trie.Convenience.{fromListWith', fromListWithL, fromListWithL'} as suggested by Ian Taylor + - Data.Trie.Convenience.{insertWith', insertWithKey', unionWith'} + * Performance: + - converted fmap, foldMap, traverse, and filterMap to worker/wrapper +0.2.2 (2010-06-10): + * Fixed: + - Corrected a major bug in mergeBy, reported by Gregory Crosswhite +0.2.1.1 (2009-12-20): + - Added a VERSION file +0.2.1 (2009-02-13): + - Most recent release before adding a VERSION file +0.1.4 (2009-01-11): + - The only previous tag diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,4 +1,4 @@ -Copyright (c) 2008--2009, wren ng thornton. +Copyright (c) 2008--2013, wren gayle romano. ALL RIGHTS RESERVED. Redistribution and use in source and binary forms, with or without diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,54 @@ +bytestring-trie +=============== +[![Hackage version](https://img.shields.io/hackage/v/bytestring-trie.svg?style=flat)](https://hackage.haskell.org/package/bytestring-trie) +[![Build Status](https://github.com/wrengr/bytestring-trie/workflows/ci/badge.svg)](https://github.com/wrengr/bytestring-trie/actions?query=workflow%3Aci) +[![Dependencies](https://img.shields.io/hackage-deps/v/bytestring-trie.svg?style=flat)](http://packdeps.haskellers.com/specific?package=bytestring-trie) + +The bytestring-trie package provides an efficient implementation +of tries mapping `ByteString` to values. The implementation is +based on Okasaki's big-endian patricia trees, à la `IntMap`. We +first trie on the elements of `ByteString` and then trie on the +big-endian bit representation of those elements. Patricia trees +have efficient algorithms for union and other merging operations, +but they're also quick for lookups and insertions. + +If you are only interested in being able to associate individual +`ByteString`s to values, then you may prefer the `hashmap` package +which is faster for those only needing a map-like structure. This +package is intended for those who need the extra capabilities that +a trie-like structure can offer (e.g., structure sharing to reduce +memory costs for highly redundant keys, taking the submap of all +keys with a given prefix, contextual mapping, extracting the minimum +and maximum keys, etc.) + + +## Install + +This is a simple package and should be easy to install. You should +be able to use the standard: + + $> cabal install bytestring-trie + + +## Portability + +The implementation is quite portable, relying only on a few basic +language extensions. The complete list of extensions used by the library is: + +* CPP +* MagicHash -- Only if using GHC +* NoImplicitPrelude + +The test suite uses a few additional extensions: + +* MultiParamTypeClasses +* FlexibleInstances +* FlexibleContexts + +## Links + +* [Website](http://wrengr.org/) +* [Blog](http://winterkoninkje.dreamwidth.org/) +* [Twitter](https://twitter.com/wrengr) +* [Hackage](http://hackage.haskell.org/package/bytestring-trie) +* [GitHub](https://github.com/wrengr/bytestring-trie) diff --git a/bytestring-trie.cabal b/bytestring-trie.cabal --- a/bytestring-trie.cabal +++ b/bytestring-trie.cabal @@ -1,38 +1,199 @@ +Cabal-Version: 2.2 +-- Cabal >=2.2 is required for: +-- +-- Since 2.1, the Cabal-Version must be the absolutely first thing +-- in the file, even before comments. Also, no longer uses ">=". +-- + ---------------------------------------------------------------- --- wren ng thornton ~ 2009.01.04 +-- wren gayle romano ~ 2025-02-11 ---------------------------------------------------------------- Name: bytestring-trie -Version: 0.1.4 -Cabal-Version: >= 1.2 +Version: 0.2.7.6 Build-Type: Simple -Stability: beta -Copyright: Copyright (c) 2008--2009 wren ng thornton -License: BSD3 +Stability: provisional +Homepage: https://wrengr.org/software/hackage.html +Bug-Reports: https://github.com/wrengr/bytestring-trie/issues +Author: wren gayle romano +Maintainer: wren@cpan.org +Copyright: 2008–2025 wren gayle romano +-- Cabal-2.2 requires us to say "BSD-3-Clause" not "BSD3" +License: BSD-3-Clause License-File: LICENSE -Author: wren ng thornton -Maintainer: wren@community.haskell.org -Homepage: http://code.haskell.org/~wren/ + Category: Data, Data Structures -Synopsis: An efficient finite map from (byte)strings to values. -Description: An efficient finite map from (byte)strings to values. - . - The implementation is based on big-endian patricia - trees, like "Data.IntMap". We first trie on the - elements of "Data.ByteString" and then trie on the - big-endian bit representation of those elements. - Patricia trees have efficient algorithms for union - and other merging operations, but they're also quick - for lookups and insertions. +Synopsis: An efficient finite map from bytestrings to values. +Description: An efficient finite map from bytestrings to values. + . + The implementation is based on big-endian patricia trees, like + "Data.IntMap". We first trie on the elements of "Data.ByteString" + and then trie on the big-endian bit representation of those + elements. Patricia trees have efficient algorithms for union + and other merging operations, but they're also quick for lookups + and insertions. + . + If you are only interested in being able to associate strings + to values, then you may prefer the @hashmap@ package which is + faster for those only needing a map-like structure. This package + is intended for those who need the extra capabilities that a + trie-like structure can offer (e.g., structure sharing to reduce + memory costs for highly redundant keys, taking the submap of + all keys with a given prefix, contextual mapping, extracting + the minimum and maximum keys, etc.) +Extra-source-files: + AUTHORS, CHANGELOG, README.md + +-- This package ought to work as far back as GHC 7.4.1; however, +-- I've tightened the lower bounds on the Build-Depends to match +-- only what we still verify via CI: +-- +-- +-- +Tested-With: + GHC ==8.0.2, + GHC ==8.2.2, + GHC ==8.4.4, + GHC ==8.6.5, + GHC ==8.8.4, + GHC ==8.10.3, + GHC ==9.0.1, + GHC ==9.2.4, + GHC ==9.4.8, + GHC ==9.6.5, + GHC ==9.8.2, + GHC ==9.10.1, + GHC ==9.12.1 + +Source-Repository head + Type: git + Location: https://github.com/wrengr/bytestring-trie.git + +---------------------------------------------------------------- +-- This stanza requires Cabal>=2.2: +-- +-- While Cabal-2.2 only ships with GHC 8.4.1, the dependencies to +-- build it have essentially the same lower bounds as we do. (They +-- require bytestring>=0.9.2.1 and deepseq>=1.3) So users of older +-- GHC should still be able to compile it; and if they can't, then +-- they already can't compile this package. +-- +-- N.B., the "import:" field must be the first thing in a stanza. +Common library-build-depends + Default-Language: Haskell2010 + Build-Depends: base >= 4.9 && < 4.22 + , bytestring >= 0.10.8 && < 0.13 + , binary >= 0.8.3 && < 0.9 + , deepseq >= 1.4.2 && < 1.6 + +-- TODO: in lieu of using CPP to expose internals to the tests/benchmarks, +-- we should consider using: +-- + Library + Import: library-build-depends Hs-Source-Dirs: src Exposed-Modules: Data.Trie - , Data.Trie.Internal , Data.Trie.Convenience - Other-Modules: Data.Trie.BitTwiddle - , Data.Trie.ByteStringInternal - Build-Depends: base, bytestring, binary + , Data.Trie.Internal + Other-Modules: Data.Trie.Internal.BitTwiddle + , Data.Trie.Internal.ByteString + , Data.Trie.Internal.Errors + +---------------------------------------------------------------- +-- See the cabal file for bytestring-lexing for more info about +-- setting up Tasty. +Test-Suite test-all + Import: library-build-depends + Type: exitcode-stdio-1.0 + Hs-Source-Dirs: dev + -- NOTE: Test-Suite Main-Is is relative to Hs-Source-Dirs; + -- unlike Executable Main-Is, which is relative to this cabal + -- file's directory instead. + Main-Is: Test/Main.hs + Other-Modules: Shared.BaseCompat + , Test.Utils + , Test.Properties + -- Test.Validity + -- Test.ByteStringInternal + -- We must include our own library for the tests to use it; but + -- we must not give a version restriction lest Cabal give warnings. + -- There's also bug : + -- if we don't pass -any, then Cabal will fill in ">= 0 && <= $ThisVersion" + -- which will also give a warning. + Build-Depends: bytestring-trie -any + , tasty >= 0.10.1.2 && < 1.6 + , tasty-smallcheck >= 0.8.0.1 && < 0.9 + , tasty-quickcheck >= 0.8.3.2 && < 0.12 + -- N.B., @tasty-hunit@ is just a partial API clone; + -- whereas @tasty-hunit-compat@ is a proper integration + -- with HUnit itself. Also, tasty-hunit-compat actually + -- depends on tasty-hunit; if you're wanting to minimize + -- dependencies. + , tasty-hunit < 0.11 + , QuickCheck >= 2.10 && < 2.16 + , smallcheck >= 1.1.1 && < 1.3 + -- lazysmallcheck >= 0.6 && < 0.7 + +---------------------------------------------------------------- +-- Can't put the "Type:" field in here; or rather even if we do, +-- it's still required in each of the Benchmark stanzas... +Common bench-common + Import: library-build-depends + if impl(ghc) + GHC-Options: -with-rtsopts=-A32m + -- TODO: this was recommended somewhere for benchmarks, but is + -- --nonmoving-gc actually good for our use case? (It's + -- concurrent-mark&sweep for the old generations; instead of + -- the default stop-the-world generational copying collector) + --if impl(ghc >= 8.10) + -- GHC-Options: -with-rtsopts=--nonmoving-gc + if impl(ghc) + GHC-Options: -Wall -O2 -rtsopts + Hs-Source-Dirs: dev + Build-Depends: bytestring-trie -any + -- TODO: try using @gauge@ instead of @criterion@, + -- to reduce the transitive dependencies. + -- BUG: @gauge@ depends on @basement@ which as of + -- version 0.0.12 doesn't support GHC 9.2; so we'll + -- have to revisit this later. + -- + -- NOTE: If you're having issues building on ghc-9.6, + -- then be sure to update your cabal info. + -- The parallel-3.2.2.0 package has metadata revisions + -- to allow building with base-4.18. + , criterion + , QuickCheck >= 2.10 && < 2.16 + +Benchmark bench-Regression + Import: bench-common + Type: exitcode-stdio-1.0 + -- NOTE: Benchmark Main-Is behaves like Test-Suite Main-Is + -- (not like Executable Main-Is). + Main-Is: Bench/Regression.hs + if impl(ghc) + GHC-Options: -main-is Bench.Regression.main + Other-Modules: Shared.BaseCompat + , Bench.Foldable + +Benchmark bench-Foldable + Import: bench-common + Type: exitcode-stdio-1.0 + Main-Is: Bench/Foldable.hs + if impl(ghc) + GHC-Options: -main-is Bench.Foldable.main + Other-Modules: Shared.BaseCompat + +Benchmark bench-MatchOne + Import: bench-common + Type: exitcode-stdio-1.0 + Main-Is: Bench/MatchOne.hs + +Benchmark bench-UnionWith + Import: bench-common + Type: exitcode-stdio-1.0 + Main-Is: Bench/UnionWith.hs ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/dev/Bench/Foldable.hs b/dev/Bench/Foldable.hs new file mode 100644 --- /dev/null +++ b/dev/Bench/Foldable.hs @@ -0,0 +1,516 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP, BangPatterns #-} + +---------------------------------------------------------------- +-- ~ 2022.03.04 +-- | +-- Module : Bench.Foldable +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : portable (with bang-patterns) +-- +-- Benchmarking definitions for 'Foldable' +---------------------------------------------------------------- + +module Bench.Foldable (main, realTrie_to_benchTrie, bgroup_Foldable) where + +import Shared.BaseCompat +import qualified Data.Trie as T +import qualified Data.Trie.Internal as TI + +import qualified Data.ByteString as S +#if MIN_VERSION_base(4,7,0) +-- [GHC 7.8.1]: "Data.Coerce" added to base. +import Data.Coerce (Coercible, coerce) +#endif +#if MIN_VERSION_base(4,9,0) +-- [GHC 8.0.1]: "Data.Semigroup" added to base. +import Data.Semigroup (Endo(..)) +#endif +-- TODO: version limits for 'Dual'? +import Data.Monoid (Dual(..)) +import Data.Word (Word8) +import Control.DeepSeq (NFData(rnf)) +import qualified Test.QuickCheck as QC +import qualified Criterion.Main as C +-- TODO: consider using the @gauge@ library instead of @criterion@. +-- It's a clone of criterion, also by BOS; with the intention of +-- having reduced dependencies. It doesn't do the HTML output, but +-- otherwise should be fine for our needs. It's what the @containers@ +-- library is using these days. +-- BUG: @gauge@ depends on @basement@ which as of version 0.0.12 +-- doesn't support GHC 9.2; so we'll have to revisit this later. +---------------------------------------------------------------- + +-- | Reference definition, to avoid needing to expose constructors +-- of the real definition. +-- +-- TODO: should maybe have a CPP macro to expose them for benchmarking, +-- to ensure this file stays in sync! +data Trie a + = Branch {-# UNPACK #-} !Word8 -- Prefix + {-# UNPACK #-} !Word8 -- Mask + !(Trie a) + !(Trie a) + | Arc {-# UNPACK #-} !S.ByteString + !(Maybe a) + !(Trie a) + | Empty + deriving Show + +-- Needed for 'C.env' +instance NFData a => NFData (Trie a) where + rnf Empty = () + rnf (Arc _ mv t) = rnf mv `seq` rnf t + rnf (Branch _ _ l r) = rnf l `seq` rnf r + +---------------------------------------------------------------- +-- | From "Data.Functor.Utils", but isn't exported. Is used heavily +-- by the default implementations, since they use so many newtype +-- wrappers. +(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c) +(#.) _f = coerce +{-# INLINE (#.) #-} + +---------------------------------------------------------------- +fold_foldMap, fold_foldrCompose, fold_foldrEta, fold_v027 + :: Monoid m => Trie m -> m + +-- | The default 'fold' definition as of base-4.16.0.0 +fold_foldMap = foldMap_v027 id + +-- | Default 'fold' via default 'foldMap' (using 'foldr_compose'), +-- but inlining the 'id' away. +-- Far worse than 'fold_foldMap', both for runtime and allocation. +fold_foldrCompose = foldr_compose mappend mempty + +-- | Default 'fold' via default 'foldMap' (using 'foldr_eta'), +-- but inlining the 'id' away. +-- This is somewhat worse than 'fold_foldrCompose' even. +fold_foldrEta = foldr_eta mappend mempty + +-- | bytestring-trie-0.2.7 definition. +-- This is a clear win over 'fold_foldMap': about half the runtime, +-- and one eighth the allocations. +fold_v027 = go + where + go Empty = mempty + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = v `mappend` go t + go (Branch _ _ l r) = go l `mappend` go r + +---------------------------------------------------------------- +foldMap_foldrCompose, foldMap_foldrEta, foldMap_v027 + :: Monoid m => (a -> m) -> Trie a -> m + +-- | The default 'fold' definition as of base-4.16.0.0 (using 'foldr_compose'). +foldMap_foldrCompose f = foldr_compose (mappend . f) mempty +-- | The default 'fold' definition as of base-4.16.0.0 (using 'foldr_eta'). +foldMap_foldrEta f = foldr_eta (mappend . f) mempty + +-- | bytestring-trie-0.2.7 definition, also used elsewhere for defaults. +-- (This was also the bytestring-trie-0.2.6 definition; the only +-- thing defined for that version.) +-- Again, this is the clear winner over the default implementation. +foldMap_v027 f = go + where + go Empty = mempty + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = f v `mappend` go t + go (Branch _ _ l r) = go l `mappend` go r + +---------------------------------------------------------------- +foldMap'_default, foldMap'_v027, foldMap'_v027_cps, foldMap'_v027_flop, foldMap'_v027_bang, foldMap'_v027_flop_bang + :: Monoid m => (a -> m) -> Trie a -> m + +-- | default definition. +foldMap'_default f = foldl'_v027 (\ acc a -> acc <> f a) mempty + +-- | bytestring-trie-0.2.7 definition. +-- Again, this is the clear winner over the default implementation. +foldMap'_v027 f = go mempty + where + go !m Empty = m + go m (Arc _ Nothing t) = go m t + go m (Arc _ (Just v) t) = go (m `mappend` f v) t + go m (Branch _ _ l r) = go (go m l) r + +-- | CPS variant, to ensure proper tail-recursion. +-- In between 'foldMap'_v027' and 'foldMap'_default'. +foldMap'_v027_cps f = \t -> go t mempty id + where + go Empty !m c = c m + go (Arc _ Nothing t) m c = go t m c + go (Arc _ (Just v) t) m c = go t (m `mappend` f v) c + go (Branch _ _ l r) m c = go l m (\m' -> go r m' c) + +-- | Based on our experience with 'foldl_v027_flop', trying the flop here. +-- On the small tries tested here, is slightly slower than +-- 'foldMap'_v027'; identical allocation. On larger tries this +-- gets to be slightly better than 'foldMap'_v027'. +foldMap'_v027_flop f = \t -> go t mempty + where + go Empty !m = m + go (Arc _ Nothing t) m = go t m + go (Arc _ (Just v) t) m = go t (m `mappend` f v) + go (Branch _ _ l r) m = go r (go l m) + +-- | This really shouldn't matter, but given our experience with 'foldl'_v027_flop_bang'... +-- And yet, is slightly better than 'foldMap'_v027_flop', slightly +-- worse than 'foldMap'_v027_flop_bang'; identical allocation for +-- all these variants. +foldMap'_v027_bang f = go mempty + where + go !m Empty = m + go m (Arc _ Nothing t) = go m t + go m (Arc _ (Just v) t) = (go $! mappend m (f v)) t + go m (Branch _ _ l r) = (go $! go m l) r + +-- | This really shouldn't matter, but given our experience with 'foldl'_v027_flop_bang'... +-- And yet, is slightly better than 'foldMap'_v027_bang', slightly +-- worse than 'foldMap'_v027'; identical allocation for all these +-- variants. For larger tries, is slightly slower than 'foldMap'_v027_flop'. +foldMap'_v027_flop_bang f = \t -> go t mempty + where + go Empty !m = m + go (Arc _ Nothing t) m = go t m + go (Arc _ (Just v) t) m = go t $! mappend m (f v) + go (Branch _ _ l r) m = go r $! go l m + +---------------------------------------------------------------- +foldr_default, foldr_compose, foldr_eta, foldr_cps_eta, foldr_cps, foldr_noClosure + :: (a -> b -> b) -> b -> Trie a -> b + +-- | The default definition as of base-4.16.0.0 +-- Actually a pretty solid baseline, for small tries. And for +-- larger tries it's far and away the clear winner. +foldr_default f z t = + appEndo (foldMap_v027 (Endo #. f) t) z + +-- | bytestring-trie-0.2.7 definition +-- Identical allocation as EndoDefault; about the same speed, or a +-- slightly slower; but only for the small tries tested here. For +-- larger tries it's far worse than 'foldr_default'. +foldr_compose f z0 = \t -> go t z0 -- eta for better inlining. + where + go Empty = id + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = f v . go t + go (Branch _ _ l r) = go l . go r + +-- | +-- Usually little faster than EndoDefault, occasionally about the +-- same, rarely slower. However, this allocates (consistently) ~43% more. +-- [I'm thinking because (a) larger thunks to store the @z@, but +-- (b) fully-saturated thunks so faster apply?] +-- For larger tries, is also far worse than 'foldr_default'. +foldr_eta f z0 = \t -> go t z0 -- eta for better inlining. + where + go Empty z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = f v (go t z) + go (Branch _ _ l r) z = go l (go r z) + +-- | Much slower; also allocates much more. (For larger tries too) +foldr_cps_eta f z0 = \t -> go t id z0 -- eta for better inlining. + where + go Empty c z = c z + go (Arc _ Nothing t) c z = go t c z + go (Arc _ (Just v) t) c z = go t (c . f v) z + go (Branch _ _ l r) c z = go r (go l c) z + +-- | Even slower; and even more allocation. (For larger tries too) +foldr_cps f z0 = \t -> go t id z0 -- eta for better inlining. + where + go Empty c = c + go (Arc _ Nothing t) c = go t c + go (Arc _ (Just v) t) c = go t (c . f v) + go (Branch _ _ l r) c = go r (go l c) + +-- | Faster than 'foldr_cps_eta', but still much slower than +-- EndoDefault/compose/eta; Also allocates about as much as +-- 'foldr_cps_eta'. +foldr_noClosure _ z Empty = z +foldr_noClosure f z (Arc _ Nothing t) = foldr_noClosure f z t +foldr_noClosure f z (Arc _ (Just v) t) = f v (foldr_noClosure f z t) +foldr_noClosure f z (Branch _ _ l r) = + foldr_noClosure f (foldr_noClosure f z r) l + +---------------------------------------------------------------- +foldr'_default, foldr'_v027, foldr'_v027_cps, foldr'_cps, foldr'_v027_flop, foldr'_v027_bang, foldr'_v027_flop_bang + :: (a -> b -> b) -> b -> Trie a -> b + +-- | The default definition as of base-4.16.0.0 (including the phrasing +-- of using a where-clause rather than a lambda). +-- The worst of the lot. +foldr'_default f z0 xs = + foldl_v027 f' id xs z0 + where f' k x z = k $! f x z + +-- | bytestring-trie-0.2.7 definition. +-- This one is the clear winner, both for time and allocation. +foldr'_v027 f z0 = go z0 -- eta for better inlining + where + go !z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = f v $! go z t + go z (Branch _ _ l r) = go (go z r) l + +-- TODO: (2022.03.05) Try a variant of 'foldr'_v027' which doesn't +-- closure over @f@. + +-- | Worse than 'foldr'_v027' but still better than the rest. +foldr'_v027_cps f z0 = \t -> go t z0 id -- eta for better inlining + where + go Empty !z c = c z + go (Arc _ Nothing t) z c = go t z c + go (Arc _ (Just v) t) z c = go t z (\ !z' -> c $! f v z') + go (Branch _ _ l r) z c = go r z (\ !z' -> go l z' c) + +-- | Worse than 'foldr'_v027_cps' but still better than 'foldr'_default' +foldr'_cps f z0 = \t -> go t id z0 -- eta for better inlining + where + go Empty c = c + go (Arc _ Nothing t) c = go t c + go (Arc _ (Just v) t) c = go t (\ !z -> c $! f v z) + go (Branch _ _ l r) c = go r (go l c) + +-- | Based on our experience with 'foldl_v027_flop', trying the flop here. +-- Actually faster than 'foldr'_v027' (but only for small tries); +-- identical allocation. For larger tries, this is a bit slower +-- than 'foldr'_v027'. +foldr'_v027_flop f z0 = \t -> go t z0 -- eta for better inlining + where + go Empty !z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = f v $! go t z + go (Branch _ _ l r) z = go l (go r z) + +-- | This really shouldn't matter, but given our experience with 'foldl'_v027_flop_bang'... +-- And yet, faster than 'foldr'_cps', slower than 'foldr'_v027_flop'. +foldr'_v027_bang f z0 = go z0 -- eta for better inlining + where + go !z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = f v $! go z t + go z (Branch _ _ l r) = (go $! go z r) l + +-- | This really shouldn't matter, but given our experience with 'foldl'_v027_flop_bang'... +-- And yet, marginally faster than 'foldr'_v027_flop', more +-- significantly slower than 'foldr'_v027'. +foldr'_v027_flop_bang f z0 = \t -> go t z0 -- eta for better inlining + where + go Empty !z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = f v $! go t z + go (Branch _ _ l r) z = go l $! go r z + +---------------------------------------------------------------- +foldl_default, foldl_default_Coerce, foldl_v027, foldl_v027_flop + :: (b -> a -> b) -> b -> Trie a -> b + +-- | The default definition as of base-4.16.0.0. +foldl_default f z t = + appEndo (getDual (foldMap_v027 (Dual . Endo . flip f) t)) z + +-- | Variant of the default that uses @('#.')@ in lieu of @('.')@. +-- This one performs massively better than 'foldl_default'. +-- TODO: send a patch to @base@ for this. +-- Hrm, for larger tries it looks like the benefit is only minimal... +foldl_default_Coerce f z t = + appEndo (getDual (foldMap_v027 (Dual #. Endo #. flip f) t)) z + +-- | bytestring-trie-0.2.7 definition, also used for defaults. +-- This one is faster than 'foldl_default_Coerce'; though it allocates +-- ~56% more than 'foldl_default_Coerce' (albeit ~11% /less/ than +-- 'foldl_default'). +foldl_v027 f z0 = go z0 -- eta for better inlining + where + go z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = go (f z v) t + go z (Branch _ _ l r) = go (go z l) r + +-- | Since 'foldr_eta' allocates ~43% more than 'foldr_default', +-- yet 'foldl_v027' allocates ~56% more: see if swapping the order +-- of arguments in the recursion changes that. +-- For small tries: +-- * Marginally slower than 'foldl_v027' (<1%). +-- * Allocates ~8% less than 'foldl_v027'; ~19% less than 'foldl_default'; +-- * albeit still ~42% more than 'foldl_default_Coerce'. +-- For larger tries: far and away the best of the lot. +foldl_v027_flop f z0 = \t -> go t z0 -- eta for better inlining + where + go Empty z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = go t (f z v) + go (Branch _ _ l r) z = go r (go l z) + +-- We could try CPSing to restore the tail-recursion; but since +-- that's worse everywhere else, there seems little point. + +---------------------------------------------------------------- +foldl'_defaultCompose, foldl'_defaultEta, foldl'_v027, foldl'_v027_flop, foldl'_v027_flop_bang + :: (b -> a -> b) -> b -> Trie a -> b + +-- | The default definition as of base-4.16.0.0 (using 'foldr_compose'). +foldl'_defaultCompose f z0 xs = + foldr_compose f' id xs z0 + where f' x k z = k $! f z x + +-- | The default definition as of base-4.16.0.0 (using 'foldr_eta'). +-- Much worse than 'foldl'_defaultCompose'. +foldl'_defaultEta f z0 xs = + foldr_eta f' id xs z0 + where f' x k z = k $! f z x + +-- | bytestring-trie-0.2.7 definition, also used for defaults. +-- Clear winner, for small tries. For large tries it's about halfway +-- between the two defaults above, and the two flops below. +foldl'_v027 f z0 = go z0 -- eta for better inlining + where + go !z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = go (f z v) t + go z (Branch _ _ l r) = go (go z l) r + +-- | For small tries: worse than 'foldl'_v027': ~10% slower, ~119% +-- more allocation. But for large tries it's much faster. +foldl'_v027_flop f z0 = \t -> go t z0 -- eta for better inlining + where + go Empty !z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = go t (f z v) + go (Branch _ _ l r) z = go r (go l z) + +-- FIXME: why on earth does this perform better?! Why aren't bang +-- patterns giving me the same thing?! +-- +-- | For small tries: ~4.4% faster than 'foldl'_v027'; but allocates +-- ~32% more still. For large tries it's only marginally faster +-- than 'foldl'_v027_flop' (which is to say massively faster than +-- 'foldl'_v027') +foldl'_v027_flop_bang f z0 = \t -> go t z0 -- eta for better inlining + where + go Empty !z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = go t $! f z v + go (Branch _ _ l r) z = go r $! go l z + +-- TODO: CPS to restore the purely tail-call for @Branch@? + +---------------------------------------------------------------- +{- +-- TODO: the base-4.16.0.0 defaults are shown +{-# INLINE toList #-} toList t = build (\ c n -> foldr c n t) +null = foldr (\_ _ -> False) True +length = foldl' (\c _ -> c+1) 0 +elem = any . (==) +-- With standalone definition: any p = getAny #. foldMap (Any #. p) + +{-# INLINEABLE maximum #-} +maximum = fromMaybe (errorWithoutStackTrace "maximum: empty structure") . + getMax . foldMap' (Max #. (Just :: a -> Maybe a)) -- -XScopedTypeVariables +{-# INLINEABLE minimum #-} +minimum = fromMaybe (errorWithoutStackTrace "minimum: empty structure") . + getMin . foldMap' (Min #. (Just :: a -> Maybe a)) -- -XScopedTypeVariables +{-# INLINEABLE sum #-} +sum = getSum #. foldMap' Sum +{-# INLINEABLE product #-} +product = getProduct #. foldMap' Product +-} + + + + +---------------------------------------------------------------- +---------------------------------------------------------------- +-- Definitely not the most efficient thing... +-- TODO: define an anamorphism function in "TI" to use here... +arbitraryTrie :: Int -> Int -> QC.Gen (Trie Int) +arbitraryTrie maxK maxL = do + k <- QC.chooseInt (0, maxK) + keys <- QC.vectorOf k $ do + l <- QC.chooseInt (0, maxL) + xs <- QC.vector l + return $ S.pack xs + return . realTrie_to_benchTrie . T.fromList $ zip keys [0..] + +-- TODO: still gotta reconstruct the 'Prefix' and 'Mask'. +-- (Even though none of the folds depend on them...) +realTrie_to_benchTrie :: T.Trie a -> Trie a +realTrie_to_benchTrie = TI.cata_ Arc (Branch 0 0) Empty + +---------------------------------------------------------------- +intToSum :: (Trie (Sum Int) -> a) -> [Trie Int] -> [a] +intToSum f = fmap f . coerce + +-- BUG: since we started generating a collection of tries, we've +-- started to see different results: each test progressively slower +-- than the last. Seems bogus, but not sure what's up since criterion +-- ought to be forcing all of @ts@ early and doing GC between each +-- benchmark... +main :: IO () +main = C.defaultMain + [ C.env (QC.generate $ QC.vectorOf 10 $ arbitraryTrie 30 10) $ \ ts -> + bgroup_Foldable ts + ] + +bgroup_Foldable :: [Trie Int] -> C.Benchmark +bgroup_Foldable ts = + C.bgroup "Foldable" + [ C.bgroup "fold" + [ C.bench "default (foldMap)" $ C.nf (intToSum fold_foldMap ) ts + , C.bench "default (foldr_compose)" $ C.nf (intToSum fold_foldrCompose) ts + , C.bench "default (foldr_eta)" $ C.nf (intToSum fold_foldrEta ) ts + , C.bench "v0.2.7" $ C.nf (intToSum fold_v027 ) ts + ] + , C.bgroup "foldMap" + [ C.bench "default (foldr_compose)" $ C.nf (intToSum (foldMap_foldrCompose id)) ts + , C.bench "default (foldr_eta)" $ C.nf (intToSum (foldMap_foldrEta id)) ts + , C.bench "v0.2.7" $ C.nf (intToSum (foldMap_v027 id)) ts + ] + , C.bgroup "foldMap'" + [ C.bench "default (foldl')" $ C.nf (intToSum (foldMap'_default id)) ts + , C.bench "v0.2.7" $ C.nf (intToSum (foldMap'_v027 id)) ts + , C.bench "v0.2.7 +cps" $ C.nf (intToSum (foldMap'_v027_cps id)) ts + , C.bench "v0.2.7 +flopped" $ C.nf (intToSum (foldMap'_v027_flop id)) ts + , C.bench "v0.2.7 +($!)" $ C.nf (intToSum (foldMap'_v027_bang id)) ts + , C.bench "v0.2.7 +flopped +($!)" $ C.nf (intToSum (foldMap'_v027_flop_bang id)) ts + ] + , C.bgroup "foldr" + [ C.bench "default (foldMap)" $ C.nf (foldr_default (+) 0 <$>) ts + , C.bench "compose" $ C.nf (foldr_compose (+) 0 <$>) ts + , C.bench "eta" $ C.nf (foldr_eta (+) 0 <$>) ts + , C.bench "cps_eta" $ C.nf (foldr_cps_eta (+) 0 <$>) ts + , C.bench "cps" $ C.nf (foldr_cps (+) 0 <$>) ts + , C.bench "noClosure" $ C.nf (foldr_noClosure (+) 0 <$>) ts + ] + , C.bgroup "foldr'" + [ C.bench "default" $ C.nf (foldr'_default (+) 0 <$>) ts + , C.bench "v0.2.7" $ C.nf (foldr'_v027 (+) 0 <$>) ts + , C.bench "v0.2.7 +cps" $ C.nf (foldr'_v027_cps (+) 0 <$>) ts + , C.bench "cps, no eta" $ C.nf (foldr'_cps (+) 0 <$>) ts + , C.bench "v0.2.7 +flopped" $ C.nf (foldr'_v027_flop (+) 0 <$>) ts + , C.bench "v0.2.7 +($!)" $ C.nf (foldr'_v027_bang (+) 0 <$>) ts + , C.bench "v0.2.7 +flopped +($!)" $ C.nf (foldr'_v027_flop_bang (+) 0 <$>) ts + ] + , C.bgroup "foldl" + [ C.bench "default (foldMap)" $ C.nf (foldl_default (+) 0 <$>) ts + , C.bench "default +Coerce" $ C.nf (foldl_default_Coerce (+) 0 <$>) ts + , C.bench "v0.2.7" $ C.nf (foldl_v027 (+) 0 <$>) ts + , C.bench "v0.2.7 +flopped" $ C.nf (foldl_v027_flop (+) 0 <$>) ts + ] + , C.bgroup "foldl'" + [ C.bench "default (foldr_compose)" $ C.nf (foldl'_defaultCompose (+) 0 <$>) ts + , C.bench "default (foldr_eta)" $ C.nf (foldl'_defaultEta (+) 0 <$>) ts + , C.bench "v0.2.7" $ C.nf (foldl'_v027 (+) 0 <$>) ts + , C.bench "v0.2.7 +flopped" $ C.nf (foldl'_v027_flop (+) 0 <$>) ts + , C.bench "v0.2.7 +flopped +($!)" $ C.nf (foldl'_v027_flop_bang (+) 0 <$>) ts + ] + ] + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Bench/MatchOne.hs b/dev/Bench/MatchOne.hs new file mode 100644 --- /dev/null +++ b/dev/Bench/MatchOne.hs @@ -0,0 +1,124 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP, BangPatterns #-} + +---------------------------------------------------------------- +-- ~ 2023.03.19 +-- | +-- Module : Bench.MatchOne +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : portable +-- +-- Benchmarking definitions for 'TI.match_' +---------------------------------------------------------------- + +module Main (main) where + +import qualified Data.Trie as T +import qualified Data.Trie.Internal as TI +import qualified Data.ByteString as S +#if !(MIN_VERSION_base(4,18,0)) +-- [base-4.18 / GHC 9.6.1]: Prelude now re-exports 'liftA2'. +import Control.Applicative (liftA2) +#endif +import Control.DeepSeq (NFData) +import qualified Test.QuickCheck as QC +import qualified Criterion.Main as C +-- TODO: consider trying . Especially so we can add tests to ensure that all these implementations are correct; but also for the comparisons across versions (since @bench-show@ and @benchgraph@ are both too heavy-weight for our needs). +---------------------------------------------------------------- + +-- Using NOINLINE to improve the reliability\/quality of the benchmarks. +match_v027, match_foldl, match_foldr, match_foldr_alt + :: T.Trie a -> S.ByteString -> Maybe (Int, a) + +-- The manual implementation. +{-# NOINLINE match_v027 #-} +match_v027 = TI.match_ + +-- This implementation is based on @base-4.16.0.0@:'GHC.List.last'. +-- N.B., the implementation uses 'foldl' in order to be a good +-- consumer, contrary to the 'foldr' we would have expected; see: +-- +-- Also, the actual definition of 'GHC.List.last' must be eta-expanded +-- in order to actually have 'foldl' make it a good consumer; see: +-- +-- (Not that that's relevant for us, since we're inlining the +-- definition anyways, and the call-site is indeed saturated.) +-- +-- TODO: since when does 'foldl' make good consumers? Did GHC +-- switch from build\/foldr to unfold\/destroy? +-- +-- Per the benchmarks below, this one is by far the slowest; thereby +-- suggesting that 'foldl' is not in fact a good consumer! +{-# NOINLINE match_foldl #-} +match_foldl t q = foldl (\_ x -> Just x) Nothing (TI.matches_ t q) + +-- This uses the definition of 'GHC.List.last' prior to #9339. Note +-- how @step@ returns the @Just@ immediately, and only lazily does +-- the case analysis. +-- +-- Per the benchmarks below, this is better than 'match_foldl' but +-- not as good as 'match_foldr_alt' +{-# NOINLINE match_foldr #-} +match_foldr t q = foldr step Nothing (TI.matches_ t q) + where + step x m = Just (case m of { Nothing -> x; Just y -> y }) + +-- And here's a version that doesn't do that @Just@-before-case... +-- +-- Per the benchmarks below, this is better than 'match_foldr' but +-- still not as good as 'match_v027'. +{-# NOINLINE match_foldr_alt #-} +match_foldr_alt t q = foldr step Nothing (TI.matches_ t q) + where + step x Nothing = Just x + step _ y = y + +-- TODO: maybe try a Codensity version of the above two, to avoid +-- redundant case analysis on the intermediate 'Maybe' + +---------------------------------------------------------------- +-- TODO: move this stuff off to the shared file. +-- TODO: should have argument to bound the 'Word8' +arbitraryBS :: Int -> QC.Gen S.ByteString +arbitraryBS maxL = do + l <- QC.chooseInt (0, maxL) + xs <- QC.vector l + return $! S.pack xs + +arbitraryTrie :: Int -> Int -> QC.Gen (T.Trie Int) +arbitraryTrie maxK maxL = do + k <- QC.chooseInt (0, maxK) + keys <- QC.vectorOf k $ arbitraryBS maxL + return $! T.fromList (zip keys [0..]) + +-- TODO: really need a better environment to work on than this... +generateEnv :: IO ([T.Trie Int], [S.ByteString]) +generateEnv = QC.generate $ do + ts <- QC.vectorOf 10 $ arbitraryTrie 30 10 + qs <- QC.vectorOf 10 $ arbitraryBS 12 + return (ts, qs) + +cartesianNF :: NFData c => (a -> b -> c) -> [a] -> [b] -> C.Benchmarkable +cartesianNF f xs ys = C.nf (liftA2 f xs) ys + +---------------------------------------------------------------- +-- TODO: +-- uses 'Control.Exception.evaluate' instead of 'C.env'. Is that +-- because of using the @gauge@ library instead of @criterion@, or +-- is there some other reason? +main :: IO () +main = C.defaultMain + [ C.env generateEnv $ \ ~(ts, qs) -> + C.bgroup "MatchOne" + [ C.bench "match_v027" $ cartesianNF match_v027 ts qs + , C.bench "match_foldl" $ cartesianNF match_foldl ts qs + , C.bench "match_foldr" $ cartesianNF match_foldr ts qs + , C.bench "match_foldr_alt" $ cartesianNF match_foldr_alt ts qs + ] + ] + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Bench/Regression.hs b/dev/Bench/Regression.hs new file mode 100644 --- /dev/null +++ b/dev/Bench/Regression.hs @@ -0,0 +1,145 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP, BangPatterns #-} + +---------------------------------------------------------------- +-- ~ 2022.04.03 +-- | +-- Module : Bench.Regression +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : portable (with bang-patterns) +-- +---------------------------------------------------------------- + +module Bench.Regression (main) where + +import Shared.BaseCompat +import qualified Data.Trie as T +-- import qualified Data.Trie.Internal as TI +import qualified Data.ByteString as S + +#if MIN_VERSION_base(4,7,0) +-- [GHC 7.8.1]: "Data.Coerce" added to base. +import Data.Coerce (Coercible, coerce) +#endif +import qualified Data.Foldable as F +import qualified Criterion.Main as C +import qualified System.IO as Sys (withFile, IOMode(..), hIsEOF) +import qualified System.Environment as Sys (getProgName, getArgs, withArgs) +import qualified System.Exit as Sys (exitFailure) +---------------------------------------------------------------- + +-- | Read a file and insert each line into a trie with its (base-0) +-- line number. +readTrieFromFile :: FilePath -> IO (T.Trie Int) +readTrieFromFile file = Sys.withFile file Sys.ReadMode (go 0 T.empty) + where + go !i !t !h = do + b <- Sys.hIsEOF h + if b + then return t + else do + line <- S.hGetLine h + go (i+1) (T.insert line i t) h + +-- TODO: maybe add some more CLI structure, for things like choosing +-- to use random inputs instead of file, or parameters of randomness, +-- etc. +main :: IO () +main = do + args <- Sys.getArgs + case args of + [] -> do + prog <- Sys.getProgName + putStrLn "ERROR: Missing file argument" + putStrLn $ "Usage: " ++ prog ++ " FILE [criterionFlags]" + Sys.exitFailure + file:rest -> + Sys.withArgs rest $ C.defaultMain + [ C.env (readTrieFromFile file) $ \ t -> + bgroup_Foldable t + -- TODO: more benchmark suites here. + ] + +---------------------------------------------------------------- +---------------------------------------------------------------- +type BenchmarkE env = env -> C.Benchmark +type BenchmarkableE env = env -> C.Benchmarkable + +bgroupE :: String -> [BenchmarkE env] -> BenchmarkE env +bgroupE name bs = \e -> C.bgroup name (($ e) <$> bs) +{-# INLINE bgroupE #-} + +benchE :: String -> BenchmarkableE env -> BenchmarkE env +benchE name b = \e -> C.bench name (b e) +{-# INLINE benchE #-} + +---------------------------------------------------------------- +#if MIN_VERSION_base(4,7,0) +{- +-- | From "Data.Functor.Utils", but isn't exported. Is used heavily +-- by the default implementations, since they use so many newtype +-- wrappers. +(#.) :: Coercible b c => (b -> c) -> (a -> b) -> (a -> c) +(#.) _ = coerce +{-# INLINE (#.) #-} +-} + +-- My own variant, ignoring the second argument instead. +(.#) :: Coercible a b => (b -> c) -> (a -> b) -> (a -> c) +(.#) f _ = coerce f +{-# INLINE (.#) #-} +#endif + +---------------------------------------------------------------- +bgroup_Foldable :: BenchmarkE (T.Trie Int) +bgroup_Foldable = + bgroupE "Foldable" + -- BUG: For these two 'fold' tests, whichever one we put first + -- runs much more slowly than the one we put second! Thus, there's + -- some sort of bug in how 'C.env' is forcing things. And adding + -- calls to 'Control.Exception.evaluate' and 'Control.DeepSeq.rnf' + -- doesn't help any. So for now we just put a nonce test first + -- to do whatever forcing is necessary. + [ benchE "IGNORE_THIS" $ C.nf (F.fold . fmap Sum) + -- Before adding rewrite rules for the 'Functor' type class, the + -- test using @(.)@ used to run monstrously slower than the version + -- using @(.#)@; indicating that those rewrites didn't automatically + -- come for free. Since defining our own rewrite rules, the two + -- versions seem about the same and it varies which one comes out + -- faster. Marking 'fold' as INLINE (instead of INLINABLE) causes + -- the @(.)@ variant to be closer to 'foldMap' than to the + -- 'fold' @(.#)@ variant. + , benchE "fold(.)" $ C.nf (F.fold . fmap Sum) +#if MIN_VERSION_base(4,7,0) + , benchE "fold(.#)" $ C.nf (F.fold .# fmap Sum) +#endif + , benchE "foldMap" $ C.nf (F.foldMap Sum) +#if MIN_VERSION_base(4,13,0) + , benchE "foldMap'" $ C.nf (F.foldMap' Sum) +#endif +-- (2021.03.12): according to this benchmark, our hand-rolled 'foldr' +-- is indeed monstrously faster than the default implementation. + , benchE "foldr" $ C.nf (F.foldr (+) 0) +#if MIN_VERSION_base(4,6,0) + , benchE "foldr'" $ C.nf (F.foldr' (+) 0) +#endif + , benchE "foldl" $ C.nf (F.foldl (+) 0) +#if MIN_VERSION_base(4,6,0) + , benchE "foldl'" $ C.nf (F.foldl' (+) 0) +#endif +#if MIN_VERSION_base(4,8,0) + , benchE "length" $ C.nf F.length + {- + , benchE "maximum" $ C.nf F.maximum -- Must first ensure non-empty! + , benchE "minimum" $ C.nf F.minimum -- Must first ensure non-empty! + -} +#else + , benchE "size" $ C.nf T.size +#endif + ] + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Bench/UnionWith.hs b/dev/Bench/UnionWith.hs new file mode 100644 --- /dev/null +++ b/dev/Bench/UnionWith.hs @@ -0,0 +1,70 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP, BangPatterns #-} + +---------------------------------------------------------------- +-- ~ 2023.03.19 +-- | +-- Module : Bench.UnionWith +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : portable +-- +-- Benchmarking the benefits of 'TI.wip_unionWith' vs 'TC.unionWith'. +---------------------------------------------------------------- + +module Main (main) where + +import qualified Data.Trie as T +import qualified Data.Trie.Internal as TI +import qualified Data.Trie.Convenience as TC +import qualified Data.ByteString as S +#if !(MIN_VERSION_base(4,18,0)) +-- [base-4.18 / GHC 9.6.1]: Prelude now re-exports 'liftA2'. +import Control.Applicative (liftA2) +#endif +import Control.DeepSeq (NFData) +import qualified Test.QuickCheck as QC +import qualified Criterion.Main as C +---------------------------------------------------------------- + + +---------------------------------------------------------------- +-- TODO: move this stuff off to the shared file. +-- TODO: should have argument to bound the 'Word8' +arbitraryBS :: Int -> QC.Gen S.ByteString +arbitraryBS maxL = do + l <- QC.chooseInt (0, maxL) + xs <- QC.vector l + return $! S.pack xs + +arbitraryTrie :: Int -> Int -> QC.Gen (T.Trie Int) +arbitraryTrie maxK maxL = do + k <- QC.chooseInt (0, maxK) + keys <- QC.vectorOf k $ arbitraryBS maxL + return $! T.fromList (zip keys [0..]) + +-- TODO: really need a better environment to work on than this... +generateEnv :: IO [T.Trie Int] +generateEnv = QC.generate $ QC.vectorOf 10 $ arbitraryTrie 30 10 + +cartesianNF :: NFData c => (a -> b -> c) -> [a] -> [b] -> C.Benchmarkable +cartesianNF f xs ys = C.nf (liftA2 f xs) ys + +---------------------------------------------------------------- +-- TODO: +-- uses 'Control.Exception.evaluate' instead of 'C.env'. Is that +-- because of using the @gauge@ library instead of @criterion@, or +-- is there some other reason? +main :: IO () +main = C.defaultMain + [ C.env generateEnv $ \ ts -> + C.bgroup "UnionWith" + [ C.bench "TI.wip_unionWith" $ cartesianNF (TI.wip_unionWith (+)) ts ts + , C.bench "TC.unionWith" $ cartesianNF (TC.unionWith (+)) ts ts + ] + ] + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Shared/BaseCompat.hs b/dev/Shared/BaseCompat.hs new file mode 100644 --- /dev/null +++ b/dev/Shared/BaseCompat.hs @@ -0,0 +1,200 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP #-} + +---------------------------------------------------------------- +-- ~ 2022.03.12 +-- | +-- Module : Shared.BaseCompat +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : portable (CPP) +-- +-- Compatibility shim for various versions of @base@. +-- +-- TODO: depending how much more we need in here, consider instead +-- just using the @base-compat-batteries@ package instead. +---------------------------------------------------------------- +module Shared.BaseCompat + ( + -- * "Data.Semigroup" + (<>) + , Sum(..) + , Product(..) + -- TODO: (?) 'Min', 'Max', 'Endo', 'Data.Monoid.Dual' + -- * "Data.Functor" + , (<&>) + -- * "Data.Functor.Compose" + , Compose(..) + -- * "Data.Functor.Identity" + , Identity(..) + ) where + +---------------------------------------------------------------- +#if MIN_VERSION_base(4,13,0) +-- [GHC 8.8.1]: @(<&>)@ added to "Data.Functor". +import Data.Functor ((<&>)) +#endif + +---------------------------------------------------------------- +#if MIN_VERSION_base(4,13,0) +-- [GHC 8.8.1]: Prelude re-exports 'Semigroup' and @('<>')@, but not 'stimes'. +#elif MIN_VERSION_base(4,9,0) +-- [GHC 8.0.1]: "Data.Semigroup" added to base. +import Data.Semigroup (Semigroup((<>))) +#elif MIN_VERSION_base(4,5,0) +-- [GHC 7.4.1]: @(<>)@ added to "Data.Monoid". +import Data.Monoid ((<>)) +#endif + +---------------------------------------------------------------- +-- For [base 4.13 / GHC 8.8.1], this module needs to explicitly +-- import 'Sum' in order to re-export it. However, for some reason +-- those versions will also cause users of this module to get an +-- \"unused import\" warning, even though the Prelude doesn't export +-- 'Data.Semigroup.Sum'. +-- TODO: figure out what the heck is going on re that warning... +-- FIXME: now getting that issue for GHC 8.0--8.6 too... +#if MIN_VERSION_base(4,9,0) +-- [GHC 8.0.1] +import Data.Semigroup (Sum(..), Product(..)) +import Data.Functor.Compose (Compose(..)) +#endif + +#if MIN_VERSION_base(4,8,0) +-- [GHC 7.10.1] +import Data.Functor.Identity (Identity(..)) +#endif + +#if MIN_VERSION_base(4,7,0) && !(MIN_VERSION_base(4,9,0)) +-- [GHC 7.8.1]: "Data.Coerce" added to base. +import Data.Coerce (coerce) +#endif + +---------------------------------------------------------------- +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,13,0)) +(<&>) :: Functor f => f a -> (a -> b) -> f b +(<&>) = flip fmap +{-# INLINE (<&>) #-} +#endif + +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,9,0)) +-- | Same as the derived instance, if the type was declared without +-- the record\/field syntax. +newtypeShowsPrec :: Show a => String -> (b -> a) -> Int -> b -> ShowS +newtypeShowsPrec name oldtype d x = + showParen (d > 10) + $ showString name + . showChar ' ' + . showsPrec 11 (oldtype x) +-- aka: +-- > newtypeShowsPrec name oldtype = +-- > Data.Functor.Classes.showsUnaryWith (\d -> showsPrec d . oldtype) name +-- But that function was also introduced in base-4.9.0.0 +#endif + +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,9,0)) +newtype Sum a = Sum { getSum :: a } + deriving (Eq, Ord, Bounded, Num) + -- We're not deriving the 'Show' instance, even though base does. + -- Also forgoing the 'Read' instance. + +instance (Show a) => Show (Sum a) where + showsPrec = newtypeShowsPrec "Sum" getSum + +instance Num a => Monoid (Sum a) where + mempty = Sum 0 +# if MIN_VERSION_base(4,7,0) + mappend = coerce ((+) :: a -> a -> a) +# else + mappend (Sum x) (Sum y) = Sum (x + y) +# endif + +-- TODO: if needed, there are also instances for: Functor, Applicative, Monad. +#endif + +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,9,0)) +newtype Product a = Product { getProduct :: a } + deriving (Eq, Ord, Bounded, Num) + -- We're not deriving the 'Show' instances, even though base does. + -- Also forgoing the 'Read' instance. + +instance (Show a) => Show (Product a) where + showsPrec = newtypeShowsPrec "Product" getProduct + +instance Num a => Monoid (Product a) where + mempty = Product 1 +# if MIN_VERSION_base(4,7,0) + mappend = coerce ((*) :: a -> a -> a) +# else + mappend (Product x) (Product y) = Product (x * y) +# endif + +-- TODO: if needed, there are also instances for: Functor, Applicative, Monad. +#endif + +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,9,0)) +newtype Compose f g a = Compose { getCompose :: f (g a) } + -- Ignoring the other instances since we don't need them. + +-- BUG: can't define 'Show' instance without 'Show1' for @f,g@; but +-- 'Show1' wasn't introduced until base-4.9.0.0. Ditto for 'Eq', +-- 'Ord'. + +instance (Functor f, Functor g) => Functor (Compose f g) where + fmap f (Compose x) = Compose (fmap (fmap f) x) +# if __GLASGOW_HASKELL__ + a <$ Compose x = Compose (fmap (a <$) x) +# endif + +instance (Foldable f, Foldable g) => Foldable (Compose f g) where + foldMap f (Compose t) = foldMap (foldMap f) t + +instance (Traversable f, Traversable g) => Traversable (Compose f g) where + traverse f (Compose t) = Compose <$> traverse (traverse f) t + +instance (Applicative f, Applicative g) => Applicative (Compose f g) where + pure x = Compose (pure (pure x)) + Compose f <*> Compose x = Compose (liftA2 (<*>) f x) +#endif + +---------------------------------------------------------------- +#if !(MIN_VERSION_base(4,8,0)) +newtype Identity a = Identity { runIdentity :: a } + deriving (Eq, Ord, Bounded, Num) + -- Ignoring the other instances since we don't need them. + +instance (Show a) => Show (Identity a) where + showsPrec = newtypeShowsPrec "Identity" runIdentity + +instance Functor Identity where +# if MIN_VERSION_base(4,7,0) + fmap = coerce +# else + fmap f (Identity x) = Identity (f x) +# endif + +instance Applicative Identity where + pure = Identity +# if MIN_VERSION_base(4,7,0) + (<*>) = coerce +# else + Identity f <*> Identity x = Identity (f x) +# endif + +-- Using 'runIdentity' here instead of pattern matching, because +-- that's what base-4.16.0.0 does. +instance Monad Identity where + m >>= k = k (runIdentity m) + +-- TODO: if needed, there's also a Foldable instance. +#endif + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Test/Main.hs b/dev/Test/Main.hs new file mode 100644 --- /dev/null +++ b/dev/Test/Main.hs @@ -0,0 +1,197 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP + , MultiParamTypeClasses + , FlexibleContexts + #-} + +---------------------------------------------------------------- +-- ~ 2022.03.05 +-- | +-- Module : Test.Main +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : semi-portable (MPTC,...) +-- +-- Testing 'Trie's. +---------------------------------------------------------------- +module Main (main) where + +#if !(MIN_VERSION_base(4,11,0)) +-- [GHC 8.4.1]: This guard is just to avoid an \"unused import\" warning. +import Shared.BaseCompat +#endif +import Test.Utils +import Test.Properties (smallcheckTests, quickcheckTests) + +import qualified Data.Trie as T +import qualified Data.ByteString as S + +import qualified System.Exit as System (exitSuccess, exitFailure) +import qualified System.IO as System (hPutStrLn, stderr) +import qualified Test.Tasty.Ingredients as Tasty (tryIngredients) +import qualified Test.Tasty.Options as Tasty (singleOption, OptionSet) +import qualified Test.Tasty.Runners as Tasty (installSignalHandlers, parseOptions) +import qualified Test.Tasty as Tasty +import qualified Test.Tasty.SmallCheck as SC +import qualified Test.Tasty.QuickCheck as QC + +---------------------------------------------------------------- +---------------------------------------------------------------- +-- We can't use 'Tasty.defaultMain' together with 'Tasty.localOption', +-- because what we want to do is to set new defaults but still allow +-- the commandline to override those defaults, and I don't see any +-- way to do that with 'Tasty.defaultMain'. +-- +-- TODO: Still need some way to remove the few remaining +-- 'Tasty.localOption' calls that aren't global... +main :: IO () +main = do + let ins = Tasty.defaultIngredients + Tasty.installSignalHandlers + opts <- Tasty.parseOptions ins tests + case Tasty.tryIngredients ins (globalOptions <> opts) tests of + Nothing -> do + System.hPutStrLn System.stderr + "No ingredients agreed to run. Something is wrong." + System.exitFailure + Just act -> do + ok <- act + if ok + then System.exitSuccess + else System.exitFailure + +-- We add the timeout for the sake of GithubActions CI, so we don't +-- accidentally blow our budget. So far, even with HPC enabled, +-- each test takes less than 1sec (almost all of them taking less +-- than 0.5sec); so allowing 10s per test is more than generous. +-- Since the full suite takes around 20sec, that means the timeout +-- will cap the suite to take no more than 200sec / 3.3min. +globalOptions :: Tasty.OptionSet +globalOptions = mconcat + [ Tasty.singleOption (Tasty.mkTimeout 10000000) -- 10sec; in microsecs + , Tasty.singleOption (QC.QuickCheckTests 500) -- QC.Args.maxSuccess + , Tasty.singleOption (QC.QuickCheckMaxSize 400) -- QC.Args.maxSize + , Tasty.singleOption (QC.QuickCheckMaxRatio 10) -- QC.Args.maxDiscardRatio + , Tasty.singleOption (SC.SmallCheckDepth 3) + ] + +-- We run the HUnit and smallcheck tests first, since they're the +-- fastest. Thus, we only pay for the QuickCheck tests if we have to. +tests :: Tasty.TestTree +tests = + Tasty.testGroup "All Tests" + [ hunitTests + , Tasty.testGroup "Properties" + [ smallcheckTests + , quickcheckTests + ] + ] + +hunitTests :: Tasty.TestTree +hunitTests = + Tasty.testGroup "HUnit" + [ test_Union + , test_Intersect + , test_Submap + , test_Insert + , test_Delete + ] + + +---------------------------------------------------------------- +---------------------------------------------------------------- +-- Because we avoid epsilons everywhere else, need to make sure +-- 'mergeBy' gets it right +test_Union :: Tasty.TestTree +test_Union = + Tasty.testGroup "mergeBy" + [ testEqual "unionL epsilon" (e1 `T.unionL` e2) e1 + , testEqual "unionR epsilon" (e1 `T.unionR` e2) e2 + , testEqual "unionL regression" (tLeft `T.unionL` tRight) tLeftResult + , testEqual "unionR regression" (tLeft `T.unionR` tRight) tRightResult + ] + where + e1 = T.singleton S.empty (4::Int) + e2 = T.singleton S.empty (2::Int) + + -- Regression test against bug filed by Gregory Crosswhite on + -- 2010.06.10 against version 0.2.1.1. + a, b :: S.ByteString + a = read "\"\231^\179\160Y\134Gr\158<)&\222\217#\156\"" + b = read "\"\172\193\GSp\222\174GE\186\151\DC1#P\213\147\SI\"" + tLeft, tRight, tRightResult, tLeftResult :: T.Trie Int + tLeft = T.fromList [(a,1), (b,0)] + tRight = T.fromList [(a,2)] + tRightResult = T.fromList [(a,2), (b,0)] + tLeftResult = T.fromList [(a,1), (b,0)] + +test_Intersect :: Tasty.TestTree +test_Intersect = + Tasty.testGroup "intersectBy" + [ testEqual "intersectL" (tLeft `T.intersectL` tRight) tLeftResult + , testEqual "intersectR" (tLeft `T.intersectR` tRight) tRightResult + ] + where + -- Trivial regression example + a, b, c :: S.ByteString + a = packC2W "a" + b = packC2W "b" + c = packC2W "c" + tLeft, tRight, tRightResult, tLeftResult :: T.Trie Int + tLeft = T.fromList [(S.empty,0), (a,1), (b,3)] + tRight = T.fromList [ (a,2), (c,4)] + tLeftResult = T.fromList [(a,1)] + tRightResult = T.fromList [(a,2)] + -- TODO: better unit test for when one string is prefix of another. + +---------------------------------------------------------------- +test_Submap :: Tasty.TestTree +test_Submap = + Tasty.testGroup "submap" + [ nullSubmap "split on arc fails" fi True + , nullSubmap "prefix of arc matches" fo False + , nullSubmap "suffix of empty fails" food True + , nullSubmap "missing branch fails" bag True + , nullSubmap "at a branch matches" ba False + ] + where + t = vocab2trie ["foo", "bar", "baz"] + fi = packC2W "fi" + fo = packC2W "fo" + food = packC2W "food" + ba = packC2W "ba" + bag = packC2W "bag" + + nullSubmap s q b = testEqual s (T.null $ T.submap q t) b + +---------------------------------------------------------------- +-- requires Eq (Trie a) and, in case it fails, Show (Trie a) +test_Insert :: Tasty.TestTree +test_Insert = + Tasty.testGroup "insert" + [ testEqual "insertion is commutative for prefix/superfix" + (T.insert aba o $ T.insert abaissed i $ T.empty) + (T.insert abaissed i $ T.insert aba o $ T.empty) + ] + where + aba = packC2W "aba" + abaissed = packC2W "abaissed" + + o = 0::Int + i = 1::Int + + +test_Delete :: Tasty.TestTree +test_Delete = + Tasty.testGroup "delete" + [ testEqual "deleting epsilon from empty trie is empty" + (T.delete epsilon T.empty) (T.empty :: T.Trie Int) + ] + where + -- TODO: why not 'S.empty'? + epsilon = packC2W "" + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Test/Properties.hs b/dev/Test/Properties.hs new file mode 100644 --- /dev/null +++ b/dev/Test/Properties.hs @@ -0,0 +1,1245 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP + , MultiParamTypeClasses + , FlexibleContexts + #-} + +-- HACK: Enable this to locally disable -Wunused-top-binds for some +-- HOF properties that I haven't figured out how to decently test +-- yet. This CPP hack allows to quickly "uncomment" all such code, +-- to make sure it still typechecks et al. +#define HideUnusedTopBonds + +---------------------------------------------------------------- +-- ~ 2023.03.19 +-- | +-- Module : Test.Properties +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : semi-portable (MPTC,...) +-- +-- Property testing 'Trie's. +---------------------------------------------------------------- +module Test.Properties (smallcheckTests, quickcheckTests) where + +import Shared.BaseCompat +import Test.Utils + +import qualified Data.Trie as T +import qualified Data.Trie.Internal as TI +import qualified Data.Trie.Convenience as TC +import qualified Data.ByteString as S + +import qualified Test.Tasty as Tasty +import qualified Test.Tasty.SmallCheck as SC +import qualified Test.Tasty.QuickCheck as QC + +#if !(MIN_VERSION_base(4,18,0)) +-- [base-4.18 / GHC 9.6.1]: Prelude now re-exports 'liftA2'. +import Control.Applicative (liftA2) +#endif +#ifndef HideUnusedTopBonds +import Control.Monad (join, (<=<)) +#endif +import Control.Monad (guard) +import Data.List (nubBy, sortBy) +import Data.Maybe (fromJust, isJust) +import Data.Ord (comparing) + +import qualified Data.Binary as B +import qualified Data.Foldable as F + +-- Repeated from "Shared.BaseCompat" because we need the class here. +#if MIN_VERSION_base(4,13,0) +-- [aka GHC 8.8.1]: Prelude re-exports 'Semigroup'. +#elif MIN_VERSION_base(4,9,0) +-- [aka GHC 8.0.1]: "Data.Semigroup" added to base. +import Data.Semigroup (Semigroup) +#endif + +#if MIN_VERSION_base(4,9,0) +import Data.Functor.Classes (Eq1) +#endif + +---------------------------------------------------------------- +---------------------------------------------------------------- +{- +TODO: see if we can't figure out some way of wrapping our properties +so that we can just write this list once and then pass in a token +for which checker to resolve to; something like: + + data Prop a = Prop String a + + data PropChecker = QuickCheck | SmallCheck + deriving (Eq, Show) + + testProp :: (QC.Testable a, SC.Testable IO a) => PropChecker -> Prop a -> Tasty.TestTree + testProp QuickCheck (Prop name a) = QC.testProperty name a + testProp SmallCheck (Prop name a) = SC.testProperty name a + +Of course, the problem with that implementation is that we need to +have the Prop remain polymorphic in the CheckGuard type, and have +testProp resolve it depending on the PropChecker. So is there a way +to do that without GADTs or impredicativity? +-} + +quickcheckTests :: Tasty.TestTree +quickcheckTests + = Tasty.testGroup "QuickCheck" + [ Tasty.testGroup "Data.Trie.ByteStringInternal" + [ QC.testProperty + "prop_breakMaximalPrefix" + (prop_breakMaximalPrefix :: WS -> WS -> Bool) + ] + , Tasty.testGroup "Trivialities (@Int)" + [ QC.testProperty + "prop_singleton" + (prop_singleton :: WS -> Int -> Bool) + , QC.testProperty + "prop_insert_lookup" + (prop_insert_lookup :: WS -> Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_delete_lookup" + (prop_delete_lookup :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_insert_size" + (prop_insert_size :: WS -> Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_delete_size" + (prop_delete_size :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_insert_delete" + (prop_insert_delete :: WS -> Int -> WTrie Int -> Bool) + ] + , Tasty.testGroup "Submap (@Int)" + [ QC.testProperty + "prop_submap_keysAreMembers" + (prop_submap_keysAreMembers :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_submap_keysHavePrefix" + (prop_submap_keysHavePrefix :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_submap_valuesEq" + (prop_submap_valuesEq :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_deleteSubmap_keysAreMembers" + (prop_deleteSubmap_keysAreMembers :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_deleteSubmap_keysLackPrefix" + (prop_deleteSubmap_keysLackPrefix :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_deleteSubmap_disunion" + (prop_deleteSubmap_disunion :: WS -> WTrie Int -> Bool) + ] + , Tasty.localOption (QC.QuickCheckMaxSize 300) + -- BUG: fix that 'Tasty.localOption' + $ Tasty.testGroup "unionWith (@Int)" + [ QC.testProperty + "prop_unionL" + (prop_unionL :: WTrie Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_unionR" + (prop_unionR :: WTrie Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_unionPlus" + (prop_unionPlus :: WTrie Int -> WTrie Int -> Bool) + ] + , Tasty.localOption (QC.QuickCheckMaxSize 300) + -- BUG: fix that 'Tasty.localOption' + $ Tasty.testGroup "intersectBy (@Int)" + [ QC.testProperty + "prop_intersectL" + (prop_intersectL :: WTrie Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_intersectR" + (prop_intersectR :: WTrie Int -> WTrie Int -> Bool) + , QC.testProperty + "prop_intersectPlus" + (prop_intersectPlus :: WTrie Int -> WTrie Int -> Bool) + ] + , Tasty.testGroup "Matching (@Int)" + [ QC.testProperty + "prop_matches_keysOrdered" + (prop_matches_keysOrdered :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_matches_keysArePrefix" + (prop_matches_keysArePrefix :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_minMatch_is_first_matches" + (prop_minMatch_is_first_matches :: WS -> WTrie Int -> Bool) + , QC.testProperty + "prop_match_is_last_matches" + (prop_match_is_last_matches :: WS -> WTrie Int -> Bool) + ] + , Tasty.testGroup "Priority-queue (@Int)" + [ QC.testProperty + "prop_minAssoc_is_first_toList" + (prop_minAssoc_is_first_toList :: WTrie Int -> Bool) + , QC.testProperty + "prop_maxAssoc_is_last_toList" + (prop_maxAssoc_is_last_toList :: WTrie Int -> Bool) + , QC.testProperty + "prop_updateMinViewBy_ident" + (prop_updateMinViewBy_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_updateMaxViewBy_ident" + (prop_updateMaxViewBy_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_updateMinViewBy_gives_minAssoc" + (prop_updateMinViewBy_gives_minAssoc :: (WS -> Int -> Maybe Int) -> WTrie Int -> Bool) + , QC.testProperty + "prop_updateMaxViewBy_gives_maxAssoc" + (prop_updateMaxViewBy_gives_maxAssoc :: (WS -> Int -> Maybe Int) -> WTrie Int -> Bool) + ] + , Tasty.testGroup "toList (@Int)" + [ QC.testProperty + "prop_toListBy_keysOrdered" + (prop_toListBy_keysOrdered :: WTrie Int -> Bool) + , QC.testProperty + "prop_keys" + (prop_keys :: WTrie Int -> Bool) + , QC.testProperty + "prop_elems" + (prop_elems :: WTrie Int -> Bool) + ] + , Tasty.testGroup "fromList (@Int)" + [ QC.testProperty + "prop_fromList_takes_first" + (prop_fromList_takes_first :: [(WS, Int)] -> Bool) + , QC.testProperty + "prop_fromListR_takes_first" + (prop_fromListR_takes_first :: [(WS, Int)] -> Bool) + , QC.testProperty + "prop_fromListL_takes_first" + (prop_fromListL_takes_first :: [(WS, Int)] -> Bool) + , QC.testProperty + "prop_fromListS_takes_first" + (prop_fromListS_takes_first :: [(WS, Int)] -> Bool) + , QC.testProperty + "prop_fromListWithConst_takes_first" + (prop_fromListWithConst_takes_first :: [(WS, Int)] -> Bool) + , QC.testProperty + "prop_fromListWithLConst_takes_first" + (prop_fromListWithLConst_takes_first :: [(WS, Int)] -> Bool) + ] + , Tasty.testGroup "Type classes" + [ Tasty.testGroup "Functor (@Int)" + [ QC.testProperty + "prop_FunctorIdentity" + (prop_FunctorIdentity :: WTrie Int -> Bool) + -- TODO: prop_FunctorCompose + {- + -- TODO: still worth it to do this test with 'undefined'? + , QC.testProperty + "prop_fmap_keys" + (prop_fmap_keys (undefined :: Int -> Int) :: WTrie Int -> Bool) + -- TODO: prop_fmap_elems + -} +#if __GLASGOW_HASKELL__ + , QC.testProperty + "prop_fmapConst" + (prop_fmapConst :: Int -> WTrie Int -> Bool) +#endif + , QC.testProperty + -- TODO: generalize to other functions. + "prop_fmap_toList" + (prop_fmap_toList (+1) :: WTrie Int -> Bool) + ] + , Tasty.testGroup "Applicative (@Int)" + [ QC.testProperty + "prop_ApplicativeIdentity" + (prop_ApplicativeIdentity :: WTrie Int -> Bool) + -- TODO: prop_ApplicativeCompose, prop_ApplicativeHom, prop_ApplicativeInterchange + ] + , Tasty.testGroup "Monad (@Int)" + [ QC.testProperty + "prop_MonadIdentityR" + (prop_MonadIdentityR :: WTrie Int -> Bool) + -- TODO: prop_MonadIdentityL, prop_MonadAssoc + ] + , Tasty.testGroup "Foldable (@Int/Sum Int)" + [ QC.testProperty + "prop_foldmapIdentity" + (prop_foldmapIdentity :: WTrie (Sum Int) -> Bool) + , QC.testProperty + "prop_foldmapFusion" + (prop_foldmapFusion Sum :: WTrie Int -> Bool) + , QC.testProperty + "prop_length" + (prop_length :: WTrie Int -> Bool) + , QC.testProperty + "prop_foldr_vs_foldrWithKey" + (prop_foldr_vs_foldrWithKey :: WTrie Int -> Bool) + , QC.testProperty + "prop_foldr'_vs_foldrWithKey'" + (prop_foldr'_vs_foldrWithKey' :: WTrie Int -> Bool) + , QC.testProperty + "prop_foldl_vs_foldlWithKey" + (prop_foldl_vs_foldlWithKey :: WTrie Int -> Bool) + , QC.testProperty + "prop_foldl'_vs_foldlWithKey'" + (prop_foldl'_vs_foldlWithKey' :: WTrie Int -> Bool) +#if MIN_VERSION_base(4,6,0) + , QC.testProperty + "prop_foldr_vs_foldr'" + (prop_foldr_vs_foldr' :: WTrie Int -> Bool) + , QC.testProperty + "prop_foldl_vs_foldl'" + (prop_foldl_vs_foldl' :: WTrie Int -> Bool) +#endif +#if MIN_VERSION_base(4,13,0) + , QC.testProperty + "prop_foldMap_vs_foldMap'" + (prop_foldMap_vs_foldMap' :: WTrie Int -> Bool) +#endif + ] + , Tasty.testGroup "Traversable (@Int)" + [ QC.testProperty + "prop_TraverseIdentity" + (prop_TraverseIdentity :: WTrie Int -> Bool) + -- TODO: prop_TraverseNaturality, prop_TraverseComposition + , QC.testProperty + "prop_SequenceaIdentity" + (prop_SequenceaIdentity :: WTrie Int -> Bool) + -- TODO: prop_SequenceaNaturality, prop_SequenceaComposition + ] + , Tasty.testGroup "Filterable (@Int)" + [ QC.testProperty + "prop_FiltermapFission" + (prop_FiltermapFission remTriple :: WTrie Int -> Bool) + , QC.testProperty + "prop_FiltermapFusion" + (prop_FiltermapFusion (`rem` 3) (not . isTriple) :: WTrie Int -> Bool) + , QC.testProperty + "prop_FiltermapConservation" + (prop_FiltermapConservation (+1) :: WTrie Int -> Bool) + -- TODO: prop_FiltermapVertComposition + , QC.testProperty + "prop_FilterDefinition" + (prop_FilterDefinition isTriple :: WTrie Int -> Bool) + , QC.testProperty + "prop_FilterVertComposition" + (prop_FilterVertComposition even isTriple :: WTrie Int -> Bool) + ] + , Tasty.testGroup "Witherable (@Int)" + [ QC.testProperty + "prop_WitherConservation" + (prop_WitherConservation (Identity . (+1)) :: WTrie Int -> Bool) + -- TODO: prop_WitherNaturality, prop_WitherPurity, prop_WitherHorizComposition + , QC.testProperty + "prop_FilteraDefinition" + (prop_FilteraDefinition (Identity . even) :: WTrie Int -> Bool) + -- TODO: prop_FilteraNaturality, prop_FilteraPurity, prop_FilteraHorizComposition, prop_FilteraHorizComposition_alt + ] +#if MIN_VERSION_base(4,9,0) + , Tasty.testGroup "Semigroup (@Sum Int)" + [ QC.testProperty + -- This one is a bit more expensive: 1~2sec instead of <=0.5sec + -- TODO: see if we can't fix that without ruining the utility of the test. + "prop_Semigroup" + (prop_Semigroup :: WTrie (Sum Int) -> WTrie (Sum Int) -> WTrie (Sum Int) -> Bool) + ] +#endif + , Tasty.testGroup "Monoid (@Sum Int)" + [ QC.testProperty + "prop_MonoidIdentityL" + (prop_MonoidIdentityL :: WTrie (Sum Int) -> Bool) + , QC.testProperty + "prop_MonoidIdentityR" + (prop_MonoidIdentityR :: WTrie (Sum Int) -> Bool) + ] + , Tasty.testGroup "Binary (@Int)" + [ QC.testProperty + "prop_Binary" + (prop_Binary :: WTrie Int -> Bool) + ] + ] + , Tasty.testGroup "Other mapping/filtering (@Int)" + [ QC.testProperty + "prop_filterMap_ident" + (prop_filterMap_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_filterMap_empty" + (prop_filterMap_empty :: WTrie Int -> Bool) + , QC.testProperty + "prop_mapBy_keys" + (prop_mapBy_keys :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualMap_ident" + (prop_contextualMap_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualMap'_ident" + (prop_contextualMap'_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualFilterMap_ident" + (prop_contextualFilterMap_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualMapBy_keys" + (prop_contextualMapBy_keys :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualMapBy_ident" + (prop_contextualMapBy_ident :: WTrie Int -> Bool) + , QC.testProperty + "prop_contextualMapBy_empty" + (prop_contextualMapBy_empty :: WTrie Int -> Bool) + ] + ] + + +---------------------------------------------------------------- +adjustSmallCheckDepth + :: (SC.SmallCheckDepth -> SC.SmallCheckDepth) + -> Tasty.TestTree -> Tasty.TestTree +adjustSmallCheckDepth = Tasty.adjustOption + +-- Throughout, we try to use 'W' or @()@ whenever we can, to reduce +-- the exponential growth problem. +smallcheckTests :: Tasty.TestTree +smallcheckTests + = Tasty.testGroup "SmallCheck" + [ Tasty.testGroup "Data.Trie.ByteStringInternal" + [ adjustSmallCheckDepth (+1) + $ SC.testProperty + "prop_breakMaximalPrefix" + (prop_breakMaximalPrefix :: WS -> WS -> Bool) + ] + , Tasty.testGroup "Trivialities (@W/@())" + [ adjustSmallCheckDepth (+2) + $ SC.testProperty + -- This one can easily handle depth=6 fine (~0.1sec), but d=7 (~4sec) + "prop_singleton" + (prop_singleton :: WS -> W -> Bool) + , SC.testProperty + -- Warning: depth 4 takes >10sec! + "prop_insert_lookup" + (prop_insert_lookup :: WS -> W -> WTrie W -> Bool) + , SC.testProperty + -- Don't waste any depth on the values! + -- Warning: depth 4 takes >10sec! + "prop_delete_lookup" + (prop_delete_lookup :: WS -> WTrie () -> Bool) + , SC.testProperty + -- Don't waste any depth on the values! + -- Warning: depth 4 takes >10sec! + "prop_insert_size" + (prop_insert_size :: WS -> () -> WTrie () -> Bool) + , SC.testProperty + -- Don't waste any depth on the values! + -- Warning: depth 4 takes >10sec! + -- FIXME: at depth 3: 388/410 (94%) did not meet the condition! + -- (That is, before rephrasing to avoid using 'CheckGuard') + "prop_delete_size" + (prop_delete_size :: WS -> WTrie () -> Bool) + , SC.testProperty + -- Note: at depth 3: 136/2120 (6.4%) did not meet the condition. + -- (That is, before rephrasing to avoid using 'CheckGuard') + "prop_insert_delete" + (prop_insert_delete :: WS -> W -> WTrie W -> Bool) + ] + , Tasty.testGroup "Submap (@()/@W)" + -- Depth=4 is at best very marginal here... + [ SC.testProperty + "prop_submap_keysAreMembers" + (prop_submap_keysAreMembers :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_submap_keysHavePrefix" + (prop_submap_keysHavePrefix :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_submap_valuesEq" + (prop_submap_valuesEq :: WS -> WTrie W -> Bool) + , SC.testProperty + "prop_deleteSubmap_keysAreMembers" + (prop_deleteSubmap_keysAreMembers :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_deleteSubmap_keysLackPrefix" + (prop_deleteSubmap_keysLackPrefix :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_deleteSubmap_disunion" + (prop_deleteSubmap_disunion :: WS -> WTrie W -> Bool) + ] + , Tasty.testGroup "unionWith (@W/@Int)" + -- Warning: Using depth=4 here is bad (the first two take about + -- 26.43sec; the last one much longer). + [ SC.testProperty + "prop_unionL" + (prop_unionL :: WTrie W -> WTrie W -> Bool) + , SC.testProperty + "prop_unionR" + (prop_unionR :: WTrie W -> WTrie W -> Bool) + , SC.testProperty + "prop_unionPlus" + (prop_unionPlus :: WTrie Int -> WTrie Int -> Bool) + ] + , Tasty.testGroup "intersectBy (@W/@Int)" + -- Warning: Using depth=4 here is bad (the first two take about + -- 26.43sec; the last one much longer). + [ SC.testProperty + "prop_intersectL" + (prop_intersectL :: WTrie W -> WTrie W -> Bool) + , SC.testProperty + "prop_intersectR" + (prop_intersectR :: WTrie W -> WTrie W -> Bool) + , SC.testProperty + "prop_intersectPlus" + (prop_intersectPlus :: WTrie Int -> WTrie Int -> Bool) + ] + , Tasty.testGroup "Matching (@()/@W)" + [ SC.testProperty + "prop_matches_keysOrdered" + (prop_matches_keysOrdered :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_matches_keysArePrefix" + (prop_matches_keysArePrefix :: WS -> WTrie () -> Bool) + , SC.testProperty + "prop_minMatch_is_first_matches" + (prop_minMatch_is_first_matches :: WS -> WTrie W -> Bool) + , SC.testProperty + "prop_match_is_last_matches" + (prop_match_is_last_matches :: WS -> WTrie W -> Bool) + ] + , Tasty.testGroup "Priority-queue (@W)" + -- Depth=4 takes about 1sec each + [ SC.testProperty + "prop_minAssoc_is_first_toList" + (prop_minAssoc_is_first_toList :: WTrie W -> Bool) + , SC.testProperty + "prop_maxAssoc_is_last_toList" + (prop_maxAssoc_is_last_toList :: WTrie W -> Bool) + , SC.testProperty + "prop_updateMinViewBy_ident" + (prop_updateMinViewBy_ident :: WTrie W -> Bool) + , SC.testProperty + "prop_updateMaxViewBy_ident" + (prop_updateMaxViewBy_ident :: WTrie W -> Bool) + -- HACK: must explicitly pass functions for these two, else they're too slow + , SC.testProperty + "prop_updateMinViewBy_gives_minAssoc" + (prop_updateMinViewBy_gives_minAssoc undefined :: WTrie W -> Bool) + , SC.testProperty + "prop_updateMaxViewBy_gives_maxAssoc" + (prop_updateMaxViewBy_gives_maxAssoc undefined :: WTrie W -> Bool) + ] + , Tasty.testGroup "toList (@()/@W)" + -- These can handle depth 4, but it takes >1sec (but still <5sec) + [ SC.testProperty + "prop_toListBy_keysOrdered" + (prop_toListBy_keysOrdered :: WTrie () -> Bool) + , SC.testProperty + "prop_keys" + (prop_keys :: WTrie () -> Bool) + , SC.testProperty + "prop_elems" + (prop_elems :: WTrie W -> Bool) + -- TODO: move these into the section for 'Foldable', to agree with QC + , SC.testProperty + "prop_foldr_vs_foldrWithKey" + (prop_foldr_vs_foldrWithKey :: WTrie W -> Bool) +#if MIN_VERSION_base(4,6,0) + , SC.testProperty + "prop_foldr_vs_foldr'" + (prop_foldr_vs_foldr' :: WTrie W -> Bool) + , SC.testProperty + "prop_foldl_vs_foldl'" + (prop_foldl_vs_foldl' :: WTrie W -> Bool) +#endif + -- TODO: prop_foldMap_vs_foldMap' requires (Num W) if we want to use W. + ] + , Tasty.adjustOption (+ (1::SC.SmallCheckDepth)) + $ Tasty.testGroup "fromList (@()/@W)" + [ SC.testProperty + "prop_fromList_takes_first" + (prop_fromList_takes_first :: [(WS, W)] -> Bool) + , SC.testProperty + "prop_fromListR_takes_first" + (prop_fromListR_takes_first :: [(WS, W)] -> Bool) + , SC.testProperty + "prop_fromListL_takes_first" + (prop_fromListL_takes_first :: [(WS, W)] -> Bool) + , SC.testProperty + "prop_fromListS_takes_first" + (prop_fromListS_takes_first :: [(WS, W)] -> Bool) + , SC.testProperty + "prop_fromListWithConst_takes_first" + (prop_fromListWithConst_takes_first :: [(WS, W)] -> Bool) + , SC.testProperty + "prop_fromListWithLConst_takes_first" + (prop_fromListWithLConst_takes_first :: [(WS, W)] -> Bool) + ] + -- TODO: do we want to smallcheck any of the "Type classes" or "Other mapping/filtering" stuff we do in quickcheck? + ] + + +---------------------------------------------------------------- +---------------------------------------------------------------- + +-- | 'TI.breakMaximalPrefix' satisfies the documented equalities. +prop_breakMaximalPrefix :: WS -> WS -> Bool +prop_breakMaximalPrefix (WS s) (WS z) = + let (pre,s',z') = TI.breakMaximalPrefix s z + in (pre <> s') == s + && (pre <> z') == z + +{- +-- FIXME: need to export the RLBS stuff if we are to test it... +prop_toStrict :: [WS] -> Bool +prop_toStrict = + (S.concat .==. (TI.toStrict . foldl' (+>) Epsilon)) . map unWS +-} + +---------------------------------------------------------------- +-- | A singleton, is. +prop_singleton :: (Eq a) => WS -> a -> Bool +prop_singleton (WS k) v = + T.singleton k v == T.insert k v T.empty + +-- | If you insert a value, you can look it up. +prop_insert_lookup :: (Eq a) => WS -> a -> WTrie a -> Bool +prop_insert_lookup (WS k) v (WT t) = + (T.lookup k . T.insert k v $ t) == Just v + +-- | If you delete a value, you can't look it up. +prop_delete_lookup :: WS -> WTrie a -> Bool +prop_delete_lookup (WS k) = + isNothing . T.lookup k . T.delete k . unWT + where + isNothing Nothing = True + isNothing (Just _) = False + +-- TODO: print/record diagnostics re what proportion of calls have +-- @n=0@ vs @n=1@, to ensure proper coverage. +prop_insert_size :: WS -> a -> WTrie a -> Bool +prop_insert_size (WS k) v (WT t) = + ((T.size . T.insert k v) .==. ((n +) . T.size)) $ t + where + n = if T.member k t then 0 else 1 + +-- TODO: print/record diagnostics re what proportion of calls have +-- @n=0@ vs @n=1@, to ensure proper coverage. +prop_delete_size :: WS -> WTrie a -> Bool +prop_delete_size (WS k) (WT t) = + ((T.size . T.delete k) .==. (subtract n . T.size)) $ t + where + n = if T.member k t then 1 else 0 + +prop_insert_delete :: (Eq a) => WS -> a -> WTrie a -> Bool +prop_insert_delete (WS k) v (WT t) + | T.member k t = ((T.delete k . T.insert k v) .==. T.delete k) $ t + | otherwise = ((T.delete k . T.insert k v) .==. id) $ t + +-- | All keys in a submap are keys in the supermap +prop_submap_keysAreMembers :: WS -> WTrie a -> Bool +prop_submap_keysAreMembers (WS q) (WT t) = + all (`T.member` t) . T.keys . T.submap q $ t + -- TODO: should we use 'QC.conjoin' (assuming another class to overload it) in lieu of 'all'? What are the actual benefits of doing so? Ditto for all the uses below. + +-- | All keys in a submap have the query as a prefix +prop_submap_keysHavePrefix :: WS -> WTrie a -> Bool +prop_submap_keysHavePrefix (WS q) = + all (q `S.isPrefixOf`) . T.keys . T.submap q . unWT + +-- | All values in a submap are the same in the supermap +prop_submap_valuesEq :: (Eq a) => WS -> WTrie a -> Bool +prop_submap_valuesEq (WS q) (WT t) = + ((`T.lookup` t') .==. (`T.lookup` t)) `all` T.keys t' + where t' = T.submap q t + +-- | All keys in the result are keys in the supermap +prop_deleteSubmap_keysAreMembers :: WS -> WTrie a -> Bool +prop_deleteSubmap_keysAreMembers (WS q) (WT t) = + all (`T.member` t) . T.keys . T.deleteSubmap q $ t + +-- | All keys in a submap lack the query as a prefix +prop_deleteSubmap_keysLackPrefix :: WS -> WTrie a -> Bool +prop_deleteSubmap_keysLackPrefix (WS q) = + all (not . S.isPrefixOf q) . T.keys . T.deleteSubmap q . unWT + +-- | 'T.submap' and 'T.deleteSubmap' partition every trie for every key. +prop_deleteSubmap_disunion :: (Eq a) => WS -> WTrie a -> Bool +prop_deleteSubmap_disunion (WS q) (WT t) = + t == (T.submap q t `TC.disunion` T.deleteSubmap q t) + +prop_unionWith :: Eq a => (a -> a -> a) -> WTrie a -> WTrie a -> Bool +prop_unionWith f (WT x) (WT y) = + TI.wip_unionWith f x y == TC.unionWith f x y + +prop_unionL :: Eq a => WTrie a -> WTrie a -> Bool +prop_unionL = prop_unionWith (\x _ -> x) + +prop_unionR :: Eq a => WTrie a -> WTrie a -> Bool +prop_unionR = prop_unionWith (\_ y -> y) + +prop_unionPlus :: (Eq a, Num a) => WTrie a -> WTrie a -> Bool +prop_unionPlus = prop_unionWith (+) + +-- TODO: other than as a helper like below, could we actually +-- generate interesting enough functions to make this worth testing +-- directly? +-- +-- | Arbitrary @x ∩ y == (x ∪ y) ⋈ (x ⋈ y)@. +prop_intersectBy :: (Eq a) => (a -> a -> Maybe a) -> WTrie a -> WTrie a -> Bool +prop_intersectBy f (WT x) (WT y) = + T.intersectBy f x y == (T.mergeBy f x y `TC.disunion` TC.disunion x y) + +-- | Left-biased @x ∩ y == (x ∪ y) ⋈ (x ⋈ y)@. +prop_intersectL :: (Eq a) => WTrie a -> WTrie a -> Bool +prop_intersectL = prop_intersectBy (\x _ -> Just x) + +-- | Right-biased @x ∩ y == (x ∪ y) ⋈ (x ⋈ y)@. +prop_intersectR :: (Eq a) => WTrie a -> WTrie a -> Bool +prop_intersectR = prop_intersectBy (\_ y -> Just y) + +-- | Additive @x ∩ y == (x ∪ y) ⋈ (x ⋈ y)@. +prop_intersectPlus :: (Eq a, Num a) => WTrie a -> WTrie a -> Bool +prop_intersectPlus = prop_intersectBy (\x y -> Just (x + y)) + +isOrdered :: (Ord a) => [a] -> Bool +isOrdered xs = and (zipWith (<=) xs (drop 1 xs)) + +-- | 'T.toListBy', 'T.toList', and 'T.keys' are ordered by keys. +prop_toListBy_keysOrdered :: WTrie a -> Bool +prop_toListBy_keysOrdered = isOrdered . T.keys . unWT + +fst3 :: (a,b,c) -> a +fst3 (a,_,_) = a + +-- | 'T.matches' is ordered by keys. +prop_matches_keysOrdered :: WS -> WTrie a -> Bool +prop_matches_keysOrdered (WS q) (WT t) = + isOrdered . map fst3 $ T.matches t q + +-- | Matching keys are a prefix of the query. +prop_matches_keysArePrefix :: WS -> WTrie a -> Bool +prop_matches_keysArePrefix (WS q) (WT t) = + all (`S.isPrefixOf` q) . map fst3 $ T.matches t q + +_eqHead :: (Eq a) => Maybe a -> [a] -> Bool +_eqHead Nothing [] = True +_eqHead (Just x) (y:_) = x == y +_eqHead _ _ = False + +_eqLast :: (Eq a) => Maybe a -> [a] -> Bool +_eqLast Nothing [] = True +_eqLast (Just x) ys@(_:_) = x == last ys +_eqLast _ _ = False + +prop_minMatch_is_first_matches :: Eq a => WS -> WTrie a -> Bool +prop_minMatch_is_first_matches (WS q) (WT t) = + _eqHead (T.minMatch t q) (T.matches t q) + +prop_match_is_last_matches :: Eq a => WS -> WTrie a -> Bool +prop_match_is_last_matches (WS q) (WT t) = + _eqLast (T.match t q) (T.matches t q) + +prop_minAssoc_is_first_toList :: Eq a => WTrie a -> Bool +prop_minAssoc_is_first_toList (WT t) = + _eqHead (TI.minAssoc t) (T.toList t) + +prop_maxAssoc_is_last_toList :: Eq a => WTrie a -> Bool +prop_maxAssoc_is_last_toList (WT t) = + _eqLast (TI.maxAssoc t) (T.toList t) + +view2assoc :: Maybe (S.ByteString, a, T.Trie a) -> Maybe (S.ByteString, a) +view2assoc Nothing = Nothing +view2assoc (Just (k,v,_)) = Just (k,v) + +-- TODO: again, can we actually generate any interesting functions here? +prop_updateMinViewBy_gives_minAssoc :: Eq a => (WS -> a -> Maybe a) -> WTrie a -> Bool +prop_updateMinViewBy_gives_minAssoc f = + ((view2assoc . TI.updateMinViewBy (f . WS)) .==. TI.minAssoc) . unWT + +prop_updateMaxViewBy_gives_maxAssoc :: Eq a => (WS -> a -> Maybe a) -> WTrie a -> Bool +prop_updateMaxViewBy_gives_maxAssoc f = + ((view2assoc . TI.updateMaxViewBy (f . WS)) .==. TI.maxAssoc) . unWT + +view2trie :: Maybe (S.ByteString, a, T.Trie a) -> T.Trie a +view2trie Nothing = T.empty +view2trie (Just (_,_,t)) = t + +prop_updateMinViewBy_ident :: Eq a => WTrie a -> Bool +prop_updateMinViewBy_ident = + ((view2trie . TI.updateMinViewBy (\_ v -> Just v)) .==. id) . unWT + +prop_updateMaxViewBy_ident :: Eq a => WTrie a -> Bool +prop_updateMaxViewBy_ident = + ((view2trie . TI.updateMaxViewBy (\_ v -> Just v)) .==. id) . unWT + +prop_keys :: WTrie a -> Bool +prop_keys = (T.keys .==. (fmap fst . T.toList)) . unWT + +prop_elems :: (Eq a) => WTrie a -> Bool +prop_elems = (T.elems .==. (fmap snd . T.toList)) . unWT + +---------------------------------------------------------------- +-- ~~~~~ 'Foldable' properties + +-- * The laws as stated in the documentation. + +prop_foldmapIdentity :: (Monoid m, Eq m) => WTrie m -> Bool +prop_foldmapIdentity = (F.fold .==. F.foldMap id) . unWT + +prop_foldmapFusion :: (Monoid m, Eq m) => (a -> m) -> WTrie a -> Bool +prop_foldmapFusion f = (F.foldMap f .==. (F.fold . fmap f)) . unWT + -- entails: foldMap f . fmap g ≡ foldMap (f . g) + +{- -- TODO: +foldr f z t ≡ appEndo (F.foldMap (Endo . f) t ) z +foldl f z t ≡ appEndo (getDual (F.foldMap (Dual . Endo . flip f) t)) z +-} + +-- By definition we already have @T.size ≡ F.length@. +prop_length :: WTrie a -> Bool +prop_length = (F.length .==. getSum . F.foldMap (Sum . const 1)) . unWT + +{- -- TODO: handle the cases of old @base@ +prop_sum :: (Num a, Eq a) => WTrie a -> Bool +prop_sum = (F.sum .==. (getSum . F.foldMap' Sum)) . unWT + +prop_product :: (Num a, Eq a) => WTrie a -> Bool +prop_product = (F.product .==. (getProduct . F.foldMap' Product)) . unWT + +-- These two laws aren't stated in the documentation, but obviously should be. +prop_minimum :: (Ord a, Eq a) => NonEmpty (WTrie a) -> Bool +prop_minimum = (F.minimum .==. (getMin . F.foldMap' Min)) . unWT + +prop_maximum :: (Ord a, Eq a) => NonEmpty (WTrie a) -> Bool +prop_maximum = (F.maximum .==. (getMax . F.foldMap' Max)) . unWT +-} + + +-- * At least make sure these at least have the same order. +-- TODO: better equivalence tests + +prop_foldr_vs_foldrWithKey :: Eq a => WTrie a -> Bool +prop_foldr_vs_foldrWithKey = + (F.foldr (:) [] .==. TI.foldrWithKey (const (:)) []) . unWT + +prop_foldr'_vs_foldrWithKey' :: Eq a => WTrie a -> Bool +prop_foldr'_vs_foldrWithKey' = + (F.foldr' (:) [] .==. TI.foldrWithKey' (const (:)) []) . unWT + +prop_foldl_vs_foldlWithKey :: Eq a => WTrie a -> Bool +prop_foldl_vs_foldlWithKey = + (F.foldl snoc [] .==. TI.foldlWithKey (const . snoc) []) . unWT + where + snoc = flip (:) + +prop_foldl'_vs_foldlWithKey' :: Eq a => WTrie a -> Bool +prop_foldl'_vs_foldlWithKey' = + (F.foldl' snoc [] .==. TI.foldlWithKey' (const . snoc) []) . unWT + where + snoc = flip (:) + +#if MIN_VERSION_base(4,6,0) +prop_foldr_vs_foldr' :: (Eq a) => WTrie a -> Bool +prop_foldr_vs_foldr' = (F.foldr (:) [] .==. F.foldr' (:) []) . unWT + +prop_foldl_vs_foldl' :: (Eq a) => WTrie a -> Bool +prop_foldl_vs_foldl' = (F.foldl snoc [] .==. F.foldl' snoc []) . unWT + where + snoc = flip (:) +#endif + +#if MIN_VERSION_base(4,13,0) +-- TODO: use a non-commutative Monoid, to ensure the order is the same. +prop_foldMap_vs_foldMap' :: (Num a, Eq a) => WTrie a -> Bool +prop_foldMap_vs_foldMap' = (F.foldMap Sum .==. F.foldMap' Sum) . unWT +#endif + +-- TODO: how can we best test that fold{l,r,Map}{,'} are sufficiently lazy\/strict? +-- See: + +-- TODO: #if MIN_VERSION_base(4,8,0), check that 'F.null' isn't cyclic definition + +---------------------------------------------------------------- +-- ~~~~~ 'Traversable' laws + +#ifndef HideUnusedTopBonds +-- | Horizontal composition of Kleisli arrows. +under :: Functor f => (b -> g c) -> (a -> f b) -> a -> Compose f g c +under g f = Compose . fmap g . f +#endif + +{- +-- TODO: We can't really autogenerate these, but maybe we could +-- have a select few interesting ones? +-- +-- | An \"applicative transformation\". +data AT f g where + AT :: (Applicative f, Applicative g) => forall a. f a -> g a + -- Such that: + -- @t (pure x) ≡ pure x@ + -- @t (f <*> x) ≡ t f <*> t x@ + +prop_TraverseNaturality + :: (Applicative f, Applicative g, Eq1 g, Eq b) + => AT f g -> (a -> f b) -> WTrie a -> Bool +prop_TraverseNaturality (AT t) f = + ((t . traverse f) .==. traverse (t . f)) . unWT +-} + +prop_TraverseIdentity :: Eq a => WTrie a -> Bool +prop_TraverseIdentity = (traverse Identity .==. Identity) . unWT + +-- TODO: how could we actually test this? Any good @f,g@ to try in particular? +#ifndef HideUnusedTopBonds +prop_TraverseComposition + :: (Applicative f, Applicative g, Eq1 g, Eq1 f, Eq c) + => (b -> g c) -> (a -> f b) -> WTrie a -> Bool +prop_TraverseComposition g f = + (traverse (g `under` f) .==. (traverse g `under` traverse f)) . unWT +#endif + +-- NOTE: The next two are redundant, since we don't provide a bespoke +-- implementation of 'sequenceA'. + +{- +prop_SequenceaNaturality + :: (Applicative f, Applicative g, Eq1 g, Eq a) + => AT f g -> WTrie (f a) -> Bool +prop_SequenceaNaturality (AT t) = + ((t . sequenceA) .==. (sequenceA . fmap t)) . unWT +-} + +prop_SequenceaIdentity :: Eq a => WTrie a -> Bool +prop_SequenceaIdentity = + ((sequenceA . fmap Identity) .==. Identity) . unWT + +-- TODO: how could we actually test this? Any good @f,g@ to try in particular? +#ifndef HideUnusedTopBonds +prop_SequenceaComposition + :: (Applicative f, Applicative g, Eq1 g, Eq1 f, Eq a) + => WTrie (g (f a)) -> Bool +prop_SequenceaComposition = + ((sequenceA . fmap Compose) .==. (sequenceA `under` sequenceA)) . unWT +#endif + +-- TODO: reiterate the laws for 'traverseWithKey'... + +---------------------------------------------------------------- +-- ~~~~~ 'Filterable' laws +-- TODO: how could we actually test these effectively? Instead of +-- just with particular example functions? + +-- | Shorthand alias for a filtering function. +type F a b = a -> Maybe b + +-- | Shorthand alias for a predicate function. +type P a = a -> Bool + +prop_FiltermapFission :: Eq b => F a b -> WTrie a -> Bool +prop_FiltermapFission f = + (T.filterMap f .==. (fmap (fromJust . f) . TI.filter (isJust . f))) . unWT + +prop_FiltermapFusion :: Eq b => (a -> b) -> P a -> WTrie a -> Bool +prop_FiltermapFusion f g = + ((fmap f . TI.filter g) .==. T.filterMap (\v -> f v <$ guard (g v))) . unWT + +prop_FiltermapConservation :: Eq b => (a -> b) -> WTrie a -> Bool +prop_FiltermapConservation f = + (T.filterMap (Just . f) .==. fmap f) . unWT + +#ifndef HideUnusedTopBonds +prop_FiltermapVertComposition :: Eq c => F b c -> F a b -> WTrie a -> Bool +prop_FiltermapVertComposition f g = + ((T.filterMap f . T.filterMap g) .==. T.filterMap (f <=< g)) . unWT +#endif + +prop_FilterDefinition :: Eq a => P a -> WTrie a -> Bool +prop_FilterDefinition f = + (TI.filter f .==. T.filterMap (\v -> v <$ guard (f v))) . unWT + +prop_FilterVertComposition :: Eq a => P a -> P a -> WTrie a -> Bool +prop_FilterVertComposition f g = + ((TI.filter f . TI.filter g) .==. TI.filter (liftA2 (&&) f g)) . unWT + +-- A handful of interesting predicates... +isTriple :: Integral a => P a +isTriple x = rem x 3 == 0 + +remTriple :: Integral a => F a a +remTriple x + | r == 0 = Nothing + | otherwise = Just r + where r = rem x 3 + +#ifndef HideUnusedTopBonds +quotTriple :: Integral a => F a a +quotTriple x + | r == 0 = Just q + | otherwise = Nothing + where (q,r) = quotRem x 3 +#endif + +---------------------------------------------------------------- +-- ~~~~~ 'Witherable' laws +-- TODO: how could we actually test these effectively? + +#ifndef HideUnusedTopBonds +-- | Shorthand alias for a withering function. +type FA a f b = a -> f (Maybe b) +#endif + +{- +prop_WitherNaturality + :: (Applicative f, Applicative g, Eq1 g, Eq b) + => AT f g -> FA a f b -> WTrie a -> Bool +prop_WitherNaturality (AT t) f = + (TI.wither (t . f) .==. (t . TI.wither f)) . unWT +-} + +{- +-- TODO: how can we nail down the functor @f@ but in a way that's portable beyond GHC? +prop_WitherPurity :: (Applicative f, Eq1 f, Eq b) => F a b -> WTrie a -> Bool +prop_WitherPurity f = + (TI.wither (pure . f) .==. (pure . T.filterMap f)) . unWT +-} + +-- HACK: adding the 'Compose' wrapper to force GHC to deduce @Eq +-- (f (T.Trie b))@ from the provided @(Eq1 f, Eq b)@. +prop_WitherConservation + :: (Applicative f, Eq1 f, Eq b) => (a -> f b) -> WTrie a -> Bool +prop_WitherConservation f = + ((Compose . TI.wither (fmap Just . f)) .==. (Compose . traverse f)) . unWT + +#ifndef HideUnusedTopBonds +prop_WitherHorizComposition + :: (Applicative f, Applicative g, Eq1 g, Eq1 f, Eq c) + => FA b g c -> FA a f b -> WTrie a -> Bool +prop_WitherHorizComposition f g = + ((TI.wither f `under` TI.wither g) .==. TI.wither (wither_Maybe f `under` g)) . unWT + +-- | Variant of wither for Maybe instead of Trie. +wither_Maybe :: Applicative f => FA a f b -> Maybe a -> f (Maybe b) +wither_Maybe f = fmap join . traverse f +#endif + +-- | Shorthand alias for an effectful predicate function. +type PA a f = a -> f Bool + +-- HACK: adding the 'Compose' wrapper to force GHC to deduce @Eq +-- (f (T.Trie a))@ from the provided @(Eq1 f, Eq a)@. +prop_FilteraDefinition :: (Applicative f, Eq1 f, Eq a) => PA a f -> WTrie a -> Bool +prop_FilteraDefinition f = + ((Compose . TI.filterA f) .==. (Compose . TI.wither (\v -> (v <$) . guard <$> f v))) . unWT + +{- +prop_FilteraNaturality + :: (Applicative f, Applicative g, Eq1 g, Eq a) + => AT f g -> PA a f -> WTrie a -> Bool +prop_FilteraNaturality (AT t) f = + (TI.filterA (t . f) .==. (t . TI.filterA f)) . unWT +-} + +{- +-- TODO: how can we nail down the @f@ but in a way that's portable beyond GHC? +prop_FilteraPurity :: (Functor f, Eq1 f, Eq a) => P a -> WTrie a -> Bool +prop_FilteraPurity f = + (TI.filterA (pure . f) .==. (pure . TI.filter f)) . unWT +-} + +#ifndef HideUnusedTopBonds +prop_FilteraHorizComposition + :: (Applicative f, Applicative g, Eq1 g, Eq1 f, Eq a) + => PA a f -> PA a g -> WTrie a -> Bool +prop_FilteraHorizComposition f g = + ((TI.filterA f `under` TI.filterA g) .==. TI.filterA (underA2 (&&) f g)) . unWT + +underA2 :: (Applicative f, Applicative g) + => (b -> c -> d) + -> (a -> g b) + -> (a -> f c) + -> a -> Compose f g d +underA2 h g f = liftA2 (liftA2 h) (g `under` pure) (pure `under` f) + +prop_FilteraHorizComposition_alt + :: (Applicative f, Applicative g, Eq1 g, Eq1 f, Eq a) + => PA a f -> PA a g -> WTrie a -> Bool +prop_FilteraHorizComposition_alt f g = + ((TI.filterA f `under` TI.filterA g) .==. TI.filterA (underF2 (&&) g f)) . unWT + +underF2 :: (Functor f, Functor g) + => (b -> c -> d) + -> (a -> f b) + -> (a -> g c) + -> a -> Compose f g d +underF2 h f g a = Compose (f a <&> ((g a <&>) . h)) +#endif + + +---------------------------------------------------------------- +-- | If there are duplicate keys in the @assocs@, then @f@ will +-- take the first value. +_takes_first :: (Eq c) => ([(S.ByteString, c)] -> T.Trie c) -> [(WS, c)] -> Bool +_takes_first f assocs = + (T.toList . f) .==. (nubBy (apFst (==)) . sortBy (comparing fst)) + $ map (first unWS) assocs + +-- | Lift a function to apply to the 'fst' of pairs, retaining the 'snd'. +first :: (a -> b) -> (a,c) -> (b,c) +first f (x,y) = (f x, y) + +-- | Lift a function to apply to the 'snd' of pairs, retaining the 'fst'. +second :: (b -> c) -> (a,b) -> (a,c) +second f (x,y) = (x, f y) + +-- | Lift a binary function to apply to the first of pairs, discarding seconds. +apFst :: (a -> b -> c) -> ((a,d) -> (b,e) -> c) +apFst f (x,_) (y,_) = f x y + +-- | 'T.fromList' takes the first value for a given key. +prop_fromList_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromList_takes_first = _takes_first T.fromList + +-- | 'T.fromListR' takes the first value for a given key. +prop_fromListR_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromListR_takes_first = _takes_first TC.fromListR + +-- | 'T.fromListL' takes the first value for a given key. +prop_fromListL_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromListL_takes_first = _takes_first TC.fromListL + +-- | 'T.fromListS' takes the first value for a given key. +prop_fromListS_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromListS_takes_first = _takes_first TC.fromListS + +-- | @('TC.fromListWith' const)@ takes the first value for a given key. +prop_fromListWithConst_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromListWithConst_takes_first = _takes_first (TC.fromListWith const) + +-- | @('TC.fromListWithL' const)@ takes the first value for a given key. +prop_fromListWithLConst_takes_first :: (Eq a) => [(WS, a)] -> Bool +prop_fromListWithLConst_takes_first = _takes_first (TC.fromListWithL const) + +---------------------------------------------------------------- +-- ~~~~~ 'Functor' properties + +prop_FunctorIdentity :: Eq a => WTrie a -> Bool +prop_FunctorIdentity = (fmap id .==. id) . unWT + +{- -- TODO: is there any way to make this remotely testable? +prop_FunctorCompose :: Eq c => (b -> c) -> (a -> b) -> WTrie a -> Bool +prop_FunctorCompose f g = (fmap (f . g) .==. (fmap f . fmap g)) . unWT +-} + +#if __GLASGOW_HASKELL__ +-- | @(<$)@ agrees with its default definition. +prop_fmapConst :: Eq a => a -> WTrie b -> Bool +prop_fmapConst v = ((v <$) .==. fmap (const v)) . unWT +#endif + +{- +-- Both of these test only a subset of what 'prop_fmap_toList' tests. I was hoping they'd help simplify the function-generation problem, but if we're testing 'prop_fmap_toList' anyways then there's no point in testing these too. + +-- | 'fmap' doesn't affect the keys. This is safe to call with an +-- undefined function, thereby proving that the function cannot +-- affect things. +prop_fmap_keys :: Eq b => (a -> b) -> WTrie a -> Bool +prop_fmap_keys f = ((T.keys . fmap f) .==. T.keys) . unWT + +prop_fmap_elems :: Eq b => (a -> b) -> WTrie a -> Bool +prop_fmap_elems f = ((T.elems . fmap f) .==. (map f . T.elems)) . unWT +-} + +-- TODO: is there any way to generate halfway useful functions for testing here? +prop_fmap_toList :: Eq b => (a -> b) -> WTrie a -> Bool +prop_fmap_toList f = + ((T.toList . fmap f) .==. (map (second f) . T.toList)) . unWT + +---------------------------------------------------------------- +prop_filterMap_ident :: Eq a => WTrie a -> Bool +prop_filterMap_ident = (T.filterMap Just .==. id) . unWT + +prop_filterMap_empty :: Eq a => WTrie a -> Bool +prop_filterMap_empty = (T.filterMap const_Nothing .==. const T.empty) . unWT + where + -- Have to fix the result type here. + const_Nothing :: a -> Maybe a + const_Nothing = const Nothing + +justConst :: a -> b -> Maybe a +justConst x _ = Just x + +prop_mapBy_keys :: WTrie a -> Bool +prop_mapBy_keys = all (uncurry (==)) . T.toList . T.mapBy justConst . unWT + +prop_contextualMap_ident :: Eq a => WTrie a -> Bool +prop_contextualMap_ident = (TI.contextualMap const .==. id) . unWT + +prop_contextualMap'_ident :: Eq a => WTrie a -> Bool +prop_contextualMap'_ident = (TI.contextualMap' const .==. id) . unWT + +prop_contextualFilterMap_ident :: Eq a => WTrie a -> Bool +prop_contextualFilterMap_ident = + (TI.contextualFilterMap justConst .==. id) . unWT + +prop_contextualMapBy_keys :: WTrie a -> Bool +prop_contextualMapBy_keys = + all (uncurry (==)) . T.toList . TI.contextualMapBy f . unWT + where + f k _ _ = Just k + +prop_contextualMapBy_ident :: Eq a => WTrie a -> Bool +prop_contextualMapBy_ident = (TI.contextualMapBy f .==. id) . unWT + where + f _ v _ = Just v + +prop_contextualMapBy_empty :: Eq a => WTrie a -> Bool +prop_contextualMapBy_empty = (TI.contextualMapBy f .==. const T.empty) . unWT + where + -- Have to fix the result type here. + f :: S.ByteString -> a -> T.Trie a -> Maybe a + f _ _ _ = Nothing + +---------------------------------------------------------------- +-- ~~~~~ 'Applicative' laws + +prop_ApplicativeIdentity :: Eq a => WTrie a -> Bool +prop_ApplicativeIdentity = ((pure id <*>) .==. id) . unWT + +{- -- (remaining, untestable) Applicative laws +prop_ApplicativeCompose = pure (.) <*> u <*> v <*> w == u <*> (v <*> w) +prop_ApplicativeHom = pure f <*> pure x == pure (f x) +prop_ApplicativeInterchange = u <*> pure y == pure ($ y) <*> u +-} + +---------------------------------------------------------------- +-- ~~~~~ 'Monad' laws + +prop_MonadIdentityR :: Eq a => WTrie a -> Bool +prop_MonadIdentityR = ((>>= return) .==. id) . unWT + +{- -- (remaining, untestable) Monad laws +prop_MonadIdentityL = (return a >>= k) == k a +prop_MonadAssoc = m >>= (\x -> k x >>= h) == (m >>= k) >>= h +-} + +---------------------------------------------------------------- +-- ~~~~~ 'Semigroup' laws + +#if MIN_VERSION_base(4,9,0) +prop_Semigroup :: (Semigroup a, Eq a) => WTrie a -> WTrie a -> WTrie a -> Bool +prop_Semigroup (WT a) (WT b) (WT c) = a <> (b <> c) == (a <> b) <> c + +-- TODO: laws for 'stimes' +#endif + +---------------------------------------------------------------- +-- ~~~~~ 'Monoid' laws + +-- N.B., base-4.11.0.0 is when Semigroup became superclass of Monoid +prop_MonoidIdentityL :: (Monoid a, Eq a) => WTrie a -> Bool +prop_MonoidIdentityL = ((mempty `mappend`) .==. id) . unWT + +prop_MonoidIdentityR :: (Monoid a, Eq a) => WTrie a -> Bool +prop_MonoidIdentityR = ((`mappend` mempty) .==. id) . unWT + +---------------------------------------------------------------- +-- ~~~~~ 'Binary' laws + +prop_Binary :: (B.Binary a, Eq a) => WTrie a -> Bool +prop_Binary = ((B.decode . B.encode) .==. id) . unWT + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/dev/Test/Utils.hs b/dev/Test/Utils.hs new file mode 100644 --- /dev/null +++ b/dev/Test/Utils.hs @@ -0,0 +1,254 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# LANGUAGE CPP + , MultiParamTypeClasses + , FlexibleInstances + , FlexibleContexts + #-} + +---------------------------------------------------------------- +-- ~ 2021.12.14 +-- | +-- Module : Test.Utils +-- Copyright : 2008--2023 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional +-- Portability : semi-portable (MPTC,...) +-- +-- Utilities for testing 'Trie's. +---------------------------------------------------------------- +module Test.Utils + ( packC2W, vocab2trie + , testEqual + , W(..), everyW + , WS(..), packWS, unpackWS + , WTrie(..) + , CheckGuard(..), (.==>.), (.==.) + ) where + +import qualified Data.Trie as T +import Data.Word (Word8) +import qualified Data.ByteString as S +import qualified Data.ByteString.Internal as S (c2w, w2c) +import Data.ByteString.Internal (ByteString(PS)) +import Control.Monad ((<=<)) + +-- N.B., "Test.Tasty.HUnit" does not in fact depend on "Test.HUnit"; +-- hence using the longer alias. +import qualified Test.Tasty as Tasty +import qualified Test.Tasty.HUnit as TastyHU +import qualified Test.QuickCheck as QC +import qualified Test.SmallCheck as SC +import qualified Test.SmallCheck.Series as SC +-- import qualified Test.LazySmallCheck as LSC +-- import qualified Test.SparseCheck as PC + +---------------------------------------------------------------- +---------------------------------------------------------------- + +-- TODO: apparently this is exported as 'Data.ByteString.Internal.packChars' +-- | Construct a bytestring from the first byte of each 'Char'. +packC2W :: String -> S.ByteString +packC2W = S.pack . map S.c2w + +-- | Construct a trie via 'packC2W' giving each key a unique value +-- (namely its position in the list). +vocab2trie :: [String] -> T.Trie Int +vocab2trie = T.fromList . flip zip [0..] . map packC2W + +---------------------------------------------------------------- +-- TODO: come up with a thing that pretty-prints the diff, instead +-- of just showing the expected\/actual. +testEqual :: (Show a, Eq a) => String -> a -> a -> Tasty.TestTree +testEqual name expected actual = + TastyHU.testCase name (TastyHU.assertEqual "" expected actual) + +---------------------------------------------------------------- +-- | A small subset of 'Word8', so that 'WS' is more likely to have +-- shared prefixes. The 'Show' instance shows it as a 'Char', for +-- better legibility and for consistency with the 'Show' instance +-- of 'WS'. +newtype W = W { unW :: Word8 } + deriving (Eq, Ord) + +instance Show W where + showsPrec p = showsPrec p . S.w2c . unW + +-- TODO: ensure that these have good bit-patterns for covering corner cases. +-- | All the possible 'W' values; or rather, all the ones generated +-- by the 'QC.Arbitrary' and 'SC.Serial' instances. +everyW :: [W] +everyW = (W . S.c2w) <$> ['a'..'m'] + +-- TODO: if we define (Enum W) then we could use 'QC.chooseEnum' +-- which is much faster than 'QC.elements'. Alternatively we might +-- consider using 'QC.growingElements' if we want something more +-- like what the SC.Serial case does. +instance QC.Arbitrary W where + arbitrary = QC.elements everyW + shrink w = takeWhile (w /=) everyW + +instance QC.CoArbitrary W where + coarbitrary = QC.coarbitrary . unW + +-- We take @(d+1)@ to match the instances for 'Char', (SC.N a), etc +instance Monad m => SC.Serial m W where + series = SC.generate (\d -> take (d+1) everyW) + +instance Monad m => SC.CoSerial m W where + coseries = fmap (. unW) . SC.coseries + +---------------------------------------------------------------- +-- TODO: we need a better instance of Arbitrary for lists to make +-- them longer than our smallcheck depth. +-- +-- | A subset of 'S.ByteString' produced by 'packWS'. +-- This newtype is to ensure that generated bytestrings are more +-- likely to have shared prefixes (and thus non-trivial tries). +newtype WS = WS { unWS :: S.ByteString } + deriving (Eq, Ord) + +instance Show WS where + showsPrec p = showsPrec p . unWS + +packWS :: [W] -> WS +packWS = WS . S.pack . map unW + +unpackWS :: WS -> [W] +unpackWS = map W . S.unpack . unWS + +-- | Like 'S.inits' but each step keeps half more, rather than just one more. +prefixes :: WS -> [WS] +prefixes (WS (PS x s l)) = + [WS (PS x s (l - k)) | k <- takeWhile (> 0) (iterate (`div` 2) l)] + +instance QC.Arbitrary WS where + arbitrary = QC.sized $ \n -> do + k <- QC.chooseInt (0,n) + xs <- QC.vector k + return $ packWS xs + shrink = QC.shrinkMap packWS unpackWS <=< prefixes + +instance QC.CoArbitrary WS where + coarbitrary = QC.coarbitrary . unpackWS + +instance Monad m => SC.Serial m WS where + series = packWS <$> SC.series + +-- TODO: While this is a perfectly valid instance, is it really the +-- most efficient one for our needs? +instance Monad m => SC.CoSerial m WS where + coseries rs = + SC.alts0 rs SC.>>- \z -> + SC.alts2 rs SC.>>- \f -> + return $ \(WS xs) -> + if S.null xs + then z + else f (W $ S.head xs) (WS $ S.tail xs) + +---------------------------------------------------------------- +-- | A subset of 'T.Trie' where all the keys are 'WS'. This newtype +-- is mainly just to avoid orphan instances. +newtype WTrie a = WT { unWT :: T.Trie a } + deriving (Eq) + +instance Show a => Show (WTrie a) where + showsPrec p = showsPrec p . unWT + +first :: (b -> c) -> (b,d) -> (c,d) +first f (x,y) = (f x, y) + +-- TODO: maybe we ought to define @T.fromListBy@ for better fusion? +fromListWT :: [(WS,a)] -> WTrie a +fromListWT = WT . T.fromList . map (first unWS) + +-- We can use 'T.toListBy' to manually fuse with the map +toListWT :: WTrie a -> [(WS,a)] +toListWT = map (first WS) . T.toList . unWT + +instance (QC.Arbitrary a) => QC.Arbitrary (WTrie a) where + arbitrary = QC.sized $ \n -> do + k <- QC.chooseInt (0,n) + labels <- QC.vector k + elems <- QC.vector k + return . fromListWT $ zip labels elems + -- Extremely inefficient, but should be effective at least. + shrink = QC.shrinkMap fromListWT toListWT + +-- TODO: instance QC.CoArbitrary (WTrie a) + +-- TODO: This instance really needs some work. The smart constructures +-- ensure only valid values are generated, but there are redundancies +-- and inefficiencies. +instance (Monad m, SC.Serial m a) => SC.Serial m (WTrie a) where + series = SC.cons0 (WT T.empty) + SC.\/ SC.cons3 arcHACK + SC.\/ SC.cons2 branch + where + arcHACK (WS k) mv (WT t) = + case mv of + Nothing -> WT (T.singleton k () >> t) + Just v -> WT (T.singleton k v >>= T.unionR t . T.singleton S.empty) + + branch (WT t0) (WT t1) = WT (t0 `T.unionR` t1) + +-- TODO: instance Monad m => SC.CoSerial m (WTrie a) + +---------------------------------------------------------------- +---------------------------------------------------------------- + +infixr 0 ==>, .==>. +infix 4 .==. + +{- +-- TODO: clean up something like this: +class ForAll src p q where + forAll :: forall a. (Show a) => src a -> (a -> p) -> q +instance (QC.Testable p) => ForAll QC.Gen p QC.Property where + forAll gen pf = QC.forAllShrink gen QC.shrink pf +instance (SC.Testable m p) => ForAll (SC.Series m) p (SC.Property m) where + forAll srs pf = SC.forAll (SC.over srs pf) + +class SuchThat src where + suchThat :: forall a. src a -> (a -> Bool) -> src a +instance SuchThat QC.Gen where + suchThat = QC.suchThat +instance SuchThat (SC.Series m) where + suchThat = flip Control.Monad.mfilter + +class Generable src a where + generate :: src a +instance (QC.Arbitrary a) => Generable QC.Gen a where + generate = QC.arbitrary +instance (SC.Serial m a) => Generable (SC.Series m) a where + generate = SC.series + +forEach :: (Forall src p q, SuchThat src, Generable src a, Show a) => (a -> Bool) -> (a -> p) -> q +forEach = forAll . suchThat generate +-} + + +-- | Deal with QC\/SC polymorphism issues because of @(==>)@. +-- Fundeps would be nice here, but @|b->a@ is undecidable, and @|a->b@ is wrong. +class CheckGuard p q where + (==>) :: Bool -> p -> q + +instance (QC.Testable p) => CheckGuard p QC.Property where + (==>) = (QC.==>) + -- TODO: might should also use 'QC.cover' with this. + -- TODO: or we may prefer to rephrase things to use 'QC.suchThat' instead (should be sufficient for our particular use case, if we can find a smallcheck analogue (probably 'SC.over')) + +instance (Monad m, SC.Testable m p) => CheckGuard p (SC.Property m) where + (==>) = (SC.==>) + +-- | Lifted implication. +(.==>.) :: CheckGuard testable prop => (a -> Bool) -> (a -> testable) -> (a -> prop) +(.==>.) p q x = p x ==> q x + +-- | Function equality / lifted equality. +(.==.) :: (Eq b) => (a -> b) -> (a -> b) -> (a -> Bool) +(.==.) f g x = f x == g x + -- TODO: should use (QC.===) or diy with QC.counterexample; assuming we can overload that for smallcheck equivalent (or for smallcheck to ignore and fall back to (==)) + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/src/Data/Trie.hs b/src/Data/Trie.hs --- a/src/Data/Trie.hs +++ b/src/Data/Trie.hs @@ -1,143 +1,213 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} - +{-# LANGUAGE NoImplicitPrelude, CPP #-} +#if __GLASGOW_HASKELL__ >= 701 +{-# LANGUAGE Safe #-} +#endif ---------------------------------------------------------------- --- ~ 2009.01.05 +-- ~ 2021.12.14 -- | -- Module : Data.Trie --- Copyright : Copyright (c) 2008--2009 wren ng thornton --- License : BSD3 --- Maintainer : wren@community.haskell.org --- Stability : beta +-- Copyright : 2008--2025 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional -- Portability : portable -- -- An efficient implementation of finite maps from strings to values. --- -- The implementation is based on /big-endian patricia trees/, like -- "Data.IntMap". We first trie on the elements of "Data.ByteString" -- and then trie on the big-endian bit representation of those -- elements. For further details on the latter, see -- -- * Chris Okasaki and Andy Gill, \"/Fast Mergeable Integer Maps/\", --- Workshop on ML, September 1998, pages 77-86, --- +-- Workshop on ML, September 1998, pages 77-86, +-- -- -- * D.R. Morrison, \"/PATRICIA -- Practical Algorithm To Retrieve/ --- /Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), --- October 1968, pages 514-534. +-- /Information Coded In Alphanumeric/\", Journal of the ACM, 15(4), +-- October 1968, pages 514-534. +-- +-- This module aims to provide an austere interface, while being +-- detailed enough for most users. For an extended interface with +-- many additional functions, see "Data.Trie.Convenience". For +-- functions that give more detailed (potentially abstraction-breaking) +-- access to the data structure, or for experimental functions +-- which aren't quite ready for the public API, see "Data.Trie.Internal". ---------------------------------------------------------------- module Data.Trie ( - -- * Data types - Trie(), KeyString, KeyElem - + -- * Data type + Trie() + -- * Basic functions , empty, null, singleton, size - + -- * Conversion functions - , fromList, toListBy, toList, keys - + , fromList, toListBy, toList, keys, elems + -- * Query functions - , lookupBy, lookup, member, submap - - -- * Single-value modification - , alterBy, insert, adjust, delete - + , lookupBy, lookup, member, submap, match, minMatch, matches + + -- * Simple modification + , insert, adjust, adjustBy, alterBy, delete, deleteSubmap + -- * Combining tries , mergeBy, unionL, unionR - + , intersectBy, intersectL, intersectR + -- * Mapping functions , mapBy, filterMap + -- TODO: re-export 'filter', 'filterA', 'wither', etc? ) where -import Prelude hiding (null, lookup) -import qualified Prelude +import Prelude hiding (null, lookup) import Data.Trie.Internal - -import Data.Maybe (isJust) -import Control.Monad (liftM) +import Data.ByteString (ByteString) +import qualified Data.ByteString as S +import Data.Maybe (isJust) ---------------------------------------------------------------- ---------------------------------------------------------------- {--------------------------------------------------------------- --- Conversion functions +-- Conversion functions ---------------------------------------------------------------} --- | Convert association list into a trie. On key conflict, values --- earlier in the list shadow later ones. -{-# INLINE fromList #-} -fromList :: [(KeyString,a)] -> Trie a -fromList = foldr (uncurry insert) empty - --- | Convert trie into association list. Keys will be in sorted order. -toList :: Trie a -> [(KeyString,a)] -toList = toListBy (,) - -- | Return all keys in the trie, in sorted order. -keys :: Trie a -> [KeyString] +-- +-- __Warning__: This function suffers . +keys :: Trie a -> [ByteString] +{-# INLINE keys #-} keys = toListBy const + {--------------------------------------------------------------- -- Query functions (just recurse) ---------------------------------------------------------------} -- | Generic function to find a value (if it exists) and the subtrie -- rooted at the prefix. +lookupBy :: (Maybe a -> Trie a -> b) -> ByteString -> Trie a -> b {-# INLINE lookupBy #-} -lookupBy :: (Maybe a -> Trie a -> b) -> KeyString -> Trie a -> b -lookupBy f = lookupBy_ f (f Nothing empty) (f Nothing) +lookupBy f = lookupBy_ (f . Just) (f Nothing) (f Nothing empty) -- | Return the value associated with a query string if it exists. +lookup :: ByteString -> Trie a -> Maybe a {-# INLINE lookup #-} -lookup :: KeyString -> Trie a -> Maybe a -lookup = lookupBy_ const Nothing (const Nothing) +lookup = lookupBy const --- TODO? move to "Data.Trie.Conventience"? -- | Does a string have a value in the trie? +member :: ByteString -> Trie a -> Bool {-# INLINE member #-} -member :: KeyString -> Trie a -> Bool member q = isJust . lookup q +getMatch :: ByteString -> (Int, a) -> (ByteString, a, ByteString) +{-# INLINE getMatch #-} +getMatch q (n,x) = + case S.splitAt n q of + (p,q') -> (p, x, q') +-- | Given a query, find the longest prefix with an associated value +-- in the trie, and return that prefix, its value, and the remainder +-- of the query. +-- +-- @since 0.2.4 +match :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString) +{-# INLINE match #-} +match t q = getMatch q <$> match_ t q + +-- | Given a query, find the shortest prefix with an associated value +-- in the trie, and return that prefix, its value, and the remainder +-- of the query. +-- +-- @since 0.2.6 +minMatch :: Trie a -> ByteString -> Maybe (ByteString, a, ByteString) +{-# INLINE minMatch #-} +minMatch t q = + case matches t q of + [] -> Nothing + x:_ -> Just x + +-- | Given a query, find all prefixes with associated values in the +-- trie, and return their (prefix, value, remainder) triples in +-- order from shortest prefix to longest. This function is a good +-- producer for list fusion. +-- +-- @since 0.2.4 +matches :: Trie a -> ByteString -> [(ByteString, a, ByteString)] +{-# INLINE matches #-} +matches t q = getMatch q <$> matches_ t q + + {--------------------------------------------------------------- --- Single-value modification functions (recurse and clone spine) +-- Simple modification functions (recurse and clone spine) ---------------------------------------------------------------} -- | Insert a new key. If the key is already present, overrides the -- old value +insert :: ByteString -> a -> Trie a -> Trie a {-# INLINE insert #-} -insert :: KeyString -> a -> Trie a -> Trie a -insert = alterBy (\_ x _ -> Just x) - --- | Apply a function to the value at a key. -{-# INLINE adjust #-} -adjust :: (a -> a) -> KeyString -> Trie a -> Trie a -adjust f q = alterBy (\_ _ -> liftM f) q undefined +insert = alterBy (\_ x _ -> Just x) +-- | Alter the value associated with a given key. If the key is not +-- present, then the trie is returned unaltered. See 'alterBy' if +-- you are interested in inserting new keys or deleting old keys. +-- Because this function does not need to worry about changing the +-- trie structure, it is somewhat faster than 'alterBy'. +-- +-- @since 0.2.6 +-- __Note__: Prior to version 0.2.6 this function was exported from +-- "Data.Trie.Internal" instead. +adjustBy :: (ByteString -> a -> a -> a) + -> ByteString -> a -> Trie a -> Trie a +{-# INLINE adjustBy #-} +adjustBy f q x = adjust (f q x) q + -- | Remove the value stored at a key. +delete :: ByteString -> Trie a -> Trie a {-# INLINE delete #-} -delete :: KeyString -> Trie a -> Trie a -delete q = alterBy (\_ _ _ -> Nothing) q undefined +delete = alterBy_ (\_ t -> (Nothing, t)) +-- | Remove all keys beginning with a prefix. +-- +-- @since 0.2.6 +deleteSubmap :: ByteString -> Trie a -> Trie a +{-# INLINE deleteSubmap #-} +deleteSubmap = alterBy_ (\_ _ -> (Nothing, empty)) + {--------------------------------------------------------------- -- Trie-combining functions ---------------------------------------------------------------} --- | Combine two tries, resolving conflicts by choosing the value --- from the left trie. -{-# INLINE unionL #-} +-- | Take the union of two tries, resolving conflicts by choosing +-- the value from the left trie. unionL :: Trie a -> Trie a -> Trie a +{-# INLINE unionL #-} unionL = mergeBy (\x _ -> Just x) --- | Combine two tries, resolving conflicts by choosing the value --- from the right trie. -{-# INLINE unionR #-} +-- | Take the union of two tries, resolving conflicts by choosing +-- the value from the right trie. unionR :: Trie a -> Trie a -> Trie a +{-# INLINE unionR #-} unionR = mergeBy (\_ y -> Just y) + +-- | Take the intersection of two tries, with values from the left trie. +-- +-- @since 0.2.6 +intersectL :: Trie a -> Trie b -> Trie a +{-# INLINE intersectL #-} +intersectL = intersectBy (\x _ -> Just x) + +-- | Take the intersection of two tries, with values from the right trie. +-- +-- @since 0.2.6 +intersectR :: Trie a -> Trie b -> Trie b +{-# INLINE intersectR #-} +intersectR = intersectBy (\_ y -> Just y) ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/src/Data/Trie/BitTwiddle.hs b/src/Data/Trie/BitTwiddle.hs deleted file mode 100644 --- a/src/Data/Trie/BitTwiddle.hs +++ /dev/null @@ -1,165 +0,0 @@ -{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-name-shadowing #-} -{-# OPTIONS_GHC -cpp -fglasgow-exts #-} - ----------------------------------------------------------------- --- ~ 2009.01.05 --- | --- Module : Data.Trie.BitTwiddle --- Copyright : Copyright (c) Daan Leijen 2002 --- License : BSD3 --- Maintainer : libraries@haskell.org, wren@community.haskell.org --- Stability : provisional --- Portability : portable (with CPP) --- --- Functions to treat 'Word' as a bit-vector for big-endian patricia --- trees. This code is duplicated from "Data.IntMap". The only --- differences are that some of the conversion functions are --- specialized to 'Word8' for bytestrings, instead of being specialized --- to 'Int'. ----------------------------------------------------------------- - -module Data.Trie.BitTwiddle - ( Prefix, Mask - , elemToNat - , zero, nomatch - , mask, shorter, branchMask - ) where - -import Data.Bits -import Data.Trie.ByteStringInternal (ByteStringElem) - -#if __GLASGOW_HASKELL__ >= 503 -import GHC.Exts ( Word(..), Int(..), shiftRL# ) -#elif __GLASGOW_HASKELL__ -import GlaExts ( Word(..), Int(..), shiftRL# ) -#else -import Data.Word (Word) -#endif - ----------------------------------------------------------------- - --- TODO: Natural word size, is 4*Word8 on my machine. Which means --- it'll be more efficient to Branch by the first 4 bytes instead --- of just one... -type KeyElem = ByteStringElem -type Prefix = KeyElem -type Mask = KeyElem - -{-# INLINE elemToNat #-} -elemToNat :: KeyElem -> Word -elemToNat i = fromIntegral i - -{-# INLINE natToElem #-} -natToElem :: Word -> KeyElem -natToElem w = fromIntegral w - -{-# INLINE shiftRL #-} -shiftRL :: Word -> Int -> Word -#if __GLASGOW_HASKELL__ --- GHC: use unboxing to get @shiftRL@ inlined. -shiftRL (W# x) (I# i) = W# (shiftRL# x i) -#else -shiftRL x i = shiftR x i -#endif - - -{--------------------------------------------------------------- --- Endian independent bit twiddling (Trie endianness, not architecture) ----------------------------------------------------------------} - --- | Is the value under the mask zero? -{-# INLINE zero #-} -zero :: KeyElem -> Mask -> Bool -zero i m = (elemToNat i) .&. (elemToNat m) == 0 - --- | Does a value /not/ match some prefix, for all the bits preceding --- a masking bit? (Hence a subtree matching the value doesn't exist.) -{-# INLINE nomatch #-} -nomatch :: KeyElem -> Prefix -> Mask -> Bool -nomatch i p m = mask i m /= p - -{-# INLINE mask #-} -mask :: KeyElem -> Mask -> Prefix -mask i m = maskW (elemToNat i) (elemToNat m) - - -{--------------------------------------------------------------- --- Big endian operations (Trie endianness, not architecture) ----------------------------------------------------------------} - --- | Get mask by setting all bits higher than the smallest bit in --- @m@. Then apply that mask to @i@. -{-# INLINE maskW #-} -maskW :: Word -> Word -> Prefix -maskW i m = natToElem (i .&. (complement (m-1) `xor` m)) - --- | Determine whether the first mask denotes a shorter prefix than --- the second. -{-# INLINE shorter #-} -shorter :: Mask -> Mask -> Bool -shorter m1 m2 = elemToNat m1 > elemToNat m2 - --- | Determine first differing bit of two prefixes. -{-# INLINE branchMask #-} -branchMask :: Prefix -> Prefix -> Mask -branchMask p1 p2 - = natToElem (highestBitMask (elemToNat p1 `xor` elemToNat p2)) - -{--------------------------------------------------------------- - Finding the highest bit (mask) in a word [x] can be done efficiently - in three ways: - * convert to a floating point value and the mantissa tells us the - [log2(x)] that corresponds with the highest bit position. The - mantissa is retrieved either via the standard C function [frexp] - or by some bit twiddling on IEEE compatible numbers (float). - Note that one needs to use at least [double] precision for an - accurate mantissa of 32 bit numbers. - * use bit twiddling, a logarithmic sequence of bitwise or's and - shifts (bit). - * use processor specific assembler instruction (asm). - - The most portable way would be [bit], but is it efficient enough? - I have measured the cycle counts of the different methods on an - AMD Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC - instruction: - - highestBitMask: method cycles - -------------- - frexp 200 - float 33 - bit 11 - asm 12 - - highestBit: method cycles - -------------- - frexp 195 - float 33 - bit 11 - asm 11 - - Wow, the bit twiddling is on today's RISC like machines even - faster than a single CISC instruction (BSR)! ----------------------------------------------------------------} - -{--------------------------------------------------------------- - [highestBitMask] returns a word where only the highest bit is - set. It is found by first setting all bits in lower positions - than the highest bit and than taking an exclusive or with the - original value. Allthough the function may look expensive, GHC - compiles this into excellent C code that subsequently compiled - into highly efficient machine code. The algorithm is derived from - Jorg Arndt's FXT library. ----------------------------------------------------------------} -{-# INLINE highestBitMask #-} -highestBitMask :: Word -> Word -highestBitMask x - = case (x .|. shiftRL x 1) of - x -> case (x .|. shiftRL x 2) of - x -> case (x .|. shiftRL x 4) of - x -> case (x .|. shiftRL x 8) of - x -> case (x .|. shiftRL x 16) of - x -> case (x .|. shiftRL x 32) of -- for 64 bit platforms - x -> (x `xor` shiftRL x 1) - ----------------------------------------------------------------- ------------------------------------------------------------ fin. diff --git a/src/Data/Trie/ByteStringInternal.hs b/src/Data/Trie/ByteStringInternal.hs deleted file mode 100644 --- a/src/Data/Trie/ByteStringInternal.hs +++ /dev/null @@ -1,175 +0,0 @@ -{-# OPTIONS_GHC -Wall -fwarn-tabs #-} - ----------------------------------------------------------------- --- ~ 2008.12.24 --- | --- Module : Data.Trie.ByteStringInternal --- Copyright : Copyright (c) 2008--2009 wren ng thornton --- License : BSD3 --- Maintainer : wren@community.haskell.org --- Stability : beta --- Portability : portable --- --- Helper functions on 'ByteString's for "Data.Trie". ----------------------------------------------------------------- - - -module Data.Trie.ByteStringInternal - ( ByteString, ByteStringElem - {-, unsafeWordHead-} - , splitMaximalPrefix - ) where - -import qualified Data.ByteString as S -import Data.ByteString.Internal (ByteString(..), inlinePerformIO) -import Data.Word - -import Control.Monad - -import Foreign.ForeignPtr (ForeignPtr, withForeignPtr) -import Foreign.Ptr (Ptr, plusPtr{-, castPtr-}) -import Foreign.Storable (Storable(..)) -import System.IO.Unsafe (unsafePerformIO) -{- -import Foreign.Marshal.Alloc (alloca) -import Data.Bits --} ----------------------------------------------------------------- - -type ByteStringElem = Word8 -- Associated type of ByteString - -{- -- Not used at present --- | The size of 'Word' in bytes. -{-# INLINE sizeOfWord #-} -sizeOfWord :: Int -sizeOfWord = sizeOf (undefined :: Word) - ----------------------------------------------------------------- - --- | Return the first natural 'Word' worth of string, padding by --- zeros as necessary. The position of elements within the word --- varies by architecture, hence this function is quasi-unsafe to --- use for trieing on the bit-vector representation. A safer version --- may be forthcoming, though trieing on multiple bytes at once --- appears impractical at the moment. -unsafeWordHead :: ByteString -> Word -unsafeWordHead (PS s o l) = inlinePerformIO $ - withForeignPtr s $ \p -> - liftM (maskInitialBytes l .&.) - (peek (p `plusPtr` o :: Ptr Word)) - -maskInitialBytes :: Int -> Word -maskInitialBytes byteCount - | isLittleEndian = case effectiveByteCount of - 0 -> 0x0000000000000000 - 1 -> 0x00000000000000FF - 2 -> 0x000000000000FFFF - 3 -> 0x0000000000FFFFFF - 4 -> 0x00000000FFFFFFFF - 5 -> 0x000000FFFFFFFFFF - 6 -> 0x0000FFFFFFFFFFFF - 7 -> 0x00FFFFFFFFFFFFFF - _ -> 0xFFFFFFFFFFFFFFFF - | otherwise = case effectiveByteCount of - 0 -> 0x0000000000000000 - 1 -> 0xFF00000000000000 - 2 -> 0xFFFF000000000000 - 3 -> 0xFFFFFF0000000000 - 4 -> 0xFFFFFFFF00000000 - 5 -> 0xFFFFFFFFFF000000 - 6 -> 0xFFFFFFFFFFFF0000 - 7 -> 0xFFFFFFFFFFFFFF00 - _ -> 0xFFFFFFFFFFFFFFFF - where - effectiveByteCount = 0 `max` (byteCount `min` sizeOfWord) - - --- TODO: How to get this to execute statically?... --- BUG? is 'alloca' safe in 'inlinePerformIO' (used in 'wordHead')? -{-# NOINLINE isLittleEndian #-} -isLittleEndian :: Bool -isLittleEndian = unsafePerformIO $ alloca $ \p -> do - poke p (0x04030201 :: Word32) - b <- peek (castPtr p :: Ptr Word8) - case b of - 0x01 -> return True - 0x04 -> return False - _ -> error ("non-standard endianness detected! " - ++ "Contact the Data.Trie maintainer.") - --- End unused code -} - ----------------------------------------------------------------- --- | Returns the longest shared prefix and the two remaining suffixes --- for a pair of strings. -splitMaximalPrefix :: ByteString -> ByteString - -> (ByteString, ByteString, ByteString) -splitMaximalPrefix - str1@(PS s1 off1 len1) - str2@(PS s2 off2 len2) - | len1 == 0 = (S.empty, S.empty, str2) - | len2 == 0 = (S.empty, str1, S.empty) - | otherwise = inlinePerformIO $ - withForeignPtr s1 $ \p1 -> - withForeignPtr s2 $ \p2 -> do - i <- indexOfDifference - (p1 `ptrElemOff` off1) - (p2 `ptrElemOff` off2) - (min len1 len2) - let pre = if off1 + len1 < off2 + len2 -- share the smaller one - then newPS s1 off1 i - else newPS s2 off2 i - let s1' = newPS s1 (off1 + i) (len1 - i) - let s2' = newPS s2 (off2 + i) (len2 - i) - - return $! (,,) !$ pre !$ s1' !$ s2' - --- | C-style pointer addition, without the liberal type of 'plusPtr'. -{-# INLINE ptrElemOff #-} -ptrElemOff :: Storable a => Ptr a -> Int -> Ptr a -ptrElemOff p i = p `plusPtr` (i * sizeOf (unsafePerformIO (peek p))) - -{-# INLINE newPS #-} -newPS :: ForeignPtr ByteStringElem -> Int -> Int -> ByteString -newPS s o l = if l <= 0 then S.empty else PS s o l - -{-# INLINE (!$) #-} -(!$) :: (a -> b) -> a -> b -(!$) = ($!) -- fix associativity bug - - --- | Calculates the first index where values differ. --- --- BUG: There has to be a smarter algorithm for this. In particular, --- it'd be nice if there was a function like C's @memcmp@, but which --- returns the index rather than the 'Ordering'. If it were agnostic --- ot architecture, that'd be even nicer. --- --- TODO: rather than revert to @goByte@, we should use xor- and --- mask-munging so @goWord@ can reuse the information it already --- has to discover the byte of difference. The trick is making that --- faster than the current version. And ensuring alignment. -{-# INLINE indexOfDifference #-} -indexOfDifference :: Ptr ByteStringElem -> Ptr ByteStringElem -> Int -> IO Int -indexOfDifference p1 p2 limit = goByte 0 - where - {- BUG: using this assumes ByteStrings are Word-aligned - goWord n = if n + sizeOfWord >= limit - then goByte n - else do w1 <- peek (p1 `plusPtr` n :: Ptr Word) - w2 <- peek (p2 `plusPtr` n :: Ptr Word) - if w1 == w2 - then goWord $! n + sizeOfWord - else goByte n - -} - - goByte n = if n >= limit - then return limit - else do c1 <- peekElemOff p1 n - c2 <- peekElemOff p2 n - if c1 == c2 - then goByte $! n+1 - else return n - ----------------------------------------------------------------- ------------------------------------------------------------ fin. diff --git a/src/Data/Trie/Convenience.hs b/src/Data/Trie/Convenience.hs --- a/src/Data/Trie/Convenience.hs +++ b/src/Data/Trie/Convenience.hs @@ -1,132 +1,271 @@ {-# OPTIONS_GHC -Wall -fwarn-tabs #-} - +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ >= 701 +{-# LANGUAGE Safe #-} +#endif ---------------------------------------------------------------- --- ~ 2009.01.05 +-- ~ 2025-02-11 -- | -- Module : Data.Trie.Convenience --- Copyright : Copyright (c) 2008--2009 wren ng thornton --- License : BSD3 --- Maintainer : wren@community.haskell.org --- Stability : beta +-- Copyright : 2008--2025 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : provisional -- Portability : portable -- --- Additional convenience versions of the generic functions. +-- Additional convenience functions. In order to keep "Data.Trie" +-- concise, non-essential and uncommonly used functions have been +-- moved here. Most of these functions simplify the generic functions +-- from "Data.Trie", following after the interface for "Data.Map" +-- and "Data.IntMap". ---------------------------------------------------------------- module Data.Trie.Convenience ( - -- * 'lookupBy' variants - lookupWithDefault - - -- * 'alterBy' variants - , insertIfAbsent, insertWith, insertWithKey + -- * Conversion functions ('fromList' variants) + -- $fromList + fromListL, fromListR, fromListS + , fromListWith, fromListWith' + , fromListWithL, fromListWithL' + + -- * Query functions ('lookupBy' variants) + , lookupWithDefault + + -- * Inserting values ('alterBy' variants) + , insertIfAbsent + , insertWith, insertWith' + , insertWithKey, insertWithKey' + + -- * Updating and adjusting values ('alterBy' and 'adjustBy' variants) , adjustWithKey , update, updateWithKey - - -- ** Conversion functions - -- $fromList - , fromListL, fromListR, fromListS - - -- * 'mergeBy' variants - , disunion, unionWith + + -- * Combining tries ('mergeBy' and 'intersectBy' variants) + , disunion + , unionWith, unionWith' + , intersectWith, intersectWith' ) where import Data.Trie -import Data.Trie.Internal (lookupBy_) -import Data.List (foldl', sortBy) -import Control.Monad (liftM) +import Data.Trie.Internal (lookupBy_, alterBy_) +import Data.ByteString (ByteString) +import Data.List (sortBy) +#if !(MIN_VERSION_base(4,20,0)) +-- [aka GHC 9.10.1]: 'foldl'' re-exported from "Prelude". +import Data.List (foldl') +#endif +import Data.Ord (comparing) ---------------------------------------------------------------- ---------------------------------------------------------------- - -- $fromList --- Just like 'fromList' both of these functions convert an association +-- Just like 'fromList' all of these functions convert an association -- list into a trie, with earlier values shadowing later ones when -- keys conflict. Depending on the order of keys in the list, there --- can be as much as 5x speed difference between the two. Yet, --- performance is about the same when matching best-case to best-case --- and worst-case to worst-case (which is which is swapped when --- reversing the list or changing which function is used). +-- can be as much as 5x speed difference between the left and right +-- variants. Yet, performance is about the same when matching +-- best-case to best-case and worst-case to worst-case (which is +-- which is swapped when reversing the list or changing which +-- function is used). + +-- | A left-fold version of 'fromList'. If you run into issues with +-- stack overflows when using 'fromList' or 'fromListR', then you +-- should use this function instead. +fromListL :: [(ByteString,a)] -> Trie a {-# INLINE fromListL #-} -fromListL :: [(KeyString,a)] -> Trie a -fromListL = foldl' (flip $ uncurry $ insertIfAbsent) empty +fromListL = foldl' (flip . uncurry $ insertIfAbsent) empty --- | This version is just an alias for 'fromList'. It is a good --- producer for list fusion. Worst-case behavior is somewhat worse --- than worst-case for 'fromListL'. + +-- | An explicitly right-fold variant of 'fromList'. It is a good +-- consumer for list fusion. Worst-case behavior is somewhat worse +-- than worst-case for 'fromListL'. The 'fromList' function is +-- currently just an alias for 'fromListR'. +fromListR :: [(ByteString,a)] -> Trie a {-# INLINE fromListR #-} -fromListR :: [(KeyString,a)] -> Trie a -fromListR = fromList +fromListR = fromList -- ≡ foldr (uncurry insert) empty --- | This version sorts the list before folding over it. This adds --- /O(n log n)/ overhead and requires the whole list be in memory --- at once, but it ensures that the list is in best-case order. The --- benefits generally outweigh the costs. + +-- TODO: compare performance against a fromListL variant, adjusting the sort appropriately +-- +-- | This variant sorts the list before folding over it. This adds +-- \(\mathcal{O}(n \log n)\) overhead and requires the whole list +-- be in memory at once, but it ensures that the list is in best-case +-- order. The benefits generally outweigh the costs. +fromListS :: [(ByteString,a)] -> Trie a {-# INLINE fromListS #-} -fromListS :: [(KeyString,a)] -> Trie a -fromListS = fromListR . sortBy (\(k,_) (q,_) -> k `compare` q) +fromListS = fromListR . sortBy (comparing fst) +-- | A variant of 'fromListR' that takes a function for combining +-- values on conflict. The first argument to the combining function +-- is the \"new\" value from the initial portion of the list; the +-- second argument is the value that has been accumulated into the +-- trie from the tail of the list (just like the first argument to +-- 'foldr'). Thus, @fromList = fromListWith const@. +fromListWith :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWith #-} +fromListWith f = foldr (uncurry $ alterBy g) empty + where + g _ v Nothing = Just v + g _ v (Just w) = Just (f v w) + + +-- | A variant of 'fromListWith' which applies the combining +-- function strictly. This function is a good consumer for list +-- fusion. If you need list fusion and are running into stack +-- overflow problems with 'fromListWith', then this function may +-- solve the problem. +-- +-- @since 0.2.3 +fromListWith' :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWith' #-} +fromListWith' f = foldr (uncurry $ alterBy g') empty + where + g' _ v Nothing = Just v + g' _ v (Just w) = Just $! f v w + + +-- | A left-fold variant of 'fromListWith'. Note that the arguments +-- to the combining function are swapped: the first is the value +-- in the trie which has been accumulated from the initial part of +-- the list; the second argument is the \"new\" value from the +-- remaining tail of the list (just like the first argument to +-- 'foldl'). Thus, @fromListL = fromListWithL const@. +-- +-- @since 0.2.3 +fromListWithL :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWithL #-} +fromListWithL f = foldl' (flip . uncurry $ alterBy flipG) empty + where + flipG _ v Nothing = Just v + flipG _ v (Just w) = Just (f w v) + + +-- | A variant of 'fromListWithL' which applies the combining +-- function strictly. +-- +-- @since 0.2.3 +fromListWithL' :: (a -> a -> a) -> [(ByteString,a)] -> Trie a +{-# INLINE fromListWithL' #-} +fromListWithL' f = foldl' (flip . uncurry $ alterBy flipG') empty + where + flipG' _ v Nothing = Just v + flipG' _ v (Just w) = Just $! f w v + ---------------------------------------------------------------- -- | Lookup a key, returning a default value if it's not found. -lookupWithDefault :: a -> KeyString -> Trie a -> a -lookupWithDefault x = lookupBy_ (\mv _ -> case mv of - Nothing -> x - Just v -> v) x (const x) +lookupWithDefault :: a -> ByteString -> Trie a -> a +lookupWithDefault def = lookupBy_ const (const def) def ---------------------------------------------------------------- -- | Insert a new key, retaining old value on conflict. -insertIfAbsent :: KeyString -> a -> Trie a -> Trie a -insertIfAbsent = alterBy $ \_ x mv -> case mv of - Nothing -> Just x - Just _ -> mv +insertIfAbsent :: ByteString -> a -> Trie a -> Trie a +insertIfAbsent = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just _ -> mv -- | Insert a new key, with a function to resolve conflicts. -insertWith :: (a -> a -> a) -> KeyString -> a -> Trie a -> Trie a -insertWith f = alterBy $ \_ x mv -> case mv of - Nothing -> Just x - Just v -> Just (f x v) +insertWith :: (a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWith f = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just v -> Just (f x v) -insertWithKey :: (KeyString -> a -> a -> a) -> KeyString -> a -> Trie a -> Trie a -insertWithKey f = alterBy $ \k x mv -> case mv of - Nothing -> Just x - Just v -> Just (f k x v) +-- | A variant of 'insertWith' which applies the combining function +-- strictly. +-- +-- @since 0.2.3 +insertWith' :: (a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWith' f = + alterBy $ \_ x mv -> + case mv of + Nothing -> Just x + Just v -> Just $! f x v +-- | A variant of 'insertWith' which also provides the key to the +-- combining function. +insertWithKey :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWithKey f = + alterBy $ \k x mv -> + case mv of + Nothing -> Just x + Just v -> Just (f k x v) + +-- | A variant of 'insertWithKey' which applies the combining +-- function strictly. +-- +-- @since 0.2.3 +insertWithKey' :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> Trie a +insertWithKey' f = + alterBy $ \k x mv -> + case mv of + Nothing -> Just x + Just v -> Just $! f k x v + {- This is a tricky one... -insertLookupWithKey :: (KeyString -> a -> a -> a) -> KeyString -> a -> Trie a -> (Maybe a, Trie a) +insertLookupWithKey :: (ByteString -> a -> a -> a) -> ByteString -> a -> Trie a -> (Maybe a, Trie a) -} +---------------------------------------------------------------- -- | Apply a function to change the value at a key. -adjustWithKey :: (KeyString -> a -> a) -> KeyString -> Trie a -> Trie a -adjustWithKey f q = alterBy (\k _ -> liftM (f k)) q undefined +adjustWithKey :: (ByteString -> a -> a) -> ByteString -> Trie a -> Trie a +adjustWithKey f q = adjust (f q) q -- | Apply a function to the value at a key, possibly removing it. -update :: (a -> Maybe a) -> KeyString -> Trie a -> Trie a -update f q = alterBy (\_ _ mx -> mx >>= f) q undefined +update :: (a -> Maybe a) -> ByteString -> Trie a -> Trie a +update f q = alterBy_ (\mx t -> (mx >>= f, t)) q -updateWithKey :: (KeyString -> a -> Maybe a) -> KeyString -> Trie a -> Trie a -updateWithKey f q = alterBy (\k _ mx -> mx >>= f k) q undefined +-- | A variant of 'update' which also provides the key to the function. +updateWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> Trie a -> Trie a +updateWithKey f q = alterBy_ (\mx t -> (mx >>= f q, t)) q {- -updateLookupWithKey :: (Key -> a -> Maybe a) -> Key -> ByteStringTrie a -> (Maybe a, ByteStringTrie a) +updateLookupWithKey :: (ByteString -> a -> Maybe a) -> ByteString -> Trie a -> (Maybe a, Trie a) -- Also tricky -} ---------------------------------------------------------------- --- | Combine two tries. If they define the same key, it is removed. +-- | Combine two tries, a la symmetric difference. If they define +-- the same key, it is removed; otherwise it is retained with the +-- value it has in whichever trie. disunion :: Trie a -> Trie a -> Trie a disunion = mergeBy (\_ _ -> Nothing) --- | Combine two tries, using a function to resolve conflicts. +-- | Take the union of two tries, using a function to resolve +-- conflicts. The resulting trie is constructed strictly, but the +-- results of the combining function are evaluated lazily. unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a unionWith f = mergeBy (\x y -> Just (f x y)) -{- TODO: (efficiently) -difference, intersection --} +-- | A variant of 'unionWith' which evaluates the combining function +-- strictly. +-- +-- @since 0.2.3 +unionWith' :: (a -> a -> a) -> Trie a -> Trie a -> Trie a +unionWith' f = mergeBy (\x y -> Just $! f x y) + +-- | Take the intersection of two tries, with a function to resolve +-- the values. The resulting trie is constructed strictly, bit the +-- results of the combining function are evaluated lazily. +-- +-- @since 0.2.6 +intersectWith :: (a -> b -> c) -> Trie a -> Trie b -> Trie c +intersectWith f = intersectBy (\x y -> Just (f x y)) + +-- | A variant of 'intersectWith' which evaluates the combining +-- function strictly. +-- +-- @since 0.2.6 +intersectWith' :: (a -> b -> c) -> Trie a -> Trie b -> Trie c +intersectWith' f = intersectBy (\x y -> Just $! f x y) + +-- TODO(github#23): add 'difference' (efficiently!) ---------------------------------------------------------------- ----------------------------------------------------------- fin. diff --git a/src/Data/Trie/Internal.hs b/src/Data/Trie/Internal.hs --- a/src/Data/Trie/Internal.hs +++ b/src/Data/Trie/Internal.hs @@ -1,617 +1,2393 @@ -{-# OPTIONS_GHC -Wall -fwarn-tabs #-} - ----------------------------------------------------------------- --- ~ 2009.01.10 --- | --- Module : Data.Trie.Internal --- Copyright : Copyright (c) 2008--2009 wren ng thornton --- License : BSD3 --- Maintainer : wren@community.haskell.org --- Stability : beta --- Portability : portable --- --- Internal definition of the 'Trie' data type and generic functions --- for manipulating them. Almost everything here is re-exported --- from "Data.Trie". ----------------------------------------------------------------- - -module Data.Trie.Internal - ( - -- * Data types - Trie(), KeyString, KeyElem, showTrie - - -- * Basic functions - , empty, null, singleton, size - - -- * Conversion functions - , toListBy - - -- * Query functions - , lookupBy_, submap - - -- * Single-value modification - , alterBy - - -- * Combining tries - , mergeBy - - -- * Mapping functions - , mapBy, filterMap - ) where - -import Prelude hiding (null, lookup) -import qualified Prelude - -import qualified Data.ByteString as S -import Data.Trie.ByteStringInternal -import Data.Trie.BitTwiddle - -import Control.Monad (liftM3, liftM4) -import Control.Applicative (Applicative(..), (<$>)) -import Data.Monoid (Monoid(..)) -import Data.Foldable (Foldable(foldMap)) -import Data.Traversable (Traversable(traverse)) -import Data.Binary ----------------------------------------------------------------- ----------------------------------------------------------------- - - -{--------------------------------------------------------------- --- ByteString Big-endian Patricia Trie ----------------------------------------------------------------} -type KeyString = ByteString -type KeyElem = ByteStringElem - -{- Idealized: -data Node a = Accept a (ArcSet a) - | Reject (Branch a) -- Invariant: Must be Branch -data Arc a = Arc KeyString (Node a) -- Invariant: never empty string -data ArcSet a = None - | One {KeyElem} (Arc a) - | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) -data Trie a = Empty - | Start KeyString (Node a) -- Maybe empty string [1] - -[1] If we maintain the invariants on how Nodes recurse, then we -can't simply have Start(Node a) because we may have a shared prefix -where the prefix itself is not Accept'ed. - - --- Squash Arc into One: --- (pure good) -data Node a = Accept a (ArcSet a) - | Reject (Branch a) -data ArcSet a = None - | Arc KeyString (Node a) - | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) -data Trie a = Empty - | Start KeyString (Node a) - - --- Squash Node together: --- (most likely good) -data Node a = Node (Maybe a) (ArcSet a) -data ArcSet a = None - | Arc KeyString (Node a) - | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) -data Trie a = Empty - | Start KeyString (Node a) - - --- Squash Empty/None and Arc/Start together: --- (Complicates invariants about non-empty strings and Node's recursion) -data Node a = Node (Maybe a) (ArcSet a) -data Trie a = Empty - | Arc KeyString (Node a) - | Branch {Prefix} {Mask} (Trie a) (Trie a) - - --- Squash Node into Arc: --- (By this point, pure good) --- Unseen invariants: --- * KeyString non-empty, unless Arc is absolute root of tree --- * If (Maybe a) is Nothing, then (Trie a) is Branch --- * With views, we could re-expand Arc into accepting and --- nonaccepting variants --- --- [2] Maybe we shouldn't unpack the KeyString. We could specialize --- or inline the splitMaximalPrefix function to prevent constructing --- a new KeyString from the parts... --} - --- | A map from 'ByteString's to @a@. For all the generic functions, --- note that tries are strict in the @Maybe@ but not in @a@. --- --- The 'Monad' instance is strange. If a key @k1@ is a prefix of --- other keys, then results from binding the value at @k1@ will --- override values from longer keys when they collide. If this is --- useful for anything, or if there's a more sensible instance, I'd --- be curious to know. - -data Trie a = Empty - | Arc {-# UNPACK #-} !KeyString - !(Maybe a) - !(Trie a) - | Branch {-# UNPACK #-} !Prefix - {-# UNPACK #-} !Mask - !(Trie a) - !(Trie a) - deriving Eq - -- Prefix/Mask should be deterministic regardless of insertion order - -- TODO: verify this is so. - - --- | Visualization fuction for debugging. -showTrie :: (Show a) => Trie a -> String -showTrie t = shows' id t "" - where - spaces f = map (const ' ') (f "") - - shows' _ Empty = (".\n"++) - shows' ss (Branch p m l r) = - let s' = ("--"++) . shows p . (","++) . shows m . ("-+"++) - ss' = ss . (tail (spaces s') ++) - in s' . shows' (ss' . ("|"++)) l - . ss' . ("|\n"++) - . ss' . ("`"++) . shows' (ss' . (" "++)) r - shows' ss (Arc k mv t') = - let s' = ("--"++) . shows k - . maybe id (\v -> ("-("++) . shows v . (")"++)) mv - . ("--"++) - in s' . shows' (ss . (spaces s' ++)) t' - - -{--------------------------------------------------------------- --- Trie instances ----------------------------------------------------------------} - -instance Binary a => Binary (Trie a) where - put Empty = put (0 :: Word8) - put (Arc k m t) = do put (1 :: Word8) - put k - put m - put t - put (Branch p m l r) = do put (2 :: Word8) - put p - put m - put l - put r - - get = do tag <- get :: Get Word8 - case tag of - 0 -> return Empty - 1 -> liftM3 Arc get get get - _ -> liftM4 Branch get get get get - - -instance Functor Trie where - fmap _ Empty = Empty - fmap f (Arc k Nothing t) = Arc k Nothing (fmap f t) - fmap f (Arc k (Just v) t) = Arc k (Just (f v)) (fmap f t) - fmap f (Branch p m l r) = Branch p m (fmap f l) (fmap f r) - - --- Does this even make sense? It's not nondeterminism like lists --- and sets. If no keys were prefixes of other keys it'd make sense --- as a decision-tree; but since keys /can/ prefix, tries formed --- from shorter keys can shadow the results from longer keys due --- to the 'unionL'. It does seem to follow the laws though... What --- computation could this possibly represent? --- --- 1. return x >>= f == f x --- 2. m >>= return == m --- 3. (m >>= f) >>= g == m >>= (\x -> f x >>= g) -instance Monad Trie where - return x = singleton S.empty x - - (>>=) Empty _ = empty - (>>=) (Branch p m l r) f = branch p m (l >>= f) (r >>= f) - (>>=) (Arc k Nothing t) f = arc k Nothing (t >>= f) - (>>=) (Arc k (Just v) t) f = arc k Nothing (f v `unionL` (t >>= f)) - where - unionL = mergeBy (\x _ -> Just x) - - -instance Monoid a => Monoid (Trie a) where - mempty = empty - mappend = mergeBy $ \x y -> Just (x `mappend` y) - - --- Not a MonadPlus for any definition I can think of. - - --- TODO: cf toListBy. We should provide foldr and foldl directly -instance Foldable Trie where - foldMap _ Empty = mempty - foldMap f (Arc _ Nothing t) = foldMap f t - foldMap f (Arc _ (Just v) t) = f v `mappend` foldMap f t - foldMap f (Branch _ _ l r) = foldMap f l `mappend` foldMap f r - - -instance Traversable Trie where - traverse _ Empty = pure Empty - traverse f (Arc k Nothing t) = Arc k Nothing <$> traverse f t - traverse f (Arc k (Just v) t) = Arc k . Just <$> f v <*> traverse f t - traverse f (Branch p m l r) = Branch p m <$> traverse f l <*> traverse f r - - -{--------------------------------------------------------------- --- Mapping functions ----------------------------------------------------------------} - --- | Apply a function to all values, potentially removing them. -filterMap :: (a -> Maybe b) -> Trie a -> Trie b -filterMap _ Empty = empty -filterMap f (Arc k Nothing t) = arc k Nothing (filterMap f t) -filterMap f (Arc k (Just v) t) = arc k (f v) (filterMap f t) -filterMap f (Branch p m l r) = branch p m (filterMap f l) (filterMap f r) - --- | Generic version of 'fmap'. This function is notably more --- expensive than 'fmap' or 'filterMap' because we have to reconstruct --- the keys. -mapBy :: (KeyString -> a -> Maybe b) -> Trie a -> Trie b -mapBy f = go S.empty - where - go _ Empty = empty - go q (Arc k Nothing t) = arc k Nothing (go q' t) where q' = S.append q k - go q (Arc k (Just v) t) = arc k (f q' v) (go q' t) where q' = S.append q k - go q (Branch p m l r) = branch p m (go q l) (go q r) - - -{--------------------------------------------------------------- --- Smart constructors and helper functions for building tries ----------------------------------------------------------------} - --- | Smart constructor to prune @Empty@ from @Branch@es. -branch :: Prefix -> Mask -> Trie a -> Trie a -> Trie a -branch _ _ Empty r = r -branch _ _ l Empty = l -branch p m l r = Branch p m l r - - --- | Smart constructor to prune @Arc@s that lead nowhere. --- N.B if mv=Just then doesn't check whether t=epsilon. It's up to callers to ensure that invariant isn't broken. -arc :: KeyString -> Maybe a -> Trie a -> Trie a -arc k mv@(Just _) t = Arc k mv t -arc _ Nothing Empty = Empty -arc k Nothing t@(Branch _ _ _ _) | S.null k = t - | otherwise = Arc k Nothing t -arc k Nothing (Arc k' mv' t') = Arc (S.append k k') mv' t' - - --- | Smart constructor to join two tries into a @Branch@ with maximal --- prefix sharing. Requires knowing the prefixes, but can combine --- either @Branch@es or @Arc@s. --- --- N.B. /do not/ use if prefixes could match entirely! -branchMerge :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a -branchMerge _ Empty _ t2 = t2 -branchMerge _ t1 _ Empty = t1 -branchMerge p1 t1 p2 t2 - | zero p1 m = Branch p m t1 t2 - | otherwise = Branch p m t2 t1 - where - m = branchMask p1 p2 - p = mask p1 m - - --- It would be better if Arc used --- Data.ByteString.TrieInternal.wordHead somehow, that way --- we can see 4/8/?*Word8 at a time instead of just one. --- But that makes maintaining invariants ...difficult :( -getPrefix :: Trie a -> Prefix -getPrefix (Branch p _ _ _) = p -getPrefix (Arc k _ _) | S.null k = 0 -- for lack of a better value - | otherwise = S.head k -getPrefix Empty = error "getPrefix: no Prefix of Empty" - - -{--------------------------------------------------------------- --- Error messages ----------------------------------------------------------------} - -errorLogHead :: String -> KeyString -> KeyElem -errorLogHead s q | S.null q = error (s ++": found null subquery") - | otherwise = S.head q - - ----------------------------------------------------------------- ----------------------------------------------------------------- - -{--------------------------------------------------------------- --- Basic functions ----------------------------------------------------------------} - --- | /O(1)/, The empty trie. -{-# INLINE empty #-} -empty :: Trie a -empty = Empty - --- | /O(1)/, Is the trie empty? -{-# INLINE null #-} -null :: Trie a -> Bool -null Empty = True -null _ = False - --- | /O(1)/, A singleton trie. -{-# INLINE singleton #-} -singleton :: KeyString -> a -> Trie a -singleton k v = Arc k (Just v) Empty --- For singletons, don't need to verify invariant on arc length >0 - --- | /O(n)/, Get count of elements in trie. -{-# INLINE size #-} -size :: Trie a -> Int -size t = size' t id 0 - --- | /O(n)/, Internal CPS accumulator function for calculating --- 'size'. -size' :: Trie a -> (Int -> Int) -> Int -> Int -size' Empty f n = f n -size' (Branch _ _ l r) f n = size' l (size' r f) n -size' (Arc _ Nothing t) f n = size' t f n -size' (Arc _ (Just _) t) f n = size' t f $! n + 1 - - -{--------------------------------------------------------------- --- Conversion functions ----------------------------------------------------------------} - --- Still rather inefficient --- --- TODO: rewrite list-catenation to be lazier (real CPS instead of --- function building? is the function building really better than --- (++) anyways?) --- --- TODO: the @q@ accumulator should be lazy ByteString and only --- forced by @f@. It's already non-strict, but we should ensure --- O(n) not O(n^2) when it's forced. --- --- | Convert a trie into a list using a function. Resulting values --- are in sorted order according to the keys. -toListBy :: (KeyString -> a -> b) -> Trie a -> [b] -toListBy f = \t -> go S.empty t [] - where - go _ Empty = id - go q (Branch _ _ l r) = go q l . go q r - go q (Arc k mv t) = case mv of - Nothing -> rest - Just v -> (f k' v :) . rest - where - rest = go k' t - k' = S.append q k - - -{--------------------------------------------------------------- --- Query functions (just recurse) ----------------------------------------------------------------} - --- | Generic function to find a value (if it exists) and the subtrie --- rooted at the prefix. The first function argument is called if and --- only if a node is exactly reachable by the query; if no node is --- exactly reachable the default value is used; if the middle of --- an arc is reached, the second function argument is used. --- --- This function is intended for internal use. For the public-facing --- version, see @lookupBy@ in "Data.Trie". -lookupBy_ :: (Maybe a -> Trie a -> b) -> b -> (Trie a -> b) - -> KeyString -> Trie a -> b -lookupBy_ f z a = lookupBy_' - where - -- | Deal with epsilon query (when there is no epsilon value) - lookupBy_' q t@(Branch _ _ _ _) | S.null q = f Nothing t - lookupBy_' q t = go q t - - -- | The main recursion - go _ Empty = z - - go q (Arc k mv t) = - let (_,k',q') = splitMaximalPrefix k q - in case (not $ S.null k', S.null q') of - (True, True) -> a (Arc k' mv t) - (True, False) -> z - (False, True) -> f mv t - (False, False) -> go q' t - - go q t_@(Branch _ _ _ _) = findArc t_ - where - qh = errorLogHead "lookupBy_" q - - -- | /O(min(m,W))/, where /m/ is number of @Arc@s in this - -- branching, and /W/ is the word size of the Prefix,Mask type. - findArc (Branch p m l r) - | nomatch qh p m = z - | zero qh m = findArc l - | otherwise = findArc r - findArc t@(Arc _ _ _) = go q t - findArc Empty = z - - --- This function needs to be here, not in "Data.Trie", because of --- 'arc' which isn't exported. We could use the monad instance --- instead, though it'd be far more circuitous. --- arc k Nothing t === singleton k () >> t --- arc k (Just v) t === singleton k v >>= unionR t . singleton S.empty --- (...except 'arc' doesn't do the invariant correction --- of (>>=) for epsilon`elem`t) --- --- | Return the subtrie containing all keys beginning with a prefix. -{-# INLINE submap #-} -submap :: KeyString -> Trie a -> Trie a -submap q = lookupBy_ (arc q) empty (arc q Nothing) q -{- -- Disable superfluous error checking. - -- @submap'@ would replace the first argument to @lookupBy_@ - where - submap' Nothing Empty = errorEmptyAfterNothing "submap" - submap' Nothing (Arc _ _ _) = errorArcAfterNothing "submap" - submap' mx t = Arc q mx t - -errorInvariantBroken :: String -> String -> a -errorInvariantBroken s e = error (s ++ ": Invariant was broken" ++ e') - where - e' = if Prelude.null e then e else ", found: " ++ e - -errorArcAfterNothing :: String -> a -errorArcAfterNothing s = errorInvariantBroken s "Arc after Nothing" - -errorEmptyAfterNothing :: String -> a -errorEmptyAfterNothing s = errorInvariantBroken s "Empty after Nothing" --- -} - - -{--------------------------------------------------------------- --- Single-value modification functions (recurse and clone spine) ----------------------------------------------------------------} - --- TODO: We should CPS on Empty to avoid cloning spine if no change. --- Difficulties arise with the calls to 'branch' and 'arc'. Will --- have to create a continuation chain, so no savings on memory --- allocation; but would have savings on held memory, if they're --- still holding the old one... --- --- | Generic function to alter a trie by one element with a function --- to resolve conflicts (or non-conflicts). -alterBy :: (KeyString -> a -> Maybe a -> Maybe a) - -> KeyString -> a -> Trie a -> Trie a -alterBy f_ q_ x_ - | S.null q_ = alterEpsilon - | otherwise = go q_ - where - f = f_ q_ x_ - nothing q = arc q (f Nothing) Empty - - alterEpsilon t_@Empty = arc q_ (f Nothing) t_ - alterEpsilon t_@(Branch _ _ _ _) = arc q_ (f Nothing) t_ - alterEpsilon t_@(Arc k mv t) | S.null k = arc q_ (f mv) t - | otherwise = arc q_ (f Nothing) t_ - - - go q Empty = nothing q - - go q t@(Branch p m l r) - | nomatch qh p m = branchMerge p t qh (nothing q) - | zero qh m = branch p m (go q l) r - | otherwise = branch p m l (go q r) - where - qh = errorLogHead "alterBy" q - - go q t_@(Arc k mv t) = - let (p,k',q') = splitMaximalPrefix k q - in case (not $ S.null k', S.null q') of - (True, True) -> -- add node to middle of arc - arc p (f Nothing) (Arc k' mv t) - (True, False) -> - case nothing q' of - Empty -> t_ -- Nothing to add, reuse old arc - l -> arc' (branchMerge (getPrefix l) l - (getPrefix r) r) - where - r = Arc k' mv t - - -- inlined version of 'arc' - arc' | S.null p = id - | otherwise = Arc p Nothing - - (False, True) -> arc k (f mv) t - (False, False) -> arc k mv (go q' t) - - -{--------------------------------------------------------------- --- Trie-combining functions ----------------------------------------------------------------} - --- TEST CASES: foldr (unionL . uncurry singleton) empty t --- foldr (uncurry insert) empty t --- where t = map (\s -> (pk s, 0)) --- ["heat","hello","hoi","apple","appa","hell","appb","appc"] --- --- | Combine two tries, using a function to resolve collisions. --- This can only define the space of functions between union and --- symmetric difference but, with those two, all set operations can --- be defined (albeit inefficiently). -mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a -mergeBy f = mergeBy' - where - -- | Deals with epsilon entries, before recursing into @go@ - mergeBy' - t0_@(Arc k0 mv0 t0) - t1_@(Arc k1 mv1 t1) - | S.null k0 && S.null k1 = arc k0 (mergeMaybe f mv0 mv1) (go t0 t1) - | S.null k0 = arc k0 mv0 (go t0 t1_) - | S.null k1 = arc k1 mv1 (go t1 t0_) - mergeBy' - (Arc k0 mv0@(Just _) t0) - t1_@(Branch _ _ _ _) - | S.null k0 = arc k0 mv0 (go t0 t1_) - mergeBy' - t0_@(Branch _ _ _ _) - (Arc k1 mv1@(Just _) t1) - | S.null k1 = arc k1 mv1 (go t1 t0_) - mergeBy' t0_ t1_ = go t0_ t1_ - - - -- | The main recursion - go Empty t1 = t1 - go t0 Empty = t0 - - -- /O(n+m)/ for this part where /n/ and /m/ are sizes of the branchings - go t0@(Branch p0 m0 l0 r0) - t1@(Branch p1 m1 l1 r1) - | shorter m0 m1 = union0 - | shorter m1 m0 = union1 - | p0 == p1 = branch p0 m0 (go l0 l1) (go r0 r1) - | otherwise = branchMerge p0 t0 p1 t1 - where - union0 | nomatch p1 p0 m0 = branchMerge p0 t0 p1 t1 - | zero p1 m0 = branch p0 m0 (go l0 t1) r0 - | otherwise = branch p0 m0 l0 (go r0 t1) - - union1 | nomatch p0 p1 m1 = branchMerge p0 t0 p1 t1 - | zero p0 m1 = branch p1 m1 (go t0 l1) r1 - | otherwise = branch p1 m1 l1 (go t0 r1) - - go t0_ t1_ = go' t0_ t1_ - where - p0 = getPrefix t0_ - p1 = getPrefix t1_ - m' = branchMask p0 p1 - p' = mask p0 m' - - go' (Arc k0 mv0 t0) - (Arc k1 mv1 t1) - | m' == 0 = - let (pre,k0',k1') = splitMaximalPrefix k0 k1 - in if S.null pre - then error "mergeBy: no mask, but no prefix string" - else let {-# INLINE arcMerge #-} - arcMerge mv' t1' t2' = arc pre mv' (go t1' t2') - in case (S.null k0', S.null k1') of - (True, True) -> arcMerge (mergeMaybe f mv0 mv1) t0 t1 - (True, False) -> arcMerge mv0 t0 (Arc k1' mv1 t1) - (False,True) -> arcMerge mv1 t1 (Arc k0' mv0 t0) - (False,False) -> arcMerge Nothing (Arc k0' mv0 t0) - (Arc k1' mv1 t1) - go' (Arc _ _ _) - (Branch _p1 m1 l r) - | nomatch p0 p1 m1 = branchMerge p1 t1_ p0 t0_ - | zero p0 m1 = branch p1 m1 (go t0_ l) r - | otherwise = branch p1 m1 l (go t0_ r) - go' (Branch _p0 m0 l r) - (Arc _ _ _) - | nomatch p1 p0 m0 = branchMerge p0 t0_ p1 t1_ - | zero p1 m0 = branch p0 m0 (go t1_ l) r - | otherwise = branch p0 m0 l (go t1_ r) - - -- Inlined branchMerge. Both tries are disjoint @Arc@s now. - go' _ _ | zero p0 m' = Branch p' m' t0_ t1_ - go' _ _ = Branch p' m' t1_ t0_ - -{-# INLINE mergeMaybe #-} -mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a -mergeMaybe _ Nothing Nothing = Nothing -mergeMaybe _ Nothing mv1@(Just _) = mv1 -mergeMaybe _ mv0@(Just _) Nothing = mv0 -mergeMaybe f (Just v0) (Just v1) = f v0 v1 - ----------------------------------------------------------------- ------------------------------------------------------------ fin. +-- Not using -Wcompat, because it wants outdated things for GHC 8.0/8.2 +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +-- Note: can also pass show-extensions to get Haddock to show +-- everything implicit/entailed by the things we actually request. +-- Apparently we're getting: MonoLocalBinds, KindSignatures, +-- ExplicitNamespaces; regardless of whether we set the default to Haskell98 or Haskell2010 +{-# OPTIONS_HADDOCK not-home #-} + +{-# LANGUAGE NoImplicitPrelude, CPP, BangPatterns #-} +#if __GLASGOW_HASKELL__ >= 708 +-- For defining the 'IsList' instance. +{-# LANGUAGE TypeFamilies #-} +#endif +#if __GLASGOW_HASKELL__ >= 701 +-- Alas, "GHC.Exts" isn't considered safe, even though 'build' and +-- 'IsList'(..) surely are. +{-# LANGUAGE Trustworthy #-} +#endif +------------------------------------------------------------ +-- ~ 2025-02-11 +-- | +-- Module : Data.Trie.Internal +-- Copyright : 2008--2025 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : experimental +-- Portability : portable (with CPP) +-- +-- Internal definition of the 'Trie' data type and generic functions +-- for manipulating them. Almost everything here is re-exported +-- from "Data.Trie", which is the preferred API for users. This +-- module is for developers who need deeper (and less stable) access +-- to the abstract type. +-- +-- @since 0.1.3 +------------------------------------------------------------ + +module Data.Trie.Internal + ( + -- * Data types + Trie() + -- BUG: can't seem to put this at the top: it'll gobble up the + -- following section name and replace it. (I'm guessing that's + -- something to do with the section not having any exported + -- entities?) + -- *** Performance Warning + -- $bug25 + + -- * Basic functions + , empty, null, singleton, size + + -- * List-conversion functions + , fromList + , toList, toListBy, elems + + -- * Query functions + , lookupBy_, submap + , match_, matches_ + + -- * Simple modification + , alterBy, alterBy_, adjust + + -- * Combining tries + , wip_unionWith + , mergeBy, intersectBy + + -- * Priority-queue functions + , minAssoc, maxAssoc + , updateMinViewBy, updateMaxViewBy + + -- * Mapping, filtering, folding, and traversing + -- ** Filterable + , filter + , filterMap + , mapBy + -- ** Witherable + , filterA + , wither + -- TODO: 'witherBy' (effectful 'mapBy') + -- ** Contextual filtering\/mapping functions + , contextualMap + , contextualMap' + , contextualFilterMap + , contextualMapBy + -- TODO: 'contextualWither' + -- TODO: 'contextualWitherBy' (effectful 'contextualMapBy') + -- ** Foldable + , foldrWithKey, foldrWithKey', foldlWithKey, foldlWithKey' + , cata_, cata + -- ** Traverse + , traverseWithKey + + -- * Internal utility functions + , showTrie + , breakMaximalPrefix + ) where + +import Prelude hiding (null, lookup, filter) + +import qualified Data.ByteString as S +import qualified Data.ByteString.Unsafe as SU +import Data.Trie.Internal.ByteString +import Data.Trie.Internal.BitTwiddle +import Data.Trie.Internal.Errors (impossible) + +import Data.Binary (Binary(..), Get, Word8) +import Data.Bits (xor) + +#if MIN_VERSION_base(4,9,0) +-- [aka GHC 8.0.1]: "Data.Semigroup" added to base. +-- +-- Note: [base-4.13.0 / GHC 8.8.1] has Prelude re-export 'Semigroup' +-- (the class name) and 'Data.Semigroup.<>'; however it does not +-- re-export 'stimes' nor (I assume) 'sconcat'! +import Data.Semigroup (Semigroup(..)) +#endif +#if !(MIN_VERSION_base(4,8,0)) +-- [aka GHC 7.10.1]: Prelude re-exports 'Monoid'. +import Data.Monoid (Monoid(..)) +#endif + +import Control.DeepSeq (NFData(rnf)) +import Control.Monad (liftM3, liftM4) + +import qualified Data.Foldable as F +-- [base-4.10.0.0 / GHC 8.2.1] moved 'liftA2' into the 'Applicative' +-- class for performance reasons; so we want to use it wherever +-- possible, for those same performance reasons. However, while +-- the Prelude re-exports all the other methods of 'Applicative' +-- since base-4.8, it does not re-export 'liftA2' (at least not up +-- through base-4.16.0.0 / GHC 9.2.1). +#if !(MIN_VERSION_base(4,18,0)) +-- [base-4.18 / GHC 9.6.1]: Prelude now re-exports 'liftA2'. +import Control.Applicative (liftA2) +#endif +#if MIN_VERSION_base(4,8,0) +-- [aka GHC 7.10.1]: Prelude re-exports 'Applicative', @(<$>)@, +-- 'Foldable', and 'Traversable'. But not 'liftA2'. +#else +import Control.Applicative (Applicative(..), (<$>)) +import Data.Foldable (Foldable()) +import Data.Traversable (Traversable(traverse)) +#endif + +#if MIN_VERSION_base(4,9,0) +import Data.Functor.Classes (Eq1(..), Ord1(..), Show1(..), Read1(..)) +import qualified Data.Functor.Classes as FC +#endif + +#ifdef __GLASGOW_HASKELL__ +import qualified Text.Read as R +import GHC.Exts (build) +#endif +#if __GLASGOW_HASKELL__ >= 708 +import qualified GHC.Exts (IsList(..)) +#endif +#if MIN_VERSION_base(4,7,0) +-- [GHC 7.8.1]: "Data.Coerce" added to base. +import Data.Coerce (coerce) +#endif + +------------------------------------------------------------ +------------------------------------------------------------ +-- $bug25 +-- #bug25# +-- Many (but not all) functions which need to reconstruct bytestrings +-- suffer from an asymptotic slowdown; see: +-- . +-- For clarity, all functions affected by this bug will have a link +-- to this section. This is not a new bug, it has affected all +-- prior versions of this library as well. However, compared to +-- older versions, /bytestring-trie-0.2.7/ mitigates the severity +-- of the bug, and in certain cases to avoids it entirely. +-- +-- In particular, this affects the \"keyed\" variants of functions +-- (for folding, traversing, filtering, etc), and anything built +-- from them, including 'toListBy' and various instances which use +-- it. +-- +-- Conversely, functions which are unaffected include: those like +-- 'alterBy' which merely pass the query back to the user as a +-- convenience; those which only need to reconstruct a single +-- bytestring (e.g., the priority-queue functions); and +-- 'Data.Trie.matches'\/'matches_'. + + +------------------------------------------------------------ +-- | Infix variant of 'uncurry'. Currently only used in 'alterBy_'. +-- The fixity-level is like @(<$>)@; but I'm making it nonassociative +-- to avoid any possible\/potential confusion. +infix 4 $$ +($$) :: (a -> b -> c) -> (a, b) -> c +($$) = uncurry +{-# INLINE ($$) #-} + +{----------------------------------------------------------- +-- ByteString Big-endian Patricia Trie datatype +-----------------------------------------------------------} +{- +In our idealized representation, we use a (directed) discrete graph +to represent our finite state machine. To organize the set of +outgoing arcs from a given Node we have ArcSet be a big-endian +patricia tree like Data.IntMap. In order to simplify things we then +go through a series of derivations. + + data Node a = Accept a (ArcSet a) + | Reject (Branch a) + data Arc a = Arc ByteString (Node a) + data ArcSet a = None + | One KeyElem (Arc a) + | Many (Branch a) + data Branch a = Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) + data Trie a = Empty + | Start (Arc a) -- [1] + +[1]: N.B., we must allow constructing @Start(Arc pre (Reject b))@ +for non-null @pre@, so that we can have a shared prefix even though +that prefix itself doesn't have an associated value. + +** Squash Arc into One and Start: +For One, this allows combining the initial KeyElem with the rest +of the ByteString, which is purely beneficial. However, it does +introduce some invariants since now we must distinguish NonEmptyBS +vs NullableBS. + + newtype NonEmptyBS = NonEmptyBS ByteString -- Invariant: never null. + newtype NullableBS = NullableBS Bytestring -- May be null. + + data Node a = Accept a (ArcSet a) + | Reject (Branch a) + data ArcSet a = None + | Arc NonEmptyBS (Node a) + | Many (Branch a) + data Branch a = Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) + data Trie a = Empty + | Start NullableBS (Node a) + +** Squash Accept and Reject together: +Maybe[2] beneficial. However, it complicates stating the invariants +about Node's recursion. + + data Node a = Node (Maybe a) (ArcSet a) + -- Invariant: if Nothing then must be Branch + data ArcSet a = None + | Arc NonEmptyBS (Node a) + | Many (Branch a) + data Branch a = Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) + data Trie a = Empty + | Start NullableBS (Node a) + +[2]: The above change is certainly beneficial from the perspective +of source-code size/repetition for 'mergeBy', 'intersectBy', and +other operations that operate on two tries in tandem; though it's +unclear whether/how much that transfers to compiled-code size/bloat. +Also since the 'Maybe' isn't unpacked, this introduces an additional +indirection to reach values. Starting at version 0.2.7 there's an +ongoing effort to try to determine whether this change is beneficial +or not, and to quantify how much it affects things. + +** Squash Branch into Many: +Purely beneficial, since there's no point in keeping them distinct anymore. + + data Node a = Node (Maybe a) (ArcSet a) + -- Invariant: if Nothing then must be Branch + data ArcSet a = None + | Arc NonEmptyBS (Node a) + | Branch {Prefix} {Mask} (ArcSet a) (ArcSet a) + data Trie a = Empty + | Start NullableBS (Node a) + +** Squash Empty/None and Arc/Start together: +Alas, this complicates the invariants about non-null strings. + + data Node a = Node (Maybe a) (ArcSet a) + -- Invariant: if Nothing then must be Branch + data Trie a = Empty + | Arc ByteString (Node a) + -- Invariant: null string only allowed if both + -- (a) the Arc is at the root, and + -- (b) the Node has a value. + | Branch {Prefix} {Mask} (Trie a) (Trie a) + +** Squash Node into Arc: +By this point, purely beneficial. However, the two unseen invariants remain. +-} + + +-- [Note: Order of constructors] +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- As discussed in "Data.IntMap.Internal", the order of constructors +-- in the definition is the order in which case analysis will always +-- test them. On GHC 9.2.1 I haven't noticed a significant difference +-- from changing the order, but that could be due to bad benchmarking +-- (i.e., reusing the test suite); so I'll trust that this hasn't +-- changed so much since GHC 7.0. The only question is whether +-- @Arc@ or @Branch@ should come first, which depends a lot on the +-- dataset being used. +-- +-- This reordering change was performed for version 0.2.7 +-- + + +-- | A map from 'ByteString's to @a@. For all the generic functions, +-- note that tries are strict in the @Maybe@ but not in @a@. +-- +-- The 'Monad' instance is strange. If a key @k1@ is a prefix of +-- other keys, then results from binding the value at @k1@ will +-- override values from longer keys when they collide. If this is +-- useful for anything, or if there's a more sensible instance, I'd +-- be curious to know. + +data Trie a + = Branch {-# UNPACK #-} !Prefix + {-# UNPACK #-} !Mask + !(Trie a) -- Must not be @Empty@. + !(Trie a) -- Must not be @Empty@. + | Arc {-# UNPACK #-} !ByteString -- Has nontrivial invariants. + !(Maybe a) + !(Trie a) -- Has complex invariants. + | Empty + deriving Eq + -- Prefix/Mask should be deterministic regardless of insertion order + -- TODO: prove this is so. + +-- [Note:Closure] aka: Local 'go' functions and capturing +-- ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +-- Copied from 'IntMap': Care must be taken when using 'go' function +-- which captures an argument. Sometimes (for example when the +-- argument is passed to a data constructor, as in insert), GHC +-- heap-allocates more than necessary. Therefore C-- code must be +-- checked for increased allocation when creating and modifying +-- such functions. + +{----------------------------------------------------------- +-- Smart constructors +-----------------------------------------------------------} + +{- +-- | A common precondition for ensuring the safety of the following +-- smart constructors. +ifJustThenNoEpsilon :: Maybe a -> Trie a -> Bool +ifJustThenNoEpsilon (Just _) (Arc k (Just _) _) = not (S.null k) +ifJustThenNoEpsilon _ _ = True +-} + +-- FIXME: [bug26] +-- We need to adjust 'arc', 'prepend', 'arcNN', and 'prependNN' to behave +-- more like a zipper, to avoid asymptotic slowdown in corner cases. + +-- | Smart constructor to prune @Arc@s that lead nowhere. +-- +-- __Preconditions__ +-- * @arc _ mv t | ifJustThenNoEpsilon mv t@ +arc :: ByteString -> Maybe a -> Trie a -> Trie a +{-# INLINE arc #-} +arc !k mv@(Just _) = Arc k mv +arc k Nothing = prepend k + +-- | Variant of 'arc' where the string is known to be non-null. +-- +-- __Preconditions__ +-- * @arcNN k _ _ | not (S.null k)@ +-- * @arcNN _ mv t | ifJustThenNoEpsilon mv t@ +arcNN :: ByteString -> Maybe a -> Trie a -> Trie a +{-# INLINE arcNN #-} +arcNN !k mv@(Just _) = Arc k mv +arcNN k Nothing = prependNN k + +-- | Prepend a possibly-null string to a trie. +-- +-- This function is only very rarely needed; most of the time you +-- already know that the string is non-null, and thus you can call +-- 'prependNN' directly. +-- +-- TODO: may actually consider exporting this one, since it could +-- be generally useful and it has no preconditions. Of course, +-- it's susceptible to [bug25][bug26] if used incorrectly... +prepend :: ByteString -> Trie a -> Trie a +{-# INLINE prepend #-} +prepend k + | S.null k = id + | otherwise = prependNN k + +-- | Prepend a non-null string to a trie. +-- +-- __Preconditions__ +-- * @prependNN k _ | not (S.null k)@ +prependNN :: ByteString -> Trie a -> Trie a +{-# INLINE prependNN #-} +prependNN !_ t@Empty = t +prependNN q t@(Branch{}) = Arc q Nothing t +prependNN q (Arc k mv s) = Arc (S.append q k) mv s + -- TODO: see [bug26]; should ensure that callers do not nest + -- calls to this function which all take this @Arc@ case. + +-- | > mayEpsilon mv ≡ arc S.empty mv +-- +-- __Preconditions__ +-- * @mayEpsilon mv t | ifJustThenNoEpsilon mv t@ +mayEpsilon :: Maybe a -> Trie a -> Trie a +{-# INLINE mayEpsilon #-} +mayEpsilon (Just v) = epsilon v +mayEpsilon Nothing = id + +-- | Canonical name for the empty arc at the top of the trie. +-- +-- > epsilon ≡ Arc S.empty . Just +-- +-- __Preconditions__ +-- * The trie argument must not already have an epsilon value. +epsilon :: a -> Trie a -> Trie a +{-# INLINE epsilon #-} +epsilon = Arc S.empty . Just + + +-- | Smart @Branch@ constructor: prunes @Empty@ on both sides. +-- This function does no other work besides pruning, so the @Prefix@, +-- @Mask@, and ordering of the 'Trie's must all be as if calling +-- the @Branch@ constructor directly. +branch :: Prefix -> Mask -> Trie a -> Trie a -> Trie a +{-# INLINE branch #-} +branch !_ !_ Empty r = r +branch _ _ l Empty = l +branch p m l r = Branch p m l r + +-- | Smart @Branch@ constructor: prunes @Empty@ on left side only. +-- +-- __Preconditions__ +-- * the right trie is not @Empty@. +branchL :: Prefix -> Mask -> Trie a -> Trie a -> Trie a +{-# INLINE branchL #-} +branchL !_ !_ Empty r = r +branchL p m l r = Branch p m l r + +-- | Smart @Branch@ constructor: prunes @Empty@ on right side only. +-- +-- __Preconditions__ +-- * the left trie is not @Empty@. +branchR :: Prefix -> Mask -> Trie a -> Trie a -> Trie a +{-# INLINE branchR #-} +branchR !_ !_ l Empty = l +branchR p m l r = Branch p m l r + +-- | Smart constructor to join two tries into a @Branch@ with maximal +-- prefix sharing, and in the correct left\/right order. Requires +-- knowing the prefixes, but can combine either @Branch@es or @Arc@s. +-- +-- __Preconditions__ +-- * Both tries must be non-@Empty@. +-- * The two prefixes /must not/ be able to match entirely! +graft :: Prefix -> Trie a -> Prefix -> Trie a -> Trie a +{-# INLINE graft #-} +graft p0 t0 p1 t1 + | zero p0 m = Branch p m t0 t1 + | otherwise = Branch p m t1 t0 + where + m = getMask p0 p1 + p = applyMask p0 m + +-- | Shorthand for prepending a non-null string to a 'graft'. This +-- is mainly useful for when @(p,k0,k1)@ came from 'breakMaximalPrefix'. +-- +-- __Preconditions__ +-- * The prefix must be non-null. +-- * Each trie must agree with their key (hence must also be non-@Empty@). +-- * The keys must not have matching prefixes. +wye :: ByteString + -> ByteString -> Trie a + -> ByteString -> Trie a + -> Trie a +wye p k0 t0 k1 t1 = + Arc p Nothing $ graft (arcPrefix k0) t0 (arcPrefix k1) t1 + +-- TODO: this is only used by 'mergeBy' (since 'intersectBy' returns +-- @Empty@ in lieu of @Branch@ for the latter cases); so maybe we +-- should move this to be a local definition there? +-- +-- | Smart constructor to join two @Arc@s into a @Branch@ when possible, +-- and to 'breakMaximalPrefix' otherwise. +-- +-- __Preconditions__ +-- * Both tries must be non-@Empty@. +arcMerge + :: ByteString -> Trie a + -> ByteString -> Trie a + -> (ByteString -> ByteString -> ByteString -> Trie a) + -> Trie a +{-# INLINE arcMerge #-} +arcMerge k0 t0 k1 t1 whenMatch + | m == 0 = + case breakMaximalPrefix k0 k1 of + (pre, k0', k1') + -- TODO: change this into an 'assert' instead. + | S.null pre -> impossible "arcMerge" -- perfect 'arcPrefix' match, yet no 'breakMaximalPrefix' prefix. + | otherwise -> whenMatch pre k0' k1' + | zero p0 m = Branch p m t0 t1 + | otherwise = Branch p m t1 t0 + where + p0 = arcPrefix k0 + p1 = arcPrefix k1 + m = getMask p0 p1 + p = applyMask p0 m + +-- It would be better if arcs used +-- 'Data.ByteString.TrieInternal.wordHead' somehow, that way +-- we can see 4/8/?*Word8 at a time instead of just one. +-- But that makes maintaining invariants ...difficult :( + +-- | Get the equivalent of the @Prefix@ stored in a @Branch@, but +-- for an @Arc@. +arcPrefix :: ByteString -> Prefix +{-# INLINE arcPrefix #-} +arcPrefix k + | S.null k = 0 -- for lack of a better value + | otherwise = SU.unsafeHead k + + +{----------------------------------------------------------- +-- Error messages +-----------------------------------------------------------} + +-- TODO: move off to "Data.Trie.Errors"? +-- TODO: shouldn't we inline the logic and just NOINLINE the string +-- constant? There are only three use sites, which themselves aren't +-- inlined... +-- TODO: this is almost identical to 'arcPrefix'; the only difference +-- is that we use this one for matching a query against a trie, +-- whereas we use 'arcPrefix' when matching two tries together. +-- That said, since our test suite never throws this error, it +-- should be safe to use 'arcPrefix' everywhere instead. Or, if +-- we want to preserve the semantic distinction, then we could start +-- using 'Control.Exception.assert' to hoist the null-check out to +-- where it belongs and still allow it to compile away. Conversely, +-- note that 'arcPrefix' is never called with a null string either +-- (since null strings are only ever allowed for epsilon values; +-- and all the use-sites of 'arcPrefix' are after handling those +-- epsilons, or otherwise guarded). +errorLogHead :: String -> ByteString -> ByteStringElem +{-# NOINLINE errorLogHead #-} +errorLogHead fn q + | S.null q = error $ "Data.Trie.Internal." ++ fn ++": found null subquery" + | otherwise = SU.unsafeHead q + +------------------------------------------------------------ +------------------------------------------------------------ + + +{----------------------------------------------------------- +-- Instances: Eq, Eq1 +-----------------------------------------------------------} + +{- +-- (2021.12.31): remove the definition of @(/=)@ for: +-- +-- +-- 'IntMap' defines their own instance so as to check the Mask +-- before the Prefix; and they have done so since at least version +-- 0.5.0.0 (2011). So I assume the performance benefits of doing +-- that are good enough to be worth it; thus, we'll do the same. +-- +-- TODO: benchmark!! +instance Eq a => Eq (Trie a) where + (==) = equal + +-- TODO: mark this INLINABLE to specialize on the Eq instance? Why doesn't IntMap? +-- TODO: Alternatively, why doesn't IntMap simply reuse the 'liftEq' implementation? +equal :: Eq a => Trie a -> Trie a -> Bool +equal (Branch p0 m0 l0 r0) + (Branch p1 m1 l1 r1) = m0 == m1 && p0 == p1 && equal l0 l1 && equal r0 r1 +equal (Arc k0 mv0 t0) + (Arc k1 mv1 t1) = k0 == k1 && mv0 == mv1 && equal t0 t1 +equal Empty Empty = True +equal _ _ = False +-} + +#if MIN_VERSION_base(4,9,0) +-- | @since 0.2.7 +instance Eq1 Trie where + liftEq = equal1 + +-- TODO: why doesn't IntMap close over @eq@? Does it really cost so much more? +-- TODO: INLINEABLE? +equal1 :: (a -> b -> Bool) -> Trie a -> Trie b -> Bool +equal1 eq (Branch p0 m0 l0 r0) (Branch p1 m1 l1 r1) = + m0 == m1 && p0 == p1 && equal1 eq l0 l1 && equal1 eq r0 r1 +equal1 eq (Arc k0 mv0 t0) (Arc k1 mv1 t1) = + k0 == k1 && liftEq eq mv0 mv1 && equal1 eq t0 t1 +equal1 _ Empty Empty = True +equal1 _ _ _ = False +#endif + +{----------------------------------------------------------- +-- Instances: Ord, Ord1 +-----------------------------------------------------------} + +-- | +-- __Warning__: This instance suffers unnecessarily from +-- . +-- +-- @since 0.2.7 +instance Ord a => Ord (Trie a) where + compare t0 t1 = compare (toList t0) (toList t1) + +#if MIN_VERSION_base(4,9,0) +-- | +-- __Warning__: This instance suffers unnecessarily from +-- . +-- +-- @since 0.2.7 +instance Ord1 Trie where + liftCompare cmp t0 t1 = + liftCompare (liftCompare cmp) (toList t0) (toList t1) +#endif + +{----------------------------------------------------------- +-- Instances: Show, Show1 +-----------------------------------------------------------} + +-- This instance does not unveil the innards of our abstract type. +-- It doesn't emit truly proper Haskell code though, since ByteStrings +-- are printed as (ASCII) Strings, but that's not our fault. +-- +-- | +-- __Warning__: This instance suffers . +-- +-- @since 0.2.2 +instance (Show a) => Show (Trie a) where + showsPrec p t = showParen (p > 10) + $ ("fromList " ++) . shows (toList t) + +#if MIN_VERSION_base(4,9,0) +-- | +-- __Warning__: This instance suffers . +-- +-- @since 0.2.7 +instance Show1 Trie where + liftShowsPrec sp sl p t = + FC.showsUnaryWith (liftShowsPrec sp' sl') "fromList" p (toList t) + where + sp' = liftShowsPrec sp sl + sl' = liftShowList sp sl +#endif + +-- | Visualization fuction for debugging. +showTrie :: (Show a) => Trie a -> String +showTrie t = shows' id t "" + where + spaces f = map (const ' ') (f "") + + shows' _ Empty = (".\n"++) + shows' ss (Branch p m l r) = + let s' = ("--"++) . shows p . (","++) . shows m . ("-+"++) + -- NOTE: GHC-9.10 warns about 'tail' being a partial function. + ss' = ss . (drop 1 (spaces s') ++) + in s' . shows' (ss' . ("|"++)) l + . ss' . ("|\n"++) + . ss' . ("`"++) . shows' (ss' . (" "++)) r + shows' ss (Arc k mv t') = + let s' = ("--"++) . shows k + . maybe id (\v -> ("-("++) . shows v . (")"++)) mv + . ("--"++) + in s' . shows' (ss . (spaces s' ++)) t' + +{----------------------------------------------------------- +-- Instances: Read, Read1 +-----------------------------------------------------------} + +-- | @since 0.2.7 +instance (Read a) => Read (Trie a) where +#ifdef __GLASGOW_HASKELL__ + readPrec = R.parens . R.prec 10 $ do + R.Ident "fromList" <- R.lexP + fromList <$> R.readPrec + + readListPrec = R.readListPrecDefault +#else + readsPrec p = readParen (p > 10) $ \ r0 -> do + ("fromList", r1) <- lex r0 + (xs, r2) <- reads r1 + return (fromList xs, r2) +#endif + +#if MIN_VERSION_base(4,9,0) +-- | @since 0.2.7 +instance Read1 Trie where + liftReadsPrec rp rl = + FC.readsData $ FC.readsUnaryWith (liftReadsPrec rp' rl') + "fromList" fromList + where + rp' = liftReadsPrec rp rl + rl' = liftReadList rp rl +#endif + +{----------------------------------------------------------- +-- Instances: Binary +-----------------------------------------------------------} + +{- +-- TODO: consider an instance more like the new one for Data.Map +-- (also for Data.IntMap), that is: +instance (Binary a) => Binary (Set.Set a) where + put s = put (size s) <> mapM_ put (toAscList s) + get = liftM fromDistinctAscList get +-- It would require redoing all the work to split bytestrings, and +-- have the overhead from storing duplicated prefixes, but is forward +-- compatible to whatever representation changes, and doesn't have +-- the invariants problem. +-- BUG: However, that would suffer from +-- , because +-- 'toList'\/@toAscList@ does. +-} +instance (Binary a) => Binary (Trie a) where + put Empty = do put (0 :: Word8) + put (Arc k mv t) = do put (1 :: Word8); put k; put mv; put t + put (Branch p m l r) = do put (2 :: Word8); put p; put m; put l; put r + + -- BUG(github#21): need to verify the invariants! + get = do tag <- get :: Get Word8 + case tag of + 0 -> return Empty + 1 -> liftM3 Arc get get get + _ -> liftM4 Branch get get get get + +-- TODO: consider adding _cereal_:'Serialize' instance. (Though that adds dependencies on: array, bytestring-builder, fail, ghc-prim). THe instances for Map/IntMap are similar to the commented ones above, though the getter uses 'fromList' rather than assuming distinct or ascending. + +-- TODO: potentially consider as well, though probably not since that has a ton of dependencies. + +-- TODO: also consider 'Storable', though that seems less appropriate for data structures than for structs and scalars... + +{----------------------------------------------------------- +-- Instances: NFData +-----------------------------------------------------------} + +-- | @since 0.2.7 +instance NFData a => NFData (Trie a) where + rnf Empty = () + rnf (Arc _ mv t) = rnf mv `seq` rnf t + rnf (Branch _ _ l r) = rnf l `seq` rnf r + +{- +-- TODO: do we want/need these? +#ifdef __GLASGOW_HASKELL__ +instance Data.Data.Data (Trie a) where ... +-- See 'IntMap' for how to do this without sacrificing abstraction. +#endif + +-- I think this macro is defined in "containers.h"? +INSTANCE_TYPEABLE(Trie) + +-- What about deriving Generic? +-} + +{----------------------------------------------------------- +-- Instances: Functor +-----------------------------------------------------------} + +instance Functor Trie where + {-# INLINE fmap #-} + fmap = fmapTrie +#ifdef __GLASGOW_HASKELL__ + -- Non-default definition since 0.2.7 + -- Leaving this definition inline and avoiding closure over @v@ + -- because that's what 'IntMap' does. + _ <$ Empty = Empty + v <$ (Arc k Nothing t) = Arc k Nothing (v <$ t) + v <$ (Arc k (Just _) t) = Arc k (Just v) (v <$ t) + v <$ (Branch p m l r) = Branch p m (v <$ l) (v <$ r) +#endif + +-- | Implementation of 'fmap', floated out so that we can give it rewrite rules. +fmapTrie :: (a -> b) -> Trie a -> Trie b +fmapTrie f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just (f v)) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) +#ifdef __GLASGOW_HASKELL__ +{-# NOINLINE [1] fmapTrie #-} +{-# RULES +"fmapTrie/fmapTrie" forall f g xs . + fmapTrie f (fmapTrie g xs) = fmapTrie (f . g) xs + #-} +-- 'IntMap' doesn't bother giving a rule for 'id'; so I guess we won't either. +#if MIN_VERSION_base(4,7,0) +{-# RULES "fmapTrie/coerce" fmapTrie coerce = coerce #-} +#endif +#endif + +-- TODO: strict version of fmap. Is there a canonical name\/class for this yet? + +{----------------------------------------------------------- +-- Instances: Traversable, Applicative, Monad +-----------------------------------------------------------} + +instance Traversable Trie where + -- TODO: 'IntMap' marks this INLINE and floats out an INLINE + -- top-level definition; no rules though, afaict. + traverse f = go + where + go Empty = pure Empty + go (Arc k Nothing t) = fmap (Arc k Nothing) (go t) + go (Arc k (Just v) t) = liftA2 (Arc k . Just) (f v) (go t) + go (Branch p m l r) = liftA2 (Branch p m) (go l) (go r) + +-- | Keyed version of 'traverse'. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.7 +traverseWithKey + :: Applicative f => (ByteString -> a -> f b) -> Trie a -> f (Trie b) +{-# INLINE traverseWithKey #-} +traverseWithKey f = go Nil + where + -- See [Note:LazyRLBS]. + go _ Empty = pure Empty + go q (Branch p m l r) = liftA2 (Branch p m) (go q l) (go q r) + go q (Arc k Nothing t) = fmap (Arc k Nothing) (go (q +>! k) t) + go q (Arc k (Just v) t) = + let q' = toStrict (q +>? k) + in liftA2 (Arc k . Just) (f q' v) (go (fromStrict q') t) + +-- [Note:LazyRLBS]: We avoid making the RLBS parameter strict, to +-- avoid incuring the cost of 'toStrict' if the user's function +-- does not force it. However, if they do force it, then we'll +-- still have the +-- problem. Using RLBS only reduces the constant factor of the +-- quadratic. + +------------------------------------------------------------ +-- TODO: would make more sense to use intersection\/zip semantics here, +-- rather than the overlaid-unionL semantics of the 'Monad'. Alas, +-- done is done. In a future major version we can try changing +-- that, and introducing newtype wrappers for this overlaid\/unionL +-- version (and the prospective underlaid\/unionR version). +-- TODO: see also +-- +-- | @since 0.2.2 +instance Applicative Trie where + pure = singleton S.empty + t0 <*> t1 = t0 >>= (<$> t1) + -- TODO: can we do better than these defaults? + -- t0 *> t1 = (id <$ t0) <*> t1 + -- t0 <* t1 = (const <$> t0) <*> t1 -- actually uses @liftA2 const t0 t1@ +#if MIN_VERSION_base(4,10,0) + -- liftA2 f t0 t1 = (f <$> t0) <*> t1 +#endif + {- + -- Inlining and case-of-case yields the following (which GHC + -- could surely derive on its own): + Empty *> _ = Empty + Branch p m l r *> t1 = branch p m (l *> t1) (r *> t1) + Arc k Nothing s *> t1 = prependNN k (s *> t1) + Arc k (Just _) s *> t1 = prepend k (t1 `unionL` (s *> t1)) + + -- This one is marginally better, since we can use @(<$)@ in the Accept case. + Empty <* _ = Empty + Branch p m l r <* t1 = branch p m (l <* t1) (r <* t1) + Arc k Nothing s <* t1 = prependNN k (s <* t1) + Arc k (Just v) s <* t1 = prepend k ((v <$ t1) `unionL` (s <* t1)) + + -- This one took a lot of inlining\/massaging, so might be worth it... + -- It's easier to see the structure if we define a closure + -- @(liftA2 f _ t1)@, but unclear if that would hurt the improvement + -- of the implementation. + liftA2 f Empty _ = Empty + liftA2 f (Branch p m l r) t1 = branch p m (liftA2 f l t1) (liftA2 f r t1) + liftA2 f (Arc k Nothing s) t1 = prependNN k (liftA2 f s t1) + liftA2 f (Arc k (Just v) s) t1 = prepend k ((f v <$> t1) `unionL` liftA2 f s t1) + -} + +------------------------------------------------------------ +-- Does this even make sense? It's not nondeterminism like lists +-- and sets. If no keys were prefixes of other keys it'd make sense +-- as a decision-tree; but since keys /can/ prefix, tries formed +-- from shorter keys can shadow the results from longer keys due +-- to the 'unionL'. It does seem to follow the laws though... What +-- computation could this possibly represent? +-- +-- 1. return x >>= f ≡ f x +-- 2. m >>= return ≡ m +-- 3. (m >>= f) >>= g ≡ m >>= (\x -> f x >>= g) +-- +-- | @since 0.2.2 +instance Monad Trie where +-- Since base-4.8 (ghc-7.10.1) we have the default @return = pure@. +-- Since ghc-9.2.1 we get a warning about providing any other +-- definition, and should instead define both 'pure' and @(*>)@ +-- directly, leaving 'return' and @(>>)@ as their defaults so they +-- can eventually be removed from the class. +-- +#if (!(MIN_VERSION_base(4,8,0))) + return = pure +#endif + -- FIXME: See [bug26]. + (>>=) Empty _ = empty + (>>=) (Branch p m l r) f = branch p m (l >>= f) (r >>= f) + (>>=) (Arc k Nothing t) f = prependNN k (t >>= f) + (>>=) (Arc k (Just v) t) f = prepend k (f v `unionL` (t >>= f)) + where unionL = mergeBy (\x _ -> Just x) + + +{----------------------------------------------------------- +-- Instances: Semigroup, Monoid +-----------------------------------------------------------} + +#if MIN_VERSION_base(4,9,0) +-- The "Data.Semigroup" module is in base since 4.9.0.0; but having +-- the 'Semigroup' superclass for the 'Monoid' instance only comes +-- into force in base 4.11.0.0. +-- | @since 0.2.5 +instance (Semigroup a) => Semigroup (Trie a) where + (<>) = mergeBy $ \x y -> Just (x <> y) + -- Non-default definition since 0.2.7 + stimes = fmap . stimes +#endif + + +-- This instance is more sensible than Data.IntMap and Data.Map's +instance (Monoid a) => Monoid (Trie a) where + mempty = empty +#if MIN_VERSION_base(4,11,0) + -- Now that the canonical instance is the default, don't define + -- 'mappend', in anticipation of Phase 4 of: + -- +{- +--if MIN_VERSION_base(4,9,0) + -- GHC 8.0/8.2 -Wnoncanonical-monoid-instances wants this + -- definition, even though 'Semigroup' isn't a superclass of + -- 'Monoid' until base-4.11 so it would require additional + -- constraints on the instance. Since we're only supporting + -- these older versions for legacy reasons, there's no reason + -- to bother adhering to the -Wcompat here. + mappend = (<>) +-} +#else + mappend = mergeBy $ \x y -> Just (x `mappend` y) +#endif + -- TODO: Can we improve `mconcat` over the default? + + +{----------------------------------------------------------- +-- Instances: Alternative, MonadPlus +-----------------------------------------------------------} + +-- Since the Monoid instance isn't natural in @a@, I can't think +-- of any other sensible instance for MonadPlus. It's as specious +-- as Maybe, IO, and STM's instances though. +-- +-- MonadPlus laws: +-- 1. forms a monoid +-- 2. mzero >>= f ≡ mzero +-- 3. m >> mzero ≡ mzero +-- 4. mplus m n >>= k ≡ mplus (m >>= k) (n >>= k) +-- 4' mplus (return a) n ≡ return a +{- +-- Follows #1, #1, and #3. But it does something like 4' instead +-- of actually doing #4 (since we'd merge the trees generated by +-- @k@ for conflicting values) +-- +-- TODO: cf Control.Applicative.Alternative (base-4, but not Hugs). +-- But (<*>) gets odd when the function is not 'pure'... maybe +-- helpful though. +instance MonadPlus Trie where + mzero = empty + mplus = unionL where unionL = mergeBy (\x _ -> Just x) +-} + + +{----------------------------------------------------------- +-- Extra mapping and filtering functions +-----------------------------------------------------------} + +{----------------------------------------------------------- +-- Pseudo-instances: Filterable, Witherable +-----------------------------------------------------------} + +-- We avoid depending on the _filterable_ package because it combines +-- too many things into its @Filterable@ class. And we avoid using +-- the _witherable_ package because it has too many dependencies, +-- and too many orphan instances. However, we go with the names +-- (mostly[1]) and laws as phrased by _witherable_. +-- +-- [1]: I'm rather not a fan of @mapMaybe@, despite its pervasiveness. +-- And similarly for @catMaybes@ etc. That's actually one of the +-- reasons I prefer the _witherable_ package over _filterable_: +-- because of the name 'wither' instead of @mapMaybeA@ :) + + +-- | Apply a function to all values, potentially removing them. +-- +-- ==== __Laws__ +-- [/Fission/] +-- @'filterMap' f ≡ 'fmap' ('Data.Maybe.fromJust' . f) . 'filter' ('Data.Maybe.isJust' . f)@ +-- +-- [/Fusion/] +-- @'fmap' f . 'filter' g ≡ 'filterMap' (\\v -> f v '<$' 'Control.Monad.guard' (g v))@ +-- +-- [/Conservation/] +-- @'filterMap' ('Just' . f) ≡ 'fmap' f@ +-- +-- [/Composition/] +-- @'filterMap' f . 'filterMap' g ≡ 'filterMap' (f 'Control.Monad.<=<' g)@ +-- +-- The fission\/fusion laws are essentially the same, they differ +-- only in which direction is more \"natural\" for use as a rewrite +-- rule. The conservation law is just a special case of fusion, +-- but it's a particularly helpful one to take note of. +-- +filterMap :: (a -> Maybe b) -> Trie a -> Trie b +filterMap f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = arc k (f v) (go t) + start t = go t + -- FIXME: See [bug26]. + go Empty = empty + go (Arc k Nothing t) = prependNN k (go t) + go (Arc k (Just v) t) = arcNN k (f v) (go t) + go (Branch p m l r) = branch p m (go l) (go r) +-- TODO: rewrite rule for the latter three laws. (The fission law +-- is unlikely to be very helpful.) +-- TODO: why not implement as @contextualFilterMap (const . f)@ ? +-- Does that actually incur additional overhead? + + +-- Some translations: +-- @'filter' f ≡ 'filterMap' (\\v -> if f v then 'Just' v else 'Nothing')@ +-- @'filter' f ≡ 'filterMap' (('<$') '<*>' 'Control.Monad.guard' . f)@ +-- @'filter' f ≡ 'filterMap' (\\v -> v '<$' 'Control.Monad.guard' (f v))@ +-- +-- | Retain only those values which satisfy some predicate. +-- +-- ==== __Laws__ +-- [/Definition/] +-- @'filter' f ≡ 'filterMap' (\\v -> v '<$' 'Control.Monad.guard' (f v))@ +-- +-- [/Composition/] +-- @'filter' f . 'filter' g ≡ 'filter' ('liftA2' ('&&') f g)@ +-- +-- The definition above is a special case of the fusion law for +-- 'filterMap'. (Also, the name just means definitional-equality; +-- it's not the actual implementation used.) +-- +-- @since 0.2.7 +filter :: (a -> Bool) -> Trie a -> Trie a +filter f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = arcB k v (f v) (go t) + start t = go t + -- FIXME: See [bug26]. + go Empty = empty + go (Arc k Nothing t) = prependNN k (go t) + go (Arc k (Just v) t) = arcNNB k v (f v) (go t) + go (Branch p m l r) = branch p m (go l) (go r) + +-- | > arcB k v b ≡ arc k (v <$ guard b) +arcB :: ByteString -> a -> Bool -> Trie a -> Trie a +{-# INLINE arcB #-} +arcB k v True = Arc k (Just v) +arcB k _ False = prepend k + +-- | > arcNNB k v b ≡ arcNN k (v <$ guard b) +arcNNB :: ByteString -> a -> Bool -> Trie a -> Trie a +{-# INLINE arcNNB #-} +arcNNB k v True = Arc k (Just v) +arcNNB k _ False = prependNN k + + +{- +#if MIN_VERSION_base(4,9,0) +import Data.Functor.Compose (Compose(Compose)) +#endif +#if MIN_VERSION_base(4,8,0) +import Data.Functor.Identity (Identity(Identity)) +#endif +-} + +-- BUG: Is there any other way to get the Haddock to force linebreaks, +-- other than using birdtracks which have a whole different +-- stylization? +-- TODO: is commutative monad sufficient, or are there other +-- requirements too? If it does in fact hold for commutative monads, +-- then we should state it as a law (like we do for naturality). +-- +-- | An effectful version of 'filterMap'. +-- +-- ==== __Laws__ +-- [/Naturality/] +-- @'wither' (t . f) ≡ t . 'wither' f@, +-- for any /applicative-transformation/ @t@ +-- +-- [/Purity/] +-- @'wither' ('pure' . f) ≡ 'pure' . 'filterMap' f@ +-- +-- [/Conservation/] +-- @'wither' ('fmap' 'Just' . f) ≡ 'traverse' f@ +-- +-- [/Horizontal Composition/] +-- @'wither' f \`under\` 'wither' g ≡ 'wither' (wither_Maybe f \`under\` g)@, +-- where: +-- +-- > under :: Functor f +-- > => (b -> g c) +-- > -> (a -> f b) +-- > -> a -> Compose f g c +-- > under g f = Compose . fmap g . f +-- > +-- > -- | Variant of wither for Maybe instead of Trie. +-- > wither_Maybe :: Applicative f +-- > => (a -> f (Maybe b)) +-- > -> Maybe a -> f (Maybe b) +-- > wither_Maybe f = fmap join . traverse f +-- +-- Note that the horizontal composition law is using two different +-- applicative functors. Conversely, a vertical composition law +-- would have the form: @'wither' f 'Control.Monad.<=<' 'wither' g ≡ ...@; +-- however, we cannot have such a law except when the applicative +-- functor is in fact a commutative monad (i.e., the order of effects +-- doesn't matter). For the curious, the terminology of +-- vs +-- +-- comes from category theory. +-- +-- Although the horizontal composition law may look baroque, it is +-- helpful to compare it to the composition law for 'traverse' +-- itself: +-- +-- @'traverse' f \`under\` 'traverse' g ≡ 'traverse' (f \`under\` g)@ +-- +-- @since 0.2.7 +wither :: Applicative f => (a -> f (Maybe b)) -> Trie a -> f (Trie b) +wither f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = liftA2 (arc k) (f v) (go t) + start t = go t + -- FIXME: See [bug26]. + go Empty = pure empty + go (Arc k Nothing t) = fmap (prependNN k) (go t) + go (Arc k (Just v) t) = liftA2 (arcNN k) (f v) (go t) + go (Branch p m l r) = liftA2 (branch p m) (go l) (go r) + +-- Some other spellings of the translation: +-- @'filterA' f ≡ 'wither' (\\v -> (\\b -> if b then 'Just' v else 'Nothing') '<$>' f v)@ +-- @'filterA' f ≡ 'wither' (\\v -> (\\b -> v '<$' 'Control.Monad.guard' b) '<$>' f v)@ +-- @'filterA' f ≡ 'wither' ('fmap' . (. 'Control.Monad.guard') . ('<$') '<*>' f)@ +-- +-- An alternative variant of 'underA2' which has a nicer type. +-- Alas, that just makes the twist show up in the law itself; also +-- the name is peculiar since it's the second argument that's run +-- under the first. Decisions decisions. +-- > underF2 :: (Functor f, Functor g) +-- > => (b -> c -> d) +-- > -> (a -> f b) +-- > -> (a -> g c) +-- > -> a -> Compose f g d +-- > underF2 h f g a = Compose (f a <&> ((g a <&>) . h)) +-- +-- +-- | An effectful version of 'filter'. +-- +-- ==== __Laws__ +-- [/Definition/] +-- @'filterA' f ≡ 'wither' (\\v -> (v '<$') . 'Control.Monad.guard' '<$>' f v)@ +-- +-- [/Naturality/] +-- @'filterA' (t . f) ≡ t . 'filterA' f@, +-- for any /applicative-transformation/ @t@ +-- +-- [/Purity/] +-- @'filterA' ('pure' . f) ≡ 'pure' . 'filter' f@ +-- +-- [/Horizontal Composition/] +-- @'filterA' f \`under\` 'filterA' g ≡ 'filterA' (underA2 ('&&') f g)@, +-- where +-- +-- > -- Like 'liftA2' for the @(a->)@ monad, but horizontal. +-- > -- The function definition should (hopefully) be straightforward +-- > -- to follow; however, do beware the oddly criss-crossed types +-- > -- for @g@ and @f@. +-- > underA2 :: (Applicative f, Applicative g) +-- > => (b -> c -> d) +-- > -> (a -> g b) +-- > -> (a -> f c) +-- > -> a -> Compose f g d +-- > underA2 h g f = liftA2 (liftA2 h) (g `under` pure) (pure `under` f) +-- +-- For the definition of @under@ and more details about horizontal +-- composition, see the laws section of 'wither'. +-- +-- @since 0.2.7 +filterA :: Applicative f => (a -> f Bool) -> Trie a -> f (Trie a) +filterA f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = liftA2 (arcB k v) (f v) (go t) + start t = go t + -- FIXME: See [bug26]. + go Empty = pure empty + go (Arc k Nothing t) = fmap (prependNN k) (go t) + go (Arc k (Just v) t) = liftA2 (arcNNB k v) (f v) (go t) + go (Branch p m l r) = liftA2 (branch p m) (go l) (go r) + + +-- | Keyed version of 'filterMap'. +-- +-- __Warning__: This function suffers . +mapBy :: (ByteString -> a -> Maybe b) -> Trie a -> Trie b +-- TODO: why not implement as @contextualMapBy (\k v _ -> f k v)@ ? +-- Does that actually incur additional overhead? +mapBy f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = arc k (f k v) (go (fromStrict k) t) + start t = go Nil t + -- FIXME: See [bug26]. + -- See [Note:LazyRLBS]. + go _ Empty = empty + go q (Branch p m l r) = branch p m (go q l) (go q r) + go q (Arc k Nothing t) = prependNN k (go (q +>! k) t) + go q (Arc k (Just v) t) = arcNN k (f q' v) (go (fromStrict q') t) + where q' = toStrict (q +>? k) + + +-- | A variant of 'fmap' which provides access to the subtrie rooted +-- at each value. +-- +-- @since 0.2.3 +contextualMap :: (a -> Trie a -> b) -> Trie a -> Trie b +contextualMap f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just (f v t)) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) + + +-- | A variant of 'contextualMap' which evaluates the function strictly. +-- +-- @since 0.2.3 +contextualMap' :: (a -> Trie a -> b) -> Trie a -> Trie b +contextualMap' f = go + where + go Empty = Empty + go (Arc k Nothing t) = Arc k Nothing (go t) + go (Arc k (Just v) t) = Arc k (Just $! f v t) (go t) + go (Branch p m l r) = Branch p m (go l) (go r) + + +-- | Contextual variant of 'filterMap'. +-- +-- @since 0.2.3 +contextualFilterMap :: (a -> Trie a -> Maybe b) -> Trie a -> Trie b +contextualFilterMap f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = arc k (f v t) (go t) + start t = go t + -- FIXME: See [bug26]. + go Empty = empty + go (Arc k Nothing t) = prependNN k (go t) + go (Arc k (Just v) t) = arcNN k (f v t) (go t) + go (Branch p m l r) = branch p m (go l) (go r) + + +-- | Contextual variant of 'mapBy', aka keyed variant of 'contextualFilterMap'. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.3 +contextualMapBy :: (ByteString -> a -> Trie a -> Maybe b) -> Trie a -> Trie b +contextualMapBy f = start + where + -- Handle epsilon values before entering the main recursion. + start (Arc k (Just v) t) = arc k (f k v t) (go (fromStrict k) t) + start t = go Nil t + -- FIXME: See [bug26]. + -- See [Note:LazyRLBS]. + go _ Empty = empty + go q (Branch p m l r) = branch p m (go q l) (go q r) + go q (Arc k Nothing t) = prependNN k (go (q +>! k) t) + go q (Arc k (Just v) t) = arcNN k (f q' v t) (go (fromStrict q') t) + where q' = toStrict (q +>? k) + + +{----------------------------------------------------------- +-- Basic functions +-----------------------------------------------------------} +-- TODO: probably want to hoist this up top + +-- | \(\mathcal{O}(1)\). Construct the empty trie. +empty :: Trie a +{-# INLINE empty #-} +empty = Empty + + +-- | \(\mathcal{O}(1)\). Is the trie empty? +null :: Trie a -> Bool +{-# INLINE null #-} +null Empty = True +null _ = False + + +-- | \(\mathcal{O}(1)\). Construct a singleton trie. +singleton :: ByteString -> a -> Trie a +{-# INLINE singleton #-} +singleton k v = Arc k (Just v) Empty +-- For singletons, don't need to verify invariant on arc length >0 + + +-- | \(\mathcal{O}(n)\). Get count of elements in trie. +size :: Trie a -> Int +{-# INLINE size #-} +size = size' 0 + +-- This is just @F.foldl' (\n _ -> n+1) 0@ manually inlined\/specialized. +-- Thus, see our implementation of 'foldl'' to see why we use this +-- particular phrasing out of the very many alternatives. +size' :: Int -> Trie a -> Int +size' !n Empty = n +size' n (Arc _ Nothing t) = size' n t +size' n (Arc _ (Just _) t) = size' (n + 1) t +size' n (Branch _ _ l r) = size' (size' n l) r + + +{----------------------------------------------------------- +-- Instances: Foldable +-----------------------------------------------------------} + +-- [Note:FoldEta]: For all the folding functions, we take only the +-- two algebra arguments on the left of the \"=\", leaving the +-- 'Trie' argument as a lambda on the right of the \"=\". This is +-- to allow the functions to be inlined when passed only the two +-- algebra arguments, rather than requiring all three arguments +-- before being inlined. + +instance Foldable Trie where + -- TODO: 'IntMap' marks 'fold' as INLINABLE; however, that seems + -- to introduce a bit of slowdown with respect to 'foldMap' + -- (which goes away when marked INLINE instead); so we should + -- do more to discern which inlinement is preferable. + {-# INLINE fold #-} + fold = go + where + go Empty = mempty + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = v `mappend` go t + go (Branch _ _ l r) = go l `mappend` go r + {-# INLINE foldMap #-} + foldMap f = go + where + go Empty = mempty + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = f v `mappend` go t + go (Branch _ _ l r) = go l `mappend` go r +#if MIN_VERSION_base(4,13,0) + -- TODO: float out this definition so folks can still use it + -- on earlier versions of base? + -- TODO: verify order of 'mappend' on some non-commutative monoid! + {-# INLINE foldMap' #-} + foldMap' f = go mempty + where + -- Benchmarking on GHC 9.2.1 indicates that for this function + -- the (m,t) argument ordering is somewhat (~3%) faster + -- than the (t,m) order; and both allocate the same. + -- This differs from the case for 'foldr'' and 'foldl'; + -- though I'm not sure why. + -- TODO: Once we disable HPC, now it's looking like the + -- flopped version is faster afterall... + go !m Empty = m + go m (Arc _ Nothing t) = go m t + go m (Arc _ (Just v) t) = go (m `mappend` f v) t + go m (Branch _ _ l r) = go (go m l) r +#endif + -- FIXME: (2022.03.06): N.B., the benchmarks are quite erratic + -- regarding whether this implementation is better than the + -- 'foldMap'-with-'Endo' definition or not; whichever one is + -- favored at some particular time is always massively favored + -- over the other. Most recently this implementation is being + -- favored. + {-# INLINE foldr #-} + foldr f z0 = \t -> go t z0 -- See [Note:FoldEta]. + where + go Empty = id + go (Arc _ Nothing t) = go t + go (Arc _ (Just v) t) = f v . go t + go (Branch _ _ l r) = go l . go r +#if MIN_VERSION_base(4,6,0) + -- TODO: float out this definition so folks can still use it + -- on earlier versions of base? + {-# INLINE foldr' #-} + foldr' f z0 = go z0 -- See [Note:FoldEta]. + where + -- Benchmarking on GHC 9.2.1 indicates that for this function: + -- for smaller tries, the (t,z) argument order is ~10% faster + -- than (z,t) and allocation is the same for both; however, + -- for larger tries the (t,z) argument order is ~6% slower + -- than (z,t). + -- + -- Also, weirdly, benchmarking indicates that the @($!)@ + -- in the Branch case slightly improved things. + -- TODO: what's going on with the @($!)@; bogus? + -- TODO: once HPC disabled, now it's saying the unflopped + -- version without the extra @($!)@ is the faster one! + -- (unflopped with @($!)@ is only marginally slower; probably + -- noise). + go !z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = f v $! go z t + go z (Branch _ _ l r) = go (go z r) l +#endif + {-# INLINE foldl #-} + foldl f z0 = \t -> go t z0 -- See [Note:FoldEta]. + where + -- Benchmarking on GHC 9.2.1 indicates that for this function + -- the (t,z) argument order is slightly faster (~0.8%) and + -- allocates ~8.4% less, compared to the (z,t) order. + -- I've no idea why the allocation would differ, especially + -- when it doesn't for 'foldr'' and 'foldMap''. + -- TODO: once HPC disabled, now it's showing the flopped + -- version is ~2x faster! bogus? + -- TODO: (2022.03.05) Rerun this benchmark on larger tries. + go Empty z = z + go (Arc _ Nothing t) z = go t z + go (Arc _ (Just v) t) z = go t (f z v) + go (Branch _ _ l r) z = go r (go l z) +#if MIN_VERSION_base(4,6,0) + -- TODO: float out this definition so folks can still use it + -- on earlier versions of base? + {-# INLINE foldl' #-} + foldl' f z0 = go z0 -- See [Note:FoldEta]. + where + -- Benchmarking on GHC 9.2.1 indicates that for this function + -- the (z,t) argument order is significantly faster (~10%) and + -- allocates half as much. + -- TODO: figure out why\/how the allocation could differ so much; bogus? + -- TODO: figure out why benchmarking indicates the \"flop_bang\" + -- version is ~4% faster (albeit ~32% more allocation); bogus? + -- TODO: once HPC disabled, the flopped version is showing ~2x faster; bogus? + -- TODO: (2022.03.05) Rerun this benchmark on larger tries. + go !z Empty = z + go z (Arc _ Nothing t) = go z t + go z (Arc _ (Just v) t) = go (f z v) t + go z (Branch _ _ l r) = go (go z l) r +#endif + -- TODO: any point in doing foldr1,foldl1? +#if MIN_VERSION_base(4,8,0) + -- TODO: float out this definition so folks can still use it + -- on earlier versions of base? + {-# INLINE length #-} + length = size + {-# INLINE null #-} + null = null -- FIXME: ensure this isn't cyclic definition! + {-# INLINE toList #-} + toList = elems -- NB: Foldable.toList /= Trie.toList + {- + -- TODO: need to move these definitions here... + -- TODO: may want to give a specialized implementation of 'member' then + {-# INLINE elem #-} -- TODO: 'IntMap' uses INLINABLE with a local definition; rather than INLINE with the top-level definition... Also, that TLD remarks to see [Note:Closure] + elem = member + -} + -- TODO: why does IntMap define these two specially, rather than using foldl' or foldl1' ? + {-# INLINABLE maximum #-} + maximum = go0 + where + go0 Empty = error "Data.Foldable.maximum @Trie: empty trie" + go0 (Arc _ Nothing t) = go0 t + go0 (Arc _ (Just v) t) = go v t + go0 (Branch _ _ l r) = go (go0 l) r + go !w Empty = w + go w (Arc _ Nothing t) = go w t + go w (Arc _ (Just v) t) = go (max w v) t + go w (Branch _ _ l r) = go (go w l) r + {-# INLINABLE minimum #-} + minimum = go0 + where + go0 Empty = error "Data.Foldable.minimum @Trie: empty trie" + go0 (Arc _ Nothing t) = go0 t + go0 (Arc _ (Just v) t) = go v t + go0 (Branch _ _ l r) = go (go0 l) r + go !w Empty = w + go w (Arc _ Nothing t) = go w t + go w (Arc _ (Just v) t) = go (min w v) t + go w (Branch _ _ l r) = go (go w l) r + {-# INLINABLE sum #-} + sum = F.foldl' (+) 0 + {-# INLINABLE product #-} + product = F.foldl' (*) 1 +#endif + +-- TODO: newtype Keys = K Trie ; instance Foldable Keys +-- TODO: newtype Assoc = A Trie ; instance Foldable Assoc + +{----------------------------------------------------------- +-- Extra folding functions +-----------------------------------------------------------} + +-- TODO: be sure to keep this in sync with whatever implementation +-- choice we use for 'F.foldr'; especially since that's the one +-- method of 'Foldable' where we can't improve substantially over +-- the default implementation. +-- +-- | Keyed variant of 'F.foldr'. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.2 +foldrWithKey :: (ByteString -> a -> b -> b) -> b -> Trie a -> b +{-# INLINE foldrWithKey #-} +foldrWithKey f z0 = \t -> go Nil t z0 -- See [Note:FoldEta]. + where + -- See [Note:LazyRLBS]. + go _ Empty = id + go q (Branch _ _ l r) = go q l . go q r + go q (Arc k Nothing t) = go (q +>! k) t + go q (Arc k (Just v) t) = f q' v . go (fromStrict q') t + where q' = toStrict (q +>? k) + +-- TODO: probably need to benchmark these separately from the +-- non-keyed variants, since the extra recursive argument will +-- surely sway things like whether to flop or not. +-- TODO: Consider just giving an +-- +-- instance, instead of naming all these separately. That adds a +-- lot of additional dependencies just to define the class, but... +-- Or maybe give an +-- instance. Again, lots of added dependencies just for the class,... +-- Then again, maybe we should just stick with doing everything +-- outside of classes; that way we could introduce a Cabal flag for +-- deciding whether the user wants either of those classes (and +-- should do the same for Witherable). + +-- | Keyed variant of 'F.foldr''. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.7 +foldrWithKey' :: (ByteString -> a -> b -> b) -> b -> Trie a -> b +{-# INLINE foldrWithKey' #-} +foldrWithKey' f z0 = go Nil z0 -- See [Note:FoldEta]. + where + -- See [Note:LazyRLBS]. + go _ !z Empty = z + go q z (Branch _ _ l r) = go q (go q z r) l + go q z (Arc k Nothing t) = go (q +>! k) z t + go q z (Arc k (Just v) t) = f q' v $! go (fromStrict q') z t + where q' = toStrict (q +>? k) + +-- | Keyed variant of 'F.foldl'. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.7 +foldlWithKey :: (b -> ByteString -> a -> b) -> b -> Trie a -> b +{-# INLINE foldlWithKey #-} +foldlWithKey f z0 = go Nil z0 -- See [Note:FoldEta]. + where + -- See [Note:LazyRLBS]. + go _ z Empty = z + go q z (Branch _ _ l r) = go q (go q z l) r + go q z (Arc k Nothing t) = go (q +>! k) z t + go q z (Arc k (Just v) t) = go (fromStrict q') (f z q' v) t + where q' = toStrict (q +>? k) + +-- | Keyed variant of 'F.foldl''. +-- +-- __Warning__: This function suffers . +-- +-- @since 0.2.7 +foldlWithKey' :: (b -> ByteString -> a -> b) -> b -> Trie a -> b +{-# INLINE foldlWithKey' #-} +foldlWithKey' f z0 = go Nil z0 -- See [Note:FoldEta]. + where + -- See [Note:LazyRLBS]. + go _ !z Empty = z + go q z (Branch _ _ l r) = go q (go q z l) r + go q z (Arc k Nothing t) = go (q +>! k) z t + go q z (Arc k (Just v) t) = go (fromStrict q') (f z q' v) t + where q' = toStrict (q +>? k) + +-- | Catamorphism for tries. Unlike most other functions ('mapBy', +-- 'contextualMapBy', 'foldrWithKey', etc), this function does /not/ +-- reconstruct the full 'ByteString' for each value; instead it +-- only returns the suffix since the previous value or branch point. +-- +-- This function is a direct\/literal catamorphism of the implementation +-- datatype, erasing only some bitmasking metadata for the branches. +-- For a more semantic catamorphism, see 'cata'. +-- +-- @since 0.2.6 +cata_ + :: (ByteString -> Maybe a -> b -> b) -- ^ Algebra for arc. + -> (b -> b -> b) -- ^ Algebra for binary branch. + -> b -- ^ Algebra for empty trie. + -> Trie a -> b +{-# INLINE cata_ #-} +cata_ a b e = go + where + go Empty = e + go (Arc k mv t) = a k mv (go t) + go (Branch _ _ l r) = b (go l) (go r) + + +-- | Catamorphism for tries. Unlike most other functions ('mapBy', +-- 'contextualMapBy', 'foldrWithKey', etc), this function does /not/ +-- reconstruct the full 'ByteString' for each value; instead it +-- only returns the suffix since the previous value or branch point. +-- +-- This function is a semantic catamorphism; that is, it tries to +-- express the invariants of the implementation, rather than exposing +-- the literal structure of the implementation. For a more literal +-- catamorphism, see 'cata_'. +-- +-- @since 0.2.6 +cata + :: (ByteString -> a -> b -> b) -- ^ Algebra for accepting arcs. + -> (ByteString -> [b] -> b) -- ^ Algebra for n-ary branch with prefix. + -> b -- ^ Algebra for empty trie. + -> Trie a -> b +cata a b e = go + where + step k (Just v) t = a k v (go t) + step k Nothing t = b k (collect t []) + go Empty = e + go (Arc k mv t) = step k mv t + go (Branch _ _ l r) = b S.empty (collect l (collect r [])) + -- TODO: would it be profitable to use 'build' for these lists? + collect Empty bs = bs + collect (Arc k mv t) bs = step k mv t : bs + collect (Branch _ _ l r) bs = collect l (collect r bs) + + + +{----------------------------------------------------------- +-- Instances: IsList +-----------------------------------------------------------} + +#if __GLASGOW_HASKELL__ >= 708 +-- | +-- __Warning__: The 'toList' method of this instance suffers +-- . +-- +-- @since 0.2.7 +instance GHC.Exts.IsList (Trie a) where + type Item (Trie a) = (ByteString, a) + fromList = fromList + toList = toList +#endif + + +-- /Moved to "Data.Trie.Internal" since 0.2.7/ +-- We define this here because 'GHC.Exts.IsList' wants it. +-- +-- | Convert association list into a trie. On key conflict, values +-- earlier in the list shadow later ones. +fromList :: [(ByteString,a)] -> Trie a +{-# INLINE fromList #-} +fromList = foldr (uncurry insert) empty + where + insert = alterBy (\_ x _ -> Just x) + + +-- /Moved to "Data.Trie.Internal" since 0.2.7/ +-- We define this here simply because so many instances want to use it. +-- TODO: would it be worth defining this directly, for optimizing +-- the case where list fusion doesn't eliminate the list? +-- +-- | Convert trie into association list. The list is ordered +-- according to the keys. +-- +-- __Warning__: This function suffers . +toList :: Trie a -> [(ByteString,a)] +{-# INLINE toList #-} +toList = toListBy (,) + + +-- cf Data.ByteString.unpack +-- +-- +-- | Convert a trie into a list using a function. Resulting values +-- are in key-sorted order. +-- +-- __Warning__: This function suffers . +toListBy :: (ByteString -> a -> b) -> Trie a -> [b] +{-# INLINE toListBy #-} +#if !defined(__GLASGOW_HASKELL__) +-- TODO: should probably inline foldrWithKey +-- TODO: compare performance of that vs both this and the GHC version +toListBy f t = foldrWithKey (((:) .) . f) [] t +#else +-- Written with 'build' to enable the build\/foldr fusion rules. +toListBy f t = build (toListByFB f t) + +-- TODO: should probably have a specialized version for strictness, +-- and a rule to rewrite generic lazy version into it. As per +-- Data.ByteString.unpack and the comments there about strictness +-- and fusion. +toListByFB :: (ByteString -> a -> b) -> Trie a -> (b -> c -> c) -> c -> c +{-# INLINE [0] toListByFB #-} +toListByFB f t cons nil = foldrWithKey ((cons .) . f) nil t +#endif + +-- /Moved to "Data.Trie.Internal" since 0.2.7/ +-- So that we can do list-fusion, and reuse the definition for Foldable +-- +-- | Return all values in the trie, in key-sorted order. +-- +-- __Note__: Prior to version 0.2.7, this function suffered +-- ; but it no longer does. +-- +-- @since 0.2.2 +elems :: Trie a -> [a] +{-# INLINE elems #-} +#ifdef __GLASGOW_HASKELL__ +elems t = build (\cons nil -> F.foldr cons nil t) +#else +elems = F.foldr (:) [] +#endif + + +------------------------------------------------------------ +------------------------------------------------------------ + + +{----------------------------------------------------------- +-- Query functions (just recurse) +-----------------------------------------------------------} + +-- TODO: All the lookup-like functions for 'IntMap' remark to see +-- [Note:Closure]; so we should double check our 'lookupBy_' adheres +-- to that performance guideline. + +-- | Generic function to find a value (if it exists) and the subtrie +-- rooted at the prefix. The first function argument is called if and +-- only if a node is exactly reachable by the query; if no node is +-- exactly reachable the default value is used; if the middle of +-- an arc is reached, the second function argument is used. +-- +-- This function is intended for internal use. For the public-facing +-- version, see 'Data.Trie.lookupBy'. +-- +-- __Note__: /Type changed in 0.2.7/ +lookupBy_ + :: (a -> Trie a -> b) -- ^ The query matches a value. + -> (Trie a -> b) -- ^ The query doesn't match, but an extension might. + -> b -- ^ The query doesn't match, nor does any extension. + -> ByteString -> Trie a -> b +lookupBy_ found missing clash = start + where + -- | Deal with epsilon query (when there is no epsilon value) + start q t@(Branch{}) | S.null q = missing t + start q t = go q t + -- | The main recursion + go _ Empty = clash + go q (Arc k mv t) = + let (_,k',q') = breakMaximalPrefix k q + in case (S.null k', S.null q') of + (False, True) -> missing (Arc k' mv t) + (False, False) -> clash + (True, True) -> + case mv of + Nothing -> missing t + Just v -> found v t + (True, False) -> go q' t + go q t_@(Branch{}) = findArc t_ + where + qh = errorLogHead "lookupBy_" q + -- | \(\mathcal{O}(\min(m,W))\), where \(m\) is number of + -- @Arc@s in this branching, and \(W\) is the word size of + -- the Prefix,Mask type. + findArc Empty = impossible "lookupBy_" -- see [Note1] + findArc t@(Arc{}) = go q t + findArc (Branch p m l r) + | nomatch qh p m = clash + | zero qh m = findArc l + | otherwise = findArc r + +-- [Note1]: Our use of the 'branch' and 'graft' smart constructors +-- ensure that 'Empty' never occurs in a 'Branch' tree ('Empty' can +-- only occur at the root, or under an 'Arc' with value); therefore +-- the @findArc Empty@ case is unreachable. If we allowed such +-- nodes, however, then this case should return the same result as +-- the 'nomatch' case. + + +-- This function needs to be here, not in "Data.Trie", because of +-- 'arc' which isn't exported. We could use the monad instance +-- instead, though it'd be far more circuitous. +-- arc k Nothing t ≡ singleton k () >> t +-- arc k (Just v) t ≡ singleton k v >>= unionR t . singleton S.empty +-- (...except 'arc' doesn't do the invariant correction +-- of (>>=) for epsilon'elem't) +-- +-- | Return the subtrie containing all keys beginning with a prefix. +submap :: ByteString -> Trie a -> Trie a +{-# INLINE submap #-} +submap q + | S.null q = id + | otherwise = lookupBy_ (Arc q . Just) (prependNN q) empty q + +{- +-- TODO: would it be worth it to define this specialization? The +-- definition is nothing special; but inlining away the first three +-- arguments to 'lookupBy_' does allow to avoid any sort of dynamic +-- dispatch or closures. +lookup :: ByteString -> Trie a -> Maybe a +lookup = start + where + -- | Deal with epsilon query (when there is no epsilon value) + start q t@(Branch{}) | S.null q = Nothing + start q t = go q t + -- | The main recursion + go _ Empty = Nothing + go q (Arc k mv t) = + let (_,k',q') = breakMaximalPrefix k q + in case (S.null k', S.null q') of + (False, _) -> Nothing + (True, True) -> mv + (True, False) -> go q' t + go q t_@(Branch{}) = findArc t_ + where + qh = errorLogHead "lookup" q + -- | \(\mathcal{O}(\min(m,W))\), where \(m\) is number of + -- @Arc@s in this branching, and \(W\) is the word size of + -- the Prefix,Mask type. + findArc Empty = impossible "lookup" -- see [Note1] + findArc t@(Arc{}) = go q t + findArc (Branch p m l r) + | nomatch qh p m = Nothing + | zero qh m = findArc l + | otherwise = findArc r +-} + + +-- TODO: would it be worth it to have a variant like 'lookupBy_' +-- which takes the three continuations? + + +-- According to our "Bench.MatchOne" benchmark, this is in fact +-- much faster than using 'matches_' and relying on list fusion. +-- +-- | Given a query, find the longest prefix with an associated value +-- in the trie, returning the length of that prefix and the associated +-- value. +-- +-- This function may not have the most useful return type. For a +-- version that returns the prefix itself as well as the remaining +-- string, see 'Data.Trie.match'. +-- +-- @since 0.2.4 +match_ :: Trie a -> ByteString -> Maybe (Int, a) +match_ = flip start + where + -- | Deal with epsilon query (when there is no epsilon value) + start q (Branch{}) | S.null q = Nothing + start q t = match1 0 q t + -- TODO: for the non-null Branch case, maybe we should jump directly to 'findArc' (i.e., inline that case of 'match1') + -- | Find the first match, or return Nothing if there isn't one. + match1 _ _ Empty = Nothing + match1 n q (Arc k mv t) = + let (p,k',q') = breakMaximalPrefix k q + !n' = n + S.length p + in case (S.null k', S.null q') of + (False, _) -> Nothing + (True, True) -> (,) n' <$> mv + (True, False) -> + case mv of + Nothing -> match1 n' q' t + Just v -> matchN n' v n' q' t + match1 n q t_@(Branch{}) = findArc t_ + where + qh = errorLogHead "match_" q + -- | \(\mathcal{O}(\min(m,W))\), where \(m\) is number of + -- @Arc@s in this branching, and \(W\) is the word size of + -- the Prefix,Mask type. + findArc Empty = impossible "match_" -- see [Note1] + findArc t@(Arc{}) = match1 n q t + findArc (Branch p m l r) + | nomatch qh p m = Nothing + | zero qh m = findArc l + | otherwise = findArc r + -- | Find the next match, or return the previous one if there are no more. + matchN n0 v0 _ _ Empty = Just (n0,v0) + matchN n0 v0 n q (Arc k mv t) = + let (p,k',q') = breakMaximalPrefix k q + !n' = n + S.length p + in case (S.null k', S.null q') of + (False, _) -> Just (n0,v0) + (True, True) -> + case mv of + Nothing -> Just (n0,v0) + Just v -> Just (n',v) + (True, False) -> + case mv of + Nothing -> matchN n0 v0 n' q' t + Just v -> matchN n' v n' q' t + matchN n0 v0 n q t_@(Branch{}) = findArc t_ + where + qh = errorLogHead "match_" q + -- | \(\mathcal{O}(\min(m,W))\), where \(m\) is number of + -- @Arc@s in this branching, and \(W\) is the word size of + -- the Prefix,Mask type. + findArc Empty = impossible "match_" -- see [Note1] + findArc t@(Arc{}) = matchN n0 v0 n q t + findArc (Branch p m l r) + | nomatch qh p m = Just (n0,v0) + | zero qh m = findArc l + | otherwise = findArc r + + +-- | Given a query, find all prefixes with associated values in the +-- trie, and return the length of each prefix with their value, in +-- order from shortest prefix to longest. This function is a good +-- producer for list fusion. +-- +-- This function may not have the most useful return type. For a +-- version that returns the prefix itself as well as the remaining +-- string, see 'Data.Trie.matches'. +-- +-- @since 0.2.4 +matches_ :: Trie a -> ByteString -> [(Int,a)] +matches_ t q = +#if !defined(__GLASGOW_HASKELL__) + matchFB_ t q (((:) .) . (,)) [] +#else + build (\cons nil -> matchFB_ t q ((cons .) . (,)) nil) +{-# INLINE matches_ #-} +#endif + +matchFB_ :: Trie a -> ByteString -> (Int -> a -> r -> r) -> r -> r +matchFB_ = \t q cons nil -> matchFB_' cons q t nil + where + matchFB_' cons = start + where + -- | Deal with epsilon query (when there is no epsilon value) + start q (Branch{}) | S.null q = id + start q t = go 0 q t + + -- | The main recursion + go _ _ Empty = id + go n q (Arc k mv t) = + let (p,k',q') = breakMaximalPrefix k q + !n' = n + S.length p + in if S.null k' + then + case mv of { Nothing -> id; Just v -> cons n' v} + . + if S.null q' then id else go n' q' t + else id + go n q t_@(Branch{}) = findArc t_ + where + qh = errorLogHead "matches_" q + -- | \(\mathcal{O}(\min(m,W))\), where \(m\) is number + -- of @Arc@s in this branching, and \(W\) is the word + -- size of the Prefix,Mask type. + findArc Empty = impossible "matches_" -- see [Note1] + findArc t@(Arc{}) = go n q t + findArc (Branch p m l r) + | nomatch qh p m = id + | zero qh m = findArc l + | otherwise = findArc r + + +{----------------------------------------------------------- +-- Simple modification functions (recurse and clone spine) +-----------------------------------------------------------} + +-- TODO: We should CPS on Empty to avoid cloning spine if no change. +-- Difficulties arise with the calls to 'branch' and 'arc'. Will +-- have to create a continuation chain, so no savings on memory +-- allocation; but would have savings on held memory, if they're +-- still holding the old one... +-- +-- | Generic function to alter a trie by one element with a function +-- to resolve conflicts (or non-conflicts). +alterBy :: (ByteString -> a -> Maybe a -> Maybe a) + -> ByteString -> a -> Trie a -> Trie a +alterBy f q x = alterBy_ (\mv t -> (f q x mv, t)) q +-- TODO: use GHC's 'inline' function so that this gets specialized away. +-- TODO: benchmark to be sure that this doesn't introduce unforseen +-- performance costs because of the uncurrying etc. +-- TODO: move to "Data.Trie" itself instead of here, since it doesn't +-- depend on any internals (unless we actually do the CPS optimization). +-- TODO: would there be any benefit in basing this off a different +-- function that captures the invariant that the subtrie is left +-- alone? + + +-- Not susceptible to [bug26] because it can only delete a single value\/subtrie. +-- +-- | A variant of 'alterBy' which also allows modifying the sub-trie. +-- If the function returns @(Just v, t)@ and @lookup 'S.empty' t == Just w@, +-- then the @w@ will be overwritten by @v@. +-- +-- @since 0.2.3 +-- __Note__: /Type changed in 0.2.6/ +alterBy_ + :: (Maybe a -> Trie a -> (Maybe a, Trie a)) + -> ByteString -> Trie a -> Trie a +alterBy_ f = start + where + start q t | not (S.null q) = go q t + start _ (Arc k mv s) | S.null k = mayEpsilon $$ f mv s + start _ t = mayEpsilon $$ f Nothing t + + -- @go@ is always called with non-null @q@, therefore @nothing@ is too. + nothing q = arcNN q $$ f Nothing Empty + + go q Empty = nothing q + go q t@(Branch p m l r) + | nomatch qh p m = + case nothing q of + Empty -> t + s -> graft p t qh s + | zero qh m = branchL p m (go q l) r + | otherwise = branchR p m l (go q r) + where qh = errorLogHead "alterBy_" q + go q t@(Arc k mv s) = + let (p,k',q') = breakMaximalPrefix k q in + case (S.null k', S.null q') of + (False, True) -> -- add node to middle of Arc + arcNN p $$ f Nothing (Arc k' mv s) + (False, False) -> + case nothing q' of + Empty -> t -- Nothing to add, reuse old Arc + Branch{} -> impossible "alterBy_" -- 'arcNN' can't Branch + l@(Arc{}) -> + -- Inlined version of @prepend p@, capturing the + -- invariant that the 'graft' must be a @Branch@. + (if S.null p then id else Arc p Nothing) + -- 'arcNN' will always have that the string in @l@ + -- must begin with @q'@, which is non-null here and + -- therefore @arcPrefix q'@ is equivalent to taking + -- the 'arcPrefix' of the string in @l@. + $ graft (arcPrefix q') l (arcPrefix k') (Arc k' mv s) + (True, True) -> arcNN k $$ f mv s + (True, False) -> arcNN k mv (go q' s) + + +-- TODO: benchmark vs the definition with alterBy\/liftM +-- TODO: add a variant that's strict in the function. +-- +-- /Since: 0.2.6/ for being exported from "Data.Trie.Internal" +-- rather than "Data.Trie" +-- +-- | Apply a function to the value at a key. If the key is not +-- present, then the trie is returned unaltered. +adjust :: (a -> a) -> ByteString -> Trie a -> Trie a +adjust f = start + where + start q t | not (S.null q) = go q t + start _ (Arc k (Just v) t) | S.null k = Arc k (Just (f v)) t + start _ t = t + + go _ Empty = Empty + go q t@(Branch p m l r) + | nomatch qh p m = t + | zero qh m = Branch p m (go q l) r + | otherwise = Branch p m l (go q r) + where qh = errorLogHead "adjust" q + go q t@(Arc k mv s) = + let (_,k',q') = breakMaximalPrefix k q in + case (S.null k', S.null q') of + (False, _) -> t + (True, True) -> Arc k (f <$> mv) s + (True, False) -> Arc k mv (go q' s) + + +{----------------------------------------------------------- +-- Trie-combining functions +-----------------------------------------------------------} + +-- TODO: it may be helpful to have a version of 'mergeBy' where the +-- function doesn't return 'Maybe' (i.e., 'Data.Trie.Convenience.unionWith'); +-- because knowing we can't delete elements would allow to use true +-- constructors directly, rather than smart constructors that patch +-- up the deletion cases. Especially since the vast majority of +-- our own uses of 'mergeBy' fall into this category. + +-- Not susceptible to [bug26] because it doesn't delete any values. +-- +-- Alas, benchmarking indicates that this gives only a very trivial +-- benefit over 'TC.unionWith' as implemented via 'mergeBy'. +-- +-- | Take the union of two tries, using a function to resolve +-- conflicts. The resulting trie is constructed strictly, but the +-- results of the combining function are evaluated lazily. +wip_unionWith :: (a -> a -> a) -> Trie a -> Trie a -> Trie a +wip_unionWith f = start + where + -- | Deals with epsilon entries, before recursing into @go@ + -- TODO: for all of these, add assertions that null bytestring entails must be Just; instead of pattern matching on it directly. + start (Arc k0 (Just v0) s0) (Arc k1 (Just v1) s1) | S.null k0 && S.null k1 + = epsilon (f v0 v1) (go s0 s1) + start (Arc k0 (Just v0) s0) t1 | S.null k0 = epsilon v0 (go s0 t1) + start t0 (Arc k1 (Just v1) s1) | S.null k1 = epsilon v1 (go t0 s1) + start t0 t1 = go t0 t1 + + -- | The main recursion + go Empty t1 = t1 + go t0 Empty = t0 + -- \(\mathcal{O}(n+m)\) for this part where \(n\) and \(m\) are + -- sizes of the branchings. + go t0@(Branch p0 m0 l0 r0) + t1@(Branch p1 m1 l1 r1) + | shorter m0 m1 = union0 + | shorter m1 m0 = union1 + | p0 == p1 = Branch p0 m0 (go l0 l1) (go r0 r1) + | otherwise = graft p0 t0 p1 t1 + where + union0 | nomatch p1 p0 m0 = graft p0 t0 p1 t1 + | zero p1 m0 = Branch p0 m0 (go l0 t1) r0 + | otherwise = Branch p0 m0 l0 (go r0 t1) + union1 | nomatch p0 p1 m1 = graft p0 t0 p1 t1 + | zero p0 m1 = Branch p1 m1 (go t0 l1) r1 + | otherwise = Branch p1 m1 l1 (go t0 r1) + -- + go t0@(Arc k0 mv0 s0) + t1@(Arc k1 mv1 s1) + = arcMerge k0 t0 k1 t1 $ \ pre k0' k1' -> + let {-# INLINE t0' #-} + t0' = Arc k0' mv0 s0 + {-# INLINE t1' #-} + t1' = Arc k1' mv1 s1 + in + case (S.null k0', S.null k1') of + (True, True) -> arcNN pre (mergeMaybe (\v0 v1 -> Just (f v0 v1)) mv0 mv1) (go s0 s1) -- TODO: if both arcs are reject, then both @s0,s1@ are branches so we can simplify the 'arcNN' to avoid the case analysis in 'prependNN'. + (True, False) -> Arc pre mv0 (go s0 t1') + (False,True) -> Arc pre mv1 (go t0' s1) + (False,False) -> wye pre k0' t0' k1' t1' + go t0@(Arc k0 _ _) + t1@(Branch p1 m1 l r) + | nomatch p0 p1 m1 = graft p1 t1 p0 t0 + | zero p0 m1 = Branch p1 m1 (go t0 l) r + | otherwise = Branch p1 m1 l (go t0 r) + where p0 = arcPrefix k0 + go t0@(Branch p0 m0 l r) + t1@(Arc k1 _ _) + | nomatch p1 p0 m0 = graft p0 t0 p1 t1 + | zero p1 m0 = Branch p0 m0 (go l t1) r + | otherwise = Branch p0 m0 l (go r t1) + where p1 = arcPrefix k1 + + +-- FIXME: See [bug26]. +-- TEST CASES: foldr (unionL . uncurry singleton) empty t +-- foldr (uncurry insert) empty t +-- where t = map (\s -> (pk s, 0)) +-- ["heat","hello","hoi","apple","appa","hell","appb","appc"] +-- +-- | Take the union of two tries, using a function to resolve collisions. +-- This can only define the space of functions between union and +-- symmetric difference but, with those two, all set operations can +-- be defined (albeit inefficiently). +mergeBy :: (a -> a -> Maybe a) -> Trie a -> Trie a -> Trie a +mergeBy f = start + where + -- | Deals with epsilon entries, before recursing into @go@ + -- TODO: for all of these, add assertions that null bytestring entails must be Just; instead of pattern matching on it directly. + start (Arc k0 (Just v0) s0) (Arc k1 (Just v1) s1) | S.null k0 && S.null k1 + = mayEpsilon (f v0 v1) (go s0 s1) + start (Arc k0 (Just v0) s0) t1 | S.null k0 = epsilon v0 (go s0 t1) + start t0 (Arc k1 (Just v1) s1) | S.null k1 = epsilon v1 (go t0 s1) + start t0 t1 = go t0 t1 + + -- | The main recursion + go Empty t1 = t1 + go t0 Empty = t0 + -- \(\mathcal{O}(n+m)\) for this part where \(n\) and \(m\) are + -- sizes of the branchings. + go t0@(Branch p0 m0 l0 r0) + t1@(Branch p1 m1 l1 r1) + | shorter m0 m1 = union0 + | shorter m1 m0 = union1 + | p0 == p1 = branch p0 m0 (go l0 l1) (go r0 r1) + | otherwise = graft p0 t0 p1 t1 + where + union0 | nomatch p1 p0 m0 = graft p0 t0 p1 t1 + | zero p1 m0 = branchL p0 m0 (go l0 t1) r0 + | otherwise = branchR p0 m0 l0 (go r0 t1) + union1 | nomatch p0 p1 m1 = graft p0 t0 p1 t1 + | zero p0 m1 = branchL p1 m1 (go t0 l1) r1 + | otherwise = branchR p1 m1 l1 (go t0 r1) + -- + go t0@(Arc k0 mv0 s0) + t1@(Arc k1 mv1 s1) + = arcMerge k0 t0 k1 t1 $ \ pre k0' k1' -> + let {-# INLINE t0' #-} + t0' = Arc k0' mv0 s0 + {-# INLINE t1' #-} + t1' = Arc k1' mv1 s1 + in + -- TODO: can be smarter than 'arcNN' here... + case (S.null k0', S.null k1') of + (True, True) -> arcNN pre (mergeMaybe f mv0 mv1) (go s0 s1) + (True, False) -> arcNN pre mv0 (go s0 t1') + (False,True) -> arcNN pre mv1 (go t0' s1) + (False,False) -> wye pre k0' t0' k1' t1' + go t0@(Arc k0 _ _) + t1@(Branch p1 m1 l r) + | nomatch p0 p1 m1 = graft p1 t1 p0 t0 + | zero p0 m1 = branchL p1 m1 (go t0 l) r + | otherwise = branchR p1 m1 l (go t0 r) + where p0 = arcPrefix k0 + go t0@(Branch p0 m0 l r) + t1@(Arc k1 _ _) + | nomatch p1 p0 m0 = graft p0 t0 p1 t1 + | zero p1 m0 = branchL p0 m0 (go l t1) r + | otherwise = branchR p0 m0 l (go r t1) + where p1 = arcPrefix k1 + + +mergeMaybe :: (a -> a -> Maybe a) -> Maybe a -> Maybe a -> Maybe a +{-# INLINE mergeMaybe #-} +mergeMaybe _ Nothing Nothing = Nothing +mergeMaybe _ Nothing mv1@(Just _) = mv1 +mergeMaybe _ mv0@(Just _) Nothing = mv0 +mergeMaybe f (Just v0) (Just v1) = f v0 v1 + + +-- FIXME: See [bug26]. +-- | Take the intersection of two tries, using a function to resolve +-- collisions. +-- +-- @since 0.2.6 +intersectBy :: (a -> b -> Maybe c) -> Trie a -> Trie b -> Trie c +intersectBy f = start + where + -- | Deals with epsilon entries, before recursing into @go@ + start (Arc k0 mv0 s0) (Arc k1 mv1 s1) | S.null k0 && S.null k1 + = mayEpsilon (intersectMaybe f mv0 mv1) (go s0 s1) + start (Arc k0 (Just _) s0) t1 | S.null k0 = go s0 t1 + start t0 (Arc k1 (Just _) s1) | S.null k1 = go t0 s1 + start t0 t1 = go t0 t1 + + -- | The main recursion + go Empty _ = Empty + go _ Empty = Empty + go t0@(Branch p0 m0 l0 r0) + t1@(Branch p1 m1 l1 r1) + | shorter m0 m1 = isect0 + | shorter m1 m0 = isect1 + | p0 == p1 = branch p0 m0 (go l0 l1) (go r0 r1) + | otherwise = Empty + where + isect0 | nomatch p1 p0 m0 = Empty + | zero p1 m0 = go l0 t1 + | otherwise = go r0 t1 + isect1 | nomatch p0 p1 m1 = Empty + | zero p0 m1 = go t0 l1 + | otherwise = go t0 r1 + go (Arc k0 mv0 s0) + (Arc k1 mv1 s1) + -- We can simplify 'getMask' to 'xor' here, avoiding the + -- cost of the 'highestBitMask'; because we don't care about + -- the actual mask itself, just the nonzero-ness. + | xor (arcPrefix k0) (arcPrefix k1) /= 0 = Empty + | otherwise = + let (pre,k0',k1') = breakMaximalPrefix k0 k1 in + if S.null pre + then error "intersectBy: no mask, but no prefix string" + else + let {-# INLINE t0' #-} + t0' = Arc k0' mv0 s0 + {-# INLINE t1' #-} + t1' = Arc k1' mv1 s1 + in + -- TODO: be smarter about the recursion and 'prependNN' + case (S.null k0', S.null k1') of + (True, True) -> arcNN pre (intersectMaybe f mv0 mv1) (go s0 s1) + (True, False) -> prependNN pre (go s0 t1') + (False,True) -> prependNN pre (go t0' s1) + (False,False) -> prependNN pre (go t0' t1') + go t0@(Arc k0 _ _) + (Branch p1 m1 l r) + | nomatch p0 p1 m1 = Empty + | zero p0 m1 = go t0 l + | otherwise = go t0 r + where p0 = arcPrefix k0 + go (Branch p0 m0 l r) + t1@(Arc k1 _ _) + | nomatch p1 p0 m0 = Empty + | zero p1 m0 = go l t1 + | otherwise = go r t1 + where p1 = arcPrefix k1 + + +intersectMaybe :: (a -> b -> Maybe c) -> Maybe a -> Maybe b -> Maybe c +{-# INLINE intersectMaybe #-} +intersectMaybe f (Just v0) (Just v1) = f v0 v1 +intersectMaybe _ _ _ = Nothing + + +-- TODO(github#23): add 'differenceBy' + + +{----------------------------------------------------------- +-- Priority-queue functions +-----------------------------------------------------------} +-- TODO: should verify that all of these are now free of the quadratic +-- slowdown from reconstructing keys. They should be, but just to +-- verify that some new quadratic hasn't accidentally crept in... + +-- TODO: Maybe we should follow 'IntMap' and separate out NOINLINE +-- versions of these functions which construct a specialized View +-- type, vs INLINE wrappers that convert the View into the +-- Maybe/pair/etc. + +-- | Return the lexicographically smallest 'ByteString' and the +-- value it's mapped to; or 'Nothing' for the empty trie. When one +-- entry is a prefix of another, the prefix will be returned. +-- +-- __Note__: Prior to version 0.2.7, this function suffered +-- ; but it no longer does. +-- +-- @since 0.2.2 +minAssoc :: Trie a -> Maybe (ByteString, a) +minAssoc = go Nil + where + go !_ Empty = Nothing + go q (Arc k (Just v) _) = Just (toStrict (q +>? k), v) + go q (Arc k Nothing t) = go (q +>! k) t + go q (Branch _ _ l _) = go q l + + +-- | Return the lexicographically largest 'ByteString' and the +-- value it's mapped to; or 'Nothing' for the empty trie. When one +-- entry is a prefix of another, the longer one will be returned. +-- +-- __Note__: Prior to version 0.2.7, this function suffered +-- ; but it no longer does. +-- +-- @since 0.2.2 +maxAssoc :: Trie a -> Maybe (ByteString, a) +maxAssoc = go Nil + where + go !_ Empty = Nothing + go q (Arc k (Just v) Empty) = Just (toStrict (q +>? k), v) + go q (Arc k (Just _) t) = go (q +>? k) t + go q (Arc k Nothing t) = go (q +>! k) t + go q (Branch _ _ _ r) = go q r + + +mapView :: (Trie a -> Trie a) + -> Maybe (ByteString, a, Trie a) -> Maybe (ByteString, a, Trie a) +{-# INLINE mapView #-} +mapView _ Nothing = Nothing +mapView f (Just (k,v,t)) = Just (k,v, f t) + + +-- Not susceptible to [bug26] because it can only delete a single value. +-- +-- | Update the 'minAssoc' and return the old 'minAssoc'. +-- +-- __Note__: Prior to version 0.2.7, this function suffered +-- ; but it no longer does. +-- +-- @since 0.2.2 +updateMinViewBy :: (ByteString -> a -> Maybe a) + -> Trie a -> Maybe (ByteString, a, Trie a) +updateMinViewBy f = go Nil + where + go !_ Empty = Nothing + go q (Arc k (Just v) t) = let q' = toStrict (q +>? k) + in Just (q',v, arc k (f q' v) t) + go q (Arc k Nothing t) = mapView (prependNN k) (go (q +>! k) t) + go q (Branch p m l r) = mapView (\l' -> branchL p m l' r) (go q l) + + +-- Not susceptible to [bug26] because it can only delete a single value. +-- +-- | Update the 'maxAssoc' and return the old 'maxAssoc'. +-- +-- __Note__: Prior to version 0.2.7, this function suffered +-- ; but it no longer does. +-- +-- @since 0.2.2 +updateMaxViewBy :: (ByteString -> a -> Maybe a) + -> Trie a -> Maybe (ByteString, a, Trie a) +updateMaxViewBy f = go Nil + where + go !_ Empty = Nothing + go q (Arc k (Just v) Empty) = let q' = toStrict (q +>? k) + in Just (q',v, arc k (f q' v) Empty) + go q (Arc k mv@(Just _) t) = mapView (Arc k mv) (go (q +>? k) t) + go q (Arc k Nothing t) = mapView (prepend k) (go (q +>! k) t) + go q (Branch p m l r) = mapView (branchR p m l) (go q r) + +------------------------------------------------------------ +------------------------------------------------------- fin. diff --git a/src/Data/Trie/Internal/BitTwiddle.hs b/src/Data/Trie/Internal/BitTwiddle.hs new file mode 100644 --- /dev/null +++ b/src/Data/Trie/Internal/BitTwiddle.hs @@ -0,0 +1,258 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs -fno-warn-name-shadowing #-} +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE CPP #-} +#ifdef __GLASGOW_HASKELL__ +{-# LANGUAGE MagicHash #-} +#endif +#if __GLASGOW_HASKELL__ >= 701 +-- Naturally, the MagicHash stuff from "GHC.Exts" isn't considered safe. +{-# LANGUAGE Trustworthy #-} +#endif +---------------------------------------------------------------- +-- ~ 2021.12.14 +-- | +-- Module : Data.Trie.Internal.BitTwiddle +-- Copyright : 2012 Clark Gaebel, 2012 Johan Tibel, 2002 Daan Leijen +-- License : BSD-3-Clause +-- Maintainer : libraries@haskell.org, wren@cpan.org +-- Stability : stable +-- Portability : portable (with CPP) +-- +-- Functions to treat 'Word' as a bit-vector for big-endian patricia +-- trees. This code is duplicated from "Data.IntMap" (or +-- "Utils.Containers.Internal.BitUtil" these days). The only +-- differences are that some of the conversion functions are +-- specialized to 'Data.Word.Word8' for bytestrings, instead of +-- being specialized to 'Int'. +---------------------------------------------------------------- + +module Data.Trie.Internal.BitTwiddle + ( + -- * Type aliases + KeyElem, Prefix, Mask + -- * Predicates + , zero, nomatch, shorter + -- * Constructors + , applyMask, getMask + ) where + +import Data.Trie.Internal.ByteString (ByteStringElem) + +-- It's too much noise to fully restrict this import, so just note +-- the requirements: +-- base 4.8.0 / GHC 7.10.1 -- 'countLeadingZeros', 'countTrailingZeros' +-- base 4.7.0 / GHC 7.8.2 -- 'FiniteBits', 'finiteBitSize' +-- base 4.5.0 / GHC 7.4.1 -- 'popCount' +import Data.Bits + +-- To make it clearer what we're really testing for. +-- TODO: make this into a Cabal flag; for easier testing if nothing else. +#define USE_CLZ_IMPLEMENTATION MIN_VERSION_base(4,8,0) + +#if __GLASGOW_HASKELL__ >= 503 +-- Before GHC 5.3 these were in "GlaExts" instead. +import GHC.Exts + ( Word(W#) + , Int(I#) +# if USE_CLZ_IMPLEMENTATION + , shiftL# +# else + , shiftRL# +# endif + ) +#else +import Data.Word (Word) +#endif + +---------------------------------------------------------------- + +-- | 'KeyElem' is what we actually use for 'Prefix' and 'Mask'. +-- For now we're using 'ByteStringElem' ('Data.Word.Word8') for +-- simplicity, but in the future we might switch to a larger word +-- size. +type KeyElem = ByteStringElem + +-- | Some prefix of the 'KeyElem', as constructed by 'applyMask'. +type Prefix = KeyElem + +-- | A single bit, signifying a mask (of all the bits preceding the +-- masking bit). +type Mask = KeyElem + +elemToNat :: KeyElem -> Word +{-# INLINE elemToNat #-} +elemToNat = fromIntegral + +natToElem :: Word -> KeyElem +{-# INLINE natToElem #-} +natToElem = fromIntegral + +-- TODO: newer versions of the containers library just use +-- 'unsafeShift{R,L}' unilaterally. So, what is the difference +-- (i.e., these days) between using the 'uncheckedShift{L,RL}#' of +-- 'unsafeShift{L,R}' vs using the 'shift{L,RL}#' of 'shift{L,R}'? +-- Also, do we no longer need to trick GHC into actually unboxing +-- and inlining these? +#if USE_CLZ_IMPLEMENTATION +shiftLL :: Word -> Int -> Word +{-# INLINE shiftLL #-} +# if __GLASGOW_HASKELL__ +-- Use unboxing to get @shiftLL@ inlined. +shiftLL (W# x) (I# i) = W# (shiftL# x i) +# else +shiftLL x i = unsafeShiftL x i +# endif +#else +shiftRL :: Word -> Int -> Word +{-# INLINE shiftRL #-} +# if __GLASGOW_HASKELL__ +-- Use unboxing to get @shiftRL@ inlined. +shiftRL (W# x) (I# i) = W# (shiftRL# x i) +# else +shiftRL x i = unsafeShiftR x i +# endif +#endif + + +{--------------------------------------------------------------- +-- Endian independent bit twiddling (Trie endianness, not architecture) +---------------------------------------------------------------} + +-- TODO: should we use the (Bits Word8) instance instead of 'elemToNat' +-- and (Bits Nat)? We need to compare Core, C--, or ASM in order +-- to decide this. The choice will apply to 'zero', 'applyMask', 'maskW',... +-- If we shouldn't, then we should probably send a patch upstream +-- to fix the (Bits Word8) instance. + +-- | Is the key zero under the masking bit? If true then whatever +-- is associated with that key should go to the left, otherwise it +-- should go to the right. +zero :: KeyElem -> Mask -> Bool +{-# INLINE zero #-} +zero i m = (elemToNat i) .&. (elemToNat m) == 0 + +-- | Does the masked key /not/ match the prefix? (Hence a subtree +-- matching the value doesn't exist.) +nomatch :: KeyElem -> Prefix -> Mask -> Bool +{-# INLINE nomatch #-} +nomatch i p m = applyMask i m /= p + +-- | Convert a masking bit to the full mask it represents, and then +-- return the prefix of the key under that mask (i.e., all the bits +-- preceding the masking bit). +applyMask :: KeyElem -> Mask -> Prefix +{-# INLINE applyMask #-} +applyMask i m = maskW (elemToNat i) (elemToNat m) + + +{--------------------------------------------------------------- +-- Big endian operations (Trie endianness, not architecture) +---------------------------------------------------------------} + +-- | Get mask by setting all bits higher than the smallest bit in +-- @m@. Then apply that mask to @i@. +maskW :: Word -> Word -> Prefix +{-# INLINE maskW #-} +maskW i m = natToElem (i .&. (complement (m-1) `xor` m)) +-- TODO: try the alternatives mentioned in the Containers paper: +-- \i m -> natToElem (i .&. (negate m - m)) +-- \i m -> natToElem (i .&. (m * complement 1)) +-- N.B. these return /all/ the low bits, and therefore they are not +-- equal functions for all m. They are, however, equal when only +-- one bit of m is set. + +-- | Determine whether the first mask denotes a shorter prefix than +-- the second. +shorter :: Mask -> Mask -> Bool +{-# INLINE shorter #-} +shorter m1 m2 = elemToNat m1 > elemToNat m2 + +-- | Determine first differing bit of two prefixes. +getMask :: Prefix -> Prefix -> Mask +{-# INLINE getMask #-} +getMask p1 p2 = natToElem (highestBitMask (elemToNat p1 `xor` elemToNat p2)) + +{--------------------------------------------------------------- + Finding the highest bit (mask) in a word [x] can be done efficiently + in three ways: + * convert to a floating point value and the mantissa tells us the + [log2(x)] that corresponds with the highest bit position. The + mantissa is retrieved either via the standard C function [frexp] + or by some bit twiddling on IEEE compatible numbers (float). + Note that one needs to use at least [double] precision for an + accurate mantissa of 32 bit numbers. + * use bit twiddling, a logarithmic sequence of bitwise or's and + shifts (bit). + * use processor specific assembler instruction (asm). + + The most portable way would be [bit], but is it efficient enough? + I have measured the cycle counts of the different methods on an + AMD Athlon-XP 1800 (~ Pentium III 1.8Ghz) using the RDTSC + instruction: + + highestBitMask: method cycles + -------------- + frexp 200 + float 33 + bit 11 + asm 12 + + highestBit: method cycles + -------------- + frexp 195 + float 33 + bit 11 + asm 11 + + Wow, the bit twiddling is on today's RISC like machines even + faster than a single CISC instruction (BSR)! +---------------------------------------------------------------} + +{--------------------------------------------------------------- + [highestBitMask] returns a word where only the highest bit is + set. It is found by first setting all bits in lower positions + than the highest bit and than taking an exclusive or with the + original value. Allthough the function may look expensive, GHC + compiles this into excellent C code that subsequently compiled + into highly efficient machine code. The algorithm is derived from + Jorg Arndt's FXT library. +---------------------------------------------------------------} + +highestBitMask :: Word -> Word +{-# INLINE highestBitMask #-} +#if USE_CLZ_IMPLEMENTATION +-- This is the implementation used in newer versions of the containers library. +-- Added this implementation here in version 0.2.7. +highestBitMask w = shiftLL 1 (wordSize - 1 - countLeadingZeros w) +#else +-- This is the classic one we used up to bytestring-trie-0.2.6.1 +-- And it's still what containers falls back to for older versions of base. +-- +-- N.B., because this is not exported and is only used by 'branchMask' +-- which operates on 'Word8' inputs, we can safely restrict the +-- algorithm to only doing the first few steps, rather than doing +-- all the steps needed for 'Word64'. +highestBitMask x + = case (x .|. shiftRL x 1) of + x -> case (x .|. shiftRL x 2) of + x -> case (x .|. shiftRL x 4) of -- for 8-bit input range. + {- + x -> case (x .|. shiftRL x 8) of -- for 16-bit + x -> case (x .|. shiftRL x 16) of -- for 32-bit + x -> case (x .|. shiftRL x 32) of -- for 64-bit platforms + -} + x -> (x `xor` shiftRL x 1) +#endif + +#if USE_CLZ_IMPLEMENTATION +wordSize :: Int +{-# INLINE wordSize #-} +# if MIN_VERSION_base(4,7,0) +wordSize = finiteBitSize (0 :: Word) +# else +wordSize = bitSize (0 :: Word) +# endif +#endif + +---------------------------------------------------------------- +----------------------------------------------------------- fin. diff --git a/src/Data/Trie/Internal/ByteString.hs b/src/Data/Trie/Internal/ByteString.hs new file mode 100644 --- /dev/null +++ b/src/Data/Trie/Internal/ByteString.hs @@ -0,0 +1,346 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE CPP, BangPatterns #-} +#if __GLASGOW_HASKELL__ >= 701 +-- Neither 'unsafeDupablePerformIO' nor 'Data.ByteString.Internal' is safe. +{-# LANGUAGE Trustworthy #-} +#endif +------------------------------------------------------------ +-- ~ 2025-02-11 +-- | +-- Module : Data.Trie.Internal.ByteString +-- Copyright : 2008--2025 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : stable +-- Portability : GHC-only +-- +-- Helper functions on 'ByteString's for "Data.Trie.Internal". +------------------------------------------------------------ + +module Data.Trie.Internal.ByteString + ( ByteString, ByteStringElem + , breakMaximalPrefix + , RevLazyByteString(..), (+>!), (+>?), fromStrict, toStrict + -- TODO: we want to export the 'Nil' constructor; but + -- do we want to export the patterns? + ) where + +import qualified Data.ByteString as S +import qualified Data.ByteString.Internal as S +#if MIN_VERSION_bytestring(0,11,5) +import qualified Foreign.Marshal.Utils as FMU +#endif +import Data.ByteString.Internal (ByteString(PS)) +import Data.Word +import Foreign.ForeignPtr (ForeignPtr) +#if MIN_VERSION_base(4,15,0) +-- [aka GHC 9.0.1]: +import GHC.ForeignPtr (unsafeWithForeignPtr) +#else +import Foreign.ForeignPtr (withForeignPtr) +#endif +import Foreign.Ptr (Ptr, plusPtr) +import Foreign.Storable (Storable(..)) +-- This module name is since @__GLASGOW_HASKELL__ >= 611@. +import GHC.IO (unsafeDupablePerformIO) + +------------------------------------------------------------ +#if !(MIN_VERSION_base(4,15,0)) +-- bytestring-0.10.12.1 and 0.11.1.0 use and export this definition; +-- however neither 0.10.12.0 nor 0.11.0.0 define nor use it.So, +-- rather than dealing with all that nonsense, we'll just do it +-- ourselves. +unsafeWithForeignPtr :: ForeignPtr a -> (Ptr a -> IO b) -> IO b +unsafeWithForeignPtr = withForeignPtr +#endif + +memcpy :: Ptr Word8 -> Ptr Word8 -> Int -> IO () +#if MIN_VERSION_bytestring(0,11,5) +-- This is the version when 'S.memcpy' got deprecated, though it +-- still exists at least through 0.12.2.0. +memcpy = FMU.copyBytes +#else +memcpy = S.memcpy +#endif + +------------------------------------------------------------ +------------------------------------------------------------ +-- | Associated type of 'ByteString' +type ByteStringElem = Word8 + + +------------------------------------------------------------ +-- The @since annotation is for when this got re-exported from +-- "Data.Trie.Internal". +-- +-- | Returns the longest shared prefix and the two remaining suffixes +-- for a pair of strings. This function performs no allocation\/copying, +-- it simply returns slices\/views of the arguments. +-- +-- * @s ≡ (\\(pre,s',z') -> pre '<>' s') ('breakMaximalPrefix' s z)@ +-- * @z ≡ (\\(pre,s',z') -> pre '<>' z') ('breakMaximalPrefix' s z)@ +-- +-- @since 0.2.2 +breakMaximalPrefix + :: ByteString + -> ByteString + -> (ByteString, ByteString, ByteString) +-- +-- [Implementation Notes] +-- +-- * We've had to define 'strictTriple' and use BangPatterns to +-- keep GHC from wrapping all the returned triples in +-- ghc-prim:'GHC.Magic.lazy'. Not sure how much this actually +-- helps performance, but it's a stepping stone towards defining +-- a custom result type which unpacks the three ByteStrings. And +-- given that GHC's worker-wrapper transform generates a worker +-- that returns an unboxed tuple and yet internally does construct +-- the tuple, this suggests that using a custom return type should +-- help performance. +-- +-- * TODO: the result of the inlined 'indexOfDifference' is still +-- being wrapped in ghc-prim:'GHC.Magic.lazy'; but nothing I can +-- do seems to change that. Is it something about the +-- 'unsafeDupablePerformIO' or what? Would changing it even help +-- performance? +-- +-- * The first two cases can safely be allowed to fall through to +-- the @i <= 0@ case. After inlining, there shouldn't be any +-- function-call overhead for letting 'goByte' do the comparison +-- instead. The only difference is that the @i <= 0@ case will +-- hold onto @s0@/@s1@ rather than replacing them by 'S.empty'. +-- Unfortunately, that difference in liveness seems to result in +-- slightly worse performance. +-- TODO: a better benchmark than just running the test suite. +-- +-- * The 'unsafeWithForeignPtr' allows for more aggressive optimization +-- than 'withForeignPtr', since it encodes the knowledge that the +-- continuation cannot diverge (loop, or throw exceptions). In +-- particular, without this, the call to 'min' will get hoisted +-- above the inner 'withForeignPtr' and the call to 'indexOfDifference' +-- will be duplicated in both branches of the 'min'; and since +-- 'indexOfDifference' will get inlined (recursive 'goBytes' and +-- all), that's a lot of code duplication. However, for whatever +-- reason the 'unsafeWithForeignPtr' version actually seems to +-- result in slightly worse performance (0.2~2% on the test suite). +-- TODO: a better benchmark than just running the test suite. +-- TODO: if that hoisting actually does help, then perhaps manually +-- lift the 'max' above both 'withForeignPtr' and manually +-- express the branch duplication. +-- TODO: Also consider whether this might be relevant: +-- +-- +-- * TODO: should we yield to the accursed call of +-- 'Data.ByteString.Internal.accursedUnutterablePerformIO'? +-- Recent versions of bytestring export it, so we wouldn't +-- even need to copy the accursed incantation itself. Regarding +-- correctness, probably the closest thing to compare against +-- are these bugs against 'S.elemIndices': +-- +-- +-- +-- * TODO: re-investigate performance of lifting the non-IO stuff +-- out of the scope of the 'unsafeDupablePerformIO', vs leaving +-- it within that scope. +-- +breakMaximalPrefix + s0@(PS fp0 off0 len0) + s1@(PS fp1 off1 len1) + | len0 <= 0 = strictTriple S.empty S.empty s1 + | len1 <= 0 = strictTriple S.empty s0 S.empty + | otherwise = + let i = unsafeDupablePerformIO $ + unsafeWithForeignPtr fp0 $ \p0 -> + unsafeWithForeignPtr fp1 $ \p1 -> + indexOfDifference + (p0 `ptrElemOff` off0) + (p1 `ptrElemOff` off1) + (len0 `min` len1) + in if i <= 0 -- can only be equal, but for paranoia's sake. + then strictTriple S.empty s0 s1 + else strictTriple + (if off0 + len0 < off1 + len1 -- share the smaller one + then PS fp0 off0 i -- TODO: assert(i<=len0) for paranoia? + else PS fp1 off1 i) -- TODO: assert(i<=len1) for paranoia? + (dropPS i fp0 off0 len0) + (dropPS i fp1 off1 len1) + +-- | Construct a triple, strict in all arguments. This helps improve +-- code generation over our previous approach. Making our own +-- datatype for this result or CPSing 'breakMaximalPrefix' may still +-- improve things further. +strictTriple :: ByteString -> ByteString -> ByteString + -> (ByteString, ByteString, ByteString) +strictTriple !p !s !z = (p,s,z) +{-# INLINE strictTriple #-} + +-- | Get the 'sizeOf' type @a@, without requiring @-XScopedTypeVariables@ +-- nor making a spurious call to 'System.IO.Unsafe.unsafePerformIO' or similar. +sizeOfElem :: Storable a => Ptr a -> Int +sizeOfElem = sizeOf . (undefined :: Ptr a -> a) +{-# INLINE sizeOfElem #-} + +-- | C-style pointer addition, without the excessively liberal type +-- of 'plusPtr'. +ptrElemOff :: Storable a => Ptr a -> Int -> Ptr a +ptrElemOff p i = p `plusPtr` (i * sizeOfElem p) +{-# INLINE [0] ptrElemOff #-} +-- This rewrite rule helps ensure that on bytestring>=0.11 we don't +-- incur any additional cost for using the 'PS' pattern synonym. +{-# RULES +"Data.Trie.ByteStringInternal ptrElemOff/0" + forall p . ptrElemOff p 0 = p + #-} + +-- For bytestring>=0.11, there's no way to improve over the 'PS' +-- constructor synonym here. After inlining, the @off=0@ from the +-- 'PS' pattern synonym will constant-propogate away, so all we'll +-- be left with is @BS (plusForeignPtr fp n) (len - n)@; which is +-- the same thing we would've written by hand. Plus, bytestring>=0.11 +-- will already define the compatibility definition of 'plusForeignPtr' +-- for use with base<4.10. +-- +-- | Unpacked version of 'S.drop', for use as a smart-constructor. +-- N.B., this assumes the @n <= 0@ case has already been handled +-- (otherwise you might as well just say @drop n (PS fp off len)@ +-- and let the compiler remove the intermediate 'PS'). +dropPS :: Int -> ForeignPtr ByteStringElem -> Int -> Int -> ByteString +dropPS !n !fp !off !len + | n >= len = S.empty + | otherwise = PS fp (off + n) (len - n) +{-# INLINE dropPS #-} + + +------------------------------------------------------------ +-- This naive algorithm doesn't depend on architecture details. We +-- could speed things up (in theory) by checking a natural word at +-- a time and then falling back to checking each byte once the +-- mismatched word is found. But in practice that doesn't seem to +-- actually speed things up. +-- +-- TODO: that's probably because of alignment issues, or because +-- we should really vectorize by the largest single load on an +-- architecture rather than by the natural word size. For more +-- details on how to do it right, see GNU glibc's implementation +-- of @memcmp@. We should be able to do a simple twist on that +-- algorithm to return the index of difference rather than the +-- ordering. That would mean requiring GPL, but unfortunately every +-- other implementations of @memcmp@ I've found (FreeBSD libc, GCC's +-- builtin,...) just uses the same naive algorithm I have below. +-- I suppose we could always fork that algorithm off into a separate +-- optional dependency of this library; where we fallback to this +-- implementation if the user doesn't want the GPL burden. +-- +-- | Calculates the first index where values differ. +indexOfDifference + :: Ptr ByteStringElem + -> Ptr ByteStringElem + -> Int + -> IO Int +indexOfDifference !p1 !p2 !limit = goByte 0 + where + goByte n + | n >= limit = return limit + | otherwise = do + c1 <- peekElemOff p1 n + c2 <- peekElemOff p2 n + if c1 == c2 + then goByte (n+1) + else return n + +-- TODO: why does bytestring-0.11 use 'peekByteOff' in lieu of +-- 'peekElemOff'? Given the definitions, the latter is more +-- direct/simpler: using @readWord8OffAddr# p# n# s@ instead of +-- @readWord8OffAddr# (plusAddr# p# n# ) 0# s@, though surely GHC +-- will optimize those to generate the same assembly. + +------------------------------------------------------------ +------------------------------------------------------------ + +-- | A \"reversed\" variant of lazy bytestrings; i.e., a snoc-list +-- of strict bytestrings. +data RevLazyByteString + = RevLazyByteString :+> {-# UNPACK #-} !S.ByteString + -- Invariant: every 'S.ByteString' is non-null. + | Nil + +-- TODO: should we add an 'assert' even though we don't check in general? +-- | \(\mathcal{O}(1)\). Unsafely\/uncheckedly append a BS to the +-- RLBS. It is up to the caller to maintain the invariant that +-- 'S.ByteString' is indeed non-null. +(+>!) :: RevLazyByteString -> S.ByteString -> RevLazyByteString +xs +>! x = xs :+> x +{-# INLINE (+>!) #-} + +-- | \(\mathcal{O}(1)\). Safely append a BS to the RLBS, maintaining +-- the invariant. +(+>?) :: RevLazyByteString -> S.ByteString -> RevLazyByteString +xs +>? PS _ _ 0 = xs +xs +>? x = xs :+> x +{-# INLINE (+>?) #-} + +-- | \(\mathcal{O}(1)\). Safely convert a strict BS to RLBS, +-- maintaining the invariant. +fromStrict :: S.ByteString -> RevLazyByteString +fromStrict = (Nil +>?) +{-# INLINE fromStrict #-} + +-- HACK: bytestring-0.10.8.1 (GHC 8.0.2) used 'S.checkedSum' (and +-- a simpler algorithm), whereas bytestring-0.10.8.2 (GHC 8.2.1) +-- introduced 'S.checkedAdd' instead; alas, those version numbers +-- cannot be differentiated by the MIN_VERSION_bytestring macro. +-- Thus, we'll simply define it ourselves. +-- TODO: since we built the trie from bytestrings that were short +-- enough to have a valid length, do we actually need to perform +-- this check at all? +-- +-- | Add two non-negative numbers. Errors out on overflow. +(+?) :: Int -> Int -> Int +x +? y + | r >= 0 = r + | otherwise = error overflowError + where r = x + y +{-# INLINE (+?) #-} + +overflowError :: String +overflowError = "Data.Trie.ByteStringInternal.toStrict: size overflow" +{-# NOINLINE overflowError #-} + +-- See commentary at LazyByteString's version of @toStrict@. This +-- implementation is from Git SHA 688f3c0887f2ca0623f2f54f78e8f675f92e31bf, +-- modulo the necessary changes for using a snoc-list in lieu of a +-- cons-list. +-- | \(\mathcal{O}(n)\). Convert the RLBS to a strict BS, by copying it. +toStrict :: RevLazyByteString -> S.ByteString +toStrict = \cs0 -> goLen0 cs0 cs0 + where + -- It's still possible that the result is empty. + goLen0 _ Nil = S.empty + goLen0 cs0 (cs :+> PS _ _ 0) = goLen0 cs0 cs + goLen0 cs0 (cs :+> c) = goLen1 cs0 c cs + -- It's still possible that the result is a single chunk. + goLen1 _ b Nil = b + goLen1 cs0 b (cs :+> PS _ _ 0) = goLen1 cs0 b cs + goLen1 cs0 (PS _ _ bl) (cs :+> PS _ _ cl) = goLen cs0 (bl +? cl) cs + -- General case, just find the total length we'll need. + goLen cs0 !total (cs :+> PS _ _ cl) = goLen cs0 (total +? cl) cs + goLen cs0 total Nil = + S.unsafeCreate total $ \ptr -> + -- TODO: this gives the correct behavior (re off-by-one + -- concerns); however, it is bad praxis to use a pointer + -- to something outside the allocated region; even if + -- it is just pointing to the first invalid byte after + -- the allocated region. + goCopy cs0 (ptr `ptrElemOff` total) + -- Copy the data + goCopy Nil !_ = return () + goCopy (cs :+> PS _ _ 0 ) !ptr = goCopy cs ptr + goCopy (cs :+> PS fp off len) !ptr = + unsafeWithForeignPtr fp $ \p -> do + let ptr' = ptr `ptrElemOff` negate len + memcpy ptr' (p `ptrElemOff` off) len + goCopy cs ptr' + +------------------------------------------------------------ +------------------------------------------------------- fin. diff --git a/src/Data/Trie/Internal/Errors.hs b/src/Data/Trie/Internal/Errors.hs new file mode 100644 --- /dev/null +++ b/src/Data/Trie/Internal/Errors.hs @@ -0,0 +1,50 @@ +{-# OPTIONS_GHC -Wall -fwarn-tabs #-} +{-# OPTIONS_HADDOCK hide #-} +{-# LANGUAGE CPP #-} +#if __GLASGOW_HASKELL__ >= 701 +{-# LANGUAGE Safe #-} +#endif +---------------------------------------------------------------- +-- ~ 2021.12.14 +-- | +-- Module : Data.Trie.Internal.Errors +-- Copyright : 2008--2025 wren romano +-- License : BSD-3-Clause +-- Maintainer : wren@cpan.org +-- Stability : internal +-- Portability : portable +-- +-- Internal convenience functions for giving error messages. +---------------------------------------------------------------- + +module Data.Trie.Internal.Errors (impossible) where + +---------------------------------------------------------------- +---------------------------------------------------------------- + +-- | The impossible happened. Use this instead of 'undefined' +-- whenever there's an unreachable case, an argument that shouldn't +-- ever get touched, etc. +impossible :: String -> a +impossible fn = error (formatMessage fn) +{-# INLINE impossible #-} +-- Inline the 'error' call itself, just not the string literals in the message. + +formatMessage :: String -> String +formatMessage fn + = "Data.Trie." ++ fn ++ ": The impossible happened." + ++ "\nThis is a bug, please report it to the maintainer.\n" +{-# NOINLINE formatMessage #-} + +-- N.B., at some point GHC adjusted 'error' to throw a +-- 'GHC.Exception.ErrorCall' which contains both the original message +-- and the location info. So we shouldn't have to resort to tricks +-- like *loch* or *placeholders* use in order to get the exact +-- location of the errors. +-- +-- For older versions of GHC, see this post for how +-- 'Control.Exception.assert' gets turned into +-- 'GHC.IO.Exception.assertError': + +---------------------------------------------------------------- +----------------------------------------------------------- fin.