B-trees & serialization with bencode

10 files changed, 539 insertions, 51 deletions

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
new file mode 100644
index 0000000..ba20abc
--- /dev/null
+++ b/.github/workflows/tests.yml
@@ -0,0 +1,17 @@
+name: Tests
+
+on:
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.head_ref }}
+
+      - uses: cachix/install-nix-action@v27
+
+      - name: Run tests
+        run: nix develop --command dune test
diff --git a/.github/workflows/uml.yml b/.github/workflows/uml.yml
deleted file mode 100644
index 2d946e8..0000000
--- a/.github/workflows/uml.yml
+++ /dev/null
@@ -1,51 +0,0 @@
-name: UML
-
-on:
-  pull_request:
-
-permissions:
-  contents: write
-
-jobs:
-  blueprint:
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-        with:
-          ref: ${{ github.head_ref }}
-
-      - uses: cachix/install-nix-action@v27
-
-      - name: Run tests
-        run: nix develop --command dune test
-
-      - name: Generate blueprint
-        run: |
-          nix shell nixpkgs#perl --command \
-            perl tools/uml_blueprint.pl --src lib --out docs/architecture.puml --strict --no-notes
-
-      - name: Render diagram
-        run: |
-          nix shell nixpkgs#plantuml --command \
-            plantuml -tsvg docs/architecture.puml
-
-      - name: Upload artifact
-        uses: actions/upload-artifact@v4
-        with:
-          name: architecture-diagram
-          path: |
-            docs/architecture.puml
-            docs/architecture.svg
-
-      - name: Commit diagram to the PR branch
-        if: github.event.pull_request.head.repo.full_name == github.repository
-        run: |
-          git config user.name  "github-actions[bot]"
-          git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
-          git add docs/architecture.puml docs/architecture.svg
-          if git diff --cached --quiet; then
-            echo "unchanged"
-          else
-            git commit -m "ci: update architecture diagram [skip ci]"
-            git push origin HEAD:${{ github.head_ref }}
-          fi
diff --git a/lib/Utilities.ml b/lib/Utilities.ml
index bd604c9..9ae0119 100644
--- a/lib/Utilities.ml
+++ b/lib/Utilities.ml
@@ -1,4 +1,12 @@
 module Result = struct
   let ( let* ) = Result.bind
   let fmap f m = Result.bind m f
+
+  let sequence ms = List.fold_right
+                      (fun m acc ->
+                        let* acc = acc in
+                        let* m = m in
+                        Ok (m::acc))
+                      ms
+                      (Ok [])
 end
