Source file store.ml

(**************************************************************************)
(* This file is part of Store.                                            *)
(*                                                                        *)
(* Copyright (C) 2023-2024 OCamlPro SAS                                   *)
(* Copyright (C) 2023-2024 Inria                                          *)
(*                                                                        *)
(* Store is distributed under the terms of the MIT license. See the       *)
(* included LICENSE file for details.                                     *)
(**************************************************************************)

(** {1 Store}

    This module implements a _store_, a bag of mutable objects
    (currently: only references). The state of the store at any point
    in time can be described as a functional _mapping_ from the store
    references to their value.

    There is a _current mapping_ of the store, on which the operations
    on store objects operate: [Store.Ref.get s r] reads the value of
    [r] in the current mapping, [Store.Ref.set s r v] updates the
    current mapping to a mapping where [r] maps to [v], etc.

    The implementation also lets user capture the current mapping of
    the store in a _snapshot_ (persistent or ephemeral), that can be
    restored in the future to become the current mapping again.

    The data about all these mappings is stored in a _store graph_,
    where each node represents a mapping; in particular, one
    distinguished node in the graph, that we call the _current node_,
    represents the current mapping.

    In slightly more details:
    - the data of the current mapping is in fact stored in the
      references directly, each reference stores its current value
    - edges in the store graph carry information on how to "go back"
      to a mapping from another.

    {b Invariant}: the store graph has a unique designated current
    node.

    {b Invariant}: The current mapping of the store, given by the
    value stored in each reference, is the mapping of the current
    node of the mapping graph.

    Let us present an example and introduce some key concepts that
    will be used to document precisely the implementation in follow-up
    comments.


    {2 Example}

    Suppose the user creates a store, with a single reference [r] of
    value [0], the corresponding graph starts with a single node A,
    and a current mapping stored in the references themselves:

         A [r = 0]

    (The node A does not store the current value of [r], it is stored
    in [r] itself.)

    Then the user may want to set [r] to [1]. This creates a new node
    B in the graph that represents this new mapping:

         A
     r=0 ↓
         B [r = 1]

    The new edge from A to B remembers that the previous value of [r]
    was [0], to be able to un-apply the change later, and the value
    inside [r] itself gets updated to [1].

    At this point, the mapping of A is [r = 0], and B is the current node
    and its mapping is [r = 1].

    It is possible to "go back" to the mapping of A by starting from
    A, and updating references until we reach the node that represents
    the "current" mapping. The updated store graph is depicted below.

         A [r = 0]
     r=1 ↑
         B

    The current node is A; when we traversed the edge from A to B, we
    updated it to have B point to A, so that following edges always
    leads us to the current node. We also updated the information
    stored on the edge to be able to re-apply the change to [r]. We
    could go back to the mapping of [B] by following the same process.

    At this point the user can set [r] again to [2], which
    creates a third node C:

           r=0
         A  →  C [r = 2]
     r=1 ↑
         B

    One can of course go back to A again at this point.

                   r=2
         A [r = 0]  ←   C
     r=1 ↑
         B

    In the examples above we depicted our store graph with directed
    edges that correspond to reachability in the implementation:
    A points to B if the book-keeping data of the node A contains B.


    {2 Graph theory reminder}

    In graph theory, an (undirected) tree is a certain kind of
    (undirected) graph: a graph that is acyclic (no cycle in
    the graph) and connected (all nodes are reachable from
    each other). In other words, an _undirected tree_ is an undirected
    graph where there exists a unique path between all pairs of nodes.

    {b Invariant}: our store graph is always an undirected tree.

    The notion of "tree" that is common in programming corresponds to
    the notion of _rooted tree_ in graph theory, a tree with
    a designated root node.

    The choice of root uniquely determines the _parent relation_ for
    the graph, the relation that sends nodes to their parent, that is,
    relates A to B when A has B as parent -- there is at most one
    parent, and the root is the only node with no parents. One can see
    this relation as a partial function defined everywhere except on
    the root.

    We will use two properties of trees and rooted trees:

    - rerooting: if we look at a given undirected trees T, and two
      different choices of root M and N, there is a simple relation
      between the parent relations of the M-rooted tree and the
      N-rooted tree: they agree on all nodes, except on the (unique)
      path from M to N where they are inverse of each other.

    - removal: to remove a node M from a rooted tree while preserving
      connectedness, we must necessarily remove at the same time all
      its ancestors or all its descendants.


    {2 Current tree}

    We call _current tree_ the rooted tree over our store graph
    whose root is the current node. We may now write _current root_
    for the current node.

    {b Invariant}: the parent relation in the current tree corresponds
    exactly to the edge direction in the (directed) store graph.

    {b Invariant}: when a new node gets created in the graph, it
    always becomes the new current root, and becomes the parent of the
    previous current root.

    In the example above, when we added C we moved from a tree
      B→A
    to
      B→A→C
    We can see that C is the new root, and the parent of the previous
    root A.


    {2 Historic tree, history of a node}

    We call _historic tree_ the rooted tree over our graph whose root
    is the first node to have been created, and where the parent of
    each node A is the node that was the current root when A was added
    to the graph.

    In our example above, we created B when A was the current root,
    at which points the historic tree was
      A
      ↑
      B
    then went back to A being the current root, which does not change
    the historic tree, and then we created C, so we had the following
    historic tree
      A←C
      ↑
      B
    for the following store graph (left) and current tree (right):
            r=0
          A  →  C [r = 2]                  A→C
      r=1 ↑                                ↑
          B                                B

    The current tree and historic tree can be compared precisely, as
    per the rerooting property:

    - On the path between the current root (C in the example) and the
      historic root (A in the example), their parent relations are
      inverse to each other.

    - Outside this path, their parent relation agree with each other
      (see B in the example).

    We call _history_ of a node the path from this node to the
    historic root.


    {2 Snapshots}

    A user can _capture_ the current mapping in a _snapshot_, and
    later _restore_ the snapshot so that its mapping becomes the
    current mapping again.

    A snapshot is implemented directly as the node that was the
    current node at the time of the capture -- the _snapshot
    node_. Restoring the snapshot is a rerooting operation, where the
    snapshot node becomes the new current root. To implement this we
    traverse the path from the snapshot node to the current root and
    reverse the edges -- as suggested by the rerooting property.

    Remark: the mapping corresponding to a node may in fact change
    over time as the graph is mutated: as long as it is not a snapshot
    node, the user cannot observe its mapping directly. Our
    implementation takes advantage of the flexibility to mutate the
    mapping of non-captured nodes to perform 'record elision'.

    Remark: there can only be a _fork_ in the historic tree, a node
    A with more than one child, if A was captured as
    a snapshot. Indeed, we had to make A the current root again to
    create its second child, and only snapshot nodes can become the
    root again.


    {2 Transactions}

    The snapshots we discussed so far are persistent; in particular,
    they can be restored several times. For efficiency reasons we also
    implement _ephemeral snapshots_ that can only be restored
    once.

    Creating an ephemeral snapshot creates a new node in the store
    graph, as explained below. Restoring an ephemeral snapshot removes
    the historic descendants of this node from the graph, which is
    faster than updating their information to allow restoring their
    mappings -- as is done with usual snapshots.

    Ephemeral snapshots are implemented using specific 'transaction'
    nodes in the graph. Taking an ephemeral snapshot when the current
    root is A creates a new transaction node T(A) that points to this
    ephemeral snapshot.

      ...→A
      becomes
      ...→A→T(A)

    Ephemeral snapshots follow a strict temporal nesting discipline:

    1. An ephemeral snapshot may only be restored while we are
       "inside" it, that is, when the current root is
       a historic children of its transaction node.

    2. Snapshots (ephemeral or not) created while "inside"
       a transaction are only valid until the transation is
       terminated. Restoring an invalid snapshot fails with an error.

    If/when the snapshot on A is restored, we remove the transaction
    node T(A) from the graph. We also remove all historic descendants of
    T(A), so in particular connectedness is preserved. A becomes the
    current root again.

      ...→A→T(A)→...→B
      becomes
      ...→A

    (Removing nodes is where the efficiency gains come from: there is
    less work to perform as we do not need to revert edges, etc.)

    This operation is only allowed if the current root B is "inside"
    the transaction T(A), that is, if the current root is a historic
    descendant of T(A).

    Removing T(A) and its historic descendants also invalidates the
    ephemeral snapshot, as well as all the snapshots
    (ephemeral or not) that are capturing T(A) or one of its historic
    descendants.

    The usual way of restoring an ephemeral snapshot changes the
    current mapping to become the mapping of the snapshot node, we
    call this 'rolling back' the transaction. We also implement
    a 'commit' variant which consumes the snapshot, but does not
    change the current mapping. All the snapshots inside the
    transaction node are invalidated, but all other operations from
    the transaction node to the current root are preserved.

    Remark: one could think of a similar design without a transaction
    node, where each node stores a reference to its ephemeral
    snapshots. But this approach makes it difficult to remember the
    historic ordering between snapshots. If we first take a persistent
    snapshot S of some node A, then an ephemeral snapshot E, consuming
    E must not invalidate S. On the other hand if we take an ephemeral
    snapshot E, then a persistent snapshot, we want consuming E to
    invalidate S; this works because capturing E creates a new
    transaction node T which becomes the root, and in particular
    S does not capture the original node A but the transaction node
    T which will get invalidated.


    {2 Memory liveness}

    Nodes in the store graph point towards the current root: the
    memory they can reach is the path from themselves to the current
    root.

    In particular, very old nodes may be collected if all snapshots
    older than them become unreachable; and conversely the nodes
    corresponding to changes in the 'future' of the current root, that
    have been un-applied by restoring an earlier snapshot, may also be
    collected if all snapshots after them become unreachable.
*)

