From 197f7eb78638195af88420c5b4f70063d92e97b1 Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Thu, 29 Aug 2013 14:30:42 +0200
Subject: [PATCH] more functions in B-encode, to compute the size of an
 expression in bytes, and write it into an already-existing string

---
 bencode.ml  | 63 +++++++++++++++++++++++++++++++++++++++++------------
 bencode.mli |  8 +++++++
 2 files changed, 57 insertions(+), 14 deletions(-)

diff --git a/bencode.ml b/bencode.ml
index e8e220bf..dc644223 100644
--- a/bencode.ml
+++ b/bencode.ml
@@ -54,24 +54,59 @@ let dict_of_list l =
 
 (** {2 Serialization (encoding)} *)
 
-let rec to_buf buf t = match t with
-  | I i -> Printf.bprintf buf "i%de" i
-  | S s ->
-    Printf.bprintf buf "%d:" (String.length s);
-    Buffer.add_string buf s
+(* length of an encoded int, in bytes *)
+let _len_int i =
+  match i with
+  | 0 -> 1
+  | _ when i < 0 -> 2 + int_of_float (log10 (float_of_int ~-i))
+  | _ -> 1 + int_of_float (log10 (float_of_int i))
+
+(* length of an encoded string, in bytes *)
+let _len_str s =
+  _len_int (String.length s) + 1 + String.length s
+
+let rec size t = match t with
+  | I i -> 2 + _len_int i
+  | S s -> _len_str s
+  | L l -> List.fold_left (fun acc i -> acc + size i) 2 l
+  | D map -> SMap.fold (fun k v acc -> acc + _len_str k + size v) map 2
+
+let write_in_string t buf o =
+  let pos = ref o in
+  let rec append t = match t with
+  | I i -> write_char 'i'; write_int i; write_char 'e'
+  | S s -> write_str s
   | L l ->
-    Buffer.add_char buf 'l';
-    List.iter (fun t' -> to_buf buf t') l;
-    Buffer.add_char buf 'e'
+    write_char 'l';
+    List.iter append l;
+    write_char 'e';
   | D m ->
-    Buffer.add_char buf 'd';
-    SMap.iter (fun key t' -> to_buf buf (S key); to_buf buf t') m;
-    Buffer.add_char buf 'e'
+    write_char 'd';
+    SMap.iter (fun key t' -> write_str key; append t') m;
+    write_char 'e'
+  and write_int i =
+    let s = string_of_int i in
+    String.blit s 0 buf !pos (String.length s);
+    pos := !pos + String.length s
+  and write_str s =
+    write_int (String.length s);
+    write_char ':';
+    String.blit s 0 buf !pos (String.length s);
+    pos := !pos + String.length s
+  and write_char c =
+    buf.[!pos] <- c;
+    incr pos
+  in
+  append t
 
 let to_string t =
-  let b = Buffer.create 25 in
-  to_buf b t;
-  Buffer.contents b
+  let len = size t in
+  let s = String.create len in
+  write_in_string t s 0;
+  s
+
+let to_buf buf t =
+  Buffer.add_string buf (to_string t)
 
 let to_chan ch t =
   let b = Buffer.create 25 in
diff --git a/bencode.mli b/bencode.mli
index fff1bf0b..4ac860d9 100644
--- a/bencode.mli
+++ b/bencode.mli
@@ -44,6 +44,14 @@ val dict_of_list : (string * t) list -> t
 
 (** {2 Serialization (encoding)} *)
 
+val size : t -> int
+  (** Size needed for serialization *)
+
+val write_in_string : t -> string -> int -> unit
+  (** [write_in_string v buf o] writes the value [v] in  the string,
+      starting at offset [o]. The portion of the string starting from [o]
+      must be big enough (ie >= [size v]) *)
+
 val to_buf : Buffer.t -> t -> unit
 val to_string : t -> string
 val to_chan : out_channel -> t -> unit