module Dict (empty,singleton,insert
            ,lookup,findWithDefault
            ,remove,member
            ,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, remove

# Query
@docs member, lookup, findWithDefault

# Combine
@docs union, intersect, diff

# Lists
@docs keys, values, toList, fromList

# Transform
@docs map, foldl, foldr

-}


import open Basics
import open Maybe
import Native.Error
import List as List
import Native.Utils

-- BBlack and NBlack should only be used during the deletion
-- algorithm. Any other occurrence is a bug and should fail an assert.
data NColor = Red | Black
            | BBlack | NBlack

data LeafColor = LBlack | LBBlack

data 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 t =
  case t of
    RBNode _ k v (RBEmpty LBlack) _ -> (k,v)
    RBNode _ _ _ l _ -> min l
    RBEmpty LBlack -> Native.Error.raise "(min Empty) is not defined"

max : Dict k v -> (k, v)
max t =
  case t of
    RBNode _ k v _ (RBEmpty _) -> (k,v)
    RBNode _ _ _ _ r -> max r
    RBEmpty _ -> Native.Error.raise "(max Empty) is not defined"

{-| Lookup the value associated with a key. -}
lookup : comparable -> Dict comparable v -> Maybe v
lookup k t =
 case t of
   RBEmpty LBlack -> Nothing
   RBNode _ k' v l r ->
    case Native.Utils.compare k k' of
      LT -> lookup k l
      EQ -> Just v
      GT -> lookup k r

{-| Find the value associated with a key. If the key is not found,
return the default value. -}
findWithDefault : v -> comparable -> Dict comparable v -> v
findWithDefault base k t =
 case t of
   RBEmpty LBlack -> base
   RBNode _ k' v l r ->
    case Native.Utils.compare k k' of
      LT -> findWithDefault base k l
      EQ -> v
      GT -> findWithDefault base k r

{-| Determine if a key is in a dictionary. -}
member : comparable -> Dict comparable v -> Bool
-- Does t contain k?
member k t = isJust <| lookup k t

rotateLeft : Dict k v -> Dict k v
rotateLeft t =
 case t of
   RBNode cy ky vy a (RBNode cz kz vz b c) -> RBNode cy kz vz (RBNode Red ky vy a b) c
   _ -> Native.Error.raise "rotateLeft of a node without enough children"

-- rotateRight -- the reverse, and
-- makes Y have Z's color, and makes Z Red.
rotateRight : Dict k v -> Dict k v
rotateRight t =
 case t of
   RBNode cz kz vz (RBNode cy ky vy a b) c -> RBNode cz ky vy a (RBNode Red kz vz b c)
   _ -> Native.Error.raise "rotateRight of a node without enough children"

rotateLeftIfNeeded : Dict k v -> Dict k v
rotateLeftIfNeeded t =
 case t of
   RBNode _ _ _ _ (RBNode Red _ _ _ _) -> rotateLeft t
   _ -> t

rotateRightIfNeeded : Dict k v -> Dict k v
rotateRightIfNeeded t =
 case t of
   RBNode _ _ _ (RBNode Red _ _ (RBNode Red _ _ _ _) _) _ -> rotateRight t
   _ -> t

otherColor c = case c of { Red -> Black ; Black -> Red }

color_flip : Dict k v -> Dict k v
color_flip t =
 case t of
   RBNode c1 bk bv (RBNode c2 ak av la ra) (RBNode c3 ck cv lc rc) ->
       RBNode (otherColor c1) bk bv
              (RBNode (otherColor c2) ak av la ra)
              (RBNode (otherColor c3) ck cv lc rc)
   _ -> Native.Error.raise "color_flip called on a Empty or Node with a Empty child"

color_flipIfNeeded : Dict k v -> Dict k v
color_flipIfNeeded t =
 case t of
   RBNode _ _ _ (RBNode Red _ _ _ _) (RBNode Red _ _ _ _) -> color_flip t
   _ -> t

fixUp t = color_flipIfNeeded (rotateRightIfNeeded (rotateLeftIfNeeded t))

ensureBlackRoot : Dict k v -> Dict k v
ensureBlackRoot t =
  case t of
    RBNode Red k v l r -> RBNode Black k v l r
    RBNode Black _ _ _ _ -> t
    RBEmpty LBlack -> t

