Utils: fix Utils.take_n in presence of duplicates

src/minutils/utils.ml

@@ -210,35 +210,40 @@ let take_n_unsorted n l =
     | x :: xs -> loop (x :: acc) (pred n) xs in
   loop [] n l
 
-module Bounded(E: Set.OrderedType) = struct
+module Bounded(E: Set.OrderedType) : sig
 
-  (* TODO one day replace list by an heap array *)
+  type t
+  val create: int -> t
+  val insert: E.t -> t -> unit
+  val get: t -> E.t list
+
+end = struct
+
+  (* TODO one day replace the list by an heap array *)
 
   type t = {
     bound : int ;
     mutable size : int ;
     mutable data : E.t list ;
   }
-  let create bound = { bound ; size = 0 ; data = [] }
+
+  let create bound =
+    if bound <= 0 then invalid_arg "Utils.Bounded(_).create" ;
+    { bound ; size = 0 ; data = [] }
 
   let rec push x = function
     | [] -> [x]
     | (y :: xs) as ys ->
-        let c = compare x y in
-        if c < 0 then x :: ys else if c = 0 then ys else y :: push x xs
-
-  let replace x xs =
-    match xs with
-    | y :: xs when compare x y > 0 ->
-        push x xs
-    | xs -> xs
+        if E.compare x y <= 0
+        then x :: ys
+        else y :: push x xs
 
   let insert x t =
     if t.size < t.bound then begin
       t.size <- t.size + 1 ;
       t.data <- push x t.data
     end else if E.compare (List.hd t.data) x < 0 then
-      t.data <- replace x t.data
+      t.data <- push x (List.tl t.data)
 
   let get { data } = data
 

test/utils/Makefile

@@ -2,6 +2,7 @@
 SRCDIR=../../src
 
 TESTS := \
+  utils \
   merkle \
   data-encoding \
   stream-data-encoding \
@@ -55,8 +56,27 @@ test-merkle: ${LIB} ${TEST_MERKLE_IMPLS:.ml=.cmx}
 clean::
 	rm -f test-merkle
 
+
 ############################################################################
-## Data_encoding
+## Utils
+
+.PHONY:run-test-utils
+run-test-utils:
+	@echo
+	./test-utils
+
+TEST_UTILS_IMPLS := \
+  test_utils.ml
+
+test-utils: ${LIB} ${TEST_UTILS_IMPLS:.ml=.cmx}
+	@echo LINK $(notdir $@)
+	@${OCAMLOPT} -linkall -linkpkg ${OCAMLFLAGS} -o $@ $^
+
+clean::
+	rm -f test-utils
+
+############################################################################
+## Utils
 
 .PHONY:run-test-data-encoding
 run-test-data-encoding:

test/utils/test_utils.ml

@@ -0,0 +1,60 @@
+(**************************************************************************)
+(*                                                                        *)
+(*    Copyright (c) 2014 - 2016.                                          *)
+(*    Dynamic Ledger Solutions, Inc. <contact@tezos.com>                  *)
+(*                                                                        *)
+(*    All rights reserved. No warranty, explicit or implicit, provided.   *)
+(*                                                                        *)
+(**************************************************************************)
+
+open Error_monad
+open Hash
+
+let rec (--) i j =
+  if j < i then []
+  else i :: (i+1) -- j
+
+let rec permut = function
+  | [] -> [[]]
+  | x :: xs ->
+      let insert xs =
+        let rec loop acc left right =
+          match right with
+          | [] -> List.rev (x :: left) :: acc
+          | y :: ys ->
+              loop
+                ((List.rev_append left (x :: right)) :: acc)
+                (y :: left) ys in
+        loop [] [] xs in
+      List.concat (List.map insert (permut xs))
+
+let test_take_n _ =
+  ListLabels.iter (permut [1;2;3;4;5;6;7;8;9]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare 1 xs) [9]
+  end ;
+  ListLabels.iter (permut [1;2;3;4;5;6;7;8;9]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare 3 xs) [7;8;9]
+  end ;
+  let inv_compare x y = compare y x in
+  ListLabels.iter (permut [1;2;3;4;5;6;7;8;9]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare:inv_compare 3 xs) [3;2;1]
+  end ;
+  (* less elements than the bound. *)
+  ListLabels.iter (permut [1;2;3;4;5;6;7;8;9]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare 12 xs) [1;2;3;4;5;6;7;8;9]
+  end ;
+  (* with duplicates. *)
+  ListLabels.iter (permut [1;2;3;3;4;5;5;5;6]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare 3 xs) [5;5;6]
+  end ;
+  ListLabels.iter (permut [1;2;3;3;4;5;5;5;6]) ~f:begin fun xs ->
+    Assert.equal ~msg:__LOC__ (take_n ~compare 5 xs) [4;5;5;5;6]
+  end ;
+  return ()
+
+let tests : (string * (string -> unit tzresult Lwt.t)) list = [
+  "take_n", test_take_n ;
+]
+
+let () =
+  Test.run "utils." tests