From fe86ffcbb84824c4293c7b2d57ecad4cde26b3ff Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Wed, 9 Oct 2013 19:01:19 +0200
Subject: [PATCH] Bencode: sequence bencode <-> sequence string; BencodeOnDisk:
 ability to write a list of values in one atomic operation

---
 bencode.ml        | 28 ++++++++++++++++++++++++++++
 bencode.mli       | 15 +++++++++++++++
 bencodeOnDisk.ml  | 13 +++++++++++--
 bencodeOnDisk.mli |  3 +++
 4 files changed, 57 insertions(+), 2 deletions(-)

diff --git a/bencode.ml b/bencode.ml
index 5fe856fd..04e64656 100644
--- a/bencode.ml
+++ b/bencode.ml
@@ -333,3 +333,31 @@ let of_string s =
   | ParseOk t -> t
   | ParsePartial -> invalid_arg "Bencode: partial parse"
   | ParseError msg -> invalid_arg msg
+
+(** {2 Iterator} *)
+
+type 'a sequence = ('a -> unit) -> unit
+
+let of_seq seq =
+  fun k ->
+    let decoder = mk_decoder () in
+    (* read a string *)
+    let rec read_chunk str =
+      match parse decoder str 0 (String.length str) with
+      | ParseOk v ->
+        k v;  (* yield, and parse the rest of the string *)
+        resume ()
+      | ParseError e -> raise (Invalid_argument e)
+      | ParsePartial -> ()  (* wait for next chunk *)
+    and resume () = match parse_resume decoder with
+      | ParseOk v ->
+        k v;
+        resume ()
+      | ParseError e -> raise (Invalid_argument e)
+      | ParsePartial -> ()  (* wait for next chunk *)
+    in
+    seq read_chunk
+
+let to_seq seq =
+  fun k -> seq (fun b -> k (to_string b))
+
diff --git a/bencode.mli b/bencode.mli
index 0bd47a80..5af81b98 100644
--- a/bencode.mli
+++ b/bencode.mli
@@ -113,3 +113,18 @@ val parse_string : string -> parse_result
 
 val of_string : string -> t
   (** Parse the string. @raise Invalid_argument if it fails to parse. *)
+
+(** {2 Iterator} *)
+
+type 'a sequence = ('a -> unit) -> unit
+
+val of_seq : string sequence -> t sequence
+  (** Given a sequence of strings into Bencode values. Strings can be
+      the result of {!Unix.read}, for instance, they don't need to be
+      valid bencode individually; Only their concatenation should
+      be a valid stream of Bencode values.
+      
+      @raise Invalid_argument if a parsing error occurs. *)
+
+val to_seq : t sequence -> string sequence
+  (** Serialize each value in the sequence of Bencode values *)
diff --git a/bencodeOnDisk.ml b/bencodeOnDisk.ml
index f327bf76..4fc55882 100644
--- a/bencodeOnDisk.ml
+++ b/bencodeOnDisk.ml
@@ -57,12 +57,11 @@ let open_out ?lock filename =
 let close_out out =
   Unix.close out.file
 
-let write out b =
+let write_string out s =
   Unix.lockf out.lock_file Unix.F_LOCK 0;
   try
     (* go to the end of the file *)
     ignore (Unix.lseek out.file 0 Unix.SEEK_END);
-    let s = Bencode.to_string b in
     (* call write() until everything is written *)
     let rec write_all n =
       if n >= String.length s
@@ -78,6 +77,16 @@ let write out b =
     Unix.lockf out.lock_file Unix.F_ULOCK 0;
     raise e
 
+let write out b =
+  let s = Bencode.to_string b in
+  write_string out s
+
+let write_batch out l =
+  let buf = Buffer.create 255 in
+  List.iter (fun b -> Bencode.to_buf buf b) l;
+  let s = Buffer.contents buf in
+  write_string out s
+
 type 'a result =
   | Ok of 'a
   | Error of string
diff --git a/bencodeOnDisk.mli b/bencodeOnDisk.mli
index 9c5eef49..b55c1ef5 100644
--- a/bencodeOnDisk.mli
+++ b/bencodeOnDisk.mli
@@ -47,6 +47,9 @@ val close_out : t -> unit
 val write : t -> Bencode.t -> unit
   (** Write "atomically" a value to the end of the file *)
 
+val write_batch : t -> Bencode.t list -> unit
+  (** Write several values at once, at the end of the file *)
+
 type 'a result =
   | Ok of 'a
   | Error of string