(** {1 Stores} *)

(** {3 Generations}

    We call "generation" of a node the number of (valid) snapshots
    capturing it or any other node in its history.

    If two nodes belonging to the same history have the same
    generation, then there is no snapshot between them.
*)
module Generation : sig
  type t = private int
  val zero : t
  val succ : t -> t
end = struct
  type t = int
  let zero = 0
  let succ n = (n + 1)
end

(** A reference in a store.

    Each reference belongs to a unique store, and it is a programming
    error to operate on a reference of the wrong store. (We could
    remember the store of the reference and check for consistency, but
    this would be costly in time and memory.)

    Note: ['a rref] is exported as ['a Ref.t] below, but it needs to
    be defined early as [data] depends on it.
*)
type 'a rref = {
  mutable value : 'a;
  (** The value of the reference in the current mapping of the
      store. *)

  mutable generation : Generation.t;
  (** The generation of the last [Set] node for this reference in the
      current history. *)
}

(** A [store] tracks the global state of the store graph. *)
type store = {
  mutable root : node;
  (** The current root of the store. *)

  mutable generation : Generation.t;
  (** The generation of the current root. *)
}

(** A [snapshot] is a persistent reference to a mapping in the store graph. *)
and snapshot = {
  store : store;
  (** The store to which the snapshot belongs.
      (This is for error checking.) *)

  root : node;
  (** The root of the current tree at the time
      the snapshot was taken. *)

  generation : Generation.t;
  (** The generation of the snapshot. *)
}

