module Dict
( Dict
, empty, singleton, insert, update
, get, remove, member
, filter
, partition
, foldl, foldr, map
, union, intersect, diff
, keys, values
, toList, fromList
) where
{-| A dictionary mapping unique keys to values. The keys can be any comparable
type. This includes `Int`, `Float`, `Time`, `Char`, `String`, and tuples or
lists of comparable types.
Insert, remove, and query operations all take *O(log n)* time.
# Build
@docs empty, singleton, insert, update, remove
# Query
@docs member, get
# Combine
@docs union, intersect, diff
# Lists
@docs keys, values, toList, fromList
# Transform
@docs map, foldl, foldr, filter, partition
-}
import Basics (..)
import Maybe (..)
import List
import List (..)
import Native.Debug
import String
-- BBlack and NBlack should only be used during the deletion
-- algorithm. Any other occurrence is a bug and should fail an assert.
type NColor
= Red
| Black
| BBlack -- Double Black, counts as 2 blacks for the invariant
| NBlack -- Negative Black, counts as -1 blacks for the invariant
showNColor : NColor -> String
showNColor c =
case c of
Red -> "Red"
Black -> "Black"
BBlack -> "BBlack"
NBlack -> "NBlack"
type LeafColor
= LBlack
| LBBlack -- Double Black, counts as 2
showLColor : LeafColor -> String
showLColor color =
case color of
LBlack -> "LBlack"
LBBlack -> "LBBlack"
type Dict k v
= RBNode NColor k v (Dict k v) (Dict k v)
| RBEmpty LeafColor
{-| Create an empty dictionary. -}
empty : Dict comparable v
empty = RBEmpty LBlack
min : Dict k v -> (k,v)
min dict =
case dict of
RBNode _ key value (RBEmpty LBlack) _ ->
(key, value)
RBNode _ _ _ left _ ->
min left
RBEmpty LBlack ->
Native.Debug.crash "(min Empty) is not defined"
max : Dict k v -> (k, v)
max dict =
case dict of
RBNode _ key value _ (RBEmpty _) ->
(key, value)
RBNode _ _ _ _ right ->
max right
RBEmpty _ ->
Native.Debug.crash "(max Empty) is not defined"
{-| Get the value associated with a key. If the key is not found, return
`Nothing`. This is useful when you are not sure if a key will be in the
dictionary.
animals = fromList [ ("Tom", Cat), ("Jerry", Mouse) ]
get "Tom" animals == Just Cat
get "Mouse" animals == Just Mouse
get "Spike" animals == Nothing
The `(?)` operator from the `Maybe` library makes it easy to give a default
value.
-}
get : comparable -> Dict comparable v -> Maybe v
get targetKey dict =
case dict of
RBEmpty LBlack ->
Nothing
RBNode _ key value left right ->
case compare targetKey key of
LT -> get targetKey left
EQ -> Just value
GT -> get targetKey right
{-| Determine if a key is in a dictionary. -}
member : comparable -> Dict comparable v -> Bool
member key dict =
case get key dict of
Just _ -> True
Nothing -> False
ensureBlackRoot : Dict k v -> Dict k v
ensureBlackRoot dict =
case dict of
RBNode Red key value left right ->
RBNode Black key value left right
RBNode Black _ _ _ _ ->
dict
RBEmpty LBlack ->
dict
{-| Insert a key-value pair into a dictionary. Replaces value when there is
a collision. -}
insert : comparable -> v -> Dict comparable v -> Dict comparable v
insert key value dict =
update key (always (Just value)) dict
{-| Remove a key-value pair from a dictionary. If the key is not found,
no changes are made. -}
remove : comparable -> Dict comparable v -> Dict comparable v
remove key dict =
update key (always Nothing) dict
type Flag = Insert | Remove | Same
showFlag : Flag -> String
showFlag f = case f of
Insert -> "Insert"
Remove -> "Remove"
Same -> "Same"
{-| Update the value of a dictionary for a specific key with a given function. -}
update : comparable -> (Maybe v -> Maybe v) -> Dict comparable v -> Dict comparable v
update k alter dict =
let up dict =
case dict of
RBEmpty LBlack ->
case alter Nothing of
Nothing -> (Same, empty)
Just v -> (Insert, RBNode Red k v empty empty)
RBNode clr key value left right ->
case compare k key of
EQ ->
case alter (Just value) of
Nothing -> (Remove, rem clr left right)
Just newValue ->
(Same, RBNode clr key newValue left right)
LT ->
let (flag, newLeft) = up left in
case flag of
Same -> (Same, RBNode clr key value newLeft right)
Insert -> (Insert, balance clr key value newLeft right)
Remove -> (Remove, bubble clr key value newLeft right)
GT ->
let (flag, newRight) = up right in
case flag of
Same -> (Same, RBNode clr key value left newRight)
Insert -> (Insert, balance clr key value left newRight)
Remove -> (Remove, bubble clr key value left newRight)
(flag, updatedDict) = up dict
in
case flag of
Same -> updatedDict
Insert -> ensureBlackRoot updatedDict
Remove -> blacken updatedDict
{-| Create a dictionary with one key-value pair. -}
singleton : comparable -> v -> Dict comparable v
singleton key value =
insert key value (RBEmpty LBlack)
isBBlack : Dict k v -> Bool
isBBlack dict =
case dict of
RBNode BBlack _ _ _ _ -> True
RBEmpty LBBlack -> True
_ -> False
moreBlack : NColor -> NColor
moreBlack color =
case color of
Black -> BBlack
Red -> Black
NBlack -> Red
BBlack -> Native.Debug.crash "Can't make a double black node more black!"
lessBlack : NColor -> NColor
lessBlack color =
case color of
BBlack -> Black
Black -> Red
Red -> NBlack
NBlack -> Native.Debug.crash "Can't make a negative black node less black!"
lessBlackTree : Dict k v -> Dict k v
lessBlackTree dict =
case dict of
RBNode c k v l r -> RBNode (lessBlack c) k v l r
RBEmpty LBBlack -> RBEmpty LBlack
reportRemBug : String -> NColor -> String -> String -> a
reportRemBug msg c lgot rgot =
Native.Debug.crash <|
String.concat
[ "Internal red-black tree invariant violated, expected "
, msg, " and got ", showNColor c, "/", lgot, "/", rgot
, "\nPlease report this bug to <https://github.com/elm-lang/Elm/issues>"
]
-- Remove the top node from the tree, may leave behind BBlacks
rem : NColor -> Dict k v -> Dict k v -> Dict k v
rem c l r =
case (l, r) of
(RBEmpty _, RBEmpty _) ->
case c of
Red -> RBEmpty LBlack
Black -> RBEmpty LBBlack
(RBEmpty cl, RBNode cr k' v' l' r') ->
case (c, cl, cr) of
(Black, LBlack, Red) ->
RBNode Black k' v' l' r'
_ ->
reportRemBug "Black/LBlack/Red" c (showLColor cl) (showNColor cr)
(RBNode cl k' v' l' r', RBEmpty cr) ->
case (c, cl, cr) of
(Black, Red, LBlack) ->
RBNode Black k' v' l' r'
_ ->
reportRemBug "Black/Red/LBlack" c (showNColor cl) (showLColor cr)
-- l and r are both RBNodes
(RBNode cl kl vl ll rl, RBNode cr kr vr lr rr) ->
let l = RBNode cl kl vl ll rl
r = RBNode cr kr vr lr rr
(k, v) = max l
l' = remove_max cl kl vl ll rl
in
bubble c k v l' r
-- Kills a BBlack or moves it upward, may leave behind NBlack
bubble : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
bubble c k v l r =
if isBBlack l || isBBlack r
then balance (moreBlack c) k v (lessBlackTree l) (lessBlackTree r)
else RBNode c k v l r
-- Removes rightmost node, may leave root as BBlack
remove_max : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
remove_max c k v l r =
case r of
RBEmpty _ ->
rem c l r
RBNode cr kr vr lr rr ->
bubble c k v l (remove_max cr kr vr lr rr)
-- generalized tree balancing act
balance : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
balance c k v l r =
balance_node (RBNode c k v l r)
blackish : Dict k v -> Bool
blackish t =
case t of
RBNode c _ _ _ _ -> c == Black || c == BBlack
RBEmpty _ -> True
balance_node : Dict k v -> Dict k v
balance_node t =
let assemble col xk xv yk yv zk zv a b c d =
RBNode (lessBlack col) yk yv (RBNode Black xk xv a b) (RBNode Black zk zv c d)
in
if blackish t
then case t of
RBNode col zk zv (RBNode Red yk yv (RBNode Red xk xv a b) c) d ->
assemble col xk xv yk yv zk zv a b c d
RBNode col zk zv (RBNode Red xk xv a (RBNode Red yk yv b c)) d ->
assemble col xk xv yk yv zk zv a b c d
RBNode col xk xv a (RBNode Red zk zv (RBNode Red yk yv b c) d) ->
assemble col xk xv yk yv zk zv a b c d
RBNode col xk xv a (RBNode Red yk yv b (RBNode Red zk zv c d)) ->
assemble col xk xv yk yv zk zv a b c d
RBNode BBlack xk xv a (RBNode NBlack zk zv (RBNode Black yk yv b c) d) ->
case d of
(RBNode Black _ _ _ _) ->
RBNode Black yk yv (RBNode Black xk xv a b) (balance Black zk zv c (redden d))
_ -> t
RBNode BBlack zk zv (RBNode NBlack xk xv a (RBNode Black yk yv b c)) d ->
case a of
(RBNode Black _ _ _ _) ->
RBNode Black yk yv (balance Black xk xv (redden a) b) (RBNode Black zk zv c d)
_ -> t
_ -> t
else t
-- make the top node black
blacken : Dict k v -> Dict k v
blacken t =
case t of
RBEmpty _ -> RBEmpty LBlack
RBNode _ k v l r -> RBNode Black k v l r
-- make the top node red
redden : Dict k v -> Dict k v
redden t =
case t of
RBEmpty _ -> Native.Debug.crash "can't make a Leaf red"
RBNode _ k v l r -> RBNode Red k v l r
{-| Apply a function to all values in a dictionary. -}
map : (comparable -> a -> b) -> Dict comparable a -> Dict comparable b
map f dict =
case dict of
RBEmpty LBlack ->
RBEmpty LBlack
RBNode clr key value left right ->
RBNode clr key (f key value) (map f left) (map f right)
{-| Fold over the key-value pairs in a dictionary, in order from lowest
key to highest key. -}
foldl : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
foldl f acc dict =
case dict of
RBEmpty LBlack -> acc
RBNode _ key value left right ->
foldl f (f key value (foldl f acc left)) right
{-| Fold over the key-value pairs in a dictionary, in order from highest
key to lowest key. -}
foldr : (comparable -> v -> b -> b) -> b -> Dict comparable v -> b
foldr f acc t =
case t of
RBEmpty LBlack -> acc
RBNode _ key value left right ->
foldr f (f key value (foldr f acc right)) left
{-| Combine two dictionaries. If there is a collision, preference is given
to the first dictionary. -}
union : Dict comparable v -> Dict comparable v -> Dict comparable v
union t1 t2 =
foldl insert t2 t1
{-| Keep a key-value pair when its key appears in the second dictionary.
Preference is given to values in the first dictionary. -}
intersect : Dict comparable v -> Dict comparable v -> Dict comparable v
intersect t1 t2 =
filter (\k _ -> k `member` t2) t1
{-| Keep a key-value pair when its key does not appear in the second dictionary.
Preference is given to the first dictionary. -}
diff : Dict comparable v -> Dict comparable v -> Dict comparable v
diff t1 t2 =
foldl (\k v t -> remove k t) t1 t2
{-| Get all of the keys in a dictionary. -}
keys : Dict comparable v -> List comparable
keys dict =
foldr (\key value keyList -> key :: keyList) [] dict
{-| Get all of the values in a dictionary. -}
values : Dict comparable v -> List v
values dict =
foldr (\key value valueList -> value :: valueList) [] dict
{-| Convert a dictionary into an association list of key-value pairs. -}
toList : Dict comparable v -> List (comparable,v)
toList dict =
foldr (\key value list -> (key,value) :: list) [] dict
{-| Convert an association list into a dictionary. -}
fromList : List (comparable,v) -> Dict comparable v
fromList assocs =
List.foldl (\(key,value) dict -> insert key value dict) empty assocs
{-| Keep a key-value pair when it satisfies a predicate. -}
filter : (comparable -> v -> Bool) -> Dict comparable v -> Dict comparable v
filter predicate dictionary =
let add key value dict =
if predicate key value
then insert key value dict
else dict
in
foldl add empty dictionary
{-| Partition a dictionary according to a predicate. The first dictionary
contains all key-value pairs which satisfy the predicate, and the second
contains the rest.
-}
partition : (comparable -> v -> Bool) -> Dict comparable v -> (Dict comparable v, Dict comparable v)
partition predicate dict =
let add key value (t1, t2) =
if predicate key value
then (insert key value t1, t2)
else (t1, insert key value t2)
in
foldl add (empty, empty) dict