Source file store.ml

(**************************************************************************)
(* This file is part of Store.                                            *)
(*                                                                        *)
(* Copyright (C) 2023-2025 OCamlPro SAS                                   *)
(* Copyright (C) 2023-2025 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 Storable types}

    In addition to basic references from the module [Store.Ref], Store
    supports working with user-defined "storable structures". Like
    references, storable structures have a current value, but
    end-users can introduce other operations than [set], and they have
    to provide specific logic to support backtracking.

    From a distance, storable structures are just like storable
    references, but the backup data that they store in the store graph
    edges is user-defined, as well as the function to unapply or re-apply
    this backup data to the current version of the structure.


    {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

(** {2 Definitions of stored data structures} *)

(** 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. *)
}

(** An absorbing structure in a store.

    Absorbing structures are used-defined structures that behave very
    much like store references, except that the logic to "capture" the
    current value of the reference and to "restore" past or future
    values is user-defined.

    The type ['s] is the (mutable) type of the absorbing structure.

    The type ['op] is the type of sequences of operations on the
    absorbing structure; each sequence goes from an {e origin} state
    [v] to an {e end} state [v'].

    The type ['anti] is the type of antioperations, undoing operations
    on the absorbing structure. Each antioperation acts like
    a snapshot that can return to a previous state of the structure,
    which we call its {e return} state [v].
*)
type ('s, 'op, 'anti) aref = {
  s : 's;
  (** The absorbing structure. *)

  mutable generation : Generation.t;
  (** The generation of the last operation node for this structure
      in the current history. *)

  descr : ('s, 'op, 'anti) absorbing_descr;
  (** User-defined functions to operate on the absorbing structure. *)
}

(** User-provided functions to operate on an absorbing structure.
    These functions are not allowed to operate on the store -- they
    must not read or modify stored references -- as they may be called
    when the store is in an inconsistent intermediary state. *)
and ('s, 'op, 'anti) absorbing_descr = {
  capture : 's -> 'anti;
  (** [capture s] captures the current value [v] of the absorbing
      structure [s], returning an anti-operation with return state [v]. *)

  rollback : 's -> 'anti -> unit;
  (** [rollback s anti] rolls back any change that happened since
      [anti] was captured, restoring [s] to its return state [v].

      The user-provided [rollback] function is used by the
      semi-persistent operations of the store
      (operating on transactions). *)

  undo : 's -> 'anti -> 'op;
  (** [undo s anti] behaves like [rollback v anti], but it also return
      a description of the operations it rolled back, a sequence of
      operation from the return state [v] of [anti] to the current
      state [v'] of the structure [s].
  
      [undo] and [redo] are used by the persistent operations of the
      store (operating on persistent snapshots). *)

  redo : 's -> 'op -> 'anti;
  (** [redo v op] re-applies operations that were previously undone.
      If [op] goes from [v] to [v'], then [s] must be in state [v]
      and [redo v op] must restore its state to [v']. The returned
      anti-operation returns to [v]. *)
}

(** A custom structure in a store.

    Custom structures are used-defined structures that behave a bit
    differently from references. They are more general, they can be
    used for data structures which cannot be presented as absorbing
    structures, but in return the API is more complex.
    
    Operations ['op] are as in absorbing structures, they have an
    origin and end state.

    Instead of "capturing" their current value before an arbitrary
    sequence of operation, each operation is logged by adding change
    description to a mutable "trail". The user provides logic to
    rollback the changes in a trail and to append the changes of
    a trail into another trail.

    A trail has the following characteristics:

    - an {b immutable} {e origin} state. The trail can only ever be
      used to go back to this origin state, which never changes.

    - a {b mutable} {e end} state. As the state of the data structure
      changes, operations are recorded on the trail so that the final
      state matches the current state of the data structure. A trail can
      only be used when the data structure is in the final state of the
      trail.

    - a {e status} which is either {e ready}, when the state can be used
      to undo operations, and {e undone}, when it has been used to 
      undo operations but may be used again in the future after redoing
      those operations.
*)
type ('s, 'op, 'trail) cref = {
  s : 's;
  (** The custom structure.
      Let [v] be its current state. *)

  mutable generation : Generation.t;
  (** The generation of the last operation node for this structure
      in the current history.
      Let [v_prev] be the state of the custom structure at this node. *)

  descr : ('s, 'op, 'trail) custom_descr;
  (** User-defined functions to operate on the custom structure. *)

  mutable trail : 'trail;
  (** The trail currently in use to record changes to the structure.
      It is ready, has origin [v_prev] and ends in [v]. *)
}

(** User-provided functions to operate on a custom structure. These
    functions are not allowed to operate on the store -- they must not
    read or modify stored references -- as they may be called when the
    store is in an inconsistent intermediary state.

    The user has to respect the obligations related to the origin
    state and end state of the trail -- it is a user mistake to provide
    a trail when the custom structure is not in the expected state.
    
    On the other hand, the guarantees that the Store interface
    provides on the trail status give implementation flexibility to
    the user, which can decide on the invariants they prefer for {e
    undone} trails. In particular, the user can decide that [undo
    s trail] leaves the ready trail unchanged (in which case the
    [trail] argument of [redo s op trail] can be ignored), or on the
    contrary it can consume its data, which [redo s op trail] must
    then reconstruct from [op].
*)
and ('s, 'op, 'trail) custom_descr = {
  record : 's -> 'trail;
  (** [record s] starts a new trail for the custom structure [s].
      If the current state of [s] is [v], then the trail [record s]
      has origin state [v] and final state [v], and is ready.
  *)

  append : 'trail -> 'trail -> unit;
  (** [append trail1 trail2] appends [trail2] at the end of [trail1].
      (Both trails must be ready.)

      Before calling [append]:
        - [trail1] has some origin state [v1] and some final state [v2].

        - [trail2] has origin state [v2] and some final state [v3].

      After calling [append]:
        - [trail1] has origin state [v1] and final state [v3].

        - [trail2] is consumed and will not be used anymore (it may be in any
          state).
  *)

  rollback : 's -> 'trail -> unit;
  (** [rollback s trail] undoes the operations recorded in [trail].

      Let us assume that [trail] has origin state [v] and final state
      [v']. Store may only call [rollback] when the [trail] is ready
      and [s] is in state [v']. When [rollback] returns, [s] must be
      in state [v] and and [trail] is consumed -- it will never be
      reused by Store afterwards.  *)

  undo : 's -> 'trail -> 'op;
  (** [undo s trail] undoes the operations recorded in [trail], and
      returns a set of operations to go back to the current state.

      Let us assume that [trail] has origin state [v] and final state
      [v']. Store may only call [undo] when [trail] is ready and [s]
      is in state [v']. When [undo] returns, [s] must be in state [v]
      and the [trail] is undone. The returned sequence of operations
      goes from [v] to [v'].
  *)

  redo : 's -> 'op -> 'trail -> unit;
  (** [redo s op trail] applies the operations recorded in [op] to the
      custom data structure.

      Let us assume that [op] goes from [v] to [v']. Store may only
      call [undo] when [s] has state [v] and [trail] is undone, has
      origin [v] and end [v']. When [redo] returns, [s] must be in
      state [v'] and [trail] is ready.  *)
}

(** {2 Definition of the store} *)

(** 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]]. *)

  | Log_op : 'a * ('a -> unit) * ('a -> unit) * node -> data
  (** [Log_op(d, op, anti, B)] represents an edge from A to B that
      logs an arbitrary operation on some mutable data [d], that was
      applied by calling [op] and can be un-applied by calling [anti].
      
      If [m] is the mapping of B and [v] is the value [m[d]], and if
      calling [anti d] in that state result in a previous value [v']
      of [d], then the mapping of [A] is [m[d ↦ v']].
      Furthermore, we have the invariant that calling [op d] in that
      state [d ↦ v'] will revert the value of [d] back to [v]. *)

  | Log_absorbing
  : ('s, 'op, 'anti) aref * ('op, 'anti) absorbing_entry * Generation.t * node -> data
  (** [Log_absorbing(s, entry, g, B)] represents an edge from A to B that
      logs an operation on an absorbing structure [s] whose generation
      was [g].

      Let us write [m[s ↦ vB]] for the mapping of B, and [m[s ↦ vA]]
      for the mapping of A. Then [entry] is an absorbing
      entry from [vB] back to [vA]. *)

  | Log_custom
  : ('a, 'op, 'trail) cref * ('op, 'trail) custom_entry * Generation.t * node -> data
  (** [Log_custom(s, entry, g, B)] represents an edge from A to B that
      logs an operation on a custom structure [s] whose generation
      was [g].

      Let us write [m[s ↦ (oB, vB)]] for the mapping of B, where [vB]
      is the state of the structure at [B] and [oB] is the origin
      state of its trail, and [m[s ↦ (oA, vA)]] for the mapping of
      A. Then we ask that [vA = oB], and [entry] is a custom entry
      over [oA, vA, vB]. *)

(** Internal data to store on edges corresponding to operations on
    absorbing structures, which keep an anti-operation or an operation
    depending on whether they point to their historic parent or away
    from it. An entry goes from a state [v_from] of the absorbing
    structure back to another state [v_to]. *)
and ('op, 'anti) absorbing_entry =
  | Undo of 'anti
    (** [Undo anti] carries an anti-operation [anti] with return
        state [v_to] -- [v_from] can be any state. *)
  | Redo of 'op
    (** [Redo op] carries a sequence of operations [op] with origin
        state [v_from] and end state [v_to]. *)

(** Internal data for edges corresponding to custom
    structures. An entry goes from a state [v_from] to
    another state [v_to]. *)
and ('ops, 'trail) custom_entry =
| Undo of { parent: 'trail }
  (** [parent] is a {e ready} trail with origin [v_to]
      that ends in [v_from]. *)
| Redo of { ops: 'ops; trail: 'trail }
  (** [ops] is an operation from [v_from] to [v_to], and [trail] is an
      {e undone} trail with origin [v_from] and end [v_to]. *)

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

let[@inline always] log_op old_root ~op ~anti t =
  let new_root = ref Mem in
  old_root := Log_op (t, op, anti, new_root);
  new_root

let[@inline always] op_and_log old_root new_root ~op ~anti t =
  old_root := Log_op (t, op, anti, new_root);
  op t;
  new_root

(* if [r] is an absorbing structure in state [v],
   and [entry] goes from [v] back to [v'], then after
   [revert_absorbing_op r entry], [r] is in state [v']
   and the returned entry goes from [v'] back to [v]. *)
let revert_absorbing_op (r : _ aref) entry =
  match entry with
  | Undo anti ->
    (* [anti] returns to [v'] *)
    Redo (r.descr.undo r.s anti)
    (* [r.descr.undo r.s anti] puts [r] in state [v'],
       and returns an operation from [v'] to [v]. *)
  | Redo op ->
    (* [op] goes from [v] to [v'] *)
    Undo (r.descr.redo r.s op)
    (* [r.descr.redo r.s op] puts [r] in state [v'],
       and returns an anti-operation that returns to [v]. *)

(* if [r] is a custom structure in state [v_cur], with a trail with
   origin state [v_prev], and [entry] is from [v_cur] to [v_cur'],
   then after [revert_absorbing_op r entry], ... *)
let revert_custom_op (r : _ cref) (entry : _ custom_entry) : _ custom_entry =
  match entry with
  | Undo { parent } ->
    (* [parent] has origin [v_parent] and ends in [v_prev] *)
    let trail = r.trail in
    (* [trail] has origin [v_prev] and ends in [v_cur]. *)
    let ops = r.descr.undo r.s trail in
    (* after [undo], [r] is in state [v_cur'], [trail] is undone,
       and [ops] goes from [v_cur'] to [v_cur]. *)
    r.trail <- parent;
    (* This write restores the invariant that the end state of
       [r.trail] must be the current state of [r]. *)
    Redo { ops; trail }
    (* [ops] goes from [v_prev] to [v_cur] as expected by [Redo],
       and [trail] is undone with origin [v_prev] and end [v_cur]. *)
  | Redo { ops; trail } ->
    (* [ops] goes from [v_cur] to [v_prev],
       and [trail] is undone with origin [v_cur] and end [v_prev]. *)
    let parent = r.trail in
    (* [parent] has origin [v_prev] and end [v_cur] *)
    r.descr.redo r.s ops trail;
    (* [r] has state [v_prev] and [trail] is ready. *)
    r.trail <- trail;
    (* This write restores the invariant that the end state of
       [r.trail] must be the current state of [r]. *)
    Undo { parent }
    (* [parent] is ready, with origin [v_cur] and end [v_prev].
       TODO: this gives an entry of type
         [v_cur, v_prev, _]
       and not [v_prev, v_prev, v_cur] as expected. *)

(** {1 Implementation of references, record elision} *)

(** {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)
    | Log_op (_t, _op, _anti, n') ->
      collect n' (n :: acc)
    | Log_absorbing (_t, _entry, _g, n') ->
      collect n' (n :: acc)
    | Log_custom (_t, _entry, _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_op (t, op, anti, n') ->
        (* [child] points to its current parent node [n],
           let us define [v := m[t]] and [v' := m'[t]];
           we have [anti(v) = v'] and [op(v') = v]. *)
        assert (n' == n);
        (* We mutate the current mapping of [t] from [v] to [anti(v) = v']. *)
        ignore (op_and_log n child ~op:anti ~anti:op t);
        (* The current mapping is now [m'], so we can conclude with [revert]. *)
        revert child rest

      | Log_set (r, v, g, n') ->
        (* Same reasoning as in [Log_op] above. *)
        assert (n' == n);
        ignore (set_and_log n child r v g);
        revert child rest

      | Log_absorbing ((r : _ aref), entry, g, n') ->
        assert (n' == n);
        (* The current mapping and the mapping of [n] are [m[r ↦ v]],
           and the mapping of [child] is [m[r ↦ v']],
           with [entry] going from [v] back to [v']. *)
        let rev_entry = revert_absorbing_op r entry in
        (* After the revert, the current mapping is
           [m[r ↦ v'] and [rev_entry] goes from [v'] back to [v]. *)
        r.generation <- g;
        n := Log_absorbing (r, rev_entry, r.generation, child);
        (* [n] now points to [child], and their mappings
           are unchanged. The current mapping is [m[r ↦ v']],
           so we can call [revert] on [child]. *)
        revert child rest

      | Log_custom ((r : _ cref), entry, g, n') ->
        (* The current mapping and the mapping of [n] are [m[r ↦ v]],
           and the mapping of [child] is [m[r ↦ v']], with [entry]
           going from [v] to [v']. *)
        assert (n' == n);
        let new_entry = revert_custom_op r entry in
        (* After the revert, the current mapping is
           [m[r ↦ v'] and [rev_entry] goes from [v'] to [v]. *)
        r.generation <- g;
        n := Log_custom (r, new_entry, r.generation, child);
        (* [n] now points to [child], and their mappings
           are unchanged. The current mapping is [m[r ↦ v']],
           so we can call [revert] on [child]. *)
        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. *)

  type anti =
    | Anti_op of (unit -> unit)
    | Anti_custom : ('a, 'op, 'trail) cref * 'trail -> anti

(** 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
    | Invalid->
      Error.invalid_rollback ()
    | 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'
    | Log_op (t, _op, anti, n') ->
      reroot n';
      anti t
    | Log_absorbing ((r : _ aref), entry, g', n') ->
      (* Same reasoning as [Log_set] above. *)
      reroot n';
      if g' <= g then (
        r.generation <- g';
        match entry with
        | Undo anti ->
          r.descr.rollback r.s anti
        | Redo _ -> assert false
      )
    | Log_custom ((r : _ cref), Undo { parent }, g', n') ->
      (* TODO *)
      reroot n';
      (* Only need to reset generation if it is the first. *)
      if g' <= g then (
        r.generation <- g';
      );
      let trail = r.trail in
      r.descr.rollback r.s trail;
      r.trail <- parent
    | Log_custom (_t, Redo _, _g', _n') ->
      (* API misuse *)
      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 acc =
    match !n with
    | Invalid->
      Error.invalid_rollback ()
    | Mem ->
      n := Invalid;
      acc
    | Transaction (_n0, n') ->
      n := Invalid;
      reroot' n' acc
    | 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' acc
    | Log_op (t, _op, anti, n') ->
      (* On the other hand, free antioperations provide no such
         unicity invariant so they must all be performed, in reverse
         order (from most to least recent). *)
      reroot' n' (Anti_op (fun () -> anti t) :: acc)
    | Log_absorbing ((r : _ aref), entry, g', n') ->
      (* Same reasoning as [Log_set] above. *)
      if g' <= g then (
        r.generation <- g';
        match entry with
        | Undo anti ->
          r.descr.rollback r.s anti
        | Redo _ -> assert false
      );
      reroot' n' acc
    | Log_custom ((r : _ cref), Undo { parent }, g', n') ->
      (* Only need to reset generation if it is the first. *)
      if g' <= g then r.generation <- g';
      reroot' n' (Anti_custom (r, parent) :: acc)
    | Log_custom (_t, Redo _, _g', _n') ->
      Error.invalid_rollback ()
  in
  (* Let [m] be the mapping of [m]. *)
  let seg = reroot' n [] in
  List.iter (fun a ->
    match a with
    | Anti_op f -> f ()
    | Anti_custom (r, parent) ->
      r.descr.rollback r.s r.trail;
      r.trail <- parent;
  ) seg;
  (* 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 antis =
    match !src with
    | Invalid ->
      Error.invalid_commit ()
    | 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]. *)
      List.iter (function
        | Anti_op _ -> assert false
        | Anti_custom (r, parent) ->
          r.descr.append parent r.trail;
          r.trail <- parent
      ) antis;
      dst
    | Transaction (_n0, src') ->
      src := Invalid;
      (* [src] and [src'] have the same mapping,
         so the invariant is preserved below. *)
      reroot g ~src:src' ~dst antis
    | 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') antis
      ) 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 antis
      )
    | Log_op (t, op, anti, src') ->
      reroot g ~src:src' ~dst:(log_op dst ~op ~anti t) antis
    | Log_absorbing ((r : _ aref), entry, g', src') ->
      (* Same reasoning as [Log_set] above. *)
      if g' < g then (
        let new_dst = ref Mem in
        dst := Log_absorbing (r, entry, g', new_dst);
        r.generation <- g;
        reroot g ~src:src' ~dst:new_dst antis
      ) else (
        if g' = g then
          r.generation <- g;
        reroot g ~src:src' ~dst antis
      )
    | Log_custom ((r : _ cref), (Undo { parent } as entry), g', src') ->
      (* TODO *)
      if g' < g then (
        let new_dst = ref Mem in
        dst := Log_custom (r, entry, g', new_dst);
        r.generation <- g;
        reroot g ~src:src' ~dst:new_dst antis
      ) else (
        if g' = g then r.generation <- g;
        reroot g ~src:src' ~dst (Anti_custom (r, parent) :: antis)
      )
    | Log_custom (_r, Redo _, _g', _src') ->
      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

(** {2 Interfaces for store types} *)

module Free : sig
  val modify :
    store -> op:('a -> unit) -> anti:('a -> unit) -> 'a -> unit 
end = struct
  let modify (s : store) ~op ~anti t = 
    s.root <- op_and_log s.root (ref Mem) ~op ~anti t
end

module Absorbing : sig
  type ('a, 'ops, 'anti) descr
    = {
      capture : 'a -> 'anti;
      rollback : 'a -> 'anti -> unit;
      undo : 'a -> 'anti -> 'ops;
      redo : 'a -> 'ops -> 'anti;
    }

  type 'a t

  val make : ('a, 'op, 'anti) descr -> 'a -> 'a t

  val get : 'a t -> 'a

  val modify : store -> _ t -> unit
end = struct
  type ('a, 'ops, 'anti) descr
    = ('a, 'ops, 'anti) absorbing_descr
    = {
      capture : 'a -> 'anti;
      rollback : 'a -> 'anti -> unit;
      undo : 'a -> 'anti -> 'ops;
      redo : 'a -> 'ops -> 'anti;
    }

  type 'a t = Absorbing : ('a, 'op, 'anti) aref -> 'a t [@@unboxed]

  let make descr s =
    let generation = Generation.zero in
    Absorbing { s; generation; descr }

  let get (Absorbing r) = r.s

  let modify (s : store) (Absorbing r) =
    let g = r.generation in
    if s.generation <> g then (
      let anti = r.descr.capture r.s in
      let new_root = ref Mem in
      s.root := Log_absorbing (r, Undo anti, g, new_root);
      r.generation <- s.generation;
      s.root <- new_root
    )
end

module Custom : sig
  type ('a, 'op, 'trail) descr
    = {
        record : 'a -> 'trail;
        append : 'trail -> 'trail -> unit;
        rollback : 'a -> 'trail -> unit;
        undo : 'a -> 'trail -> 'op;
        redo : 'a -> 'op -> 'trail -> unit;
      }

  type ('a, 'trail) t
  val make : ('a, 'op, 'trail) descr -> 'a -> ('a, 'trail) t
  val get : ('a, 'trail) t -> 'a
  val modify : store -> ('a, 'trail) t -> 'trail
  (** This function must be called before modifying the state
      of a custom reference, and any modification must be recorded
      in the provided trail. *)
 end = struct
  type ('a, 'op, 'trail) descr = ('a, 'op, 'trail) custom_descr
    = {
        record : 'a -> 'trail;
        append : 'trail -> 'trail -> unit;
        rollback : 'a -> 'trail -> unit;
        undo : 'a -> 'trail -> 'op;
        redo : 'a -> 'op -> 'trail -> unit;
      }

  type ('a, 'trail) t = Custom : ('a, 'op, 'trail) cref -> ('a, 'trail) t [@@unboxed]

  let make descr s =
    let trail = descr.record s in
    Custom {
      s;
      generation = Generation.zero;
      descr;
      trail;
    }

  let get (Custom r) = r.s

  let modify (s : store) (Custom r) =
   let g = r.generation in
   let current_trail = r.trail in
   if s.generation <> g then (
     let new_trail = r.descr.record r.s in
     let new_root = ref Mem in
     s.root := Log_custom (r, Undo { parent = current_trail }, g, new_root);
     r.trail <- new_trail;
     r.generation <- s.generation;
     s.root <- new_root;
     new_trail
   ) else
     current_trail
end