and transaction = {
  snapshot: snapshot;
  node : node;
}

and node = data ref

and data =
  | Mem
  (** [Mem] indicates that the node is the current root.

      The mapping of this node is given by the value stored in each
      reference. *)

  | Invalid
  (** [Invalid] indicates that the node was removed from the graph.

      It does not represent a mapping anymore. *)

  | Transaction of node * node
  (** [Transaction (n0, n)] records a transaction started from the node [n0].

      The node has the same mapping as [n]. *)

  | Log_set : 'a rref * 'a * Generation.t * node -> data
  (** If the node is A, the data [Log_set (r, v, g, B)] represents an edge from
      A to B that logs a write to the reference [r] whose value before the
      write was [v], and whose generation was [g].

      If [m] is the mapping of B, then the mapping of A is [m[r ↦ v]]. *)

type t = store

module Error = struct
  let[@inline] invalid_operation ~op ~obj =
    Printf.ksprintf invalid_arg
      "Store.%s: invalid %s" op obj

  let[@inline never] invalid_restore () =
    invalid_operation ~op:"restore" ~obj:"snapshot"

  let[@inline never] invalid_rollback () =
    invalid_operation ~op:"rollback" ~obj:"transaction"

  let[@inline never] invalid_commit () =
    invalid_operation ~op:"commit" ~obj:"transaction"

  let[@inline] wrong_store ~op =
    Printf.ksprintf invalid_arg
      "Store.%s: wrong store" op
