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 map : 'a t -> f:('a -> 'b) -> 'b t
  16
  17   val filter : 'a t -> f:('a -> bool) -> 'a t
  18
  19   val concat : ('a t) list -> 'a t
  20 end = struct
  21   module S = Stream
  22
  23   type 'a t =
  24     {mutable streams : ('a S.t) list}
  25
  26   let create f =
  27     {streams = [S.from (fun _ -> f ())]}
  28
  29   let rec next t =
  30     match t.streams with
  31     | [] ->
  32         None
  33     | s :: streams ->
  34         (match S.next s with
  35         | exception Stream.Failure ->
  36             t.streams <- streams;
  37             next t
  38         | x ->
  39             Some x
  40         )
  41
  42   let map t ~f =
  43     create (fun () ->
  44       match next t with
  45       | None   -> None
  46       | Some x -> Some (f x)
  47     )
  48
  49   let filter t ~f =
  50     let rec filter () =
  51       match next t with
  52       | None ->
  53           None
  54       | Some x when f x ->
  55           Some x
  56       | Some _ ->
  57           filter ()
  58     in
  59     create filter
  60
  61   let iter t ~f =
  62     List.iter t.streams ~f:(S.iter f)
  63
  64   let concat ts =
  65     {streams = List.concat (List.map ts ~f:(fun {streams} -> streams))}
  66 end
  67
  68 module In_channel : sig
  69   val lines : in_channel -> string Stream.t
  70 end = struct
  71   let lines ic =
  72     Stream.create (fun () ->
  73       match input_line ic with
  74       | exception End_of_file ->
  75           None
  76       | line ->
  77           Some line
  78     )
  79 end
  80
  81 module File : sig
  82   type t =
  83     { path : string
  84     ; size : int
  85     }
  86
  87   val find : string -> t Stream.t
  88   (** Find all files in the directory tree, starting from the given root path *)
  89
  90   val lookup : string Stream.t -> t Stream.t
  91   (** Lookup file info for given paths *)
  92
  93   module Set : sig include Set.S with type elt := t end
  94 end = struct
  95   type t =
  96     { path : string
  97     ; size : int
  98     }
  99
 100   let compare {path=p1; _} {path=p2; _} =
 101     Stdlib.compare p1 p2
 102
 103   let lookup paths =
 104     Stream.map paths ~f:(fun path ->
 105       let {Unix.st_size = size; _} = Unix.lstat path in
 106       {path; size}
 107     )
 108
 109   let find root =
 110     let dirs  = Queue.create () in
 111     let files = Queue.create () in
 112     let explore parent =
 113       Array.iter (Sys.readdir parent) ~f:(fun child ->
 114         let path = Filename.concat parent child in
 115         let {Unix.st_kind = file_kind; st_size; _} = Unix.lstat path in
 116         match file_kind with
 117         | Unix.S_REG ->
 118             let file = {path; size = st_size} in
 119             Queue.add file files
 120         | Unix.S_DIR ->
 121             Queue.add path dirs
 122         | Unix.S_CHR
 123         | Unix.S_BLK
 124         | Unix.S_LNK
 125         | Unix.S_FIFO
 126         | Unix.S_SOCK ->
 127             ()
 128       )
 129     in
 130     explore root;
 131     let rec next () =
 132       match Queue.is_empty files, Queue.is_empty dirs with
 133       | false, _     -> Some (Queue.take files)
 134       | true , true  -> None
 135       | true , false ->
 136           explore (Queue.take dirs);
 137           next ()
 138     in
 139     Stream.create next
 140
 141   module Set = Set.Make(struct
 142     type elt = t
 143     type t = elt
 144     let compare = compare
 145   end)
 146 end
 147
 148 type input =
 149   | Stdin
 150   | Directories of string list
 151
 152 type output =
 153   | Stdout
 154   | Directory of string
 155
 156 type opt =
 157   { input  : input
 158   ; output : output
 159   ; ignore : Str.regexp option
 160   ; sample : int
 161   }
 162
 163 type count =
 164   { considered_files    : int ref
 165   ; considered_bytes    : int ref
 166   ; empty               : int ref
 167   ; ignored_files       : int ref
 168   ; ignored_bytes       : int ref
 169   ; unique_size_files   : int ref
 170   ; unique_size_bytes   : int ref
 171   ; unique_sample_files : int ref
 172   ; unique_sample_bytes : int ref
 173   ; sampled_files       : int ref
 174   ; sampled_bytes       : int ref
 175   ; hashed_files        : int ref
 176   ; hashed_bytes        : int ref
 177   ; digests             : int ref
 178   }
 179
 180 let add sum addend =
 181   sum := !sum + addend
 182
 183 let make_input_stream input ignore count =
 184   let input =
 185     match input with
 186     | Stdin ->
 187         File.lookup (In_channel.lines stdin)
 188     | Directories paths ->
 189         let paths = StrSet.elements (StrSet.of_list paths) in
 190         Stream.concat (List.map paths ~f:File.find)
 191   in
 192   Stream.filter input ~f:(fun {File.path; size} ->
 193     incr count.considered_files;
 194     add count.considered_bytes size;
 195     let empty = size = 0 in
 196     let ignored =
 197       match ignore with
 198       | Some regexp when (Str.string_match regexp path 0) ->
 199           incr count.ignored_files;
 200           add count.ignored_bytes size;
 201           true
 202       | Some _ | None ->
 203           false
 204     in
 205     if empty   then incr count.empty;
 206     (not empty) && (not ignored)
 207   )
 208
 209 let make_output_fun = function
 210   | Stdout ->
 211       fun digest n_files files ->
 212         printf "%s %d\n%!" (Digest.to_hex digest) n_files;
 213         List.iter (File.Set.elements files) ~f:(fun {File.path; _} ->
 214           printf "    %S\n%!" path
 215         )
 216   | Directory dir ->
 217       fun digest _ files ->
 218         let digest = Digest.to_hex digest in
 219         let dir = Filename.concat dir (String.sub digest 0 2) in
 220         Unix.mkdir dir ~perm:0o700;
 221         let oc = open_out (Filename.concat dir digest) in
 222         List.iter (File.Set.elements files) ~f:(fun {File.path; _} ->
 223           output_string oc (sprintf "%S\n%!" path)
 224         );
 225         close_out oc
 226
 227 let sample path ~len ~count =
 228   let buf = Bytes.make len ' ' in
 229   let ic = open_in_bin path in
 230   let rec read pos len =
 231     assert (len >= 0);
 232     if len = 0 then
 233       ()
 234     else begin
 235       let chunk_size = input ic buf pos len in
 236       add count.sampled_bytes chunk_size;
 237       if chunk_size = 0 then (* EOF *)
 238         ()
 239       else
 240         read (pos + chunk_size) (len - chunk_size)
 241     end
 242   in
 243   read 0 len;
 244   close_in ic;
 245   Bytes.to_string buf
 246
 247 let main {input; output; ignore; sample = sample_len} =
 248   let t0 = Sys.time () in
 249   let count =
 250     { considered_files    = ref 0
 251     ; considered_bytes    = ref 0
 252     ; empty               = ref 0
 253     ; ignored_files       = ref 0
 254     ; ignored_bytes       = ref 0
 255     ; unique_size_files   = ref 0
 256     ; unique_size_bytes   = ref 0
 257     ; sampled_files       = ref 0
 258     ; sampled_bytes       = ref 0
 259     ; hashed_files        = ref 0
 260     ; hashed_bytes        = ref 0
 261     ; unique_sample_files = ref 0
 262     ; unique_sample_bytes = ref 0
 263     ; digests             = ref 0
 264     }
 265   in
 266   let output = make_output_fun  output in
 267   let input  = make_input_stream input ignore count in
 268   let files_by_size = Hashtbl.create 1_000_000 in
 269   let files_by_sample = Hashtbl.create 1_000_000 in
 270   let files_by_digest = Hashtbl.create 1_000_000 in
 271   let process tbl ~group ~file =
 272     let count, files =
 273       match Hashtbl.find_opt tbl group with
 274       | None ->
 275           (0, File.Set.empty)
 276       | Some (n, files) ->
 277           (n, files)
 278     in
 279     Hashtbl.replace tbl group (count + 1, File.Set.add file files)
 280   in
 281   (* TODO: Make a nice(r) abstraction to re-assemble pieces in the pipeline:
 282    *
 283    * from input           to files_by_size
 284    * from files_by_size   to files_by_sample
 285    * from files_by_sample to files_by_digest
 286    * from files_by_digest to output
 287    *
 288    * input |> files_by_size |> files_by_sample |> files_by_digest |> output
 289    *)
 290   let t0_group_by_size = Sys.time () in
 291   Stream.iter input ~f:(fun ({File.size; _} as file) ->
 292     process files_by_size ~group:size ~file
 293   );
 294   let t1_group_by_size = Sys.time () in
 295   let t0_group_by_sample = Sys.time () in
 296   Hashtbl.iter
 297     (fun _ (n, files) ->
 298       (* Skip files with unique sizes *)
 299       if n > 1 then
 300         File.Set.iter
 301           (fun ({File.path; _} as file) ->
 302             incr count.sampled_files;
 303             process
 304               files_by_sample
 305               ~group:(sample path ~len:sample_len ~count)
 306               ~file
 307           )
 308           files
 309       else
 310         File.Set.iter
 311           (fun {File.size; _} ->
 312             incr count.unique_size_files;
 313             add count.unique_size_bytes size
 314           )
 315           files
 316     )
 317     files_by_size;
 318   let t1_group_by_sample = Sys.time () in
 319   let t0_group_by_digest = Sys.time () in
 320   Hashtbl.iter
 321     (fun _ (n, files) ->
 322       (* Skip files with unique samples *)
 323       if n > 1 then
 324         File.Set.iter
 325           (fun ({File.path; size} as file) ->
 326             incr count.hashed_files;
 327             add count.hashed_bytes size;
 328             process files_by_digest ~group:(Digest.file path) ~file
 329           )
 330           files
 331       else
 332         File.Set.iter
 333           (fun {File.size; _} ->
 334             incr count.unique_sample_files;
 335             add count.unique_sample_bytes size;
 336           )
 337           files
 338     )
 339     files_by_sample;
 340   let t1_group_by_digest = Sys.time () in
 341   Hashtbl.iter
 342     (fun d (n, files) ->
 343       incr count.digests;
 344       if n > 1 then
 345         output d n files
 346     )
 347     files_by_digest;
 348   let t1 = Sys.time () in
 349   let b_to_mb b = (float_of_int b) /. 1024. /. 1024. in
 350   let b_to_gb b = (b_to_mb b) /. 1024. in
 351   eprintf "Time                         : %8.2f seconds\n%!" (t1 -. t0);
 352   eprintf "Considered                   : %8d files  %6.2f Gb\n%!"
 353     !(count.considered_files)
 354     (b_to_gb !(count.considered_bytes));
 355   eprintf "Sampled                      : %8d files  %6.2f Gb\n%!"
 356     !(count.sampled_files)
 357     (b_to_gb !(count.sampled_bytes));
 358   eprintf "Hashed                       : %8d files  %6.2f Gb  %6.2f seconds\n%!"
 359     !(count.hashed_files)
 360     (b_to_gb !(count.hashed_bytes))
 361     (t1_group_by_digest -. t0_group_by_digest);
 362   eprintf "Digests                      : %8d\n%!"
 363     !(count.digests);
 364   eprintf "Duplicates (Hashed - Digests): %8d\n%!"
 365     (!(count.hashed_files) - !(count.digests));
 366   eprintf "Skipped due to 0      size   : %8d files\n%!" !(count.empty);
 367   eprintf "Skipped due to unique size   : %8d files  %6.2f Gb  %6.2f seconds\n%!"
 368     !(count.unique_size_files)
 369     (b_to_gb !(count.unique_size_bytes))
 370     (t1_group_by_size -. t0_group_by_size);
 371   eprintf "Skipped due to unique sample : %8d files  %6.2f Gb  %6.2f seconds\n%!"
 372     !(count.unique_sample_files)
 373     (b_to_gb !(count.unique_sample_bytes))
 374     (t1_group_by_sample -. t0_group_by_sample);
 375   eprintf "Ignored due to regex match   : %8d files  %6.2f Gb\n%!"
 376     !(count.ignored_files)
 377     (b_to_gb !(count.ignored_bytes))
 378
 379 let get_opt () : opt =
 380   let assert_ test x msg =
 381     if not (test x) then begin
 382       eprintf "%s\n%!" msg;
 383       exit 1
 384     end
 385   in
 386   let assert_file_exists path =
 387     assert_ Sys.file_exists path (sprintf "File does not exist: %S" path)
 388   in
 389   let assert_file_is_dir path =
 390     assert_ Sys.is_directory path (sprintf "File is not a directory: %S" path)
 391   in
 392   let input  = ref Stdin in
 393   let output = ref Stdout in
 394   let ignore = ref None in
 395   let sample = ref 256 in
 396   let spec =
 397     [ ( "-out"
 398       , Arg.String (fun path ->
 399           assert_file_exists path;
 400           assert_file_is_dir path;
 401           output := Directory path
 402         )
 403       , " Output to this directory instead of stdout."
 404       )
 405     ; ( "-ignore"
 406       , Arg.String (fun regexp -> ignore := Some (Str.regexp regexp))
 407       , " Ignore file paths which match this regexp pattern (see Str module)."
 408       )
 409     ; ( "-sample"
 410       , Arg.Set_int sample
 411       , (sprintf " Byte size of file samples to use. Default: %d" !sample)
 412       )
 413     ]
 414   in
 415   Arg.parse
 416     (Arg.align spec)
 417     (fun path ->
 418       assert_file_exists path;
 419       assert_file_is_dir path;
 420       match !input with
 421       | Stdin ->
 422           input := Directories [path]
 423       | Directories paths ->
 424           input := Directories (path :: paths)
 425     )
 426     "";
 427   assert_
 428     (fun x -> x > 0)
 429     !sample
 430     (sprintf "Sample size cannot be negative: %d" !sample);
 431   { input  = !input
 432   ; output = !output
 433   ; ignore = !ignore
 434   ; sample = !sample
 435   }
 436
 437 let () =
 438   main (get_opt ())