-------------------------------------------------------------
-- |
-- Module        : Data.IntTrie
-- Copyright     : (c) Luke Palmer 2010
-- License       : BSD3
--
-- Maintainer    : Luke Palmer <lrpalmer@gmail.com>
-- Stability     : experimental
-- Portability   : Haskell 2010
--
-- Provides a minimal infinite, lazy trie for integral types.
-- It intentionally leaves out ideas such as delete and
-- emptiness so that it can be used lazily, eg. as the target
-- of an infinite foldr.  Essentially its purpose is to be an
-- efficient implementation of a function from integral type,
-- given point-at-a-time modifications.
-------------------------------------------------------------

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(..))

-- | A trie from integers to values of type a. 
-- 
-- Semantics: [[IntTrie a]] = Integer -> a
data IntTrie a = IntTrie (BitTrie a) a (BitTrie a)  -- negative, 0, positive

data BitTrie a = BitTrie a (BitTrie a) (BitTrie a)

instance Functor BitTrie where
    fmap :: (a -> b) -> BitTrie a -> BitTrie b
fmap a -> b
f ~(BitTrie a
x BitTrie a
l BitTrie a
r) = b -> BitTrie b -> BitTrie b -> BitTrie b
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> b
f a
x) ((a -> b) -> BitTrie a -> BitTrie b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
l) ((a -> b) -> BitTrie a -> BitTrie b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
r)

instance Applicative BitTrie where
    pure :: a -> BitTrie a
pure a
x = (BitTrie a -> BitTrie a) -> BitTrie a
forall a. (a -> a) -> a
fix (\BitTrie a
g -> a -> BitTrie a -> BitTrie a -> BitTrie a
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) <*> :: BitTrie (a -> b) -> BitTrie a -> BitTrie b
<*> ~(BitTrie a
x BitTrie a
xl BitTrie a
xr) = b -> BitTrie b -> BitTrie b -> BitTrie b
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> b
f a
x) (BitTrie (a -> b)
fl BitTrie (a -> b) -> BitTrie a -> BitTrie b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> BitTrie a
xl) (BitTrie (a -> b)
fr BitTrie (a -> b) -> BitTrie a -> BitTrie b
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
(<>) = (a -> a -> a) -> BitTrie a -> BitTrie a -> BitTrie a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

instance Monoid a => Monoid (BitTrie a) where
    mempty :: BitTrie a
mempty = a -> BitTrie a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
    mappend :: BitTrie a -> BitTrie a -> BitTrie a
mappend = (a -> a -> a) -> BitTrie a -> BitTrie a -> BitTrie a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Monoid a => a -> a -> a
mappend

instance Functor IntTrie where
    fmap :: (a -> b) -> IntTrie a -> IntTrie b
fmap a -> b
f ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) = BitTrie b -> b -> BitTrie b -> IntTrie b
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie ((a -> b) -> BitTrie a -> BitTrie b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
neg) (a -> b
f a
z) ((a -> b) -> BitTrie a -> BitTrie b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> b
f BitTrie a
pos)

instance Applicative IntTrie where
    pure :: a -> IntTrie a
pure a
x = BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (a -> BitTrie a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x) a
x (a -> BitTrie a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
x)
    IntTrie BitTrie (a -> b)
