<commit_before>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<commit_msg>Add efficient update function for Dict<commit_after>module Dict (empty,singleton,insert,update
            ,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, update, 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 = let updater _ = v 
               in update k updater t

{-| Update the value of a dictionary for a specific key with a given function. -}
update : comparable -> (Maybe v -> v) -> Dict comparable v -> Dict comparable v
update k u t = -- Invariant: t is a valid left-leaning rb tree
  let up t =
      case t of
        -- Insert case
        RBEmpty LBlack -> RBNode Red k (u Nothing) (RBEmpty LBlack) (RBEmpty LBlack)
        RBNode c k' v' l r ->
          let h = case Native.Utils.compare k k' of
                    LT -> RBNode c k' v' (up l) r
                    EQ -> RBNode c k' (u (Just v')) l r -- update case
                    GT -> RBNode c k' v' l (up r)
          in  fixUp h
  in  ensureBlackRoot (up 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
{'fraction_non_alphanumeric': 0.0850716488300381, 'fraction_numerical': 0.0010883366588064574, 'mean_word_length': 2.3014970059880238, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 151, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2179603', 'n_tokens_mistral': 3713, 'n_tokens_neox': 3585, 'n_words': 1590}