end

let create () : store = {
  root = ref Mem;
  generation = Generation.zero;
}

(** [log A r v g] logs the antioperation [r := v] by adding a new node to the
    current tree. It transforms
      ...→A
    into
      ...→A→B
    where the edge from A to B logs [r := v], and returns the new root B.

    In formal terms, if the mapping of the current root A is [m]
    before the call, then [log A r v g] creates and returns a new current
    root B with mapping [m], and changes the mapping of A to [m[r := v]].
*)
let[@inline always] log_set old_root r v g =
  let new_root = ref Mem in
  old_root := Log_set (r, v, g, new_root);
  new_root

let[@inline always] set r v g =
  r.value <- v;
  r.generation <- g

let[@inline always] set_and_log old_root new_root r v g =
  old_root := Log_set (r, r.value, r.generation, new_root);
  set r v g;
  new_root

(** {1 Store references} *)

(** {3 Record elision}

    Suppose we want to set a reference [r] to a new value [v].
    Let S be the latest [Set] node for this [r] in the current
    history. If the generation of [r] is the same as the current
    generation, it means that no snapshot was taken between S and
    the current root. In this case there is no need to record an
    additional [Set] antioperation, as restoring any snapshot
    will un-apply S and get its intended value for [r].

    We call this the _record elision_ optimization. It mutates
    the mapping of all nodes on the segment from the current root
    to S excluded. For each such node, let its previous mapping
    be [m], its new mapping is [m[r ↦ v]]. Because there is
    no snapshot on this segment, mutating the mapping of those
    nodes is not observable.

    (Note: a priori the nodes on this segment may have other
    descendants outside the segment, but such a "fork" is
    impossible if none of the nodes of the segment have been
    taken as snapshots.)

    Note: {!get} and {!set} are the performance-critial operations in
    Store, and this is a very important fast path in the common
    workload where snapshot captures are infrequent and many writes to
    the same references occur between them. With record elision, the
    performance of {!set} is basically as fast as with OCaml's raw
    references -- the generation check is faster than the write, at
    least for non-immediate values.

    {b Invariant} (_record unicity_): for a given reference, along any
    history, there is at most one [Set] node on that reference per
    generation.
*)