fneg a -> b
fz BitTrie (a -> b)
fpos <*> :: IntTrie (a -> b) -> IntTrie a -> IntTrie b
<*> IntTrie BitTrie a
xneg a
xz BitTrie a
xpos = 
        BitTrie b -> b -> BitTrie b -> IntTrie b
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (BitTrie (a -> b)
fneg BitTrie (a -> b) -> BitTrie a -> BitTrie b
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 BitTrie (a -> b) -> BitTrie a -> BitTrie b
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
(<>) = (a -> a -> a) -> IntTrie a -> IntTrie a -> IntTrie a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Semigroup a => a -> a -> a
(<>)

instance Monoid a => Monoid (IntTrie a) where
    mempty :: IntTrie a
mempty = a -> IntTrie a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
forall a. Monoid a => a
mempty
    mappend :: IntTrie a -> IntTrie a -> IntTrie a
mappend = (a -> a -> a) -> IntTrie a -> IntTrie a -> IntTrie a
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 a -> a -> a
forall a. Monoid a => a -> a -> a
mappend

-- | Apply the trie to an argument.  This is the semantic map.
apply :: (Ord b, Num b, Bits b) => IntTrie a -> b -> a
apply :: IntTrie a -> b -> a
apply (IntTrie BitTrie a
neg a
z BitTrie a
pos) b
x =
    case b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
        Ordering
LT -> BitTrie a -> b -> a
forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
neg (-b
x)
        Ordering
EQ -> a
z
        Ordering
GT -> BitTrie a -> b -> a
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 :: BitTrie a -> b -> a
applyPositive (BitTrie a
one BitTrie a
even BitTrie a
odd) b
x
    | b
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
1 = a
one
    | b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = BitTrie a -> b -> a
forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
odd (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)
    | Bool
otherwise   = BitTrie a -> b -> a
forall b a. (Num b, Bits b) => BitTrie a -> b -> a
applyPositive BitTrie a
even (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)

-- | The identity trie.  
--
-- > apply identity = id
identity :: (Num a, Bits a) => IntTrie a
identity :: IntTrie a
identity = BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie ((a -> a) -> BitTrie a -> BitTrie a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap a -> a
forall a. Num a => a -> a
negate BitTrie a
forall a. (Num a, Bits a) => BitTrie a
identityPositive) a
0 BitTrie a
forall a. (Num a, Bits a) => BitTrie a
identityPositive

identityPositive :: (Num a, Bits a) => BitTrie a
identityPositive :: BitTrie a
identityPositive = BitTrie a
go
    where
    go :: BitTrie a
go = a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
1 ((a -> a) -> BitTrie a -> BitTrie a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) BitTrie a
go) ((a -> a) -> BitTrie a -> BitTrie a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (\a
n -> (a
n a -> Int -> a
forall a. Bits a => a -> Int -> a
`shiftL` Int
1) a -> a -> a
forall a. Bits a => a -> a -> a
.|. a
1) BitTrie a
go)

-- | Modify the function at one point
--
-- > apply (modify x f t) i | i == x = f (apply t i)
-- >                        | otherwise = apply t i
modify :: (Ord b, Num b, Bits b) => b -> (a -> a) -> IntTrie a -> IntTrie a
modify :: b -> (a -> a) -> IntTrie a -> IntTrie a
modify b
x a -> a
f ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) =
    case b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
        Ordering
