GenericDict.elm

module Dict
    ( Dict
    , empty, singleton, insert, update
    , isEmpty, get, remove, member, size
    , 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 k
type. This includes `Int`, `Float`, `Time`, `Char`, `String`, and tuples or
lists of k types.

Insert, remove, and query operations all take *O(log n)* time. Dictionary
equality with `(==)` is unreliable and should not be used.

# Dictionaries
@docs Dict

# Build
@docs empty, singleton, insert, update, remove

# Query
@docs isEmpty, member, get, size

# Combine
@docs union, intersect, diff

# Lists
@docs keys, values, toList, fromList

# Transform
@docs map, foldl, foldr, filter, partition

-}


import Basics exposing (..)
import Maybe exposing (..)
import List exposing (..)
import Native.Debug
import String

type alias Comparable k rest = { rest | compare : k -> k -> Order}


-- 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 k rest -> k -> Dict k v -> Maybe v
get inst targetKey dict =
  case dict of
    RBEmpty_elm_builtin _ ->
      Nothing

    RBNode_elm_builtin _ key value left right ->
      case inst.compare targetKey key of
        LT ->
          get inst targetKey left

        EQ ->
          Just value

        GT ->
          get inst targetKey right


{-| Determine if a key is in a dictionary. -}
member : Comparable k rest -> k -> Dict k v -> Bool
member inst key dict =
  case get inst 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 k rest -> k -> v -> Dict k v -> Dict k v
insert inst key value dict =
  update inst 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 k rest -> k -> Dict k v -> Dict k v
remove inst key dict =
  update inst 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 k rest -> k -> (Maybe v -> Maybe v) -> Dict k v -> Dict k v
update inst 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 inst.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 k rest -> k -> v -> Dict k v
singleton inst key value =
  insert inst key value empty


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 c l r =
  case (l, r) of
    (RBEmpty_elm_builtin _, RBEmpty_elm_builtin _) ->
      case c 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 (c, cl, cr) of
        (Black, LBlack, Red) ->
          RBNode_elm_builtin Black k' v' l' r'

        _ ->
          reportRemBug "Black/LBlack/Red" c (toString cl) (toString cr)

    (RBNode_elm_builtin cl k' v' l' r', RBEmpty_elm_builtin cr) ->
      case (c, cl, cr) of
        (Black, Red, LBlack) ->
          RBNode_elm_builtin Black k' v' l' r'

        _ ->
          reportRemBug "Black/Red/LBlack" c (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

        l' =
          removeMax 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_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


{-| Apply a function to all values in a dictionary.
-}
map : (k -> a -> b) -> Dict k a -> Dict k 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 : (k -> v -> b -> b) -> b -> Dict k 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 : (k -> v -> b -> b) -> b -> Dict k 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


{-| Combine two dictionaries. If there is a collision, preference is given
to the first dictionary.
-}
union : Comparable k rest -> Dict k v -> Dict k v -> Dict k v
union inst t1 t2 =
  foldl (insert inst) 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 : Comparable k rest -> Dict k v -> Dict k v -> Dict k v
intersect inst t1 t2 =
  filter inst (\k _ -> member inst k t2) t1


{-| Keep a key-value pair when its key does not appear in the second dictionary.
-}
diff : Comparable k rest -> Dict k v -> Dict k v -> Dict k v
diff inst t1 t2 =
  foldl (\k v t -> remove inst k t) t1 t2


{-| Get all of the keys in a dictionary, sorted from lowest to highest.

    keys (fromList [(0,"Alice"),(1,"Bob")]) == [0,1]
-}
keys : Dict k v -> List k
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 k 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 k v -> List (k,v)
toList dict =
  foldr (\key value list -> (key,value) :: list) [] dict


{-| Convert an association list into a dictionary. -}
fromList : Comparable k rest -> List (k,v) -> Dict k v
fromList inst assocs =
  List.foldl (\(key,value) dict -> insert inst key value dict) empty assocs


{-| Keep a key-value pair when it satisfies a predicate. -}
filter : Comparable k rest -> (k -> v -> Bool) -> Dict k v -> Dict k v
filter inst predicate dictionary =
  let
    add key value dict =
      if predicate key value then
        insert inst 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 k rest -> (k -> v -> Bool) -> Dict k v -> (Dict k v, Dict k v)
partition inst predicate dict =
  let
    add key value (t1, t2) =
      if predicate key value then
        (insert inst key value t1, t2)

      else
        (t1, insert inst key value t2)
  in
    foldl add (empty, empty) dict