module Data.IntTrie
( IntTrie, identity, apply, modify, modify', overwrite,
mirror, modifyAscList, modifyDescList )
where
import Control.Applicative
import Control.Arrow (first, second)
import Data.Bits
import Data.Function (fix)
import Data.Monoid (Monoid(..))
import Data.Semigroup (Semigroup(..))
data IntTrie a = IntTrie (BitTrie a) a (BitTrie a)
data BitTrie a = BitTrie a (BitTrie a) (BitTrie a)
instance Functor BitTrie where
fmap :: forall a b. (a -> b) -> BitTrie a -> BitTrie b
fmap a -> b
f ~(BitTrie a
x BitTrie a
l BitTrie a
r) = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> b
f a
x) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
l) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
r)
instance Applicative BitTrie where
pure :: forall a. a -> BitTrie a
pure a
x = forall a. (a -> a) -> a
fix (\BitTrie a
g -> forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
x BitTrie a
g BitTrie a
g)
~(BitTrie a -> b
f BitTrie (a -> b)
fl BitTrie (a -> b)
fr) <*> :: forall a b. BitTrie (a -> b) -> BitTrie a -> BitTrie b
<*> ~(BitTrie a
x BitTrie a
xl BitTrie a
xr) = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> b
f a
x) (BitTrie (a -> b)
fl forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BitTrie a
xl) (BitTrie (a -> b)
fr forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BitTrie a
xr)
instance Semigroup a => Semigroup (BitTrie a) where
<> :: BitTrie a -> BitTrie a -> BitTrie a
(<>) = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Semigroup a => a -> a -> a
(<>)
instance Monoid a => Monoid (BitTrie a) where
mempty :: BitTrie a
mempty = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
mappend :: BitTrie a -> BitTrie a -> BitTrie a
mappend = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Monoid a => a -> a -> a
mappend
instance Functor IntTrie where
fmap :: forall a b. (a -> b) -> IntTrie a -> IntTrie b
fmap a -> b
f ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) = forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
neg) (a -> b
f a
z) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
pos)
instance Applicative IntTrie where
pure :: forall a. a -> IntTrie a
pure a
x = forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x) a
x (forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
IntTrie BitTrie (a -> b)
fneg a -> b
fz BitTrie (a -> b)
fpos <*> :: forall a b. IntTrie (a -> b) -> IntTrie a -> IntTrie b
<*> IntTrie BitTrie a
xneg a
xz BitTrie a
xpos =
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (BitTrie (a -> b)
fneg forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BitTrie a
xneg) (a -> b
fz a
xz) (BitTrie (a -> b)
fpos forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BitTrie a
xpos)
instance Semigroup a => Semigroup (IntTrie a) where
<> :: IntTrie a -> IntTrie a -> IntTrie a
(<>) = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Semigroup a => a -> a -> a
(<>)
instance Monoid a => Monoid (IntTrie a) where
mempty :: IntTrie a
mempty = forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a. Monoid a => a
mempty
mappend :: IntTrie a -> IntTrie a -> IntTrie a
mappend = forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 forall a. Monoid a => a -> a -> a
mappend
apply :: (Ord b, Num b, Bits b) => IntTrie a -> b -> a
apply :: forall b a. (Ord b, Num b, Bits b) => IntTrie a -> b -> a
apply (IntTrie BitTrie a
neg a
z BitTrie a
pos) b
x =
case forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
Ordering
LT -> forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
neg (-b
x)
Ordering
EQ -> a
z
Ordering
GT -> forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
pos b
x
applyPositive :: (Num b, Bits b) => BitTrie a -> b -> a
applyPositive :: forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive (BitTrie a
one BitTrie a
even BitTrie a
odd) b
x
| b
x forall a. Eq a => a -> a -> Bool
== b
1 = a
one
| forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
odd (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
| Bool
otherwise = forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
even (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
identity :: (Num a, Bits a) => IntTrie a
identity :: forall a. (Num a, Bits a) => IntTrie a
identity = forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap forall a. Num a => a -> a
negate forall a. (Num a, Bits a) => BitTrie a
identityPositive) a
0 forall a. (Num a, Bits a) => BitTrie a
identityPositive
identityPositive :: (Num a, Bits a) => BitTrie a
identityPositive :: forall a. (Num a, Bits a) => BitTrie a
identityPositive = BitTrie a
go
where
go :: BitTrie a
go = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
1 (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (forall a. Bits a => a -> Int -> a
`shiftL` Int
1) BitTrie a
go) (forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
n -> (a
n forall a. Bits a => a -> Int -> a
`shiftL` Int
1) forall a. Bits a => a -> a -> a
.|. a
1) BitTrie a
go)
modify :: (Ord b, Num b, Bits b) => b -> (a -> a) -> IntTrie a -> IntTrie a
modify :: forall b a.
(Ord b, Num b, Bits b) =>
b -> (a -> a) -> IntTrie a -> IntTrie a
modify b
x a -> a
f ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) =
case forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
Ordering
LT -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive (-b
x) a -> a
f BitTrie a
neg) a
z BitTrie a
pos
Ordering
EQ -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg (a -> a
f a
z) BitTrie a
pos
Ordering
GT -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg a
z (forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive b
x a -> a
f BitTrie a
pos)
modifyPositive :: (Num b, Bits b) => b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive :: forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive b
x a -> a
f ~(BitTrie a
one BitTrie a
even BitTrie a
odd)
| b
x forall a. Eq a => a -> a -> Bool
== b
1 = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> a
f a
one) BitTrie a
even BitTrie a
odd
| forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one BitTrie a
even (forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
odd)
| Bool
otherwise = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one (forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
even) BitTrie a
odd
modify' :: (Ord b, Num b, Bits b) => b -> (a -> a) -> IntTrie a -> IntTrie a
modify' :: forall b a.
(Ord b, Num b, Bits b) =>
b -> (a -> a) -> IntTrie a -> IntTrie a
modify' b
x a -> a
f (IntTrie BitTrie a
neg a
z BitTrie a
pos) =
case forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
Ordering
LT -> (forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie forall a b. (a -> b) -> a -> b
$! forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' (-b
x) a -> a
f BitTrie a
neg) a
z BitTrie a
pos
Ordering
EQ -> (forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg forall a b. (a -> b) -> a -> b
$! a -> a
f a
z) BitTrie a
pos
Ordering
GT -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg a
z forall a b. (a -> b) -> a -> b
$! forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' b
x a -> a
f BitTrie a
pos
modifyPositive' :: (Num b, Bits b) => b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' :: forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' b
x a -> a
f (BitTrie a
one BitTrie a
even BitTrie a
odd)
| b
x forall a. Eq a => a -> a -> Bool
== b
1 = (forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie forall a b. (a -> b) -> a -> b
$! a -> a
f a
one) BitTrie a
even BitTrie a
odd
| forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one BitTrie a
even forall a b. (a -> b) -> a -> b
$! forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
odd
| Bool
otherwise = (forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one forall a b. (a -> b) -> a -> b
$! forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' (b
x forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
even) BitTrie a
odd
overwrite :: (Ord b, Num b, Bits b) => b -> a -> IntTrie a -> IntTrie a
overwrite :: forall b a.
(Ord b, Num b, Bits b) =>
b -> a -> IntTrie a -> IntTrie a
overwrite b
i a
x = forall b a.
(Ord b, Num b, Bits b) =>
b -> (a -> a) -> IntTrie a -> IntTrie a
modify b
i (forall a b. a -> b -> a
const a
x)
mirror :: IntTrie a -> IntTrie a
mirror :: forall a. IntTrie a -> IntTrie a
mirror ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) = forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
pos a
z BitTrie a
neg
modifyAscList :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> IntTrie a -> IntTrie a
modifyAscList :: forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> IntTrie a -> IntTrie a
modifyAscList [(b, a -> a)]
ifs ~t :: IntTrie a
t@(IntTrie BitTrie a
neg a
z BitTrie a
pos) =
case forall a. (a -> Bool) -> [a] -> ([a], [a])
break ((forall a. Ord a => a -> a -> Bool
>= b
0) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a, b) -> a
fst) [(b, a -> a)]
ifs of
([], []) -> IntTrie a
t
([(b, a -> a)]
nifs, (b
0, a -> a
f):[(b, a -> a)]
pifs) -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall {a}. [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListNegative [(b, a -> a)]
nifs BitTrie a
neg) (a -> a
f a
z)
(forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [(b, a -> a)]
pifs BitTrie a
pos)
([(b, a -> a)]
nifs, [(b, a -> a)]
pifs) -> forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (forall {a}. [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListNegative [(b, a -> a)]
nifs BitTrie a
neg) a
z
(forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [(b, a -> a)]
pifs BitTrie a
pos)
where modifyAscListNegative :: [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListNegative = forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall a. Num a => a -> a
negate) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse
modifyDescList :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> IntTrie a -> IntTrie a
modifyDescList :: forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> IntTrie a -> IntTrie a
modifyDescList [(b, a -> a)]
ifs = forall a. IntTrie a -> IntTrie a
mirror forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> IntTrie a -> IntTrie a
modifyAscList (forall a b. (a -> b) -> [a] -> [b]
map (forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first forall a. Num a => a -> a
negate) [(b, a -> a)]
ifs) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IntTrie a -> IntTrie a
mirror
modifyAscListPositive :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive :: forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [] BitTrie a
t = BitTrie a
t
modifyAscListPositive ((b
0, a -> a
_):[(b, a -> a)]
_) BitTrie a
_ =
forall a. HasCallStack => [Char] -> a
error [Char]
"modifyAscList: expected strictly monotonic indices"
modifyAscListPositive ifs :: [(b, a -> a)]
ifs@((b
i, a -> a
f):[(b, a -> a)]
_) ~(BitTrie a
one BitTrie a
even BitTrie a
odd) = forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one' BitTrie a
even' BitTrie a
odd' where
(a
one', [(b, a -> a)]
ifs') = if b
i forall a. Eq a => a -> a -> Bool
== b
1 then (a -> a
f a
one, forall a. [a] -> [a]
tail [(b, a -> a)]
ifs) else (a
one, [(b, a -> a)]
ifs)
even' :: BitTrie a
even' = forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [(b, a -> a)]
ifsEven BitTrie a
even
odd' :: BitTrie a
odd' = forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [(b, a -> a)]
ifsOdd BitTrie a
odd
([(b, a -> a)]
ifsOdd, [(b, a -> a)]
ifsEven) = forall {t} {b}. (t -> b) -> (t, t) -> (b, b)
both (forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a -> b) -> a -> b
$ forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (forall a. Bits a => a -> Int -> a
`shiftR` Int
1)) forall a b. (a -> b) -> a -> b
$ forall b a.
(Num b, Bits b) =>
[(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices [(b, a -> a)]
ifs'
both :: (t -> b) -> (t, t) -> (b, b)
both t -> b
f (t
x, t
y) = (t -> b
f t
x, t -> b
f t
y)
partitionIndices :: (Num b, Bits b) => [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices :: forall b a.
(Num b, Bits b) =>
[(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices [] = ([], [])
partitionIndices [(b, a -> a)
x] = if forall a. Bits a => a -> Int -> Bool
testBit (forall a b. (a, b) -> a
fst (b, a -> a)
x) Int
0 then ([(b, a -> a)
x], []) else ([], [(b, a -> a)
x])
partitionIndices ((b, a -> a)
x:xs :: [(b, a -> a)]
xs@((b, a -> a)
y:[(b, a -> a)]
_)) = case forall a. Bits a => a -> Int -> Bool
testBit (forall a b. (a, b) -> a
fst (b, a -> a)
x) Int
0 of
Bool
False -> (if forall a. Bits a => a -> Int -> Bool
testBit (forall a b. (a, b) -> a
fst (b, a -> a)
y) Int
0 then [(b, a -> a)]
odd else forall {a}. (b, a -> a)
padforall a. a -> [a] -> [a]
:[(b, a -> a)]
odd, (b, a -> a)
xforall a. a -> [a] -> [a]
:[(b, a -> a)]
even)
Bool
True -> ((b, a -> a)
xforall a. a -> [a] -> [a]
:[(b, a -> a)]
odd, if forall a. Bits a => a -> Int -> Bool
testBit (forall a b. (a, b) -> a
fst (b, a -> a)
y) Int
0 then forall {a}. (b, a -> a)
padforall a. a -> [a] -> [a]
:[(b, a -> a)]
even else [(b, a -> a)]
even)
where ~([(b, a -> a)]
odd, [(b, a -> a)]
even) = forall b a.
(Num b, Bits b) =>
[(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices [(b, a -> a)]
xs
pad :: (b, a -> a)
pad = (forall a b. (a, b) -> a
fst (b, a -> a)
y forall a. Num a => a -> a -> a
- b
1, forall a. a -> a
id)