diff --git a/lib/Utilities.mli b/lib/Utilities.mli
index c9852f0..cdb53df 100644
--- a/lib/Utilities.mli
+++ b/lib/Utilities.mli
@@ -1,4 +1,5 @@
 module Result : sig
   val ( let* ) : ('a, 'e) result -> ('a -> ('b, 'e) result) -> ('b, 'e) result
   val fmap : ('a -> ('b, 'e) result) -> ('a, 'e) result -> ('b, 'e) result
+  val sequence : ('a, 'b) result list -> ('a list, 'b) result
 end
diff --git a/lib/concepts/Codec.ml b/lib/concepts/Codec.ml
new file mode 100644
index 0000000..758d718
--- /dev/null
+++ b/lib/concepts/Codec.ml
@@ -0,0 +1,182 @@
+ module Error = struct
+  open Condition
+
+  let unexpected_end =
+    condition "codec-unexpected-end" "Unexpected end of input while decoding" empty
+
+  let malformed_integer digits =
+    condition "codec-malformed-integer" "A byte run is not a valid integer"
+      ("digits" |=| Value.String digits)
+
+  let unexpected_byte c =
+    condition "codec-unexpected-byte" "Unexpected byte while decoding"
+      ("byte" |=| Value.String (String.make 1 c))
+
+  let string_out_of_bounds length =
+    condition "codec-string-out-of-bounds" "A string length runs past the end of input"
+      ("length" |=| Value.Integer length)
+
+  let trailing_bytes =
+    condition "codec-trailing-bytes" "Bytes remain after a complete value" empty
+
+  let tag_mismatch ~expected ~actual =
+    condition "codec-tag-mismatch" "A value with a different tag was expected"
+      ("expected" |=| Value.String (String.make 1 expected) &
+       "actual" |=| Value.String (String.make 1 actual))
+
+  let type_mismatch ~expected =
+    condition "codec-type-mismatch" "A Bencode value of a different shape was expected"
+      ("expected" |=| Value.String expected)
+
+  let missing_field key =
+    condition "codec-missing-field" "A required field is absent"
+      ("field" |=| Value.String key)
+
+  let unexpected_object ~expected ~actual =
+    condition "codec-unexpected-object" "The decoded object is not of the expected kind"
+      ("expected" |=| Value.String expected & "actual" |=| Value.String actual)
+end
+
+(* A self-describing serialization used as the uniform wire format for
+   durable objects. Chosen over an ad-hoc key=value scheme because it
+   nests (a multigroup holds a map of relation hashes) and has a
+   single byte form for a given value (dict keys are emitted sorted),
+   which matters once these bytes are content-addressed.
+
+   Grammar:
+   - Int    "i" <decimal> "e" (negatives allowed)
+   - String <length> ":" <raw bytes>
+   - List   "l" <element>* "e"
+   - Dict   "d" (<String key> <value>)* "e" keys sorted on encode *)
+module Bencode = struct
+  type t = Tagged of char * t | Int of int | String of string | List of t list | Dict of (string * t) list
+
+  let rec encode buf = function
+    | Tagged (tag, value) ->
+       Buffer.add_char buf 't';
+       Buffer.add_char buf tag;
+       encode buf value
+    | Int n ->
+       Buffer.add_char buf 'i';
+       Buffer.add_string buf (string_of_int n);
+       Buffer.add_char buf 'e'
+    | String s ->
+       Buffer.add_string buf (string_of_int (String.length s));
+       Buffer.add_char buf ':';
+       Buffer.add_string buf s
+    | List xs ->
+       Buffer.add_char buf 'l';
+       List.iter (encode buf) xs;
+       Buffer.add_char buf 'e'
+    | Dict kvs ->
+       Buffer.add_char buf 'd';
+       List.iter
+         (fun (k, v) ->
+           encode buf (String k);
+           encode buf v )
+         (List.sort (fun (a, _) (b, _) -> String.compare a b) kvs);
+       Buffer.add_char buf 'e'
+
+  let to_bytes (value : t) : bytes =
+    let buf = Buffer.create 256 in
+    encode buf value;
+    Buffer.to_bytes buf
+
+  let as_string = function String s -> Ok s | _ -> Error (Error.type_mismatch ~expected:"string")
+  let as_int = function Int n -> Ok n | _ -> Error (Error.type_mismatch ~expected:"integer")
+  let as_list = function List l -> Ok l | _ -> Error (Error.type_mismatch ~expected:"list")
+  let as_dict = function Dict d -> Ok d | _ -> Error (Error.type_mismatch ~expected:"dict")
+
+  let of_bytes (bytes : bytes) : (t, Condition.condition) result =
+    let open Utilities.Result in
+    let s = Bytes.to_string bytes in
+    let n = String.length s in
+    let pos = ref 0 in
+    let peek () = if !pos >= n then Error (Error.unexpected_end) else Ok (s.[!pos]) in
+    let advance () = incr pos in
+    (* Read the decimal run up to (and consuming) [term]. *)
+    let read_int_from start term =
+      let rec scan () =
+        let* c = peek () in
+        if c = term then begin
+          let digits = String.sub s start (!pos - start) in
+          advance ();
+          match int_of_string_opt digits with
+          | Some i -> Ok i
+          | None -> Error (Error.malformed_integer digits)
+        end
+        else begin
+          advance ();
+          scan ()
+        end
+      in
+      scan ()
+    in
+    let read_int term = read_int_from !pos term in
+    let read_char () =
+      let* c = peek () in
+      advance ();
+      Ok c
+    in
+    let rec parse () =
+      let* c = peek () in
+      advance ();
+      match c with
+      | 't' ->
+         let* tag = read_char () in
+         let* data = parse () in
+         Ok (Tagged (tag, data))
+      | 'i' ->
+         let* n = read_int 'e' in
+         Ok (Int n)
+      | 'l' -> parse_list []
+      | 'd' -> parse_dict []
+      | c when c >= '0' && c <= '9' -> parse_string_from (!pos - 1)
+      | c -> Error (Error.unexpected_byte c)
+    and parse_string () =
+      parse_string_from !pos
+    and parse_string_from start =
+      let* len = read_int_from start ':' in
+      if len < 0 || !pos + len > n then
+        Error (Error.string_out_of_bounds len)
+      else
+        let str = String.sub s !pos len in
+        pos := !pos + len;
+        Ok (String str)
+    and parse_list acc =
+      let* c = peek () in
+      if c = 'e' then (advance (); Ok (List (List.rev acc)))
+      else
+        let* result = parse () in
+        parse_list (result :: acc)
+    and parse_dict acc =
+      let* c = peek () in
+      if c = 'e' then (advance (); Ok (Dict (List.rev acc)))
+      else
+        let* key = parse_string () in
+        let* key = as_string key in
+        let* value = parse () in
+        parse_dict ((key, value) :: acc)
+    in
+    let* value = parse () in
+    if !pos <> n then Error Error.trailing_bytes else
+      Ok value
+
+  let to_blob x = to_bytes x |> Representation.blob_of_bytes
+  let of_blob b = Representation.bytes_of_blob b |> of_bytes
+
+  let with_tag t = function
+    | Tagged (t', v) ->
+       if t = t' then
+         Ok v
+       else
+         Error (Error.tag_mismatch ~expected:t ~actual:t')
+    | _ -> Error (Error.type_mismatch ~expected:"tagged")
+
+  let field key = function
+    | Dict kvs -> (
+      match List.assoc_opt key kvs with
+      | Some v -> Ok v
+      | None -> Error (Error.missing_field key) )
+    | _ -> Error (Error.type_mismatch ~expected:"dict")
+end
diff --git a/lib/concepts/Hash.ml b/lib/concepts/Hash.ml
index f467ee7..aac7428 100644
--- a/lib/concepts/Hash.ml
+++ b/lib/concepts/Hash.ml
@@ -2,8 +2,31 @@
 
 type hash = Digestif.SHA256.t
 
+let size = 256/8
+
+let compare h1 h2 = Digestif.SHA256.unsafe_compare h1 h2
+                    |> Ordering.of_int
+
 let hash_of_value (v : Value.value) =
   Representation.(blob_of_value v |> bytes_of_blob) |> Digestif.SHA256.digest_bytes
 
+let hash_of_bytes bytes = Digestif.SHA256.digest_bytes bytes
+
+let hash_of_blob blob = Representation.bytes_of_blob blob |> Digestif.SHA256.digest_bytes
+
+let to_raw_string hash = Digestif.SHA256.to_raw_string hash
+
+let of_raw_string str = Digestif.SHA256.of_raw_string str
+
 let hash_equals = Digestif.SHA256.equal
 let bytes_of_hash (h : hash) = Bytes.of_string (Digestif.SHA256.to_raw_string h)
+
+let power_of_two p =
+  let b = Bytes.make 32 '\x00' in
+  let byte = 31 - (p / 8) in
+  Bytes.set b byte (Char.chr (1 lsl (p mod 8)));
+  Bytes.unsafe_to_string b
+
+let pick n hash =
+  let threshold = Digestif.SHA256.of_raw_string (power_of_two (256 - n)) in
+  compare hash threshold = Ordering.Smaller
diff --git a/lib/concepts/Hash.mli b/lib/concepts/Hash.mli
index b100cd7..44fb3d0 100644
--- a/lib/concepts/Hash.mli
+++ b/lib/concepts/Hash.mli
@@ -1,11 +1,32 @@
 (** an abstract value representing a hash *)
 type hash
 
+(** the size of a hash, in bytes *)
+val size : int
+
 (** given a value, compute it's hash *)
 val hash_of_value : Value.value -> hash
 
+(** given a bag of bytes, compute it's hash *)
+val hash_of_bytes : bytes -> hash
+
+(** given a blob, compute it's hash *)
+val hash_of_blob : Representation.blob -> hash
+
+(** given a hash, return a representation of it as a string (/not/ the usual hex representation!) *)
+val to_raw_string : hash -> string
+
+(** given a string containing a representation of a hash (as per `to_raw_string`), return the corresponding hash *)
+val of_raw_string : string -> hash
+
+(** compare two hashes lexicographically *)
+val compare : hash -> hash -> Ordering.ordering
+
 (** given two hashes, are they the same? *)
 val hash_equals : hash -> hash -> bool
 
 (** given a hash, return a `bytes` representation of it *)
 val bytes_of_hash : hash -> bytes
+
+(** pick one every 1/2^n hashes *)
+val pick : int -> hash -> bool
diff --git a/lib/concepts/Ordering.ml b/lib/concepts/Ordering.ml
new file mode 100644
index 0000000..42b20e5
--- /dev/null
+++ b/lib/concepts/Ordering.ml
@@ -0,0 +1,4 @@
+type ordering = Equal | Smaller | Greater
+
+let of_int x =
+  if x < 0 then Smaller else if x > 0 then Greater else Equal
diff --git a/lib/concepts/Ordering.mli b/lib/concepts/Ordering.mli
new file mode 100644
index 0000000..ab8a3f4
--- /dev/null
+++ b/lib/concepts/Ordering.mli
@@ -0,0 +1,3 @@
+type ordering = Equal | Smaller | Greater
+
+val of_int : int -> ordering
diff --git a/lib/kernel/Merkle.ml b/lib/kernel/Merkle.ml
new file mode 100644
index 0000000..1fd364f
--- /dev/null
+++ b/lib/kernel/Merkle.ml
@@ -0,0 +1,280 @@
+module type KEY = sig
+  type t
+
+  val encode : t -> Concepts.Representation.blob
+  val compare : t -> t -> Concepts.Ordering.ordering
+  val decode : Concepts.Representation.blob -> (t, Concepts.Condition.condition) result
+end
+
+module type TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> sig
+  type address = Concepts.Hash.hash
+
+  type 'a node
+
+  val find : S.transaction -> address -> ('a node option, Concepts.Condition.condition) result
+
+  val empty : 'a node
+
+  val hash_of : 'a node -> address
+
+  val insert : S.transaction -> 'a node -> K.t -> address -> ('a node, Concepts.Condition.condition) result
+  val remove : S.transaction -> 'a node -> K.t -> ('a node, Concepts.Condition.condition) result
+  val lookup : S.transaction -> 'a node -> K.t -> (address option, Concepts.Condition.condition) result
+end
+
+module Make : TREE = functor (S : Abstract.Storage.STORAGE) (K : KEY) -> struct
+  type address = Concepts.Hash.hash
+
+  module Error = struct
+    open Concepts.Condition
+
+    let malformed_node () = condition "malformed-node" "A node representation did not conform to what was expected" empty
+  end
+
+  type 'a node =
+    Leaf of
+      { keys : K.t BatFingerTree.t;
+        values : address BatFingerTree.t }
+  | Trunk of
+      { keys : K.t BatFingerTree.t;
+        children : address BatFingerTree.t }
+
+  let empty = Leaf { keys = BatFingerTree.empty; values = BatFingerTree.empty }
+
+  let rec lookup1' keys key bottom top =
+    let open Concepts.Ordering in
+    if bottom >= top then (bottom, false)
+    else
+      let mid = bottom + (top - bottom) / 2 in
+      let node = BatFingerTree.get keys mid in
+      match K.compare node key with
+      | Equal -> (mid, true)
+      | Greater -> lookup1' keys key bottom mid
+      | Smaller -> lookup1' keys key (mid + 1) top
+
+  let lookup1 keys key = lookup1' keys key 0 (BatFingerTree.size keys)
+
+  let keys_of = function
+    | Leaf { keys; _ } -> keys
+    | Trunk { keys; _ } -> keys
+
+  let of_bencode =
+    let open Concepts.Codec.Bencode in
+    let open Utilities.Result in
+    (* This is terrible *)
+    let decode_key t = as_string t
+                       |> Result.map String.to_bytes
+                       |> Result.map Concepts.Representation.blob_of_bytes
+                       |> fmap K.decode in
+    let decode_value t = as_string t |> Result.map Concepts.Hash.of_raw_string in
+    let decode_list f data = data
+                             |> fmap as_list
+                             |> Result.map (List.map f)
+                             |> fmap sequence
+                             |> Result.map BatFingerTree.of_list in
+    function
+    | Tagged ('!', (Dict _ as data)) ->
+       let* keys = field "keys" data |> decode_list decode_key in
+       let* values = field "values" data |> decode_list decode_value in
+       Ok (Leaf { keys; values })
+    | Tagged ('#', (Dict _ as data)) ->
+       let* keys = field "keys" data |> decode_list decode_key in
+       let* children = field "children" data |> decode_list decode_value in
+       Ok (Trunk { keys; children })
+    | _ -> Error (Error.malformed_node ())
+
+  let from_blob blob = Concepts.Codec.Bencode.of_blob blob
+                       |> Utilities.Result.fmap of_bencode
+
+  let to_bencode =
+    let open Concepts.Codec.Bencode in
+    (* Ditto. *)
+    let bencode_key k = K.encode k
+                        |> Concepts.Representation.bytes_of_blob
+                        |> String.of_bytes
+                        |> (fun s -> String s) in
+    let bencode_hash h = String (Concepts.Hash.to_raw_string h) in
+    function
+    | Leaf { keys; values } ->
+       Tagged ('!', Dict [("keys", List (BatFingerTree.to_list keys |> List.map bencode_key));
+                          ("values", List (BatFingerTree.to_list values |> List.map bencode_hash))])
+    | Trunk { keys; children } ->
+       Tagged ('#', Dict [("keys", List (BatFingerTree.to_list keys |> List.map bencode_key));
+                          ("children", List (BatFingerTree.to_list children |> List.map bencode_hash))])
+
+  let to_blob node = to_bencode node |> Concepts.Codec.Bencode.to_blob
+
+  (* FIXME *)
+  let hash_of node = to_blob node |> Concepts.Hash.hash_of_blob
+
+  let find tx node =
+    let open Utilities.Result in
+    let* data = S.get tx node in
+    match data with
+    | None -> Ok None
+    | Some data ->
+       let* node = from_blob data in
+       Ok (Some node)
+
+  let find' tx node =
+    let open Utilities.Result in
+    let* child = find tx node in
+    match child with
+    | Some child -> Ok child
+    | None -> failwith "A child node was not found on the underlying storage. Either your database is corrupted, or this is a bug on RNT!"
+
+  let rec lookup tx node key =
+    let open Utilities.Result in
+    let i, found = lookup1 (keys_of node) key in
+    match node with
+    | Leaf { values; _ } ->
+       if found then Ok (Some (BatFingerTree.get values i)) else Ok None
+    | Trunk { children; _ } ->
+       let* child = BatFingerTree.get children (if found then i+1 else i) |> find' tx in
+       lookup tx child key
+
+  let persist tx node =
+    let open Utilities.Result in
+    let data = to_blob node in
+    let addr = Concepts.Hash.hash_of_blob data in
+    let* () = S.put tx addr data in
+    Ok addr
+
+  type 'a op = Update of address * 'a node | Split of address * K.t * address
+
+  let emplace i v ft =
+    let size = BatFingerTree.size ft in
+    if i < 0 || i > size then
+      invalid_arg "Merkle.emplace: index out of bounds"
+    else if size = 0 then
+      BatFingerTree.singleton v
+    else if i = size then
+      BatFingerTree.snoc ft v
+    else
+      let fl, fr = BatFingerTree.split_at ft i in
+      BatFingerTree.singleton v
+      |> BatFingerTree.append fl
+      |> (Fun.flip BatFingerTree.append) fr
+
+  let insert1 node key value =
+    let i, present = lookup1 (keys_of node) key in
+    match node with
+    | Leaf ({ keys; values } as leaf) ->
+       if present then
+         Leaf { leaf with values = BatFingerTree.set values i value }
+       else
+         Leaf { keys = emplace i key keys; values = emplace i value values }
+    (* When in danger or in doubt: run in circles, scream and shout. *)
+    | Trunk _ -> failwith "Cannot insert a value on a trunk node. This is a bug in RNT!"
+
+  let pivot_at ft i =
+    let l, r' = BatFingerTree.split_at ft i in
+    let r, p = BatFingerTree.front_exn r' in
+    l, p, r
+
+  let split_at pivot = function
+    | Leaf { keys; values } ->
+       let kl, kr = BatFingerTree.split_at keys pivot in
+       let vl, vr = BatFingerTree.split_at values pivot in
+       let p = BatFingerTree.get kr 0 in
+       Leaf { keys = kl; values = vl },
+       p,
+       Leaf { keys = kr; values = vr }
+    | Trunk { keys; children } ->
+       let kl, kp, kr = pivot_at keys pivot in
+       let cl, cr = BatFingerTree.split_at children (pivot + 1) in
+       Trunk { keys = kl; children = cl },
+       kp,
+       Trunk { keys = kr; children = cr }
+
+  (* FIXME: this works out to a geometric distribution with roughly 64
+     keys per node. Some of the consequences of this are that:
+     - most of our nodes will end up with one or a few elements, which
+       would make our trees taller than they would normally need to be
+     - some nodes will end up much larger than usual, which is bad for
+       lookup performance
+
+     What we want is a normal distribution centered around our desired
+     node size. See
+     https://www.dolthub.com/docs/architecture/storage-engine/prolly-tree/#controlling-chunk-size
+   *)
+  let should_split = Concepts.Hash.pick 6
+
+  let hash_of_keys ks =
+    let buffer = Buffer.create 1024 in
+    BatFingerTree.iter
+      (fun k ->
+        K.encode k
+        |> Concepts.Representation.bytes_of_blob
+        |> Buffer.add_bytes buffer)
+      ks;
+    Concepts.Hash.hash_of_bytes (Buffer.to_bytes buffer)
+
+  let split_position keys =
+    let rec loop left right =
+      match BatFingerTree.front right with
+      | None -> None
+      | Some (tail, head) ->
+         (* TODO: rolling hash *)
+         let left' = BatFingerTree.snoc left head in
+         let hash = hash_of_keys left' in
+         if should_split hash then
+           Some (BatFingerTree.size left)
+         else
+           loop left' tail
+    in
+    loop BatFingerTree.empty keys
+
+  let commit_node tx node =
+    let open Utilities.Result in
+    match split_position (keys_of node) with
+    | Some i ->
+       let l, p, r = split_at i node in
+       let* l_addr = persist tx l in
+       let* r_addr = persist tx r in
+       Ok (Split (l_addr, p, r_addr))
+    | None ->
+       let* addr = persist tx node in
+       Ok (Update (addr, node))
+
+  let rec insert' tx node key value =
+    let open Utilities.Result in
+    match node with
+    | Leaf _ ->
+       insert1 node key value
+       |> commit_node tx
+    | Trunk { keys; children } ->
+       let i, found = lookup1 keys key in
+       let i = if found then i+1 else i in
+       let* child = BatFingerTree.get children i |> find' tx in
+       let* r = insert' tx child key value in
+       match r with
+       | Update (addr, _) ->
+          Trunk { keys; children = BatFingerTree.set children i addr }
+          |> commit_node tx
+       | Split (l, p, r) ->
+          Trunk { keys = emplace i p keys;
+                  children = BatFingerTree.set children i l |> emplace (i + 1) r }
+          |> commit_node tx
+
+  let insert tx node key value =
+    let open Utilities.Result in
+    let* op = insert' tx node key value in
+    match op with
+    | Update (_, node) -> Ok node
+    | Split (l, p, r) ->
+       let new_root =
+         Trunk { keys = BatFingerTree.singleton p;
+                 children = BatFingerTree.empty
+                            |> (Fun.flip BatFingerTree.snoc) l
+                            |> (Fun.flip BatFingerTree.snoc) r } in
+       let* _ = persist tx new_root in
+       Ok new_root
+
+  let remove tx node key = failwith "TODO" [@@warning "-27"]
+end
+
+(* module type UTILS = functor (T : TREE) (V : VALUE) -> sig *)
+(*   val insert : S.transaction -> 'a node -> 'a key -> V.t -> ('a node, Concepts.Condition.condition) result *)
+(*   val lookup : S.transaction -> 'a node -> 'a key -> (V.t option, Concepts.Condition.condition) result *)
+(* end *)