5a09124c8bf920cedff189b038c58aaedd46600e
[dups.git] / 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 ())
This page took 0.052826 seconds and 3 git commands to generate.