diff --git a/LICENSE b/LICENSE --- a/LICENSE +++ b/LICENSE @@ -1,27 +1,30 @@ -Copyright 2012 Tony Morris +Copyright Andrew Martin, chessai (c) 2018 All rights reserved. Redistribution and use in source and binary forms, with or without -modification, are permitted provided that the following conditions -are met: -1. Redistributions of source code must retain the above copyright - notice, this list of conditions and the following disclaimer. -2. Redistributions in binary form must reproduce the above copyright - notice, this list of conditions and the following disclaimer in the - documentation and/or other materials provided with the distribution. -3. Neither the name of the author nor the names of his contributors - may be used to endorse or promote products derived from this software - without specific prior written permission. +modification, are permitted provided that the following conditions are met: -THIS SOFTWARE IS PROVIDED BY THE REGENTS AND CONTRIBUTORS ``AS IS'' AND -ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE -IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE -ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHORS OR CONTRIBUTORS BE LIABLE -FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL -DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS -OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) -HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT -LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY -OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF -SUCH DAMAGE. + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + * Neither the name of Andrew Martin nor the names of other + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +"AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR +A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT +OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, +SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT +LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, +DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY +THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/Setup.hs b/Setup.hs --- a/Setup.hs +++ b/Setup.hs @@ -1,3 +1,2 @@ import Distribution.Simple main = defaultMain - diff --git a/diet.cabal b/diet.cabal --- a/diet.cabal +++ b/diet.cabal @@ -1,38 +1,63 @@ -Name: diet -Version: 0.0.1 -License: BSD3 -License-File: LICENSE -Author: Tony Morris <ʇǝu˙sıɹɹoɯʇ@ןןǝʞsɐɥ> -Maintainer: Tony Morris -Copyright: Tony Morris -Synopsis: Discrete Interval Encoding Tree -Category: Development -Description: Discrete Interval Encoding Tree described by Martin Erwig in /Diets for Fat Sets, January 1993/. -Homepage: https://github.com/tonymorris/diet -Cabal-Version: >= 1.6 -Build-Type: Simple - -Source-Repository head - Type: git - Location: git@github.com:tonymorris/diet.git - -Flag small_base - Description: Choose the new, split-up base package. - -Library - Build-Depends: - base < 5 && >= 3 - - GHC-Options: - -Wall - - Hs-Source-Dirs: - src +name: + diet +version: + 0.1.0.0 +synopsis: + Discrete Interval Encoding Trees +description: + Discrete Interval Encoding Tree described by Martin Erwig in /Diets for Fat Sets, January 1993/. +homepage: + https://github.com/andrewthad/primitive-containers/ +bug-reports: + https://github.com/andrewthad/primitive-containers/issues +author: + Andrew Martin, chessai +maintainer: + andrew.thaddeus@gmail.com, chessai1996@gmail.com +copyright: + 2018 Andrew Martin, 2018 chessai +license: + BSD3 +license-file: + LICENSE +build-type: + Simple +cabal-version: 2.0 - Exposed-Modules: - Data.Set.Diet +source-repository head + type: git + location: https://github.com/andrewthad/primitive-containers - GHC-Options: - -Wall - -funbox-strict-fields +library + hs-source-dirs: + src + build-depends: + base >=4.9 && <5 + , primitive >= 0.6.4 + , contiguous + exposed-modules: + Data.Diet.Set.Lifted + Data.Diet.Set.Unboxed + Data.Diet.Map.Lifted.Lifted + Data.Diet.Map.Unboxed.Lifted + other-modules: + Data.Diet.Map.Internal + Data.Diet.Set.Internal + ghc-options: -Wall -O2 + default-language: Haskell2010 +test-suite test + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + build-depends: + base + , QuickCheck + , containers + , diet + , primitive + , quickcheck-classes >= 0.4.11.1 + , tasty + , tasty-quickcheck + ghc-options: -Wall -O2 + default-language: Haskell2010 diff --git a/src/Data/Diet/Map/Internal.hs b/src/Data/Diet/Map/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Map/Internal.hs @@ -0,0 +1,401 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE MagicHash #-} + +{-# OPTIONS_GHC -O2 -Wall #-} +module Data.Diet.Map.Internal + ( Map + , empty + , singleton + , map + , append + , lookup + , concat + , equals + , showsPrec + , liftShowsPrec2 + -- list conversion + , fromListN + , fromList + , fromListAppend + , fromListAppendN + , toList + ) where + +import Prelude hiding (lookup,showsPrec,concat,map) + +import Control.Applicative (liftA2) +import Control.Monad.ST (ST,runST) +import Data.Semigroup (Semigroup) +import Data.Foldable (foldl') +import Text.Show (showListWith) +import Data.Primitive.Contiguous (Contiguous,Element,Mutable) +import qualified Data.List as L +import qualified Data.Semigroup as SG +import qualified Prelude as P +import qualified Data.Primitive.Contiguous as I + +-- The key array is twice as long as the value array since +-- everything is stored as a range. Also, figure out how to +-- unpack these two arguments at some point. +data Map karr varr k v = Map !(karr k) !(varr v) + +empty :: (Contiguous karr, Contiguous varr) => Map karr varr k v +empty = Map I.empty I.empty + +map :: (Contiguous karr, Element karr k, Contiguous varr, Element varr v, Element varr w) => (v -> w) -> Map karr varr k v -> Map karr varr k w +map f (Map k v) = Map k (I.map f v) + +equals :: (Contiguous karr, Element karr k, Eq k, Contiguous varr, Element varr v, Eq v) => Map karr varr k v -> Map karr varr k v -> Bool +equals (Map k1 v1) (Map k2 v2) = I.equals k1 k2 && I.equals v1 v2 + +fromListN :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) => Int -> [(k,k,v)] -> Map karr varr k v +fromListN = fromListWithN (\_ a -> a) + +fromList :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) => [(k,k,v)] -> Map karr varr k v +fromList = fromListN 1 + +fromListAppendN :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Semigroup v, Eq v) => Int -> [(k,k,v)] -> Map karr varr k v +fromListAppendN = fromListWithN (SG.<>) + +fromListAppend :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Semigroup v, Eq v) => [(k,k,v)] -> Map karr varr k v +fromListAppend = fromListAppendN 1 + +fromListWithN :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) => (v -> v -> v) -> Int -> [(k,k,v)] -> Map karr varr k v +fromListWithN combine _ xs = + concatWith combine (P.map (\(lo,hi,v) -> singleton lo hi v) xs) + +concat :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Semigroup v, Eq v) => [Map karr varr k v] -> Map karr varr k v +concat = concatWith (SG.<>) + +singleton :: forall karr varr k v. (Contiguous karr, Element karr k,Ord k,Contiguous varr, Element varr v) => k -> k -> v -> Map karr varr k v +singleton !lo !hi !v = if lo <= hi + then Map + ( runST $ do + !(arr :: Mutable karr s k) <- I.new 2 + I.write arr 0 lo + I.write arr 1 hi + I.unsafeFreeze arr + ) + ( runST $ do + !(arr :: Mutable varr s v) <- I.new 1 + I.write arr 0 v + I.unsafeFreeze arr + ) + else empty + +lookup :: forall karr varr k v. (Contiguous karr, Element karr k, Ord k, Contiguous varr, Element varr v) => k -> Map karr varr k v -> Maybe v +lookup a (Map keys vals) = go 0 (I.size vals - 1) where + go :: Int -> Int -> Maybe v + go !start !end = if end <= start + then if end == start + then + let !valLo = I.index keys (2 * start) + !valHi = I.index keys (2 * start + 1) + in if a >= valLo && a <= valHi + then case I.index# vals start of + (# v #) -> Just v + else Nothing + else Nothing + else + let !mid = div (end + start + 1) 2 + !valLo = I.index keys (2 * mid) + in case P.compare a valLo of + LT -> go start (mid - 1) + EQ -> case I.index# vals mid of + (# v #) -> Just v + GT -> go mid end +{-# INLINEABLE lookup #-} + + +append :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Semigroup v, Eq v) => Map karr varr k v -> Map karr varr k v -> Map karr varr k v +append (Map ksA vsA) (Map ksB vsB) = + case unionArrWith (SG.<>) ksA vsA ksB vsB of + (k,v) -> Map k v + +appendWith :: (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) => (v -> v -> v) -> Map karr varr k v -> Map karr varr k v -> Map karr varr k v +appendWith combine (Map ksA vsA) (Map ksB vsB) = + case unionArrWith combine ksA vsA ksB vsB of + (k,v) -> Map k v + + +unionArrWith :: forall karr varr k v. (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) + => (v -> v -> v) + -> karr k -- keys a + -> varr v -- values a + -> karr k -- keys b + -> varr v -- values b + -> (karr k, varr v) +unionArrWith combine keysA valsA keysB valsB + | I.size valsA < 1 = (keysB,valsB) + | I.size valsB < 1 = (keysA,valsA) + | otherwise = runST action + where + action :: forall s. ST s (karr k, varr v) + action = do + let !szA = I.size valsA + !szB = I.size valsB + !(keysDst :: Mutable karr s k) <- I.new (max szA szB * 8) + !(valsDst :: Mutable varr s v) <- I.new (max szA szB * 4) + let writeKeyRange :: Int -> k -> k -> ST s () + writeKeyRange !ix !lo !hi = do + I.write keysDst (2 * ix) lo + I.write keysDst (2 * ix + 1) hi + writeDstHiKey :: Int -> k -> ST s () + writeDstHiKey !ix !hi = I.write keysDst (2 * ix + 1) hi + writeDstValue :: Int -> v -> ST s () + writeDstValue !ix !v = I.write valsDst ix v + readDstHiKey :: Int -> ST s k + readDstHiKey !ix = I.read keysDst (2 * ix + 1) + readDstVal :: Int -> ST s v + readDstVal !ix = I.read valsDst ix + indexLoKeyA :: Int -> k + indexLoKeyA !ix = I.index keysA (ix * 2) + indexLoKeyB :: Int -> k + indexLoKeyB !ix = I.index keysB (ix * 2) + indexHiKeyA :: Int -> k + indexHiKeyA !ix = I.index keysA (ix * 2 + 1) + indexHiKeyB :: Int -> k + indexHiKeyB !ix = I.index keysB (ix * 2 + 1) + indexValueA :: Int -> v + indexValueA !ix = I.index valsA ix + indexValueB :: Int -> v + indexValueB !ix = I.index valsB ix + -- In the go functon, ixDst is always at least one. Similarly, + -- all key arguments are always greater than minBound. + let go :: Int -> k -> k -> v -> Int -> k -> k -> v -> Int -> ST s Int + go !ixA !loA !hiA !valA !ixB !loB !hiB !valB !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + prevVal <- readDstVal (ixDst - 1) + case compare loA loB of + LT -> do + let (upper,ixA') = if hiA < loB + then (hiA,ixA + 1) + else (pred loB,ixA) + ixDst' <- if pred loA == prevHi && valA == prevVal + then do + writeDstHiKey (ixDst - 1) upper + return ixDst + else do + writeKeyRange ixDst loA upper + writeDstValue ixDst valA + return (ixDst + 1) + if ixA' < szA + then do + let (loA',hiA') = if hiA < loB + then (indexLoKeyA ixA',indexHiKeyA ixA') + else (loB,hiA) + go ixA' loA' hiA' (indexValueA ixA') ixB loB hiB valB ixDst' + else copyB ixB loB hiB valB ixDst' + GT -> do + let (upper,ixB') = if hiB < loA + then (hiB,ixB + 1) + else (pred loA,ixB) + ixDst' <- if pred loB == prevHi && valB == prevVal + then do + writeDstHiKey (ixDst - 1) upper + return ixDst + else do + writeKeyRange ixDst loB upper + writeDstValue ixDst valB + return (ixDst + 1) + if ixB' < szB + then do + let (loB',hiB') = if hiB < loA + then (indexLoKeyB ixB',indexHiKeyB ixB') + else (loA,hiB) + go ixA loA hiA valA ixB' loB' hiB' (indexValueB ixB') ixDst' + else copyA ixA loA hiA valA ixDst' + EQ -> do + let valCombination = combine valA valB + case compare hiA hiB of + LT -> do + ixDst' <- if pred loA == prevHi && valCombination == prevVal + then do + writeDstHiKey (ixDst - 1) hiA + return ixDst + else do + writeKeyRange ixDst loA hiA + writeDstValue ixDst valCombination + return (ixDst + 1) + let ixA' = ixA + 1 + loB' = succ hiA + if ixA' < szA + then go ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') (indexValueA ixA') ixB loB' hiB valB ixDst' + else copyB ixB loB' hiB valB ixDst' + GT -> do + ixDst' <- if pred loB == prevHi && valCombination == prevVal + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + writeDstValue ixDst valCombination + return (ixDst + 1) + let ixB' = ixB + 1 + loA' = succ hiB + if ixB' < szB + then go ixA loA' hiA valA ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') (indexValueB ixB') ixDst' + else copyA ixA loA' hiA valA ixDst' + EQ -> do + ixDst' <- if pred loB == prevHi && valCombination == prevVal + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + writeDstValue ixDst valCombination + return (ixDst + 1) + let ixA' = ixA + 1 + ixB' = ixB + 1 + if ixA' < szA + then if ixB' < szB + then go ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') (indexValueA ixA') ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') (indexValueB ixB') ixDst' + else copyA ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') (indexValueA ixA') ixDst' + else if ixB' < szB + then copyB ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') (indexValueB ixB') ixDst' + else return ixDst' + copyB :: Int -> k -> k -> v -> Int -> ST s Int + copyB !ixB !loB !hiB !valB !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + prevVal <- readDstVal (ixDst - 1) + ixDst' <- if pred loB == prevHi && valB == prevVal + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + writeDstValue ixDst valB + return (ixDst + 1) + let ixB' = ixB + 1 + remaining = szB - ixB' + I.copy keysDst (ixDst' * 2) keysB (ixB' * 2) (remaining * 2) + I.copy valsDst ixDst' valsB ixB' remaining + return (ixDst' + remaining) + copyA :: Int -> k -> k -> v -> Int -> ST s Int + copyA !ixA !loA !hiA !valA !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + prevVal <- readDstVal (ixDst - 1) + ixDst' <- if pred loA == prevHi && valA == prevVal + then do + writeDstHiKey (ixDst - 1) hiA + return ixDst + else do + writeKeyRange ixDst loA hiA + writeDstValue ixDst valA + return (ixDst + 1) + let ixA' = ixA + 1 + remaining = szA - ixA' + I.copy keysDst (ixDst' * 2) keysA (ixA' * 2) (remaining * 2) + I.copy valsDst ixDst' valsA ixA' remaining + return (ixDst' + remaining) + let !loA0 = indexLoKeyA 0 + !loB0 = indexLoKeyB 0 + !hiA0 = indexHiKeyA 0 + !hiB0 = indexHiKeyB 0 + !valA0 = indexValueA 0 + !valB0 = indexValueB 0 + total <- case compare loA0 loB0 of + LT -> if hiA0 < loB0 + then do + writeKeyRange 0 loA0 hiA0 + writeDstValue 0 valA0 + if 1 < szA + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) (indexValueA 1) 0 loB0 hiB0 valB0 1 + else copyB 0 loB0 hiB0 valB0 1 + else do + -- here we know that hiA > loA + let !upperA = pred loB0 + writeKeyRange 0 loA0 upperA + writeDstValue 0 valA0 + go 0 loB0 hiA0 valA0 0 loB0 hiB0 valB0 1 + EQ -> case compare hiA0 hiB0 of + LT -> do + writeKeyRange 0 loA0 hiA0 + writeDstValue 0 (combine valA0 valB0) + if 1 < szA + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) (indexValueA 1) 0 (succ hiA0) hiB0 valB0 1 + else copyB 0 (succ hiA0) hiB0 valB0 1 + GT -> do + writeKeyRange 0 loB0 hiB0 + writeDstValue 0 (combine valA0 valB0) + if 1 < szB + then go 0 (succ hiB0) hiA0 valA0 1 (indexLoKeyB 1) (indexHiKeyB 1) (indexValueB 1) 1 + else copyA 0 (succ hiB0) hiA0 valA0 1 + EQ -> do + writeKeyRange 0 loA0 hiA0 + writeDstValue 0 (combine valA0 valB0) + if 1 < szA + then if 1 < szB + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) (indexValueA 1) 1 (indexLoKeyB 1) (indexHiKeyB 1) (indexValueB 1) 1 + else copyA 1 (indexLoKeyA 1) (indexHiKeyA 1) (indexValueA 1) 1 + else if 1 < szB + then copyB 1 (indexLoKeyB 1) (indexHiKeyB 1) (indexValueB 1) 1 + else return 1 + GT -> if hiB0 < loA0 + then do + writeKeyRange 0 loB0 hiB0 + writeDstValue 0 valB0 + if 1 < szB + then go 0 loA0 hiA0 valA0 1 (indexLoKeyB 1) (indexHiKeyB 1) (indexValueB 1) 1 + else copyA 0 loA0 hiA0 valA0 1 + else do + let !upperB = pred loA0 + writeKeyRange 0 loB0 upperB + writeDstValue 0 valB0 + go 0 loA0 hiA0 valA0 0 loA0 hiB0 valB0 1 + !keysFinal <- I.resize keysDst (total * 2) + !valsFinal <- I.resize valsDst total + liftA2 (,) (I.unsafeFreeze keysFinal) (I.unsafeFreeze valsFinal) + +concatWith :: forall karr varr k v. (Contiguous karr, Element karr k, Ord k, Enum k, Contiguous varr, Element varr v, Eq v) => (v -> v -> v) -> [Map karr varr k v] -> Map karr varr k v +concatWith combine = go [] where + go :: [Map karr varr k v] -> [Map karr varr k v] -> Map karr varr k v + go !stack [] = foldl' (appendWith combine) empty (L.reverse stack) + go !stack (x : xs) = if size x > 0 + then go (pushStack x stack) xs + else go stack xs + pushStack :: Map karr varr k v -> [Map karr varr k v] -> [Map karr varr k v] + pushStack x [] = [x] + pushStack x (s : ss) = if size x >= size s + then pushStack (appendWith combine s x) ss + else x : s : ss + +size :: (Contiguous varr, Element varr v) => Map karr varr k v -> Int +size (Map _ vals) = I.size vals + +toList :: (Contiguous karr, Element karr k, Contiguous varr, Element varr v) => Map karr varr k v -> [(k,k,v)] +toList = foldrWithKey (\lo hi v xs -> (lo,hi,v) : xs) [] + +foldrWithKey :: (Contiguous karr, Element karr k, Contiguous varr, Element varr v) => (k -> k -> v -> b -> b) -> b -> Map karr varr k v -> b +foldrWithKey f z (Map keys vals) = + let !sz = I.size vals + go !i + | i == sz = z + | otherwise = + let !lo = I.index keys (i * 2) + !hi = I.index keys (i * 2 + 1) + !v = I.index vals i + in f lo hi v (go (i + 1)) + in go 0 + +showsPrec :: (Contiguous karr, Element karr k, Show k, Contiguous varr, Element varr v, Show v) => Int -> Map karr varr k v -> ShowS +showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) + +liftShowsPrec2 :: (Contiguous karr, Element karr k, Contiguous varr, Element varr v) => (Int -> k -> ShowS) -> ([k] -> ShowS) -> (Int -> v -> ShowS) -> ([v] -> ShowS) -> Int -> Map karr varr k v -> ShowS +liftShowsPrec2 showsPrecK _ showsPrecV _ p xs = showParen (p > 10) $ + showString "fromList " . showListWith (\(a,b,c) -> show_tuple [showsPrecK 0 a, showsPrecK 0 b, showsPrecV 0 c]) (toList xs) + +-- implementation copied from GHC.Show +show_tuple :: [ShowS] -> ShowS +show_tuple ss = id + . showChar '(' + . foldr1 (\s r -> s . showChar ',' . r) ss + . showChar ')' + + diff --git a/src/Data/Diet/Map/Lifted/Lifted.hs b/src/Data/Diet/Map/Lifted/Lifted.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Map/Lifted/Lifted.hs @@ -0,0 +1,76 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +{-# OPTIONS_GHC -O2 #-} +module Data.Diet.Map.Lifted.Lifted + ( Map + , singleton + , lookup + -- * List Conversion + , fromList + , fromListAppend + , fromListN + , fromListAppendN + ) where + +import Prelude hiding (lookup,map) + +import Data.Semigroup (Semigroup) +import Data.Primitive (Array) +import qualified GHC.Exts as E +import qualified Data.Semigroup as SG +import qualified Data.Diet.Map.Internal as I + +newtype Map k v = Map (I.Map Array Array k v) + +-- | /O(1)/ Create a diet map with a single element. +singleton :: Ord k + => k -- ^ inclusive lower bound + -> k -- ^ inclusive upper bound + -> v -- ^ value + -> Map k v +singleton lo hi v = Map (I.singleton lo hi v) + +-- | /O(log n)/ Lookup the value at a key in the map. +lookup :: Ord k => k -> Map k v -> Maybe v +lookup a (Map s) = I.lookup a s + +instance (Show k, Show v) => Show (Map k v) where + showsPrec p (Map m) = I.showsPrec p m + +instance (Eq k, Eq v) => Eq (Map k v) where + Map x == Map y = I.equals x y + +instance (Ord k, Enum k, Semigroup v, Eq v) => Semigroup (Map k v) where + Map x <> Map y = Map (I.append x y) + +instance (Ord k, Enum k, Semigroup v, Eq v) => Monoid (Map k v) where + mempty = Map I.empty + mappend = (SG.<>) + mconcat = Map . I.concat . E.coerce + +instance (Ord k, Enum k, Eq v) => E.IsList (Map k v) where + type Item (Map k v) = (k,k,v) + fromListN n = Map . I.fromListN n + fromList = Map . I.fromList + toList (Map s) = I.toList s + +fromList :: (Ord k, Enum k, Eq v) => [(k,k,v)] -> Map k v +fromList = Map . I.fromList + +fromListN :: (Ord k, Enum k, Eq v) + => Int -- ^ expected size of resulting 'Map' + -> [(k,k,v)] -- ^ key-value pairs + -> Map k v +fromListN n = Map . I.fromListN n + +fromListAppend :: (Ord k, Enum k, Semigroup v, Eq v) => [(k,k,v)] -> Map k v +fromListAppend = Map . I.fromListAppend + +fromListAppendN :: (Ord k, Enum k, Semigroup v, Eq v) + => Int -- ^ expected size of resulting 'Map' + -> [(k,k,v)] -- ^ key-value pairs + -> Map k v +fromListAppendN n = Map . I.fromListAppendN n diff --git a/src/Data/Diet/Map/Unboxed/Lifted.hs b/src/Data/Diet/Map/Unboxed/Lifted.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Map/Unboxed/Lifted.hs @@ -0,0 +1,78 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +{-# OPTIONS_GHC -O2 #-} +module Data.Diet.Map.Unboxed.Lifted + ( Map + , singleton + , lookup + -- * List Conversion + , fromList + , fromListAppend + , fromListN + , fromListAppendN + ) where + +import Prelude hiding (lookup,map) + +import Data.Semigroup (Semigroup) +import Data.Primitive.Types (Prim) +import Data.Primitive.PrimArray (PrimArray) +import Data.Primitive.Array (Array) +import qualified GHC.Exts as E +import qualified Data.Semigroup as SG +import qualified Data.Diet.Map.Internal as I + +newtype Map k v = Map (I.Map PrimArray Array k v) + +-- | /O(1)/ Create a diet map with a single element. +singleton :: (Prim k,Ord k) + => k -- ^ inclusive lower bound + -> k -- ^ inclusive upper bound + -> v -- ^ value + -> Map k v +singleton lo hi v = Map (I.singleton lo hi v) + +-- | /O(log n)/ Lookup the value at a key in the map. +lookup :: (Prim k, Ord k) => k -> Map k v -> Maybe v +lookup a (Map s) = I.lookup a s + +instance (Prim k, Show k, Show v) => Show (Map k v) where + showsPrec p (Map m) = I.showsPrec p m + +instance (Prim k, Eq k, Eq v) => Eq (Map k v) where + Map x == Map y = I.equals x y + +instance (Prim k, Ord k, Enum k, Semigroup v, Eq v) => Semigroup (Map k v) where + Map x <> Map y = Map (I.append x y) + +instance (Prim k, Ord k, Enum k, Semigroup v, Eq v) => Monoid (Map k v) where + mempty = Map I.empty + mappend = (SG.<>) + mconcat = Map . I.concat . E.coerce + +instance (Prim k, Ord k, Enum k, Eq v) => E.IsList (Map k v) where + type Item (Map k v) = (k,k,v) + fromListN n = Map . I.fromListN n + fromList = Map . I.fromList + toList (Map s) = I.toList s + +fromList :: (Ord k, Enum k, Prim k, Eq v) => [(k,k,v)] -> Map k v +fromList = Map . I.fromList + +fromListN :: (Ord k, Enum k, Prim k, Eq v) + => Int -- ^ expected size of resulting 'Map' + -> [(k,k,v)] -- ^ key-value pairs + -> Map k v +fromListN n = Map . I.fromListN n + +fromListAppend :: (Ord k, Enum k, Prim k, Semigroup v, Eq v) => [(k,k,v)] -> Map k v +fromListAppend = Map . I.fromListAppend + +fromListAppendN :: (Ord k, Enum k, Prim k, Semigroup v, Eq v) + => Int -- ^ expected size of resulting 'Map' + -> [(k,k,v)] -- ^ key-value pairs + -> Map k v +fromListAppendN n = Map . I.fromListAppendN n diff --git a/src/Data/Diet/Set/Internal.hs b/src/Data/Diet/Set/Internal.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Set/Internal.hs @@ -0,0 +1,488 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +{-# OPTIONS_GHC -O2 -Wall #-} +module Data.Diet.Set.Internal + ( Set + , empty + , singleton + , append + , member + , concat + , equals + , showsPrec + , difference + , foldr + -- splitting + , aboveInclusive + , belowInclusive + , betweenInclusive + -- list conversion + , fromListN + , fromList + , toList + ) where + +import Prelude hiding (lookup,showsPrec,concat,map,foldr) + +import Control.Monad.ST (ST,runST) +import Data.Primitive.Contiguous (Contiguous,Element,Mutable) +import qualified Data.Foldable as F +import qualified Prelude as P +import qualified Data.Primitive.Contiguous as I + +-- The key array is twice as long as the value array since +-- everything is stored as a range. Also, figure out how to +-- unpack these two arguments at some point. +newtype Set arr a = Set (arr a) + +empty :: Contiguous arr => Set arr a +empty = Set I.empty + +equals :: (Contiguous arr, Element arr a, Eq a) => Set arr a -> Set arr a -> Bool +equals (Set x) (Set y) = I.equals x y + +fromListN :: (Contiguous arr, Element arr a, Ord a, Enum a) => Int -> [(a,a)] -> Set arr a +fromListN _ xs = concat (P.map (uncurry singleton) xs) + +fromList :: (Contiguous arr, Element arr a, Ord a, Enum a) => [(a,a)] -> Set arr a +fromList = fromListN 1 + +concat :: forall arr a. (Contiguous arr, Element arr a, Ord a, Enum a) + => [Set arr a] + -> Set arr a +concat = go [] where + go :: [Set arr a] -> [Set arr a] -> Set arr a + go !stack [] = F.foldl' append empty stack + go !stack (x : xs) = if size x > 0 + then go (pushStack x stack) xs + else go stack xs + pushStack :: Set arr a -> [Set arr a] -> [Set arr a] + pushStack x [] = [x] + pushStack x (s : ss) = if size x >= size s + then pushStack (append x s) ss + else x : s : ss + +singleton :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a -- ^ lower inclusive bound + -> a -- ^ upper inclusive bound + -> Set arr a +singleton !lo !hi = if lo <= hi + then uncheckedSingleton lo hi + else empty + +-- precondition: lo must be less than or equal to hi +uncheckedSingleton :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a -- ^ lower inclusive bound + -> a -- ^ upper inclusive bound + -> Set arr a +uncheckedSingleton lo hi = runST $ do + !(arr :: Mutable arr s a) <- I.new 2 + I.write arr 0 lo + I.write arr 1 hi + r <- I.unsafeFreeze arr + return (Set r) + +member :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a + -> Set arr a + -> Bool +member a (Set arr) = go 0 ((div (I.size arr) 2) - 1) where + go :: Int -> Int -> Bool + go !start !end = if end <= start + then if end == start + then + let !valLo = I.index arr (2 * start) + !valHi = I.index arr (2 * start + 1) + in a >= valLo && a <= valHi + else False + else + let !mid = div (end + start + 1) 2 + !valLo = I.index arr (2 * mid) + in case P.compare a valLo of + LT -> go start (mid - 1) + EQ -> True + GT -> go mid end +{-# INLINEABLE member #-} + +-- not exported since it provides an index into the underlying set +-- Right means that the needle was found. The index provided is the +-- index of the range that contains it [O,n). Left means that the needle +-- was not contained by any of the ranges. The index provided is +-- the index of the range to its right [0,n] +locate :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a + -> Set arr a + -> Either Int Int +locate a (Set arr) = go 0 ((div (I.size arr) 2) - 1) where + go :: Int -> Int -> Either Int Int + go !start !end = if end <= start + then if end == start + then + let !valLo = I.index arr (2 * start) + !valHi = I.index arr (2 * start + 1) + in if (a >= valLo) + then if a <= valHi + then Right start + else Left (start + 1) + else Left start + else Left 0 + else + let !mid = div (end + start + 1) 2 + !valLo = I.index arr (2 * mid) + in case P.compare a valLo of + LT -> go start (mid - 1) + EQ -> Right mid + GT -> go mid end + +betweenInclusive :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a -- ^ inclusive lower bound + -> a -- ^ inclusive upper bound + -> Set arr a + -> Set arr a +betweenInclusive lo hi (Set arr) + | hi < lo = empty + | I.size arr > 0 && I.index arr 0 >= lo && I.index arr (I.size arr - 1) <= hi = Set arr + | otherwise = case locate lo (Set arr) of + Left ixLo -> case locate hi (Set arr) of + Left ixHi -> Set (I.clone arr (ixLo * 2) ((ixHi - ixLo) * 2)) + Right ixHi -> runST $ do + let len = ixHi - ixLo + 1 + res <- I.new (len * 2) + rightLo <- I.indexM arr (ixHi * 2) + I.copy res 0 arr (ixLo * 2) (len * 2 - 2) + I.write res (len * 2 - 2) rightLo + I.write res (len * 2 - 1) hi + r <- I.unsafeFreeze res + return (Set r) + Right ixLo -> case locate hi (Set arr) of + Left ixHi -> runST $ do + let len = ixHi - ixLo + res <- I.new (len * 2) + leftHi <- I.indexM arr (ixLo * 2 + 1) + I.write res 0 lo + I.write res 1 leftHi + I.copy res 2 arr (ixLo * 2 + 2) (len * 2 - 2) + r <- I.unsafeFreeze res + return (Set r) + Right ixHi -> if ixLo == ixHi + then uncheckedSingleton lo hi + else runST $ do + let len = ixHi - ixLo + 1 + res <- I.new (len * 2) + leftHi <- I.indexM arr (ixLo * 2 + 1) + I.write res 0 lo + I.write res 1 leftHi + I.copy res 2 arr (ixLo * 2 + 2) (len * 2 - 4) + rightLo <- I.indexM arr (ixHi * 2) + I.write res (len * 2 - 2) rightLo + I.write res (len * 2 - 1) hi + r <- I.unsafeFreeze res + return (Set r) + + +aboveInclusive :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a -- ^ inclusive lower bound + -> Set arr a + -> Set arr a +aboveInclusive x (Set arr) = case locate x (Set arr) of + Left ix -> if ix == 0 + then Set arr + else Set (I.clone arr (ix * 2) (I.size arr - ix * 2)) + Right ix -> + let lo = I.index arr (ix * 2) + hi = I.index arr (ix * 2 + 1) + in if lo == x + then if ix == 0 + then Set arr + else Set (I.clone arr (ix * 2) (I.size arr - ix * 2)) + else runST $ do + result <- I.new (I.size arr - ix * 2) + I.write result 0 x + I.write result 1 hi + I.copy result 2 arr ((ix + 1) * 2) (I.size arr - ix * 2 - 2) + r <- I.unsafeFreeze result + return (Set r) + +belowInclusive :: forall arr a. (Contiguous arr, Element arr a, Ord a) + => a -- ^ inclusive upper bound + -> Set arr a + -> Set arr a +belowInclusive x (Set arr) = case locate x (Set arr) of + Left ix -> if ix * 2 == I.size arr + then Set arr + else Set (I.clone arr 0 (ix * 2)) + Right ix -> + let lo = I.index arr (ix * 2) + hi = I.index arr (ix * 2 + 1) + in if hi == x + then if ix * 2 == I.size arr + then Set arr + else Set (I.clone arr 0 ((ix + 1) * 2)) + else runST $ do + result <- I.new ((ix + 1) * 2) + I.copy result 0 arr 0 (ix * 2) + I.write result (ix * 2) lo + I.write result (ix * 2 + 1) x + r <- I.unsafeFreeze result + return (Set r) + +append :: forall arr a. (Contiguous arr, Element arr a, Ord a, Enum a) + => Set arr a + -> Set arr a + -> Set arr a +append (Set keysA) (Set keysB) + | szA < 1 = Set keysB + | szB < 1 = Set keysA + | otherwise = runST action + where + !szA = div (I.size keysA) 2 + !szB = div (I.size keysB) 2 + action :: forall s. ST s (Set arr a) + action = do + !(keysDst :: Mutable arr s a) <- I.new (max szA szB * 8) + let writeKeyRange :: Int -> a -> a -> ST s () + writeKeyRange !ix !lo !hi = do + I.write keysDst (2 * ix) lo + I.write keysDst (2 * ix + 1) hi + writeDstHiKey :: Int -> a -> ST s () + writeDstHiKey !ix !hi = I.write keysDst (2 * ix + 1) hi + readDstHiKey :: Int -> ST s a + readDstHiKey !ix = I.read keysDst (2 * ix + 1) + indexLoKeyA :: Int -> a + indexLoKeyA !ix = I.index keysA (ix * 2) + indexLoKeyB :: Int -> a + indexLoKeyB !ix = I.index keysB (ix * 2) + indexHiKeyA :: Int -> a + indexHiKeyA !ix = I.index keysA (ix * 2 + 1) + indexHiKeyB :: Int -> a + indexHiKeyB !ix = I.index keysB (ix * 2 + 1) + -- In the go functon, ixDst is always at least one. Similarly, + -- all key arguments are always greater than minBound. + let go :: Int -> a -> a -> Int -> a -> a -> Int -> ST s Int + go !ixA !loA !hiA !ixB !loB !hiB !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + case compare loA loB of + LT -> do + let (upper,ixA') = if hiA < loB + then (hiA,ixA + 1) + else (pred loB,ixA) + ixDst' <- if pred loA == prevHi + then do + writeDstHiKey (ixDst - 1) upper + return ixDst + else do + writeKeyRange ixDst loA upper + return (ixDst + 1) + if ixA' < szA + then do + let (loA',hiA') = if hiA < loB + then (indexLoKeyA ixA',indexHiKeyA ixA') + else (loB,hiA) + go ixA' loA' hiA' ixB loB hiB ixDst' + else copyB ixB loB hiB ixDst' + GT -> do + let (upper,ixB') = if hiB < loA + then (hiB,ixB + 1) + else (pred loA,ixB) + ixDst' <- if pred loB == prevHi + then do + writeDstHiKey (ixDst - 1) upper + return ixDst + else do + writeKeyRange ixDst loB upper + return (ixDst + 1) + if ixB' < szB + then do + let (loB',hiB') = if hiB < loA + then (indexLoKeyB ixB',indexHiKeyB ixB') + else (loA,hiB) + go ixA loA hiA ixB' loB' hiB' ixDst' + else copyA ixA loA hiA ixDst' + EQ -> do + case compare hiA hiB of + LT -> do + ixDst' <- if pred loA == prevHi + then do + writeDstHiKey (ixDst - 1) hiA + return ixDst + else do + writeKeyRange ixDst loA hiA + return (ixDst + 1) + let ixA' = ixA + 1 + loB' = succ hiA + if ixA' < szA + then go ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') ixB loB' hiB ixDst' + else copyB ixB loB' hiB ixDst' + GT -> do + ixDst' <- if pred loB == prevHi + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + return (ixDst + 1) + let ixB' = ixB + 1 + loA' = succ hiB + if ixB' < szB + then go ixA loA' hiA ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') ixDst' + else copyA ixA loA' hiA ixDst' + EQ -> do + ixDst' <- if pred loB == prevHi + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + return (ixDst + 1) + let ixA' = ixA + 1 + ixB' = ixB + 1 + if ixA' < szA + then if ixB' < szB + then go ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') ixDst' + else copyA ixA' (indexLoKeyA ixA') (indexHiKeyA ixA') ixDst' + else if ixB' < szB + then copyB ixB' (indexLoKeyB ixB') (indexHiKeyB ixB') ixDst' + else return ixDst' + copyB :: Int -> a -> a -> Int -> ST s Int + copyB !ixB !loB !hiB !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + ixDst' <- if pred loB == prevHi + then do + writeDstHiKey (ixDst - 1) hiB + return ixDst + else do + writeKeyRange ixDst loB hiB + return (ixDst + 1) + let ixB' = ixB + 1 + remaining = szB - ixB' + I.copy keysDst (ixDst' * 2) keysB (ixB' * 2) (remaining * 2) + return (ixDst' + remaining) + copyA :: Int -> a -> a -> Int -> ST s Int + copyA !ixA !loA !hiA !ixDst = do + prevHi <- readDstHiKey (ixDst - 1) + ixDst' <- if pred loA == prevHi + then do + writeDstHiKey (ixDst - 1) hiA + return ixDst + else do + writeKeyRange ixDst loA hiA + return (ixDst + 1) + let ixA' = ixA + 1 + remaining = szA - ixA' + I.copy keysDst (ixDst' * 2) keysA (ixA' * 2) (remaining * 2) + return (ixDst' + remaining) + let !loA0 = indexLoKeyA 0 + !loB0 = indexLoKeyB 0 + !hiA0 = indexHiKeyA 0 + !hiB0 = indexHiKeyB 0 + total <- case compare loA0 loB0 of + LT -> if hiA0 < loB0 + then do + writeKeyRange 0 loA0 hiA0 + if 1 < szA + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) 0 loB0 hiB0 1 + else copyB 0 loB0 hiB0 1 + else do + -- here we know that hiA > loA + let !upperA = pred loB0 + writeKeyRange 0 loA0 upperA + go 0 loB0 hiA0 0 loB0 hiB0 1 + EQ -> case compare hiA0 hiB0 of + LT -> do + writeKeyRange 0 loA0 hiA0 + if 1 < szA + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) 0 (succ hiA0) hiB0 1 + else copyB 0 (succ hiA0) hiB0 1 + GT -> do + writeKeyRange 0 loB0 hiB0 + if 1 < szB + then go 0 (succ hiB0) hiA0 1 (indexLoKeyB 1) (indexHiKeyB 1) 1 + else copyA 0 (succ hiB0) hiA0 1 + EQ -> do + writeKeyRange 0 loA0 hiA0 + if 1 < szA + then if 1 < szB + then go 1 (indexLoKeyA 1) (indexHiKeyA 1) 1 (indexLoKeyB 1) (indexHiKeyB 1) 1 + else copyA 1 (indexLoKeyA 1) (indexHiKeyA 1) 1 + else if 1 < szB + then copyB 1 (indexLoKeyB 1) (indexHiKeyB 1) 1 + else return 1 + GT -> if hiB0 < loA0 + then do + writeKeyRange 0 loB0 hiB0 + if 1 < szB + then go 0 loA0 hiA0 1 (indexLoKeyB 1) (indexHiKeyB 1) 1 + else copyA 0 loA0 hiA0 1 + else do + let !upperB = pred loA0 + writeKeyRange 0 loB0 upperB + go 0 loA0 hiA0 0 loA0 hiB0 1 + !keysFinal <- I.resize keysDst (total * 2) + fmap Set (I.unsafeFreeze keysFinal) + +difference :: forall a arr. (Contiguous arr, Element arr a, Ord a, Enum a) + => Set arr a + -> Set arr a + -> Set arr a +difference setA@(Set arrA) setB@(Set arrB) + | szA == 0 = empty + | szB == 0 = setA + | otherwise = + let inners :: Int -> [Set arr a] + inners !ix = if ix < szB - 1 + then + let inner = betweenInclusive + (succ (I.index arrB (2 * ix + 1))) + (pred (I.index arrB (2 * ix + 2))) + (Set arrA) + in inner : inners (ix + 1) + else [] + lowestA = I.index arrA 0 + highestA = I.index arrA (szA * 2 - 1) + lowestB = I.index arrB 0 + highestB = I.index arrB (szB * 2 - 1) + -- TODO: if we ever add exclusive variants of below + -- and above, we should switch to using them here. + lowFragment = if lowestA < lowestB + then [belowInclusive (pred lowestB) (Set arrA)] + else [] + highFragment = if highestA > highestB + then [aboveInclusive (succ highestB) (Set arrA)] + else [] + -- we should use a more efficient concat since + -- we know everything is ordered. + in concat (lowFragment ++ inners 0 ++ highFragment) + where + !szA = size setA + !szB = size setB + +size :: (Contiguous arr, Element arr a) => Set arr a -> Int +size (Set arr) = quot (I.size arr) 2 + +toList :: (Contiguous arr, Element arr a) => Set arr a -> [(a,a)] +toList = foldr (\lo hi xs -> (lo,hi) : xs) [] + +foldr :: (Contiguous arr, Element arr a) => (a -> a -> b -> b) -> b -> Set arr a -> b +foldr f z (Set arr) = + let !sz = div (I.size arr) 2 + go !i + | i == sz = z + | otherwise = + let !lo = I.index arr (i * 2) + !hi = I.index arr (i * 2 + 1) + in f lo hi (go (i + 1)) + in go 0 +{-# INLINABLE foldr #-} + +showsPrec :: (Contiguous arr, Element arr a, Show a) + => Int + -> Set arr a + -> ShowS +showsPrec p xs = showParen (p > 10) $ + showString "fromList " . shows (toList xs) + diff --git a/src/Data/Diet/Set/Lifted.hs b/src/Data/Diet/Set/Lifted.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Set/Lifted.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +{-# OPTIONS_GHC -O2 #-} +module Data.Diet.Set.Lifted + ( Set + , singleton + , member + , difference + -- * Split + , aboveInclusive + , belowInclusive + , betweenInclusive + -- * Folds + , foldr + -- * List Conversion + , fromList + , fromListN + ) where + +import Prelude hiding (lookup,map,foldr) + +import Data.Semigroup (Semigroup) +import Data.Primitive (Array) +import qualified GHC.Exts as E +import qualified Data.Semigroup as SG +import qualified Data.Diet.Set.Internal as I + +newtype Set a = Set (I.Set Array a) + +-- | /O(1)/ Create a diet set with a single element. +singleton :: Ord a + => a -- ^ inclusive lower bound + -> a -- ^ inclusive upper bound + -> Set a +singleton lo hi = Set (I.singleton lo hi) + +-- | /O(log n)/ Lookup the value at a key in the map. +member :: Ord a => a -> Set a -> Bool +member a (Set s) = I.member a s + +instance Show a => Show (Set a) where + showsPrec p (Set s) = I.showsPrec p s + +instance Eq a => Eq (Set a) where + Set x == Set y = I.equals x y + +instance Ord a => Ord (Set a) where + compare (Set xs) (Set ys) = compare (I.toList xs) (I.toList ys) + +instance (Ord a, Enum a) => Semigroup (Set a) where + Set x <> Set y = Set (I.append x y) + +instance (Ord a, Enum a) => Monoid (Set a) where + mempty = Set I.empty + mappend = (SG.<>) + mconcat = Set . I.concat . E.coerce + +instance (Ord a, Enum a) => E.IsList (Set a) where + type Item (Set a) = (a,a) + fromListN n = Set . I.fromListN n + fromList = Set . I.fromList + toList (Set s) = I.toList s + +fromList :: (Ord a, Enum a) => [(a,a)] -> Set a +fromList = Set . I.fromList + +fromListN :: (Ord a, Enum a) + => Int -- ^ expected size of resulting diet 'Set' + -> [(a,a)] -- ^ key-value pairs + -> Set a +fromListN n = Set . I.fromListN n + +-- | /O(n + m*log n)/ Subtract the subtrahend of size @m@ from the +-- minuend of size @n@. It should be possible to improve the improve +-- the performance of this to /O(n + m)/. Anyone interested in doing +-- this should open a PR. +difference :: (Ord a, Enum a) + => Set a -- ^ minuend + -> Set a -- ^ subtrahend + -> Set a +difference (Set x) (Set y) = Set (I.difference x y) + +foldr :: (a -> a -> b -> b) -> b -> Set a -> b +foldr f z (Set arr) = I.foldr f z arr + +-- | /O(n)/ The subset where all elements are greater than +-- or equal to the given value. +aboveInclusive :: (Ord a) + => a -- ^ inclusive lower bound + -> Set a + -> Set a +aboveInclusive x (Set s) = Set (I.aboveInclusive x s) + +-- | /O(n)/ The subset where all elements are less than +-- or equal to the given value. +belowInclusive :: (Ord a) + => a -- ^ inclusive upper bound + -> Set a + -> Set a +belowInclusive x (Set s) = Set (I.belowInclusive x s) + +-- | /O(n)/ The subset where all elements are greater than +-- or equal to the lower bound and less than or equal to +-- the upper bound. +betweenInclusive :: (Ord a) + => a -- ^ inclusive lower bound + -> a -- ^ inclusive upper bound + -> Set a + -> Set a +betweenInclusive x y (Set s) = Set (I.betweenInclusive x y s) + diff --git a/src/Data/Diet/Set/Unboxed.hs b/src/Data/Diet/Set/Unboxed.hs new file mode 100644 --- /dev/null +++ b/src/Data/Diet/Set/Unboxed.hs @@ -0,0 +1,73 @@ +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE TypeFamilies #-} + +{-# OPTIONS_GHC -O2 #-} +module Data.Diet.Set.Unboxed + ( Set + , singleton + , member + -- * List Conversion + , toList + , fromList + , fromListN + ) where + +import Prelude hiding (lookup,map) + +import Data.Semigroup (Semigroup) +import Data.Primitive.Types (Prim) +import Data.Primitive.PrimArray (PrimArray) +import qualified GHC.Exts as E +import qualified Data.Semigroup as SG +import qualified Data.Diet.Set.Internal as I + +newtype Set a = Set (I.Set PrimArray a) + +-- | /O(1)/ Create a diet set with a single element. +singleton :: (Ord a, Prim a) + => a -- ^ inclusive lower bound + -> a -- ^ inclusive upper bound + -> Set a +singleton lo hi = Set (I.singleton lo hi) + +-- | /O(log n)/ Lookup the value at a key in the map. +member :: (Ord a, Prim a) => a -> Set a -> Bool +member a (Set s) = I.member a s + +instance (Show a, Prim a) => Show (Set a) where + showsPrec p (Set s) = I.showsPrec p s + +instance (Eq a, Prim a) => Eq (Set a) where + Set x == Set y = I.equals x y + +instance (Ord a, Prim a) => Ord (Set a) where + compare (Set xs) (Set ys) = compare (I.toList xs) (I.toList ys) + +instance (Ord a, Enum a, Prim a) => Semigroup (Set a) where + Set x <> Set y = Set (I.append x y) + +instance (Ord a, Enum a, Prim a) => Monoid (Set a) where + mempty = Set I.empty + mappend = (SG.<>) + mconcat = Set . I.concat . E.coerce + +instance (Ord a, Enum a, Prim a) => E.IsList (Set a) where + type Item (Set a) = (a,a) + fromListN n = Set . I.fromListN n + fromList = Set . I.fromList + toList (Set s) = I.toList s + +toList :: Prim a => Set a -> [(a,a)] +toList (Set x) = I.toList x + +fromList :: (Ord a, Enum a, Prim a) => [(a,a)] -> Set a +fromList = Set . I.fromList + +fromListN :: (Ord a, Enum a, Prim a) + => Int -- ^ expected size of resulting diet 'Set' + -> [(a,a)] -- ^ key-value pairs + -> Set a +fromListN n = Set . I.fromListN n + diff --git a/src/Data/Set/Diet.hs b/src/Data/Set/Diet.hs deleted file mode 100644 --- a/src/Data/Set/Diet.hs +++ /dev/null @@ -1,299 +0,0 @@ --- | Discrete Interval Encoding Tree described by Martin Erwig in /Diets for Fat Sets, January 1993/. -module Data.Set.Diet( - Interval -, point -, interval -, intervalMin -, intervalMax -, mergeI -, isPointed -, mapI -, Diet -, member -, notMember -, insert -, delete -, empty -, single -, singleI -, size -, diet -, toList -, fromList -, mapD -) where - -import Data.Ix -import Data.Foldable(foldl', Foldable) - --- | An interval with discrete values between. -data Interval a = - Interval a a - deriving (Eq, Ord) - -instance (Eq a, Show a) => Show (Interval a) where - show (Interval a1 a2) = - '[' : show a1 ++ (if a1 == a2 then [] else '|' : show a2) ++ "]" - --- | An interval with the same minimum and maximum. -point :: - a - -> Interval a -point a = - Interval a a - --- | Construct an interval ensuring that the minimum is less than or equal to maximum. -interval :: - Ord a => - a - -> a - -> Interval a -interval a1 a2 = - if a1 <= a2 - then - Interval a1 a2 - else - Interval a2 a1 - --- | The minimum of the interval. -intervalMin :: - Interval a - -> a -intervalMin (Interval a _) = - a - --- | The maximum of the interval. -intervalMax :: - Interval a - -> a -intervalMax (Interval _ a) = - a - --- | Merge two intervals if they are overlapping or adjacent. -mergeI :: - (Ord a, Enum a) => - Interval a - -> Interval a - -> Maybe (Interval a) -mergeI (Interval a1 a2) (Interval aa1 aa2) = - if (a1 <= aa2 && succ a2 >= aa1) || (aa1 <= a2 && succ aa2 >= a1) - then - Just $ Interval (min a1 aa1) (max a2 aa2) - else - Nothing - --- | Returns whether or not the interval has the same minimum and maximum. -isPointed :: - Eq a => - Interval a - -> Bool -isPointed (Interval a1 a2) = - a1 == a2 - --- | Map a function across the minimum and maximum of the interval. -mapI :: - Ord b => - (a -> b) - -> Interval a - -> Interval b -mapI f (Interval a1 a2) = - interval (f a1) (f a2) - --- | A Discrete Interval Encoding Tree. -data Diet a = - Empty - | Node (Diet a) (Interval a) (Diet a) - deriving (Eq, Ord) - -instance (Eq a, Show a) => Show (Diet a) where - showsPrec _ Empty = - id - showsPrec n (Node l i r) = - showsPrec n l . shows i . showsPrec n r - --- | Test for membership in the interval tree. -member :: - Ix a => - a - -> Diet a - -> Bool -member _ Empty = - False -member x (Node l (Interval a1 a2) r) = - inRange (a1, a2) x || member x (if x < a1 then l else r) - --- | Test for non-membership in the interval tree. -notMember :: - Ix a => - a - -> Diet a - -> Bool -notMember a = - not . member a - --- | Insert an element into the interval tree. -insert :: - (Ord a, Enum a) => - a - -> Diet a - -> Diet a -insert x Empty = - Node Empty (point x) Empty -insert x d@(Node l i@(Interval a1 a2) r) - | x < a1 = - if succ x == a1 - then - let joinLeft md@(Node Empty _ _) = - md - joinLeft (Node ml mi@(Interval ma1 ma2) mr) = - let (ml', Interval ml1 ml2) = splitMax ml - in if succ ml2 == ma1 - then - Node ml' (Interval ml1 ma2) mr - else - Node ml mi mr - joinLeft Empty = - error "Broken invariant @ Data.Set.Diet#joinLeft" - in joinLeft (Node l (Interval x a2) r) - else - Node (insert x l) i r - | x > a2 = - if succ a2 == x - then - let splitMin (Node Empty mi mr) = - (mr, mi) - splitMin (Node ml mi mr) = - let (md, mi') = splitMin ml - in (Node md mi mr, mi') - splitMin Empty = - error "Broken invariant @ Data.Set.Diet#splitMin" - joinRight jd@(Node _ _ Empty) = - jd - joinRight (Node jl ji@(Interval ja1 ja2) jr) = - let (jr', Interval jr1 jr2) = splitMin jr - in if succ ja2 == jr1 - then - Node jl (Interval ja1 jr2) jr' - else - Node jl ji jr - joinRight Empty = - error "Broken invariant @ Data.Set.Diet#joinRight" - in joinRight (Node l (Interval a1 x) r) - else - Node l i (insert x r) - | otherwise = - d - --- | Delete an element from the interval tree. -delete :: - (Ord a, Enum a) => - a - -> Diet a - -> Diet a -delete _ Empty = - Empty -delete x (Node l i@(Interval a1 a2) r) - | x < a1 = - Node (delete x l) i r - | x > a2 = - Node l i (delete x r) - | x == a1 = - let merge ml Empty = - ml - merge Empty mr = - mr - merge ml mr = - let (ml', mi) = splitMax ml - in Node ml' mi mr - in if isPointed i - then - merge l r - else - Node l (Interval (succ a1) a2) r - | x == a2 = - Node l (Interval a1 (pred a2)) r - | otherwise = - Node l (Interval a1 (pred x)) (Node Empty (Interval (succ x) a2) r) - --- | Construct an interval tree with no elements. -empty :: - Diet a -empty = - Empty - --- | Construct an interval tree with a single element. -single :: - a - -> Diet a -single a = - Node Empty (point a) Empty - --- | Construct an interval tree with a single interval. -singleI :: - Interval a - -> Diet a -singleI a = - Node Empty a Empty - --- | Return the number of elements in the interval tree. -size :: - Ix a => - Diet a - -> Int -size Empty = - 0 -size (Node l (Interval a1 a2) r) = - sum [size l, rangeSize (a1, a2), size r] - --- | Fold on the interval tree. -diet :: - (b -> Interval a -> b -> b) - -> b - -> Diet a - -> b -diet _ z Empty = - z -diet f z (Node l i r) = - f (diet f z l) i (diet f z r) - --- | Return all elements of the interval tree as a list. -toList :: - Ix a => - Diet a - -> [a] -toList = - diet (\l (Interval a1 a2) r -> concat [l, range (a1, a2), r]) [] - --- | Construct an interval tree with the elements of the list. -fromList :: - (Foldable t, Ord a, Enum a) => - t a - -> Diet a -fromList = - foldl' (flip insert) Empty - --- | Map a function across the interval tree. -mapD :: - Ord b => - (a -> b) - -> Diet a - -> Diet b -mapD _ Empty = - Empty -mapD f (Node l i r) = - Node (mapD f l) (mapI f i) (mapD f r) - --- BEGIN not exported - -splitMax :: - Diet a - -> (Diet a, Interval a) -splitMax (Node l i Empty) = - (l, i) -splitMax (Node l i r) = - let (d, i') = splitMax r - in (Node l i d, i') -splitMax Empty = - error "Broken invariant @ Data.Set.Diet#splitMax" - --- END not exported diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,297 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE CPP #-} +{-# LANGUAGE KindSignatures #-} +{-# LANGUAGE MagicHash #-} +{-# LANGUAGE UnboxedTuples #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TypeApplications #-} + +{-# OPTIONS_GHC -fno-warn-orphans #-} + +import Data.Primitive +import Data.Primitive.UnliftedArray (PrimUnlifted) +import Data.Word +import Data.Proxy (Proxy(..)) +import Data.Int + +import Test.Tasty (defaultMain,testGroup,TestTree) +import Test.QuickCheck (Arbitrary,Gen,(===),(==>)) +import Data.List.NonEmpty (NonEmpty((:|))) +import Control.Monad.ST (ST) +import qualified Test.Tasty.QuickCheck as TQC +import qualified Test.QuickCheck as QC +import qualified Test.QuickCheck.Classes as QCC +import qualified Data.Semigroup as SG +import qualified Data.Map as M +import qualified Data.Set as S +import qualified Data.Foldable as F +import qualified GHC.Exts as E +import qualified Test.QuickCheck.Classes.IsList as QCCL + +import qualified Data.Diet.Map.Unboxed.Lifted as DMUL +import qualified Data.Diet.Map.Lifted.Lifted as DMLL +import qualified Data.Diet.Set.Lifted as DSL + +main :: IO () +main = defaultMain $ testGroup "Data" + [ testGroup "Diet" + [ testGroup "Set" + [ testGroup "Lifted" + [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (DSL.Set Word16))) + , lawsToTest (QCC.ordLaws (Proxy :: Proxy (DSL.Set Word16))) + , lawsToTest (QCC.commutativeMonoidLaws (Proxy :: Proxy (DSL.Set Word16))) + , lawsToTest (QCC.isListLaws (Proxy :: Proxy (DSL.Set Word16))) + , TQC.testProperty "member" (dietMemberProp @Word8 E.fromList DSL.member) + , TQC.testProperty "difference" dietSetDifferenceProp + , TQC.testProperty "aboveInclusive" dietSetAboveProp + , testGroup "belowInclusive" + [ TQC.testProperty "basic" dietSetBelowProp + , TQC.testProperty "lowest" dietSetBelowLowestProp + , TQC.testProperty "highest" dietSetBelowHighestProp + ] + , testGroup "betweenInclusive" + [ TQC.testProperty "basic" dietSetBetweenProp + , TQC.testProperty "border" dietSetBetweenBorderProp + , TQC.testProperty "inside" dietSetBetweenBorderNearProp + ] + ] + ] + , testGroup "Map" + [ testGroup "Lifted" + [ testGroup "Lifted" + [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (DMLL.Map Word8 Integer))) + , lawsToTest (QCC.semigroupLaws (Proxy :: Proxy (DMLL.Map Word8 Word))) + , lawsToTest (QCC.commutativeMonoidLaws (Proxy :: Proxy (DMLL.Map Word8 Int))) + , lawsToTest (QCC.isListLaws (Proxy :: Proxy (DMLL.Map Word8 Integer))) + , TQC.testProperty "lookup" (dietLookupPropA @Word8 @Int E.fromList DMLL.lookup) + --, TQC.testProperty "doubleton" dietDoubletonProp + , TQC.testProperty "valid" dietValidProp + ] + ] + , testGroup "Unboxed" + [ testGroup "Lifted" + [ lawsToTest (QCC.eqLaws (Proxy :: Proxy (DMUL.Map Word8 Integer))) + , lawsToTest (QCC.semigroupLaws (Proxy :: Proxy (DMUL.Map Word8 Word))) + , lawsToTest (QCC.commutativeMonoidLaws (Proxy :: Proxy (DMUL.Map Word8 Int))) + , lawsToTest (QCC.isListLaws (Proxy :: Proxy (DMUL.Map Word8 Integer))) + , TQC.testProperty "lookup" (dietLookupPropA @Word32 @Int E.fromList DMUL.lookup) + ] + ] + ] + ] + ] + +int16 :: Proxy Int16 +int16 = Proxy + +int32 :: Proxy Int32 +int32 = Proxy + +dietSetDifferenceProp :: QC.Property +dietSetDifferenceProp = QC.property $ \(xs :: DSL.Set Word8) (ys :: DSL.Set Word8) -> + let xs' = dietSetToSet xs + ys' = dietSetToSet ys + in DSL.difference xs ys === DSL.fromList (map (\x -> (x,x)) (S.toList (S.difference xs' ys'))) + +dietSetAboveProp :: QC.Property +dietSetAboveProp = QC.property $ \(y :: Word8) (ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + (_,isMember,c) = S.splitMember y ys' + r = if isMember then S.insert y c else c + in DSL.aboveInclusive y ys === DSL.fromList (map (\x -> (x,x)) (S.toList r)) + +dietSetBelowProp :: QC.Property +dietSetBelowProp = QC.property $ \(y :: Word8) (ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + (c,isMember,_) = S.splitMember y ys' + r = if isMember then S.insert y c else c + in DSL.belowInclusive y ys === DSL.fromList (map (\x -> (x,x)) (S.toList r)) + +dietSetBelowLowestProp :: QC.Property +dietSetBelowLowestProp = QC.property $ \(ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + in case S.lookupMin ys' of + Nothing -> QC.property QC.Discard + Just y -> + let (c,isMember,_) = S.splitMember y ys' + r = if isMember then S.insert y c else c + in QC.property (DSL.belowInclusive y ys === DSL.fromList (map (\x -> (x,x)) (S.toList r))) + +dietSetBelowHighestProp :: QC.Property +dietSetBelowHighestProp = QC.property $ \(ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + in case S.lookupMax ys' of + Nothing -> QC.property QC.Discard + Just y -> + let (c,isMember,_) = S.splitMember y ys' + r = if isMember then S.insert y c else c + in QC.property (DSL.belowInclusive y ys === DSL.fromList (map (\x -> (x,x)) (S.toList r))) + +dietSetBetweenProp :: QC.Property +dietSetBetweenProp = QC.property $ \(x :: Word8) (y :: Word8) (ys :: DSL.Set Word8) -> + (x <= y) + ==> + ( let ys' = dietSetToSet ys + r = S.filter (\e -> e >= x && e <= y) ys' + in DSL.betweenInclusive x y ys === DSL.fromList (map (\z -> (z,z)) (S.toList r)) + ) + +dietSetBetweenBorderProp :: QC.Property +dietSetBetweenBorderProp = QC.property $ \(ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + in case S.lookupMax ys' of + Nothing -> QC.property QC.Discard + Just hi -> case S.lookupMin ys' of + Nothing -> QC.property QC.Discard + Just lo -> + let r = S.filter (\e -> e >= lo && e <= hi) ys' + in DSL.betweenInclusive lo hi ys === DSL.fromList (map (\z -> (z,z)) (S.toList r)) + +dietSetBetweenBorderNearProp :: QC.Property +dietSetBetweenBorderNearProp = QC.property $ \(ys :: DSL.Set Word8) -> + let ys' = dietSetToSet ys + in ( S.size ys' > 1 + ==> + ( let hi = pred (S.findMax ys') + lo = succ (S.findMin ys') + r = S.filter (\e -> e >= lo && e <= hi) ys' + in DSL.betweenInclusive lo hi ys === DSL.fromList (map (\z -> (z,z)) (S.toList r)) + ) + ) + +-- This enumerates all of the element contained by all ranges +-- in the diet set. +dietSetToSet :: (Enum a, Ord a) => DSL.Set a -> S.Set a +dietSetToSet = DSL.foldr + (\lo hi s -> S.fromList (enumFromTo lo hi) SG.<> s) + mempty + +memberProp :: forall a t. (Arbitrary a, Show a) => ([a] -> t a) -> (a -> t a -> Bool) -> QC.Property +memberProp containerFromList containerMember = QC.property $ \(xs :: [a]) -> + let c = containerFromList xs + in all (\x -> containerMember x c) xs === True + +lookupProp :: forall k v t. (Arbitrary k, Show k, Ord k, Arbitrary v, Show v, Eq v) => ([(k,v)] -> t k v) -> (k -> t k v -> Maybe v) -> QC.Property +lookupProp containerFromList containerLookup = QC.property $ \(xs :: [(k,v)]) -> + let ys = M.fromList xs + c = containerFromList xs + in all (\(x,_) -> containerLookup x c == M.lookup x ys) xs === True + +dietMemberProp :: forall a t. (Arbitrary a, Show a, Ord a, Arbitrary a, Show (t a)) => ([(a,a)] -> t a) -> (a -> t a -> Bool) -> QC.Property +dietMemberProp containerFromList containerLookup = QC.property $ \(xs :: [a]) -> + let c = containerFromList (map (\a -> (a,a)) xs) + in QC.counterexample ("original list: " ++ show xs ++ "; diet set: " ++ show c) (all (\x -> containerLookup x c == True) xs === True) + +dietLookupPropA :: forall k v t. (Arbitrary k, Show k, Ord k, Arbitrary v, Show v, Eq v, Show (t k v)) => ([(k,k,v)] -> t k v) -> (k -> t k v -> Maybe v) -> QC.Property +dietLookupPropA containerFromList containerLookup = QC.property $ \(xs :: [(k,v)]) -> + let ys = M.fromList xs + c = containerFromList (map (\(k,v) -> (k,k,v)) xs) + in QC.counterexample ("original list: " ++ show xs ++ "; diet map: " ++ show c) (all (\(x,_) -> containerLookup x c == M.lookup x ys) xs === True) + +--dietDoubletonProp :: QC.Property +--dietDoubletonProp = QC.property $ \(loA :: Word8) (hiA :: Word8) (valA :: Int) (loB :: Word8) (hiB :: Word8) (valB :: Int) -> +-- (hiA >= loA && hiB >= loB) +-- ==> +-- (simpleDoubletonToList loA hiA valA loB hiB valB === E.toList (DMLL.singleton loA hiA valA SG.<> DMLL.singleton loB hiB valB)) + +dietValidProp :: QC.Property +dietValidProp = QC.property $ \(xs :: DMLL.Map Word8 Int) -> + True === validDietTriples (E.toList xs) + +simpleDoubletonToList :: (Ord k, Enum k, SG.Semigroup v, Eq v) => k -> k -> v -> k -> k -> v -> [(k,k,v)] +simpleDoubletonToList key1A key2A valA key1B key2B valB = + let loA = min key1A key2A + hiA = max key1A key2A + loB = min key1B key2B + hiB = max key1B key2B + in deduplicate $ case compare loA loB of + LT -> case compare hiA loB of + LT -> [(loA,hiA,valA),(loB,hiB,valB)] + EQ -> case compare hiA hiB of + LT -> [(loA,pred loB,valA),(loB,hiA,valA SG.<> valB),(succ hiA,hiB,valB)] + EQ -> [(loA,pred loB,valA),(loB,hiA,valA SG.<> valB)] + GT -> error "simpleDoubletonToList: invariant violated" + GT -> case compare hiA hiB of + LT -> [(loA,pred loB,valA),(loB,hiA,valA SG.<> valB),(succ hiA,hiB,valB)] + EQ -> [(loA,pred loB,valA),(loB,hiA,valA SG.<> valB)] + GT -> [(loA,pred loB,valA),(loB,hiB,valA SG.<> valB),(succ hiB,hiA,valA)] + EQ -> case compare hiA hiB of + LT -> [(loA,hiA,valA SG.<> valB),(succ hiA, hiB, valB)] + GT -> [(loB,hiB,valA SG.<> valB),(succ hiB, hiA, valA)] + EQ -> [(loA,hiA,valA SG.<> valB)] + GT -> case compare hiB loA of + LT -> [(loB,hiB,valB),(loA,hiA,valA)] + EQ -> case compare hiB hiA of + LT -> [(loB,pred loA,valB),(loA,hiB,valA SG.<> valB),(succ hiB,hiA,valA)] + EQ -> [(loB,pred loA,valB),(loA,hiB,valA SG.<> valB)] + GT -> error "simpleDoubletonToList: invariant violated" + GT -> case compare hiB hiA of + LT -> [(loB,pred loA,valB),(loA,hiB,valA SG.<> valB),(succ hiB,hiA,valA)] + EQ -> [(loB,pred loA,valB),(loA,hiB,valA SG.<> valB)] + GT -> [(loB,pred loA,valB),(loA,hiA,valA SG.<> valB),(succ hiA,hiB,valB)] + +validDietTriples :: (Enum k,Eq k,Eq v) => [(k,k,v)] -> Bool +validDietTriples xs = deduplicate xs == xs + +deduplicate :: (Enum k,Eq k, Eq v) => [(k,k,v)] -> [(k,k,v)] +deduplicate [] = [] +deduplicate (x : xs) = F.toList (deduplicateNonEmpty (x :| xs)) + +deduplicateNonEmpty :: (Enum k, Eq k, Eq v) => NonEmpty (k,k,v) -> NonEmpty (k,k,v) +deduplicateNonEmpty ((lo,hi,v) :| xs) = case xs of + y : ys -> case deduplicateNonEmpty (y :| ys) of + (lo',hi',v') :| xs' -> if v == v' && pred lo' == hi + then (lo,hi',v) :| xs' + else (lo,hi,v) :| ((lo',hi',v') : xs') + [] -> (lo,hi,v) :| [] + +lawsToTest :: QCC.Laws -> TestTree +lawsToTest (QCC.Laws name pairs) = testGroup name (map (uncurry TQC.testProperty) pairs) + +instance (Arbitrary a, Prim a) => Arbitrary (PrimArray a) where + arbitrary = fmap E.fromList QC.arbitrary + +instance (Arbitrary k, Ord k, Enum k, Bounded k, Arbitrary v, SG.Semigroup v, Eq v) => Arbitrary (DMLL.Map k v) where + arbitrary = DMLL.fromListAppend <$> QC.vectorOf 10 arbitraryOrderedPairValue + shrink x = map E.fromList (QC.shrink (E.toList x)) + +instance (Arbitrary k, Prim k, Ord k, Enum k, Bounded k, Arbitrary v, SG.Semigroup v, Eq v) => Arbitrary (DMUL.Map k v) where + arbitrary = DMUL.fromListAppend <$> QC.vectorOf 10 arbitraryOrderedPairValue + shrink x = map E.fromList (QC.shrink (E.toList x)) + +instance (Arbitrary a, Ord a, Enum a, Bounded a) => Arbitrary (DSL.Set a) where + arbitrary = DSL.fromList <$> QC.vectorOf 7 arbitraryOrderedPair + shrink x = map E.fromList (QC.shrink (E.toList x)) + +arbitraryOrderedPair :: (Ord k, Enum k, Bounded k, Arbitrary k) => Gen (k,k) +arbitraryOrderedPair = do + a0 <- QC.arbitrary + let a1 = if a0 < maxBound then succ a0 else a0 + a2 = if a1 < maxBound then succ a1 else a1 + a3 = if a2 < maxBound then succ a2 else a2 + a' <- QC.elements [a0,a1,a2,a3] + return (a0,a') + +arbitraryOrderedPairValue :: (Ord k, Enum k, Bounded k, Arbitrary k, Arbitrary v) => Gen (k,k,v) +arbitraryOrderedPairValue = do + (lo,hi) <- arbitraryOrderedPair + v <- QC.arbitrary + return (lo,hi,v) + +instance SG.Semigroup Word where + w <> _ = w + +instance SG.Semigroup Int where + (<>) = (+) + +instance Monoid Int where + mempty = 0 + mappend = (SG.<>) + +instance SG.Semigroup Integer where + (<>) = (+) + +instance Monoid Integer where + mempty = 0 + mappend = (SG.<>) +