module Ref = struct
  type 'a t = 'a rref

  let make (_ : store) value =
    (* We use the generation 0, which is the same as the generation of
       the initial store right after its creation and before the first
       snapshot is taken. This means that writes happening at this
       point will not be recorded (they benefit from
       record elision). This is the desired behavior, as there is no
       way for users to 'restore' an earlier mapping. *)
    { value; generation = Generation.zero }

  let[@inline always] get (_ : store) (r : 'a rref) : 'a =
    r.value

  (** If the current mapping is [m],
      [slow_set s r v] changes it to [m[r ↦ v]].
      (See {!set} below first.) *)
  let slow_set (s : store) (r : 'a rref) (v : 'a) : unit =
    (* Let A be the current root with mapping [m]. *)
    s.root <- set_and_log s.root (ref Mem) r v s.generation;
    (* Now [r] points to [v], and A has a new mapping [m']
       with [op(m') = m].

       At this point, the root is a new node B with mapping [m'].
       The mapping of A changed to [op(m')], that is,
       it is [m] as before. *)
    ()

  (** If the current mapping is [m],
      [set s r v] changes it to [m[r ↦ v]]. *)
  let[@inline always] set (s : store) (r : 'a rref) (v : 'a) : unit =
    if s.generation = r.generation
    then
      (* record elision *)
      r.value <- v
    else
      (* The slow path where an antioperation needs to be recorded. *)
      slow_set s r v

  let eq (_ : store) (x : 'a t) (y : 'a t) : bool =
    x == y
end


(** {1 Snapshots} *)

let snapshot (s: store) : snapshot = {
  store = s;
  root = s.root;
  generation = s.generation;
}

let capture s =
  let snap = snapshot s in
  (* The new root is reachable by a new snapshot,
     so we increment the current generation. *)
  s.generation <- Generation.succ s.generation;
  snap

(** [reroot_restore n] makes [n] the current root
    while preserving the mapping of all nodes in the graph. *)
let reroot_restore new_root =
  let rec collect n acc =
    (* Return the current root of G, and its descendants along the
       path to [n] (included) ordered from child to parent. *)
    match !n with
    | Mem ->
      n, acc
    | Invalid ->
      Error.invalid_restore ()
    | Transaction (_tr, n') ->
      collect n' (n :: acc)
    | Log_set (_r, _v, _g, n') ->
      collect n' (n :: acc)
  in
  let old_root, seg = collect new_root [] in
  let rec revert n seg =
    (* Precondition of [revert n seg]:
       - [n] has mapping [m]
       - the current mapping is also [m]
       - [seg] is a segment from [n] excluded to [new_root] included,
         ordered from parent to child
       - if all the nodes in [seg] are changed so that they have the
         same mapping as they had in G, then all nodes in the store
         graph will have the same mappings as in G.

       [revert] mutates [n] and the nodes in [seg],
       but no other nodes.

       Postcondition:
       - all nodes have the same mapping as in G
       - [new_root] is the current root
    *)
    match seg with
    | child :: rest ->
      (* [child] is a child of [n],
         let [m'] be its mapping. *)
      begin match !child with
      | Mem | Invalid ->
        (* Cannot be the data of a child node in the graph. *)
        assert false
      | Transaction (tr, n') ->
        assert (n' == n);
        (* [Transaction] node: [m] and [m'] are the same mapping. *)
        n := Transaction(tr, child);
        (* Checking the preconditions of [revert child rest]:
           - the current mapping is [m] and thus also [m']
           - If [child] recovers its mapping [m'], then [n] will have the
             same mapping, which is also the expected mapping [m]. *)
        revert child rest
      | Log_set (r, v, g, n') ->
        assert (n' == n);
        (* [child] points to its current parent node [n]
           with antioperation [op], so we have [m' = op(m)].

           We mutate the current mapping from [m] to [op(m) = m']. *)
        ignore (set_and_log n child r v g);
        (* Checking the preconditions of [revert child rest]:
           - the current mapping is [m']
           - If [child] recovers its mapping [m'], then [n] will have
             mapping [rev_op(m') = m], as expected. *)
        revert child rest
      end
    | [] ->
      (* If the segment is empty, then [n] must be [new_root]. *)
      let new_root = n in
      new_root := Mem;
      (* Now [new_root] has the mapping [m] has expected;
         [seg] is empty so all the nodes have recovered the mapping
         they had in G. *)
  in
  revert old_root seg

let restore s snap =
  if s != snap.store then Error.wrong_store ~op:"restore";
  if !(snap.root) = Mem then
    (* Fast path if we are restoring the current root.

       (This comes up naturally if one implements an
       observationally-pure interface on top of Store, where snapshots
       are restored before each operation and captured after each
       operation. In almost all cases the snapshot restored will be on
       the current mapping.) *)
    ()
  else begin
   reroot_restore snap.root;
   (* [snap.root] is now the current root. *)
   s.root <- snap.root;
   (* The new root is reachable by a snapshot, at generation
      [snap.generation], that was not present in the history
      of the snapshot root. *)
   s.generation <- Generation.succ snap.generation
 end


(** {1 Transactions} *)

let transaction (s : store) : transaction =
  let snap = snapshot s in
  let old_root = s.root in
  let new_root = ref Mem in
  assert (!old_root = Mem);
  old_root := Transaction (old_root, new_root);
  s.root <- new_root;
  s.generation <- Generation.succ s.generation;
  {
    snapshot = snap;
    node = new_root;
  }

(** {3 Selective invalidation}

    Instead of invalidating all the nodes below a consumed
    transaction, we invalidate only:
    - the final [Mem] node
    - each [Transaction] node encountered along the way

    Invalidating the final [Mem] node guarantees that any
    operation trying to restore a snapshot, persistent or
    ephemeral (so restore, rollback, commit) will fail when
    it reaches the 'end' of the segment.

    However, this is not quite enough due to a combination of:
    1. our current implementations of 'commit' and 'rollback',
       which perform mutations on nodes before they
       reach the end of the segment (unlike 'restore')
    2. our exceptations that functions fail cleanly on invalid inputs,
       before performing any mutation.

    To preserve both (1) and (2) -- both of which are cute choices
    that are not part of our user-facing specification -- we must ensure
    that commit/rollback on an invalidated transaction fail before performing
    any mutation. This is done by invalidating the [Transaction] node to
    ensure that they fail immediately.

    In contrast, 'restore' is not necessarily called on a transaction
    node, but its implementation does not perform any update before
    reaching the final Mem node. *)

(** If the mapping of [n] is [m], then [reroot_rollback n] invalidates
    the current ancestors of [n], and makes it the current root,
    still at mapping [m]. *)
let reroot_rollback g n =
  (* If the node [n] has mapping [m] before the call, then
     [reroot n] ensures that:

     1. The current mapping of the graph becomes [m].

     2. [n] and its current ancestors become [Invalid].
        In particular there is no [Mem] node anymore -- temporarily.
        Other nodes are unchanged. *)
  let rec reroot n =
    match !n with
    | Mem ->
      n := Invalid;
    | Transaction (_n0, n') ->
      n := Invalid;
      reroot n';
    | Log_set (r, v, g', n') ->
      reroot n';
      (* If [g' > g], then this is not the last write on this
         reference in history, we can elide it. *)
      if g' <= g then
        set r v g'
    | Invalid->
      Error.invalid_rollback ()
  in
  ignore reroot;
  (* [reroot' n] is a tail-recursive variant of [reroot n] that has
     the same external specification. The specification is not strong
     enough to verify the recursive calls, but one can reason about
     the function by comparing it to [reroot]. *)
  let rec reroot' n =
    match !n with
    | Mem ->
      n := Invalid;
    | Transaction (_n0, n') ->
      n := Invalid;
      reroot' n';
    | Log_set (r, v, g', n') ->
      (* Writing to [r] can be moved *before* the recursive call
         (compared to [reroot] above), because the generation
         condition, combined with the 'record unicity' invariant,
         guarantees that we will write only once per reference, so
         writing earlier or later does not change the result. *)
      if g' <= g then
        set r v g';
      reroot' n';
    | Invalid->
      Error.invalid_rollback ()
  in
  (* Let [m] be the mapping of [m]. *)
  reroot' n;
  (* At this point the current mapping is [m], and [n] and its current
     ancestors are invalid, there is no [Mem] node anymore. *)
  n := Mem;
  (* The store graph is well-formed again, and [n] is the current
     root at mapping [m] as expected. *)
  ()

let get_transaction_snapshot ~invalid (trans : transaction) : snapshot =
  let snap = trans.snapshot in
  match !(snap.root) with
    | Transaction (n0, n) ->
      (* A transaction may only be terminated from 'inside', when the
         current root is a historic children of the transaction
         node. We are 'inside' when the Transaction edge goes from the
         snapshot root to the transaction node, and not the other way
         around -- and not to some other transaction node. *)
      if n0 == snap.root && n == trans.node
      then snap
      else invalid ()
    | _ -> invalid ()

(** [rollback s trans] reverts the store to the root and mapping
    of [trans], and invalidates [trans] and any other transaction
    and snapshot taken on the historic descendants of [trans]. *)
let rollback (s : store) (trans : transaction) : unit =
  let snap =
    get_transaction_snapshot ~invalid:Error.invalid_rollback trans in
  if s != snap.store then Error.wrong_store ~op:"rollback";
  (* Invalidating the current ancestors of [snap.root] will in
     particular invalidate the [Transaction] node and in turn all its
     historic descendants. *)
  reroot_rollback snap.generation snap.root;
  (* At this point [snap.root] is the current root, and its
     mapping is unchanged. Its ancestors in the previous current
     tree have been invalidated, which invalidates all its historic
     descendants. *)
  s.root <- snap.root;
  (* [trans] has been invalidated, so the number
     of snapshots in history is back to [snap.generation]. *)
  s.generation <- snap.generation;
  ()

(** Preconditions for [reroot_commit g n]:
    - the store graph is [G]
    - the current mapping is [mcur],
    - [n] has generation [g]
    - [n] is a historic ancestor of the current root
    - [s] is the first snapshot node from [n] to the historic root

    Postconditions on returning a node [n']:
    - the current ancestors of [n] in [G] are now invalidated
    - [s] and its current descendants are unchanged from [G]
    - [n'] is the new current root
    - [n'] has generation [g]
    - [n'] has mapping [mcur]
*)
let reroot_commit g n =
  (* Invariants of [reroot]:

     - [dst] has generation [g]

     - setting [dst := !src] results in a valid store graph whose
       current mapping is [mcur], and where [s] and its current
       descendants are unchanged from [G].

     Its postcondition is the one of [reroot_commit] above. *)
  let rec reroot g ~src ~dst =
    match !src with
    | Mem ->
      src := Invalid;
      dst := Mem;
      (* From our invariant we now know that the current mapping is
         [mcur], and therefore [dst] which is the current node must
         have mapping [mcur]. *)
      dst
    | Transaction (_n0, src') ->
      src := Invalid;
      (* [src] and [src'] have the same mapping,
         so the invariant is preserved below. *)
      reroot g ~src:src' ~dst
    | Log_set (r, v, g', src') ->
      (* The generation [g'] is the generation of the reference [r] at
         the time this [set] operation was performed. It counts the
         number of snapshots before the last write to that reference
         on the path from the history root to this node. *)
      if g' < g then (
        (* If [g' < g], then the last [Set] node occurred historically
           before [s] in the graph. We must log the new write
           (at generation [g]), otherwise we could not restore [s]
           correctly -- formally, the mapping of [s] would change. *)
        r.generation <- g;
        reroot g ~src:src' ~dst:(log_set dst r v g')
      ) else (
        (* If [g' >= g], then there is a [Set] node after [s] already.
           We must not log the new write, as this would break the
           'record unicity' invariant. *)
        begin
          (* The current write may have happened at a higher generation
             than [g]. We need to fix the generation of [r] to be no
             higher than [g]. *)
          if g' = g then
            r.generation <- g
          else
            (* If [g' > g], then we have already encountered a [Set]
               node at generation [g'] during traversal, and we have
               already fixed the generation of [r]. *)
            ()
        end;
        reroot g ~src:src' ~dst
      )
    | Invalid ->
      Error.invalid_commit ()
  in
  (* [n] has generation [g], setting [n := !n] trivially gives a valid
     store graph of current mapping [mcur]. *)
  reroot g ~src:n ~dst:n

(** [commit s trans] preserves the current mapping of the store, but
    moves its to a root at the same generation as [trans]. This
    invalidates [trans] and any other transaction and snapshot taken
    on the historic descendants of [trans]. *)
let commit (s : store) (trans : transaction) : unit =
  let snap =
    get_transaction_snapshot ~invalid:Error.invalid_commit trans in
  if s != snap.store then Error.wrong_store ~op:"commit";
  s.root <- reroot_commit snap.generation snap.root;
  (* [trans] has been invalidated, so the number
     of snapshots in history is back to [snap.generation]. *)
  s.generation <- snap.generation;
  ()

(** {2 High-level wrappers} *)

let with_transaction ~on_success (s : store) (f : unit -> 'a) : 'a =
  let t = transaction s in
  match f () with
  | result ->
    on_success s t;
    result
  | exception e ->
    let b = Printexc.get_raw_backtrace () in
    rollback s t;
    Printexc.raise_with_backtrace e b

let temporarily s f =
  with_transaction ~on_success:rollback s f

let tentatively s f =
  with_transaction ~on_success:commit s f