Permalink
662 lines (480 sloc) 16 KB
module Dict exposing
( Dict
, empty, singleton, insert, update
, isEmpty, get, remove, member, size
, filter
, partition
, foldl, foldr, map
, union, intersect, diff, merge
, keys, values
, toList, fromList
)
{-| 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.
# Dictionaries
@docs Dict
# Build
@docs empty, singleton, insert, update, remove
# Query
@docs isEmpty, member, get, size
# Lists
@docs keys, values, toList, fromList
# Transform
@docs map, foldl, foldr, filter, partition
# Combine
@docs union, intersect, diff, merge
-}
import Basics exposing (..)
import Maybe exposing (..)
import List exposing (..)
import Native.Debug
import String
-- DICTIONARIES
-- 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
type LeafColor
= LBlack
| LBBlack -- Double Black, counts as 2
{-| A dictionary of keys and values. So a `(Dict String User)` is a dictionary
that lets you look up a `String` (such as user names) and find the associated
`User`.
-}
type Dict k v
= RBNode_elm_builtin NColor k v (Dict k v) (Dict k v)
| RBEmpty_elm_builtin LeafColor
{-| Create an empty dictionary. -}
empty : Dict k v
empty =
RBEmpty_elm_builtin LBlack
maxWithDefault : k -> v -> Dict k v -> (k, v)
maxWithDefault k v r =
case r of
RBEmpty_elm_builtin _ ->
(k, v)
RBNode_elm_builtin _ kr vr _ rr ->
maxWithDefault kr vr rr
{-| 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 "Jerry" animals == Just Mouse
get "Spike" animals == Nothing
-}
get : comparable -> Dict comparable v -> Maybe v
get targetKey dict =
case dict of
RBEmpty_elm_builtin _ ->
Nothing
RBNode_elm_builtin _ 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
{-| Determine the number of key-value pairs in the dictionary. -}
size : Dict k v -> Int
size dict =
sizeHelp 0 dict
sizeHelp : Int -> Dict k v -> Int
sizeHelp n dict =
case dict of
RBEmpty_elm_builtin _ ->
n
RBNode_elm_builtin _ _ _ left right ->
sizeHelp (sizeHelp (n+1) right) left
{-| Determine if a dictionary is empty.
isEmpty empty == True
-}
isEmpty : Dict k v -> Bool
isEmpty dict =
dict == empty
{- The actual pattern match here is somewhat lax. If it is given invalid input,
it will do the wrong thing. The expected behavior is:
red node => black node
black node => same
bblack node => xxx
nblack node => xxx
black leaf => same
bblack leaf => xxx
-}
ensureBlackRoot : Dict k v -> Dict k v
ensureBlackRoot dict =
case dict of
RBNode_elm_builtin Red key value left right ->
RBNode_elm_builtin Black key value left right
_ ->
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
{-| 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
-- expecting only black nodes, never double black nodes here
RBEmpty_elm_builtin _ ->
case alter Nothing of
Nothing ->
(Same, empty)
Just v ->
(Insert, RBNode_elm_builtin Red k v empty empty)
RBNode_elm_builtin 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_elm_builtin clr key newValue left right)
LT ->
let (flag, newLeft) = up left in
case flag of
Same ->
(Same, RBNode_elm_builtin 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_elm_builtin 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 empty
-- HELPERS
isBBlack : Dict k v -> Bool
isBBlack dict =
case dict of
RBNode_elm_builtin BBlack _ _ _ _ ->
True
RBEmpty_elm_builtin 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!"
{- The actual pattern match here is somewhat lax. If it is given invalid input,
it will do the wrong thing. The expected behavior is:
node => less black node
bblack leaf => black leaf
black leaf => xxx
-}
lessBlackTree : Dict k v -> Dict k v
lessBlackTree dict =
case dict of
RBNode_elm_builtin c k v l r ->
RBNode_elm_builtin (lessBlack c) k v l r
RBEmpty_elm_builtin _ ->
RBEmpty_elm_builtin 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 ", toString c, "/", lgot, "/", rgot
, "\nPlease report this bug to <https://github.com/elm-lang/core/issues>"
]
-- Remove the top node from the tree, may leave behind BBlacks
rem : NColor -> Dict k v -> Dict k v -> Dict k v
rem color left right =
case (left, right) of
(RBEmpty_elm_builtin _, RBEmpty_elm_builtin _) ->
case color of
Red ->
RBEmpty_elm_builtin LBlack
Black ->
RBEmpty_elm_builtin LBBlack
_ ->
Native.Debug.crash "cannot have bblack or nblack nodes at this point"
(RBEmpty_elm_builtin cl, RBNode_elm_builtin cr k v l r) ->
case (color, cl, cr) of
(Black, LBlack, Red) ->
RBNode_elm_builtin Black k v l r
_ ->
reportRemBug "Black/LBlack/Red" color (toString cl) (toString cr)
(RBNode_elm_builtin cl k v l r, RBEmpty_elm_builtin cr) ->
case (color, cl, cr) of
(Black, Red, LBlack) ->
RBNode_elm_builtin Black k v l r
_ ->
reportRemBug "Black/Red/LBlack" color (toString cl) (toString cr)
-- l and r are both RBNodes
(RBNode_elm_builtin cl kl vl ll rl, RBNode_elm_builtin _ _ _ _ _) ->
let
(k, v) =
maxWithDefault kl vl rl
newLeft =
removeMax cl kl vl ll rl
in
bubble color k v newLeft right
-- 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_elm_builtin c k v l r
-- Removes rightmost node, may leave root as BBlack
removeMax : NColor -> k -> v -> Dict k v -> Dict k v -> Dict k v
removeMax c k v l r =
case r of
RBEmpty_elm_builtin _ ->
rem c l r
RBNode_elm_builtin cr kr vr lr rr ->
bubble c k v l (removeMax 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 =
let
tree =
RBNode_elm_builtin c k v l r
in
if blackish tree then
balanceHelp tree
else
tree
blackish : Dict k v -> Bool
blackish t =
case t of
RBNode_elm_builtin c _ _ _ _ ->
c == Black || c == BBlack
RBEmpty_elm_builtin _ ->
True
balanceHelp : Dict k v -> Dict k v
balanceHelp tree =
case tree of
-- double red: left, left
RBNode_elm_builtin col zk zv (RBNode_elm_builtin Red yk yv (RBNode_elm_builtin Red xk xv a b) c) d ->
balancedTree col xk xv yk yv zk zv a b c d
-- double red: left, right
RBNode_elm_builtin col zk zv (RBNode_elm_builtin Red xk xv a (RBNode_elm_builtin Red yk yv b c)) d ->
balancedTree col xk xv yk yv zk zv a b c d
-- double red: right, left
RBNode_elm_builtin col xk xv a (RBNode_elm_builtin Red zk zv (RBNode_elm_builtin Red yk yv b c) d) ->
balancedTree col xk xv yk yv zk zv a b c d
-- double red: right, right
RBNode_elm_builtin col xk xv a (RBNode_elm_builtin Red yk yv b (RBNode_elm_builtin Red zk zv c d)) ->
balancedTree col xk xv yk yv zk zv a b c d
-- handle double blacks
RBNode_elm_builtin BBlack xk xv a (RBNode_elm_builtin NBlack zk zv (RBNode_elm_builtin Black yk yv b c) (RBNode_elm_builtin Black _ _ _ _ as d)) ->
RBNode_elm_builtin Black yk yv (RBNode_elm_builtin Black xk xv a b) (balance Black zk zv c (redden d))
RBNode_elm_builtin BBlack zk zv (RBNode_elm_builtin NBlack xk xv (RBNode_elm_builtin Black _ _ _ _ as a) (RBNode_elm_builtin Black yk yv b c)) d ->
RBNode_elm_builtin Black yk yv (balance Black xk xv (redden a) b) (RBNode_elm_builtin Black zk zv c d)
_ ->
tree
balancedTree : NColor -> k -> v -> k -> v -> k -> v -> Dict k v -> Dict k v -> Dict k v -> Dict k v -> Dict k v
balancedTree col xk xv yk yv zk zv a b c d =
RBNode_elm_builtin
(lessBlack col)
yk
yv
(RBNode_elm_builtin Black xk xv a b)
(RBNode_elm_builtin Black zk zv c d)
-- make the top node black
blacken : Dict k v -> Dict k v
blacken t =
case t of
RBEmpty_elm_builtin _ ->
RBEmpty_elm_builtin LBlack
RBNode_elm_builtin _ k v l r ->
RBNode_elm_builtin Black k v l r
-- make the top node red
redden : Dict k v -> Dict k v
redden t =
case t of
RBEmpty_elm_builtin _ ->
Native.Debug.crash "can't make a Leaf red"
RBNode_elm_builtin _ k v l r ->
RBNode_elm_builtin Red k v l r
-- COMBINE
{-| 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 _ -> member k t2) t1
{-| Keep a key-value pair when its key does not appear in the second dictionary.
-}
diff : Dict comparable v -> Dict comparable v -> Dict comparable v
diff t1 t2 =
foldl (\k v t -> remove k t) t1 t2
{-| The most general way of combining two dictionaries. You provide three
accumulators for when a given key appears:
1. Only in the left dictionary.
2. In both dictionaries.
3. Only in the right dictionary.
You then traverse all the keys from lowest to highest, building up whatever
you want.
-}
merge
: (comparable -> a -> result -> result)
-> (comparable -> a -> b -> result -> result)
-> (comparable -> b -> result -> result)
-> Dict comparable a
-> Dict comparable b
-> result
-> result
merge leftStep bothStep rightStep leftDict rightDict initialResult =
let
stepState rKey rValue (list, result) =
case list of
[] ->
(list, rightStep rKey rValue result)
(lKey, lValue) :: rest ->
if lKey < rKey then
stepState rKey rValue (rest, leftStep lKey lValue result)
else if lKey > rKey then
(list, rightStep rKey rValue result)
else
(rest, bothStep lKey lValue rValue result)
(leftovers, intermediateResult) =
foldl stepState (toList leftDict, initialResult) rightDict
in
List.foldl (\(k,v) result -> leftStep k v result) intermediateResult leftovers
-- TRANSFORM
{-| 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_elm_builtin _ ->
RBEmpty_elm_builtin LBlack
RBNode_elm_builtin clr key value left right ->
RBNode_elm_builtin 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_elm_builtin _ ->
acc
RBNode_elm_builtin _ 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_elm_builtin _ ->
acc
RBNode_elm_builtin _ key value left right ->
foldr f (f key value (foldr f acc right)) left
{-| 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
-- LISTS
{-| Get all of the keys in a dictionary, sorted from lowest to highest.
keys (fromList [(0,"Alice"),(1,"Bob")]) == [0,1]
-}
keys : Dict comparable v -> List comparable
keys dict =
foldr (\key value keyList -> key :: keyList) [] dict
{-| Get all of the values in a dictionary, in the order of their keys.
values (fromList [(0,"Alice"),(1,"Bob")]) == ["Alice", "Bob"]
-}
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, sorted by keys. -}
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