dups.ml

   1 open Printf
   2
   3 module Array = ArrayLabels
   4 module List  = ListLabels
   5 module StrSet= Set.Make(String)
   6 module Unix  = UnixLabels
   7
   8 module Stream : sig
   9   type 'a t
  10
  11   val create : (unit -> 'a option) -> 'a t
  12
  13   val iter : 'a t -> f:('a -> unit) -> unit
  14
  15   val concat : ('a t) list -> 'a t
  16 end = struct
  17   module S = Stream
  18
  19   type 'a t =
  20     ('a S.t) list
  21
  22   let create f =
  23     [S.from (fun _ -> f ())]
  24
  25   let iter t ~f =
  26     List.iter t ~f:(S.iter f)
  27
  28   let concat ts =
  29     List.concat ts
  30 end
  31
  32 module In_channel : sig
  33   val lines : in_channel -> string Stream.t
  34 end = struct
  35   let lines ic =
  36     Stream.create (fun () ->
  37       match input_line ic with
  38       | exception End_of_file ->
  39           None
  40       | line ->
  41           Some line
  42     )
  43 end
  44
  45 module Directory_tree : sig
  46   val find_files : string -> string Stream.t
  47 end = struct
  48   let find_files root =
  49     let dirs  = Queue.create () in
  50     let files = Queue.create () in
  51     let explore parent =
  52       Array.iter (Sys.readdir parent) ~f:(fun child ->
  53         let path = Filename.concat parent child in
  54         let {Unix.st_kind = file_kind; _} = Unix.lstat path in
  55         match file_kind with
  56         | Unix.S_REG ->
  57             Queue.add path files
  58         | Unix.S_DIR ->
  59             Queue.add path dirs
  60         | Unix.S_CHR
  61         | Unix.S_BLK
  62         | Unix.S_LNK
  63         | Unix.S_FIFO
  64         | Unix.S_SOCK ->
  65             ()
  66       )
  67     in
  68     explore root;
  69     let rec next () =
  70       match Queue.is_empty files, Queue.is_empty dirs with
  71       | false, _     -> Some (Queue.take files)
  72       | true , true  -> None
  73       | true , false ->
  74           explore (Queue.take dirs);
  75           next ()
  76     in
  77     Stream.create next
  78 end
  79
  80 type input =
  81   | Root_paths of string list
  82   | Paths_on_stdin
  83
  84 type output =
  85   | Stdout
  86   | Directory of string
  87
  88 let make_input_stream = function
  89   | Paths_on_stdin ->
  90       In_channel.lines stdin
  91   | Root_paths paths ->
  92       let paths = StrSet.elements (StrSet.of_list paths) in
  93       Stream.concat (List.map paths ~f:Directory_tree.find_files)
  94
  95 let make_output_fun = function
  96   | Stdout ->
  97       fun digest n_paths paths ->
  98         printf "%s %d\n%!" (Digest.to_hex digest) n_paths;
  99         List.iter (StrSet.elements paths) ~f:(printf "    %S\n%!")
 100   | Directory dir ->
 101       fun digest _ paths ->
 102         let digest = Digest.to_hex digest in
 103         let dir = Filename.concat dir (String.sub digest 0 2) in
 104         Unix.mkdir dir ~perm:0o700;
 105         let oc = open_out (Filename.concat dir digest) in
 106         List.iter (StrSet.elements paths) ~f:(fun path ->
 107           output_string oc (sprintf "%S\n%!" path)
 108         );
 109         close_out oc
 110
 111 let main input output =
 112   let output = make_output_fun  output in
 113   let input  = make_input_stream input in
 114   let paths_by_digest = Hashtbl.create 1_000_000 in
 115   let path_count = ref 0 in
 116   let t0 = Sys.time () in
 117   Stream.iter input ~f:(fun path ->
 118     incr path_count;
 119     try
 120       let digest = Digest.file path in
 121       let count, paths =
 122         match Hashtbl.find_opt paths_by_digest digest with
 123         | None ->
 124             (0, StrSet.empty)
 125         | Some (n, paths) ->
 126             (n, paths)
 127       in
 128       Hashtbl.replace paths_by_digest digest (count + 1, StrSet.add path paths)
 129     with Sys_error e ->
 130       eprintf "WARNING: Failed to process %S: %S\n%!" path e
 131   );
 132   Hashtbl.iter (fun d (n, ps) -> if n > 1 then output d n ps) paths_by_digest;
 133   let t1 = Sys.time () in
 134   eprintf "Processed %d files in %f seconds.\n%!" !path_count (t1 -. t0)
 135
 136 let () =
 137   let input  = ref Paths_on_stdin in
 138   let output = ref Stdout in
 139   let assert_file_exists path =
 140     if Sys.file_exists path then
 141       ()
 142     else begin
 143       eprintf "File does not exist: %S\n%!" path;
 144       exit 1
 145     end
 146   in
 147   let assert_file_is_dir path =
 148     if Sys.is_directory path then
 149       ()
 150     else begin
 151       eprintf "File is not a directory: %S\n%!" path;
 152       exit 1
 153     end
 154   in
 155   let spec =
 156     [ ( "-out"
 157       , Arg.String (fun path ->
 158           assert_file_exists path;
 159           assert_file_is_dir path;
 160           output := Directory path
 161         )
 162       , " Output to this directory instead of stdout."
 163       )
 164     ]
 165   in
 166   Arg.parse
 167     (Arg.align spec)
 168     (fun path ->
 169       assert_file_exists path;
 170       assert_file_is_dir path;
 171       match !input with
 172       | Paths_on_stdin ->
 173           input := Root_paths [path]
 174       | Root_paths paths ->
 175           input := Root_paths (path :: paths)
 176     )
 177     "";
 178   main !input !output