Handle multiple root directories
[dups.git] / dups.ml
diff --git a/dups.ml b/dups.ml
index a1dca83..941d8fa 100644 (file)
--- a/dups.ml
+++ b/dups.ml
@@ -2,6 +2,7 @@ open Printf
 
 module Array = ArrayLabels
 module List  = ListLabels
+module StrSet= Set.Make(String)
 
 module Stream : sig
   type 'a t
@@ -9,17 +10,22 @@ module Stream : sig
   val create : (unit -> 'a option) -> 'a t
 
   val iter : 'a t -> f:('a -> unit) -> unit
+
+  val concat : ('a t) list -> 'a t
 end = struct
   module S = Stream
 
   type 'a t =
-    'a S.t
+    ('a S.t) list
 
   let create f =
-    S.from (fun _ -> f ())
+    [S.from (fun _ -> f ())]
 
   let iter t ~f =
-    S.iter f t
+    List.iter t ~f:(S.iter f)
+
+  let concat ts =
+    List.concat ts
 end
 
 module In_channel : sig
@@ -35,7 +41,7 @@ end = struct
     )
 end
 
-module Directory : sig
+module Directory_tree : sig
   val find_files : string -> string Stream.t
 end = struct
   let find_files root =
@@ -58,36 +64,30 @@ end = struct
             ()
       )
     in
-    let next_dir () =
-      match Queue.take dirs with
-      | exception Queue.Empty ->
-          ()
-      | dir ->
-          explore dir
-    in
-    let next_file () =
-      match Queue.take files with
-      | exception Queue.Empty ->
-          None
-      | file_path ->
-          Some file_path
-    in
     explore root;
-    Stream.create (fun () ->
-      next_dir ();
-      next_file ()
-    )
+    let rec next () =
+      match Queue.is_empty files, Queue.is_empty dirs with
+      | false, _     -> Some (Queue.take files)
+      | true , true  -> None
+      | true , false ->
+          explore (Queue.take dirs);
+          next ()
+    in
+    Stream.create next
 end
 
 type input =
-  | Root_path of string
+  | Root_paths of string list
   | Paths_on_stdin
 
 let main input =
   let paths =
     match input with
-    | Paths_on_stdin -> In_channel.lines stdin
-    | Root_path root -> Directory.find_files root
+    | Paths_on_stdin ->
+        In_channel.lines stdin
+    | Root_paths paths ->
+        let paths = StrSet.elements (StrSet.of_list paths) in
+        Stream.concat (List.map paths ~f:Directory_tree.find_files)
   in
   let paths_by_digest = Hashtbl.create 1_000_000 in
   let path_count = ref 0 in
@@ -99,20 +99,20 @@ let main input =
       let paths =
         match Hashtbl.find_opt paths_by_digest digest with
         | None ->
-            []
+            StrSet.empty
         | Some paths ->
             paths
       in
-      Hashtbl.replace paths_by_digest digest (path :: paths)
+      Hashtbl.replace paths_by_digest digest (StrSet.add path paths)
     with Sys_error e ->
       eprintf "WARNING: Failed to process %S: %S\n%!" path e
   );
   Hashtbl.iter
     (fun digest paths ->
-      let n_paths = List.length paths in
+      let n_paths = StrSet.cardinal paths in
       if n_paths > 1 then begin
         printf "%s %d\n%!" (Digest.to_hex digest) n_paths;
-        List.iter paths ~f:(fun path -> printf "    %s\n%!" path)
+        List.iter (StrSet.elements paths) ~f:(printf "    %s\n%!")
       end
     )
     paths_by_digest;
@@ -121,12 +121,19 @@ let main input =
 
 let () =
   let input = ref Paths_on_stdin in
-  Arg.parse [] (fun path ->
-    if Sys.file_exists path then
-      input := Root_path path
-    else begin
-      eprintf "File does not exist: %S\n%!" path;
-      exit 1
-    end
-  ) "";
+  Arg.parse
+    []
+    (function
+    | path when Sys.file_exists path ->
+        (match !input with
+        | Paths_on_stdin ->
+            input := Root_paths [path]
+        | Root_paths paths ->
+            input := Root_paths (path :: paths)
+        )
+    | path ->
+        eprintf "File does not exist: %S\n%!" path;
+        exit 1
+    )
+    "";
   main !input
This page took 0.023122 seconds and 4 git commands to generate.