{-| 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 k v t =  -- Invariant: t is a valid left-leaning rb tree
  let ins t =
      case t of
        RBEmpty LBlack -> RBNode Red k v (RBEmpty LBlack) (RBEmpty LBlack)
        RBNode c k' v' l r ->
          let h = case Native.Utils.compare k k' of
                    LT -> RBNode c k' v' (ins l) r
                    EQ -> RBNode c k' v  l r  -- replace
                    GT -> RBNode c k' v' l (ins r)
          in  fixUp h
  in  ensureBlackRoot (ins t)

{-| Create a dictionary with one key-value pair. -}
singleton : comparable -> v -> Dict comparable v
singleton k v = insert k v (RBEmpty LBlack)

{- Remove helpers: everything from here to remove should only be used
   internally by remove as they would otherwise break rb-invariants -}
                
isBBlack : Dict k v -> Bool
isBBlack t = case t of
  RBNode c _ _ _ _ -> case c of
    BBlack -> True
    _      -> False
  RBEmpty LBBlack -> True
  _               -> False

moreBlack : NColor -> NColor
moreBlack c = case c of
  BBlack -> BBlack
  Black  -> BBlack
  Red    -> Black
  NBlack -> Red
lessBlack : NColor -> NColor
lessBlack c = case c of
  BBlack -> Black
  Black  -> Red
  Red    -> NBlack
  NBlack -> NBlack

moreBlackTree : Dict k v -> Dict k v
moreBlackTree t = case t of
  RBNode c k v l r -> RBNode (moreBlack c) k v l r
  RBEmpty _        -> RBEmpty LBBlack

lessBlackTree : Dict k v -> Dict k v
lessBlackTree t = case t of
  RBNode c k v l r -> RBNode (lessBlack c) k v l r
  RBEmpty _ -> RBEmpty LBlack

-- Finds and deletes k in t
del : comparable -> Dict comparable v -> Dict comparable v
del k t = case t of
  RBEmpty _        -> t
  RBNode c k' v l r -> case Native.Utils.compare k k' of
    LT -> bubble c k' v (del k l) r
    EQ -> rem t
    GT -> bubble c k' v l (del k r)

-- Remove the top node from the tree, may leave behind BBlacks
rem : Dict k v -> Dict k v
rem t = case t of
  RBNode c k v (RBEmpty _) (RBEmpty _) -> case c of
    Red -> RBEmpty LBlack
    Black -> RBEmpty LBBlack
  RBNode Black _ _ (RBEmpty _) (RBNode _ k v l r) ->
    RBNode Black k v l r
  RBNode Black _ _ (RBNode _ k v l r) (RBEmpty _) ->
    RBNode Black k v l r
  -- l and r are both RBNodes
  RBNode c _ _ l r ->
    let (k, v) = max l
        l'     = remove_max l
    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 : Dict k v -> Dict k v
remove_max t = case t of
  RBNode c k v l (RBEmpty _) -> rem t
  RBNode c k v l r           -> bubble c k v l (remove_max r)

-- 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 (RBNode c _ _ _ _) = c == Black || c == BBlack

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.Error.raise "can't make a Leaf red"
  RBNode _ k v l r -> RBNode Red k v l r

{-| 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 k t = blacken <| del k t

{-| Apply a function to all values in a dictionary. -}
map : (a -> b) -> Dict comparable a -> Dict comparable b
map f t =
  case t of
    RBEmpty LBlack -> RBEmpty LBlack
    RBNode c k v l r -> RBNode c k (f v) (map f l) (map f r)

{-| 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 t =
  case t of
    RBEmpty LBlack -> acc
    RBNode _ k v l r -> foldl f (f k v (foldl f acc l)) r

{-| 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 _ k v l r -> foldr f (f k v (foldr f acc r)) l

{-| 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 =
 let combine k v t = if k `member` t2 then insert k v t else t
 in  foldl combine empty 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 -> [comparable]
keys t   = foldr (\k v acc -> k :: acc) [] t

{-| Get all of the values in a dictionary. -}
values : Dict comparable v -> [v]
values t = foldr (\k v acc -> v :: acc) [] t

{-| Convert a dictionary into an association list of key-value pairs. -}
toList : Dict comparable v -> [(comparable,v)]
toList t = foldr (\k v acc -> (k,v) :: acc) [] t

{-| Convert an association list into a dictionary. -}
fromList : [(comparable,v)] -> Dict comparable v
fromList assocs = List.foldl (\(k,v) d -> insert k v d) empty assocs