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 end = struct
  93   type t =
  94     { path : string
  95     ; size : int
  96     }
  97
  98   let lookup paths =
  99     Stream.map paths ~f:(fun path ->
 100       let {Unix.st_size = size; _} = Unix.lstat path in
 101       {path; size}
 102     )
 103
 104   let find root =
 105     let dirs  = Queue.create () in
 106     let files = Queue.create () in
 107     let explore parent =
 108       Array.iter (Sys.readdir parent) ~f:(fun child ->
 109         let path = Filename.concat parent child in
 110         let {Unix.st_kind = file_kind; st_size; _} = Unix.lstat path in
 111         match file_kind with
 112         | Unix.S_REG ->
 113             let file = {path; size = st_size} in
 114             Queue.add file files
 115         | Unix.S_DIR ->
 116             Queue.add path dirs
 117         | Unix.S_CHR
 118         | Unix.S_BLK
 119         | Unix.S_LNK
 120         | Unix.S_FIFO
 121         | Unix.S_SOCK ->
 122             ()
 123       )
 124     in
 125     explore root;
 126     let rec next () =
 127       match Queue.is_empty files, Queue.is_empty dirs with
 128       | false, _     -> Some (Queue.take files)
 129       | true , true  -> None
 130       | true , false ->
 131           explore (Queue.take dirs);
 132           next ()
 133     in
 134     Stream.create next
 135 end
 136
 137 type input =
 138   | Stdin
 139   | Directories of string list
 140
 141 type output =
 142   | Stdout
 143   | Directory of string
 144
 145 type count =
 146   { considered  : int ref
 147   ; empty       : int ref
 148   ; ignored     : int ref
 149   ; unique_size : int ref
 150   ; hashed      : int ref
 151   }
 152
 153 let make_input_stream input ignore count =
 154   let input =
 155     match input with
 156     | Stdin ->
 157         File.lookup (In_channel.lines stdin)
 158     | Directories paths ->
 159         let paths = StrSet.elements (StrSet.of_list paths) in
 160         Stream.concat (List.map paths ~f:File.find)
 161   in
 162   Stream.filter input ~f:(fun {File.path; size} ->
 163     incr count.considered;
 164     let empty = size = 0 in
 165     let ignored =
 166       match ignore with
 167       | Some regexp when (Str.string_match regexp path 0) ->
 168           true
 169       | Some _ | None ->
 170           false
 171     in
 172     if empty   then incr count.empty;
 173     if ignored then incr count.ignored;
 174     (not empty) && (not ignored)
 175   )
 176
 177 let make_output_fun = function
 178   | Stdout ->
 179       fun digest n_paths paths ->
 180         printf "%s %d\n%!" (Digest.to_hex digest) n_paths;
 181         List.iter (StrSet.elements paths) ~f:(printf "    %S\n%!")
 182   | Directory dir ->
 183       fun digest _ paths ->
 184         let digest = Digest.to_hex digest in
 185         let dir = Filename.concat dir (String.sub digest 0 2) in
 186         Unix.mkdir dir ~perm:0o700;
 187         let oc = open_out (Filename.concat dir digest) in
 188         List.iter (StrSet.elements paths) ~f:(fun path ->
 189           output_string oc (sprintf "%S\n%!" path)
 190         );
 191         close_out oc
 192
 193 let main input output ignore =
 194   let t0 = Sys.time () in
 195   let count =
 196     { considered  = ref 0
 197     ; empty       = ref 0
 198     ; ignored     = ref 0
 199     ; unique_size = ref 0
 200     ; hashed      = ref 0
 201     }
 202   in
 203   let output = make_output_fun  output in
 204   let input  = make_input_stream input ignore count in
 205   let paths_by_size = Hashtbl.create 1_000_000 in
 206   let paths_by_digest = Hashtbl.create 1_000_000 in
 207   let process tbl path ~f =
 208     let key = f path in
 209     let count, paths =
 210       match Hashtbl.find_opt tbl key with
 211       | None ->
 212           (0, StrSet.empty)
 213       | Some (n, paths) ->
 214           (n, paths)
 215     in
 216     Hashtbl.replace tbl key (count + 1, StrSet.add path paths)
 217   in
 218   Stream.iter input ~f:(fun {File.path; size} ->
 219     process paths_by_size path ~f:(fun _ -> size)
 220   );
 221   Hashtbl.iter
 222     (fun _ (n, paths) ->
 223       (* Skip files with unique sizes *)
 224       if n > 1 then
 225         StrSet.iter
 226           (fun path ->
 227             incr count.hashed;
 228             process paths_by_digest path ~f:Digest.file
 229           )
 230           paths
 231       else
 232         incr count.unique_size;
 233     )
 234     paths_by_size;
 235   Hashtbl.iter (fun d (n, ps) -> if n > 1 then output d n ps) paths_by_digest;
 236   let t1 = Sys.time () in
 237   eprintf "Time                       : %f seconds\n%!" (t1 -. t0);
 238   eprintf "Considered                 : %d\n%!" !(count.considered);
 239   eprintf "Hashed                     : %d\n%!" !(count.hashed);
 240   eprintf "Skipped due to 0      size : %d\n%!" !(count.empty);
 241   eprintf "Skipped due to unique size : %d\n%!" !(count.unique_size);
 242   eprintf "Ignored due to regex match : %d\n%!" !(count.ignored)
 243
 244 let () =
 245   let input  = ref Stdin in
 246   let output = ref Stdout in
 247   let ignore = ref None in
 248   let assert_file_exists path =
 249     if Sys.file_exists path then
 250       ()
 251     else begin
 252       eprintf "File does not exist: %S\n%!" path;
 253       exit 1
 254     end
 255   in
 256   let assert_file_is_dir path =
 257     if Sys.is_directory path then
 258       ()
 259     else begin
 260       eprintf "File is not a directory: %S\n%!" path;
 261       exit 1
 262     end
 263   in
 264   let spec =
 265     [ ( "-out"
 266       , Arg.String (fun path ->
 267           assert_file_exists path;
 268           assert_file_is_dir path;
 269           output := Directory path
 270         )
 271       , " Output to this directory instead of stdout."
 272       )
 273     ; ( "-ignore"
 274       , Arg.String (fun regexp -> ignore := Some (Str.regexp regexp))
 275       , " Ignore file paths which match this regexp pattern (see Str module)."
 276       )
 277     ]
 278   in
 279   Arg.parse
 280     (Arg.align spec)
 281     (fun path ->
 282       assert_file_exists path;
 283       assert_file_is_dir path;
 284       match !input with
 285       | Stdin ->
 286           input := Directories [path]
 287       | Directories paths ->
 288           input := Directories (path :: paths)
 289     )
 290     "";
 291   main !input !output !ignore