From d689a686b922050ee94a042b04e2d49afce30cd7 Mon Sep 17 00:00:00 2001
From: iguer <mohamed.iguernelala@gmail.com>
Date: Wed, 22 Apr 2015 15:18:25 +0200
Subject: [PATCH] Map.add x y m returns the original map if x is already bound
 to a value that is physically equal to y in m

---
 stdlib/map.ml                      | 10 ++++++----
 stdlib/map.mli                     |  5 ++++-
 testsuite/tests/lib-set/testmap.ml | 27 +++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/stdlib/map.ml b/stdlib/map.ml
index b5d9a08ca..fd05e30a8 100644
--- a/stdlib/map.ml
+++ b/stdlib/map.ml
@@ -103,14 +103,16 @@ module Make(Ord: OrderedType) = struct
     let rec add x data = function
         Empty ->
           Node(Empty, x, data, Empty, 1)
-      | Node(l, v, d, r, h) ->
+      | Node(l, v, d, r, h) as m ->
           let c = Ord.compare x v in
           if c = 0 then
-            Node(l, x, data, r, h)
+            if d == data then m else Node(l, x, data, r, h)
           else if c < 0 then
-            bal (add x data l) v d r
+            let ll = add x data l in
+            if l == ll then m else bal ll v d r
           else
-            bal l v d (add x data r)
+            let rr = add x data r in
+            if r == rr then m else bal l v d rr
 
     let rec find x = function
         Empty ->
diff --git a/stdlib/map.mli b/stdlib/map.mli
index 4a6ced5c8..038531dc7 100644
--- a/stdlib/map.mli
+++ b/stdlib/map.mli
@@ -77,7 +77,10 @@ module type S =
     val add: key -> 'a -> 'a t -> 'a t
     (** [add x y m] returns a map containing the same bindings as
        [m], plus a binding of [x] to [y]. If [x] was already bound
-       in [m], its previous binding disappears. *)
+       in [m] to a value that is physically equal to [y],
+       [m] is returned unchanged (the result of the function is
+       then physically equal to [m]). Otherwise, the previous binding
+       of [x] in [m] disappears. *)
 
     val singleton: key -> 'a -> 'a t
     (** [singleton x y] returns the one-element map that contains a binding [y]
diff --git a/testsuite/tests/lib-set/testmap.ml b/testsuite/tests/lib-set/testmap.ml
index 8792ae7f0..4cbf36c62 100644
--- a/testsuite/tests/lib-set/testmap.ml
+++ b/testsuite/tests/lib-set/testmap.ml
@@ -155,3 +155,30 @@ let () =
   for i = 1 to 10 do m1 := M.add i (float i) !m1 done;
   let m2 = M.filter (fun e _ -> e >= 0) !m1 in
   assert (m2 == !m1)
+
+let () =
+  (* check that adding a binding "x -> y" to a map that already
+     contains it doesn't allocate and return the original map. *)
+  let m1 = ref M.empty in
+  let tmp = ref None in
+  for i = 1 to 10 do
+    tmp := Some (float i);
+    m1 := M.add i !tmp !m1
+  done;
+  let m2 = ref !m1 in
+
+  let a0 = Gc.allocated_bytes () in
+  let a1 = Gc.allocated_bytes () in
+
+  (* 10 |-> !tmp is already present in !m2 *)
+  m2 := M.add 10 !tmp !m2;
+
+  let a2 = Gc.allocated_bytes () in
+
+  assert (!m2 == !m1);
+  assert(a2 -. a1 = a1 -. a0);
+
+  (* 4 |-> Some 84. is not present in !m2 *)
+  m2 := M.add 4 (Some 84.) !m2;
+
+  assert (not (!m2 == !m1));