diff --git a/sorted_map/map.mbt b/sorted_map/map.mbt index 09197c0e8..fb2766f43 100644 --- a/sorted_map/map.mbt +++ b/sorted_map/map.mbt @@ -40,7 +40,12 @@ pub fn of[K : Compare, V](entries : Array[(K, V)]) -> T[K, V] { pub fn add[K : Compare, V](self : T[K, V], key : K, value : V) -> Unit { let parent = loop self.root, self.root { Some(curr), _ => - continue if key < curr.key { curr.left } else { curr.right }, Some(curr) + if key == curr.key { + curr.value = value + return + } else { + continue if key < curr.key { curr.left } else { curr.right }, Some(curr) + } None, parent => parent } let new = new_node(key, value, ~parent) @@ -51,7 +56,6 @@ pub fn add[K : Compare, V](self : T[K, V], key : K, value : V) -> Unit { } else { parent.unwrap().right = Some(new) } - new.parent = parent self.balance(new) self.size += 1 } diff --git a/sorted_map/map_test.mbt b/sorted_map/map_test.mbt index 6dddaa1f0..c36f98695 100644 --- a/sorted_map/map_test.mbt +++ b/sorted_map/map_test.mbt @@ -41,6 +41,13 @@ test "add" { inspect(map.size(), content="8")? } +test "add duplicate key" { + let map = of([(3, "c"), (2, "b"), (1, "a")]) + inspect(map.debug_tree(), content="([1]2,b,([0]1,a,_,_),([0]3,c,_,_))")? + map.add(1, "x") + inspect(map.debug_tree(), content="([1]2,b,([0]1,x,_,_),([0]3,c,_,_))")? +} + test "RL rotate" { let map = new() map.add(5, "a")