LT -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (b -> (a -> a) -> BitTrie a -> BitTrie a
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 -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg (a -> a
f a
z) BitTrie a
pos
        Ordering
GT -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg a
z (b -> (a -> a) -> BitTrie a -> BitTrie a
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 :: b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive b
x a -> a
f ~(BitTrie a
one BitTrie a
even BitTrie a
odd)
    | b
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
1      = a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> a
f a
one) BitTrie a
even BitTrie a
odd
    | b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one BitTrie a
even (b -> (a -> a) -> BitTrie a -> BitTrie a
forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
odd)
    | Bool
otherwise   = a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one (b -> (a -> a) -> BitTrie a -> BitTrie a
forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
even) BitTrie a
odd


-- | Modify the function at one point (strict version)
modify' :: (Ord b, Num b, Bits b) => b -> (a -> a) -> IntTrie a -> IntTrie a
modify' :: b -> (a -> a) -> IntTrie a -> IntTrie a
modify' b
x a -> a
f (IntTrie BitTrie a
neg a
z BitTrie a
pos) =
    case b -> b -> Ordering
forall a. Ord a => a -> a -> Ordering
compare b
x b
0 of
        Ordering
LT -> (BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie (BitTrie a -> a -> BitTrie a -> IntTrie a)
-> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a b. (a -> b) -> a -> b
$! b -> (a -> a) -> BitTrie a -> BitTrie a
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 -> (BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg (a -> BitTrie a -> IntTrie a) -> a -> BitTrie a -> IntTrie a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
z) BitTrie a
pos
        Ordering
GT -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
neg a
z (BitTrie a -> IntTrie a) -> BitTrie a -> IntTrie a
forall a b. (a -> b) -> a -> b
$! b -> (a -> a) -> BitTrie a -> BitTrie a
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' :: b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' b
x a -> a
f (BitTrie a
one BitTrie a
even BitTrie a
odd)
    | b
x b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
1      = (a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie (a -> BitTrie a -> BitTrie a -> BitTrie a)
-> a -> BitTrie a -> BitTrie a -> BitTrie a
forall a b. (a -> b) -> a -> b
$! a -> a
f a
one) BitTrie a
even BitTrie a
odd
    | b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit b
x Int
0 = a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one BitTrie a
even (BitTrie a -> BitTrie a) -> BitTrie a -> BitTrie a
forall a b. (a -> b) -> a -> b
$! b -> (a -> a) -> BitTrie a -> BitTrie a
forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
odd
    | Bool
otherwise   = (a -> BitTrie a -> BitTrie a -> BitTrie a
forall a. a -> BitTrie a -> BitTrie a -> BitTrie a
BitTrie a
one (BitTrie a -> BitTrie a -> BitTrie a)
-> BitTrie a -> BitTrie a -> BitTrie a
forall a b. (a -> b) -> a -> b
$! b -> (a -> a) -> BitTrie a -> BitTrie a
forall b a.
(Num b, Bits b) =>
b -> (a -> a) -> BitTrie a -> BitTrie a
modifyPositive' (b
x b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1) a -> a
f BitTrie a
even) BitTrie a
odd


-- | Overwrite the function at one point
--
-- > overwrite i x = modify i (const x)
overwrite :: (Ord b, Num b, Bits b) => b -> a -> IntTrie a -> IntTrie a
overwrite :: b -> a -> IntTrie a -> IntTrie a
overwrite b
i a
x = b -> (a -> a) -> IntTrie a -> IntTrie a
forall b a.
(Ord b, Num b, Bits b) =>
b -> (a -> a) -> IntTrie a -> IntTrie a
modify b
i (a -> a -> a
forall a b. a -> b -> a
const a
x)


-- | Negate the domain of the function
--
-- > apply (mirror t) i = apply t (-i)
-- > mirror . mirror = id
mirror :: IntTrie a -> IntTrie a
mirror :: IntTrie a -> IntTrie a
mirror ~(IntTrie BitTrie a
neg a
z BitTrie a
pos) = BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie BitTrie a
pos a
z BitTrie a
neg


-- | Modify the function at a (potentially infinite) list of points in ascending order
--
-- > modifyAscList [(i0, f0)..(iN, fN)] = modify i0 f0 . ... . modify iN fN
modifyAscList :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> IntTrie a -> IntTrie a
modifyAscList :: [(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 ((b, a -> a) -> Bool)
-> [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
break ((b -> b -> Bool
forall a. Ord a => a -> a -> Bool
>= b
0) (b -> Bool) -> ((b, a -> a) -> b) -> (b, a -> a) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (b, a -> a) -> b
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) -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie ([(b, a -> a)] -> BitTrie a -> BitTrie a
forall a. [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListNegative [(b, a -> a)]
nifs BitTrie a
neg) (a -> a
f a
z)
                                       ([(b, a -> a)] -> BitTrie a -> BitTrie a
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)        -> BitTrie a -> a -> BitTrie a -> IntTrie a
forall a. BitTrie a -> a -> BitTrie a -> IntTrie a
IntTrie ([(b, a -> a)] -> BitTrie a -> BitTrie a
forall a. [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListNegative [(b, a -> a)]
nifs BitTrie a
neg) a
z
                                       ([(b, a -> a)] -> BitTrie a -> BitTrie a
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 = [(b, a -> a)] -> BitTrie a -> BitTrie a
forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive ([(b, a -> a)] -> BitTrie a -> BitTrie a)
-> ([(b, a -> a)] -> [(b, a -> a)])
-> [(b, a -> a)]
-> BitTrie a
-> BitTrie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((b, a -> a) -> (b, a -> a)) -> [(b, a -> a)] -> [(b, a -> a)]
forall a b. (a -> b) -> [a] -> [b]
map ((b -> b) -> (b, a -> a) -> (b, a -> a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first b -> b
forall a. Num a => a -> a
negate) ([(b, a -> a)] -> [(b, a -> a)])
-> ([(b, a -> a)] -> [(b, a -> a)])
-> [(b, a -> a)]
-> [(b, a -> a)]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(b, a -> a)] -> [(b, a -> a)]
forall a. [a] -> [a]
reverse

-- | Modify the function at a (potentially infinite) list of points in descending order
modifyDescList :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> IntTrie a -> IntTrie a
modifyDescList :: [(b, a -> a)] -> IntTrie a -> IntTrie a
modifyDescList [(b, a -> a)]
ifs = IntTrie a -> IntTrie a
forall a. IntTrie a -> IntTrie a
mirror (IntTrie a -> IntTrie a)
-> (IntTrie a -> IntTrie a) -> IntTrie a -> IntTrie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [(b, a -> a)] -> IntTrie a -> IntTrie a
forall b a.
(Ord b, Num b, Bits b) =>
[(b, a -> a)] -> IntTrie a -> IntTrie a
modifyAscList (((b, a -> a) -> (b, a -> a)) -> [(b, a -> a)] -> [(b, a -> a)]
forall a b. (a -> b) -> [a] -> [b]
map ((b -> b) -> (b, a -> a) -> (b, a -> a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first b -> b
forall a. Num a => a -> a
negate) [(b, a -> a)]
ifs) (IntTrie a -> IntTrie a)
-> (IntTrie a -> IntTrie a) -> IntTrie a -> IntTrie a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. IntTrie a -> IntTrie a
forall a. IntTrie a -> IntTrie a
mirror

modifyAscListPositive :: (Ord b, Num b, Bits b) => [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive :: [(b, a -> a)] -> BitTrie a -> BitTrie a
modifyAscListPositive [] BitTrie a
t = BitTrie a
t
modifyAscListPositive ((b
0, a -> a
_):[(b, a -> a)]
_) BitTrie a
_ =
    [Char] -> 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) = a -> BitTrie a -> BitTrie a -> BitTrie a
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 b -> b -> Bool
forall a. Eq a => a -> a -> Bool
== b
1 then (a -> a
f a
one, [(b, a -> a)] -> [(b, a -> a)]
forall a. [a] -> [a]
tail [(b, a -> a)]
ifs) else (a
one, [(b, a -> a)]
ifs)
    even' :: BitTrie a
even'             = [(b, a -> a)] -> BitTrie a -> BitTrie a
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'              = [(b, a -> a)] -> BitTrie a -> BitTrie a
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) = ([(b, a -> a)] -> [(b, a -> a)])
-> ([(b, a -> a)], [(b, a -> a)]) -> ([(b, a -> a)], [(b, a -> a)])
forall t b. (t -> b) -> (t, t) -> (b, b)
both (((b, a -> a) -> (b, a -> a)) -> [(b, a -> a)] -> [(b, a -> a)]
forall a b. (a -> b) -> [a] -> [b]
map (((b, a -> a) -> (b, a -> a)) -> [(b, a -> a)] -> [(b, a -> a)])
-> ((b, a -> a) -> (b, a -> a)) -> [(b, a -> a)] -> [(b, a -> a)]
forall a b. (a -> b) -> a -> b
$ (b -> b) -> (b, a -> a) -> (b, a -> a)
forall (a :: * -> * -> *) b c d.
Arrow a =>
a b c -> a (b, d) (c, d)
first (b -> Int -> b
forall a. Bits a => a -> Int -> a
`shiftR` Int
1)) (([(b, a -> a)], [(b, a -> a)]) -> ([(b, a -> a)], [(b, a -> a)]))
-> ([(b, a -> a)], [(b, a -> a)]) -> ([(b, a -> a)], [(b, a -> a)])
forall a b. (a -> b) -> a -> b
$ [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
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)

-- Like `partition (flip testBit 0 . fst)`, except that this version addresses the
-- problem of infinite lists of only odd or only even indices by injecting an `id`
-- into the other result list wherever there are two evens or two odds in a row.
-- This allows `modifyAscListPositive` to return a value as soon as the next index is
-- higher than the current location in the trie instead of scanning for the end of
-- the list, which for infinite lists may never be reached.
partitionIndices :: (Num b, Bits b) => [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices :: [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
partitionIndices []           = ([], [])
partitionIndices [(b, a -> a)
x]          = if b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit ((b, a -> a) -> b
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 b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit ((b, a -> a) -> b
forall a b. (a, b) -> a
fst (b, a -> a)
x) Int
0 of
    Bool
False -> (if b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit ((b, a -> a) -> b
forall a b. (a, b) -> a
fst (b, a -> a)
y) Int
0 then [(b, a -> a)]
odd else (b, a -> a)
forall a. (b, a -> a)
pad(b, a -> a) -> [(b, a -> a)] -> [(b, a -> a)]
forall a. a -> [a] -> [a]
:[(b, a -> a)]
odd, (b, a -> a)
x(b, a -> a) -> [(b, a -> a)] -> [(b, a -> a)]
forall a. a -> [a] -> [a]
:[(b, a -> a)]
even)
    Bool
True  -> ((b, a -> a)
x(b, a -> a) -> [(b, a -> a)] -> [(b, a -> a)]
forall a. a -> [a] -> [a]
:[(b, a -> a)]
odd, if b -> Int -> Bool
forall a. Bits a => a -> Int -> Bool
testBit ((b, a -> a) -> b
forall a b. (a, b) -> a
fst (b, a -> a)
y) Int
0 then (b, a -> a)
forall a. (b, a -> a)
pad(b, a -> a) -> [(b, a -> a)] -> [(b, a -> a)]
forall a. a -> [a] -> [a]
:[(b, a -> a)]
even else [(b, a -> a)]
even)
    where ~([(b, a -> a)]
odd, [(b, a -> a)]
even) = [(b, a -> a)] -> ([(b, a -> a)], [(b, a -> a)])
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 = ((b, a -> a) -> b
forall a b. (a, b) -> a
fst (b, a -> a)
y b -> b -> b
forall a. Num a => a -> a -> a
- b
1, a -> a
forall a. a -> a
id)