From 9a2249292fb5fd6467bb26b43fbe3bc583c99e6d Mon Sep 17 00:00:00 2001
From: Simon Cruanes <simon.cruanes.2007@m4x.org>
Date: Mon, 19 Sep 2022 21:50:13 -0400
Subject: [PATCH] feat: add sidekick.bencode for serialization

---
 src/bencode/Sidekick_bencode.ml  | 119 +++++++++++++++++++++++++++++++
 src/bencode/Sidekick_bencode.mli |  14 ++++
 src/bencode/dune                 |   6 ++
 3 files changed, 139 insertions(+)
 create mode 100644 src/bencode/Sidekick_bencode.ml
 create mode 100644 src/bencode/Sidekick_bencode.mli
 create mode 100644 src/bencode/dune

diff --git a/src/bencode/Sidekick_bencode.ml b/src/bencode/Sidekick_bencode.ml
new file mode 100644
index 00000000..254d2334
--- /dev/null
+++ b/src/bencode/Sidekick_bencode.ml
@@ -0,0 +1,119 @@
+module V = Ser_value
+
+type t = Ser_value.t
+
+module Encode = struct
+  let to_buffer (buf : Buffer.t) (v : t) : unit =
+    let[@inline] char c = Buffer.add_char buf c in
+    let[@inline] str s = Buffer.add_string buf s in
+    let[@inline] int i = Printf.bprintf buf "%d" i in
+
+    let rec enc_v (v : t) : unit =
+      match v with
+      | Int i ->
+        char 'i';
+        int i;
+        char 'e'
+      | Bool true -> str "i1e"
+      | Bool false -> str "i0e"
+      | Str s | Bytes s ->
+        int (String.length s);
+        char ':';
+        str s
+      | List l ->
+        char 'l';
+        List.iter (fun v -> enc_v v) l;
+        char 'e'
+      | Dict l ->
+        char 'd';
+        Util.Str_map.iter
+          (fun k v ->
+            enc_v (V.string k);
+            enc_v v)
+          l;
+        char 'e'
+    in
+    enc_v v
+
+  let to_string v : string =
+    let buf = Buffer.create 16 in
+    to_buffer buf v;
+    Buffer.contents buf
+end
+
+module Decode = struct
+  exception Fail
+
+  let of_string s =
+    let i = ref 0 in
+
+    let[@inline] check_not_eof () =
+      if !i >= String.length s then raise_notrace Fail
+    in
+
+    let rec top () : t =
+      check_not_eof ();
+      match String.unsafe_get s !i with
+      | 'l' ->
+        incr i;
+        read_list []
+      | 'd' ->
+        incr i;
+        read_map Util.Str_map.empty
+      | 'i' ->
+        incr i;
+        let n = read_int 'e' true 0 in
+        V.int n
+      | '0' .. '9' -> V.string (parse_str_len ())
+      | _ -> raise_notrace Fail
+    (* read integer until char [stop] is met, consume [stop], return int *)
+    and read_int stop sign n : int =
+      check_not_eof ();
+      match String.unsafe_get s !i with
+      | c when c == stop ->
+        incr i;
+        if sign then
+          n
+        else
+          -n
+      | '-' when stop == 'e' && sign && n = 0 ->
+        incr i;
+        read_int stop false n
+      | '0' .. '9' as c ->
+        incr i;
+        read_int stop sign (Char.code c - Char.code '0' + (10 * n))
+      | _ -> raise_notrace Fail
+    and parse_str_len () : string =
+      let n = read_int ':' true 0 in
+      if !i + n > String.length s then raise_notrace Fail;
+      let s = String.sub s !i n in
+      i := !i + n;
+      s
+    and read_list acc =
+      check_not_eof ();
+      match String.unsafe_get s !i with
+      | 'e' ->
+        incr i;
+        V.list (List.rev acc)
+      | _ ->
+        let x = top () in
+        read_list (x :: acc)
+    and read_map acc =
+      check_not_eof ();
+      match String.unsafe_get s !i with
+      | 'e' ->
+        incr i;
+        V.dict acc
+      | _ ->
+        let k = parse_str_len () in
+        let v = top () in
+        read_map (Util.Str_map.add k v acc)
+    in
+
+    try Some (top ()) with Fail -> None
+
+  let of_string_exn s =
+    match of_string s with
+    | Some x -> x
+    | None -> failwith "bencode.decode: invalid string"
+end
diff --git a/src/bencode/Sidekick_bencode.mli b/src/bencode/Sidekick_bencode.mli
new file mode 100644
index 00000000..1a377b67
--- /dev/null
+++ b/src/bencode/Sidekick_bencode.mli
@@ -0,0 +1,14 @@
+type t = Ser_value.t
+
+module Encode : sig
+  val to_buffer : Buffer.t -> t -> unit
+  val to_string : t -> string
+end
+
+module Decode : sig
+  val of_string : string -> t option
+
+  val of_string_exn : string -> t
+  (** Parse string.
+      @raise Error.Error if the string is not valid bencode. *)
+end
diff --git a/src/bencode/dune b/src/bencode/dune
new file mode 100644
index 00000000..ddcd6ca3
--- /dev/null
+++ b/src/bencode/dune
@@ -0,0 +1,6 @@
+(library
+ (name sidekick_bencode)
+ (public_name sidekick.bencode)
+ (synopsis "basic Bencode serialization")
+ (flags :standard -open Sidekick_util)
+ (libraries containers sidekick.util))