elixir-mst / lib / mst / tree.ex

defmodule MST.Tree do
  @moduledoc """
  An in-memory Merkle Search Tree.

  `MST.Tree` is the primary interface for building and querying MSTs. It pairs
  a root CID (or `nil` for an empty tree) with a block store that maps CIDs to
  decoded `MST.Node` structs.

  All mutation operations (`put/3`, `delete/3`) return a new `MST.Tree` —
  the data structure is persistent/immutable. The underlying store accumulates
  all written nodes across mutations; stale nodes are not removed (no GC).

  ## Example

      store = MST.Store.Memory.new()
      tree  = MST.Tree.new(store)

      val = DASL.CID.compute("my record data")
      {:ok, tree} = MST.Tree.put(tree, "collection/record-key", val)
      {:ok, ^val} = MST.Tree.get(tree, "collection/record-key")

      {:ok, tree} = MST.Tree.delete(tree, "collection/record-key")
      {:error, :not_found} = MST.Tree.get(tree, "collection/record-key")

  """

  use TypedStruct
  import Kernel, except: [length: 1]

  alias DASL.CID
  alias MST.{Height, Node, Node.Entry, Store}

  @type store() :: Store.t()
  @type tree_error() :: {:error, atom()}

  typedstruct enforce: true do
    field :root, CID.t() | nil
    field :store, store()
  end

  # ---------------------------------------------------------------------------
  # Construction
  # ---------------------------------------------------------------------------

  @doc """
  Returns a new, empty tree backed by the given store.

  ## Examples

      iex> tree = MST.Tree.new(MST.Store.Memory.new())
      iex> tree.root
      nil

  """
  @spec new(store()) :: t()
  def new(store), do: %__MODULE__{root: nil, store: store}

  @doc """
  Returns a tree referencing an existing root node CID in the given store.

  Use this to wrap an already-populated store (e.g. after loading from a CAR
  file).

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> node = MST.Node.empty()
      iex> {:ok, cid} = MST.Node.cid(node)
      iex> store = MST.Store.put(store, cid, node)
      iex> tree = MST.Tree.from_root(cid, store)
      iex> tree.root == cid
      true

  """
  @spec from_root(CID.t() | nil, store()) :: t()
  def from_root(root, store), do: %__MODULE__{root: root, store: store}

  # ---------------------------------------------------------------------------
  # Lookup
  # ---------------------------------------------------------------------------

  @doc """
  Looks up `key` in the tree.

  Returns `{:ok, value_cid}` if found, `{:error, :not_found}` otherwise.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> MST.Tree.get(tree, "col/key")
      {:error, :not_found}

  """
  @spec get(t(), binary()) :: {:ok, CID.t()} | {:error, :not_found} | tree_error()
  def get(%__MODULE__{root: nil}, _key), do: {:error, :not_found}

  def get(%__MODULE__{root: root, store: store}, key) do
    search(store, root, key)
  end

  # ---------------------------------------------------------------------------
  # Mutation
  # ---------------------------------------------------------------------------

  @doc """
  Inserts or updates `key` → `value` in the tree.

  Returns `{:ok, new_tree}` on success. The new tree shares the store with the
  old tree, but both may be used independently — nodes are append-only.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> val = DASL.CID.compute("data")
      iex> {:ok, tree2} = MST.Tree.put(tree, "col/key", val)
      iex> MST.Tree.get(tree2, "col/key")
      {:ok, val}

  """
  @spec put(t(), binary(), CID.t()) :: {:ok, t()} | tree_error()
  def put(%__MODULE__{root: nil, store: store}, key, value) do
    # Empty tree — create a leaf node. No intermediate layers needed for a
    # single-key tree (the spec says empty nodes at the top must be pruned).
    node = leaf_node(key, value)

    with {:ok, root, store} <- write_node(store, node) do
      {:ok, %__MODULE__{root: root, store: store}}
    end
  end

  def put(%__MODULE__{root: root, store: store}, key, value) do
    with {:ok, root_node} <- fetch_node(store, root) do
      if root_node.entries == [] and root_node.left == nil do
        # Empty root from CAR import — treat as fresh empty tree.
        put(%__MODULE__{root: nil, store: store}, key, value)
      else
        key_height = Height.for_key(key)
        tree_height = require_height(store, root_node)

        with {:ok, new_root, store} <-
               do_insert(store, root, key, value, key_height, tree_height) do
          {:ok, %__MODULE__{root: new_root, store: store}}
        end
      end
    end
  end

  @doc """
  Removes `key` from the tree.

  Returns `{:ok, new_tree}` on success, `{:error, :not_found}` if the key
  does not exist.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> val = DASL.CID.compute("data")
      iex> {:ok, tree2} = MST.Tree.put(tree, "col/key", val)
      iex> {:ok, tree3} = MST.Tree.delete(tree2, "col/key")
      iex> MST.Tree.get(tree3, "col/key")
      {:error, :not_found}

  """
  @spec delete(t(), binary()) :: {:ok, t()} | {:error, :not_found} | tree_error()
  def delete(%__MODULE__{root: nil}, _key), do: {:error, :not_found}

  def delete(%__MODULE__{root: root, store: store}, key) do
    with {:ok, root_node} <- fetch_node(store, root) do
      key_height = Height.for_key(key)
      tree_height = require_height(store, root_node)

      with {:ok, new_root, store} <-
             do_remove(store, root, key, key_height, tree_height) do
        # Trim empty wrappers from the top after deletion.
        {:ok, new_root, store} = trim_top(store, new_root)
        {:ok, %__MODULE__{root: new_root, store: store}}
      end
    end
  end

  # ---------------------------------------------------------------------------
  # Traversal
  # ---------------------------------------------------------------------------

  @doc """
  Returns all key-value pairs in the tree as a sorted list.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> val = DASL.CID.compute("data")
      iex> {:ok, tree} = MST.Tree.put(tree, "col/b", val)
      iex> {:ok, tree} = MST.Tree.put(tree, "col/a", val)
      iex> {:ok, pairs} = MST.Tree.to_list(tree)
      iex> Enum.map(pairs, &elem(&1, 0))
      ["col/a", "col/b"]

  """
  @spec to_list(t()) :: {:ok, [{binary(), CID.t()}]} | tree_error()
  def to_list(%__MODULE__{root: nil}), do: {:ok, []}

  def to_list(%__MODULE__{root: root, store: store}) do
    with {:ok, reversed} <- walk(store, root, []) do
      {:ok, Enum.reverse(reversed)}
    end
  end

  @doc """
  Returns a lazy stream of `{key, value_cid}` pairs in sorted order.

  The stream reads nodes from the store on demand. Raises on missing nodes
  (consistent with lazy stream semantics).

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> val = DASL.CID.compute("x")
      iex> {:ok, tree} = MST.Tree.put(tree, "col/a", val)
      iex> tree |> MST.Tree.stream() |> Enum.to_list()
      [{"col/a", val}]

  """
  @spec stream(t()) :: Enumerable.t()
  def stream(%__MODULE__{root: nil}), do: []

  def stream(%__MODULE__{root: root, store: store}) do
    Stream.resource(
      fn -> [{:visit, root}] end,
      fn
        [] ->
          {:halt, []}

        [{:yield, k, v} | rest] ->
          {[{k, v}], rest}

        [{:visit, cid} | rest] ->
          node = fetch_node!(store, cid)
          full_keys = Node.keys(node)
          {[], in_order_items(node, full_keys) ++ rest}
      end,
      fn _ -> :ok end
    )
  end

  @doc """
  Returns the number of key-value pairs in the tree.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> {:ok, 0} = MST.Tree.length(tree)
      iex> val = DASL.CID.compute("x")
      iex> {:ok, tree} = MST.Tree.put(tree, "col/a", val)
      iex> MST.Tree.length(tree)
      {:ok, 1}

  """
  @spec length(t()) :: {:ok, non_neg_integer()} | tree_error()
  def length(tree) do
    with {:ok, pairs} <- to_list(tree) do
      {:ok, Kernel.length(pairs)}
    end
  end

  # ---------------------------------------------------------------------------
  # Block collection
  # ---------------------------------------------------------------------------

  @doc """
  Collects all MST nodes reachable from the root into a map of CID → encoded bytes.

  Useful for serialising the tree to a CAR file.

  ## Examples

      iex> store = MST.Store.Memory.new()
      iex> tree = MST.Tree.new(store)
      iex> val = DASL.CID.compute("x")
      iex> {:ok, tree} = MST.Tree.put(tree, "col/a", val)
      iex> {:ok, blocks} = MST.Tree.collect_blocks(tree)
      iex> map_size(blocks) >= 1
      true

  """
  @spec collect_blocks(t()) :: {:ok, %{CID.t() => binary()}} | tree_error()
  def collect_blocks(%__MODULE__{root: nil}), do: {:ok, %{}}

  def collect_blocks(%__MODULE__{root: root, store: store}) do
    collect_reachable(store, root, %{})
  end

  # ---------------------------------------------------------------------------
  # Private — search
  # ---------------------------------------------------------------------------

  @spec search(store(), CID.t(), binary()) ::
          {:ok, CID.t()} | {:error, :not_found} | tree_error()
  defp search(store, cid, key) do
    with {:ok, node} <- fetch_node(store, cid) do
      full_keys = Node.keys(node)
      search_node(store, node, full_keys, key)
    end
  end

  @spec search_node(store(), Node.t(), [binary()], binary()) ::
          {:ok, CID.t()} | {:error, :not_found} | tree_error()
  defp search_node(store, node, full_keys, key) do
    case locate(full_keys, key) do
      {:found, idx} ->
        {:ok, Enum.at(node.entries, idx).value}

      {:left} ->
        descend(store, node.left, key)

      {:right, idx} ->
        descend(store, Enum.at(node.entries, idx).right, key)
    end
  end

  @spec descend(store(), CID.t() | nil, binary()) ::
          {:ok, CID.t()} | {:error, :not_found} | tree_error()
  defp descend(_store, nil, _key), do: {:error, :not_found}
  defp descend(store, cid, key), do: search(store, cid, key)

  # ---------------------------------------------------------------------------
  # Private — insert
  # ---------------------------------------------------------------------------

  # Recursive insert into the subtree rooted at `cid`.
  # `tree_height` is the known height of the node (threaded from the parent).
  @spec do_insert(store(), CID.t(), binary(), CID.t(), non_neg_integer(), non_neg_integer()) ::
          {:ok, CID.t(), store()} | tree_error()
  defp do_insert(store, cid, key, value, key_height, tree_height) do
    with {:ok, node} <- fetch_node(store, cid) do
      cond do
        key_height > tree_height ->
          # Key belongs at a higher layer. Wrap the current node in an empty
          # parent and recurse at tree_height + 1.
          wrapper = %Node{left: cid, entries: []}

          with {:ok, wrapper_cid, store} <- write_node(store, wrapper) do
            do_insert(store, wrapper_cid, key, value, key_height, tree_height + 1)
          end

        key_height < tree_height ->
          # Descend into the appropriate subtree.
          {kv_pairs, subtrees} = node_to_arrays(node)
          keys = Enum.map(kv_pairs, &elem(&1, 0))
          idx = lower_bound(keys, key)
          subtree_cid = Enum.at(subtrees, idx)

          with {:ok, new_sub, store} <-
                 insert_into_subtree(
                   store,
                   subtree_cid,
                   key,
                   value,
                   key_height,
                   tree_height - 1
                 ) do
            new_subtrees = List.replace_at(subtrees, idx, new_sub)
            write_node(store, arrays_to_node(kv_pairs, new_subtrees))
          end

        true ->
          # key_height == tree_height — insert at this level.
          put_here(store, node, key, value)
      end
    end
  end

  # Insert into a subtree that may be nil. When nil, creates a new leaf and
  # wraps it in as many empty intermediate layers as needed.
  @spec insert_into_subtree(
          store(),
          CID.t() | nil,
          binary(),
          CID.t(),
          non_neg_integer(),
          non_neg_integer()
        ) :: {:ok, CID.t(), store()} | tree_error()
  defp insert_into_subtree(store, nil, key, value, key_height, expected_height) do
    leaf = leaf_node(key, value)

    with {:ok, leaf_cid, store} <- write_node(store, leaf) do
      wrap_with_empty_layers(store, leaf_cid, expected_height - key_height)
    end
  end

  defp insert_into_subtree(store, cid, key, value, key_height, expected_height) do
    do_insert(store, cid, key, value, key_height, expected_height)
  end

  # Insert a key at the current level (key_height == tree_height).
  # Splits the subtree at the insertion point recursively.
  @spec put_here(store(), Node.t(), binary(), CID.t()) ::
          {:ok, CID.t(), store()} | tree_error()
  defp put_here(store, node, key, value) do
    {kv_pairs, subtrees} = node_to_arrays(node)
    keys = Enum.map(kv_pairs, &elem(&1, 0))
    idx = lower_bound(keys, key)

    if idx < Kernel.length(keys) and Enum.at(keys, idx) == key do
      # Overwrite existing key.
      new_kv = List.replace_at(kv_pairs, idx, {key, value})
      write_node(store, arrays_to_node(new_kv, subtrees))
    else
      # Split the subtree at the insertion point recursively.
      with {:ok, lsub, rsub, store} <- split_on_key(store, Enum.at(subtrees, idx), key) do
        new_kv = List.insert_at(kv_pairs, idx, {key, value})

        new_subtrees =
          List.replace_at(subtrees, idx, lsub) |> List.insert_at(idx + 1, rsub)

        write_node(store, arrays_to_node(new_kv, new_subtrees))
      end
    end
  end

  # Recursively splits the subtree at `key`. Returns {left_cid, right_cid}
  # where left contains all keys < `key` and right contains all keys >= `key`.
  # Either side may be nil if empty.
  @spec split_on_key(store(), CID.t() | nil, binary()) ::
          {:ok, CID.t() | nil, CID.t() | nil, store()} | tree_error()
  defp split_on_key(store, nil, _key), do: {:ok, nil, nil, store}

  defp split_on_key(store, cid, key) do
    with {:ok, node} <- fetch_node(store, cid) do
      {kv_pairs, subtrees} = node_to_arrays(node)
      keys = Enum.map(kv_pairs, &elem(&1, 0))
      idx = lower_bound(keys, key)

      # Recursively split the subtree at the boundary position.
      with {:ok, inner_l, inner_r, store} <-
             split_on_key(store, Enum.at(subtrees, idx), key) do
        left_kv = Enum.take(kv_pairs, idx)
        left_subs = Enum.take(subtrees, idx) ++ [inner_l]

        right_kv = Enum.drop(kv_pairs, idx)
        right_subs = [inner_r | Enum.drop(subtrees, idx + 1)]

        with {:ok, left_cid, store} <-
               write_node_to_nullable(store, arrays_to_node(left_kv, left_subs)),
             {:ok, right_cid, store} <-
               write_node_to_nullable(store, arrays_to_node(right_kv, right_subs)) do
          {:ok, left_cid, right_cid, store}
        end
      end
    end
  end

  # ---------------------------------------------------------------------------
  # Private — delete
  # ---------------------------------------------------------------------------

  # Recursive delete. `tree_height` is the known height of the node at `cid`.
  @spec do_remove(store(), CID.t(), binary(), non_neg_integer(), non_neg_integer()) ::
          {:ok, CID.t() | nil, store()} | {:error, :not_found} | tree_error()
  defp do_remove(store, cid, key, key_height, tree_height) do
    with {:ok, node} <- fetch_node(store, cid) do
      cond do
        key_height > tree_height ->
          {:error, :not_found}

        key_height < tree_height ->
          {kv_pairs, subtrees} = node_to_arrays(node)
          keys = Enum.map(kv_pairs, &elem(&1, 0))
          idx = lower_bound(keys, key)

          case Enum.at(subtrees, idx) do
            nil ->
              {:error, :not_found}

            sub_cid ->
              with {:ok, new_sub, store} <-
                     do_remove(store, sub_cid, key, key_height, tree_height - 1) do
                new_subtrees = List.replace_at(subtrees, idx, new_sub)
                write_node_to_nullable(store, arrays_to_node(kv_pairs, new_subtrees))
              end
          end

        true ->
          # key_height == tree_height — key must be at this level if it exists.
          {kv_pairs, subtrees} = node_to_arrays(node)
          keys = Enum.map(kv_pairs, &elem(&1, 0))
          idx = lower_bound(keys, key)

          if idx < Kernel.length(keys) and Enum.at(keys, idx) == key do
            # Found! Merge the adjacent subtrees that flanked the deleted key.
            with {:ok, merged_sub, store} <-
                   do_merge(store, Enum.at(subtrees, idx), Enum.at(subtrees, idx + 1)) do
              new_kv = List.delete_at(kv_pairs, idx)

              new_subtrees =
                Enum.take(subtrees, idx) ++ [merged_sub | Enum.drop(subtrees, idx + 2)]

              write_node_to_nullable(store, arrays_to_node(new_kv, new_subtrees))
            end
          else
            {:error, :not_found}
          end
      end
    end
  end

  # Recursively merges two adjacent subtree pointers. The boundary subtrees
  # (rightmost of left, leftmost of right) are merged recursively.
  @spec do_merge(store(), CID.t() | nil, CID.t() | nil) ::
          {:ok, CID.t() | nil, store()} | tree_error()
  defp do_merge(store, nil, right_cid), do: {:ok, right_cid, store}
  defp do_merge(store, left_cid, nil), do: {:ok, left_cid, store}

  defp do_merge(store, left_cid, right_cid) do
    with {:ok, left_node} <- fetch_node(store, left_cid),
         {:ok, right_node} <- fetch_node(store, right_cid) do
      {left_kv, left_subs} = node_to_arrays(left_node)
      {right_kv, right_subs} = node_to_arrays(right_node)

      with {:ok, merged_boundary, store} <-
             do_merge(store, List.last(left_subs), hd(right_subs)) do
        new_kv = left_kv ++ right_kv
        new_subs = Enum.slice(left_subs, 0..-2//1) ++ [merged_boundary | tl(right_subs)]
        write_node_to_nullable(store, arrays_to_node(new_kv, new_subs))
      end
    end
  end

  # Strips empty wrapper nodes from the top of the tree. Only called after
  # a top-level delete — intermediate empty nodes are preserved during
  # recursive descent.
  @spec trim_top(store(), CID.t() | nil) :: {:ok, CID.t() | nil, store()} | tree_error()
  defp trim_top(store, nil), do: {:ok, nil, store}

  defp trim_top(store, cid) do
    with {:ok, node} <- fetch_node(store, cid) do
      cond do
        node.entries != [] -> {:ok, cid, store}
        node.left == nil -> {:ok, nil, store}
        true -> trim_top(store, node.left)
      end
    end
  end

  # ---------------------------------------------------------------------------
  # Private — in-order traversal (to_list)
  # ---------------------------------------------------------------------------

  @spec walk(store(), CID.t(), [{binary(), CID.t()}]) ::
          {:ok, [{binary(), CID.t()}]} | tree_error()
  defp walk(store, cid, acc) do
    with {:ok, node} <- fetch_node(store, cid) do
      full_keys = Node.keys(node)
      walk_node(store, node, full_keys, acc)
    end
  end

  @spec walk_node(store(), Node.t(), [binary()], [{binary(), CID.t()}]) ::
          {:ok, [{binary(), CID.t()}]} | tree_error()
  defp walk_node(store, node, full_keys, acc) do
    # Walk in-order: left subtree, then entries interleaved with right subtrees.
    # We collect in reverse for efficiency, then reverse at the end.
    with {:ok, acc} <- walk_subtree(store, node.left, acc) do
      walk_entries(store, node.entries, full_keys, acc)
    end
  end

  @spec walk_subtree(store(), CID.t() | nil, [{binary(), CID.t()}]) ::
          {:ok, [{binary(), CID.t()}]} | tree_error()
  defp walk_subtree(_store, nil, acc), do: {:ok, acc}
  defp walk_subtree(store, cid, acc), do: walk(store, cid, acc)

  @spec walk_entries(store(), [Entry.t()], [binary()], [{binary(), CID.t()}]) ::
          {:ok, [{binary(), CID.t()}]} | tree_error()
  # Accumulates pairs in reverse order; to_list/1 reverses once at the top.
  defp walk_entries(_store, [], [], acc), do: {:ok, acc}

  defp walk_entries(store, [entry | rest_e], [key | rest_k], acc) do
    acc = [{key, entry.value} | acc]

    with {:ok, acc} <- walk_subtree(store, entry.right, acc) do
      walk_entries(store, rest_e, rest_k, acc)
    end
  end

  # ---------------------------------------------------------------------------
  # Private — stream helpers
  # ---------------------------------------------------------------------------

  # Expand a node into an in-order sequence of {:visit, cid} and {:yield, k, v}
  # items. When prepended to a DFS stack, this produces correct left-to-right
  # traversal: left subtree is visited first, then each entry interleaved with
  # its right subtree.
  @spec in_order_items(Node.t(), [binary()]) :: list()
  defp in_order_items(node, full_keys) do
    left_items = if node.left, do: [{:visit, node.left}], else: []

    entry_items =
      Enum.zip(node.entries, full_keys)
      |> Enum.flat_map(fn {e, k} ->
        right_items = if e.right, do: [{:visit, e.right}], else: []
        [{:yield, k, e.value} | right_items]
      end)

    left_items ++ entry_items
  end

  # ---------------------------------------------------------------------------
  # Private — block collection
  # ---------------------------------------------------------------------------

  @spec collect_reachable(store(), CID.t(), %{CID.t() => binary()}) ::
          {:ok, %{CID.t() => binary()}} | tree_error()
  defp collect_reachable(store, cid, acc) do
    if Map.has_key?(acc, cid) do
      {:ok, acc}
    else
      with {:ok, node} <- fetch_node(store, cid),
           {:ok, bytes} <- Node.encode(node) do
        acc = Map.put(acc, cid, bytes)
        collect_children(store, node, acc)
      else
        {:error, :not_found} -> {:error, :missing_node}
        {:error, :encode, reason} -> {:error, reason}
      end
    end
  end

  @spec collect_children(store(), Node.t(), %{CID.t() => binary()}) ::
          {:ok, %{CID.t() => binary()}} | tree_error()
  defp collect_children(store, node, acc) do
    subtrees =
      if(node.left, do: [node.left], else: []) ++
        Enum.flat_map(node.entries, fn e -> if e.right, do: [e.right], else: [] end)

    Enum.reduce_while(subtrees, {:ok, acc}, fn cid, {:ok, acc} ->
      case collect_reachable(store, cid, acc) do
        {:ok, acc} -> {:cont, {:ok, acc}}
        err -> {:halt, err}
      end
    end)
  end

  # ---------------------------------------------------------------------------
  # Private — node I/O
  # ---------------------------------------------------------------------------

  @spec fetch_node(store(), CID.t()) :: {:ok, Node.t()} | tree_error()
  defp fetch_node(store, cid) do
    case Store.get(store, cid) do
      {:ok, node} -> {:ok, node}
      {:error, :not_found} -> {:error, :missing_node}
    end
  end

  @spec fetch_node!(store(), CID.t()) :: Node.t()
  defp fetch_node!(store, cid) do
    case Store.get(store, cid) do
      {:ok, node} -> node
      {:error, :not_found} -> raise "MST node not found: #{CID.encode(cid)}"
    end
  end

  @spec write_node(store(), Node.t()) :: {:ok, CID.t(), store()} | tree_error()
  defp write_node(store, node) do
    case Node.cid(node) do
      {:ok, cid} -> {:ok, cid, Store.put(store, cid, node)}
      {:error, :encode, reason} -> {:error, reason}
    end
  end

  # Write a node unless it is truly empty (no entries, no left). Returns nil
  # for empty leaf-level nodes; preserves empty intermediate nodes that have
  # a left subtree pointer.
  @spec write_node_to_nullable(store(), Node.t()) ::
          {:ok, CID.t() | nil, store()} | tree_error()
  defp write_node_to_nullable(store, %Node{left: nil, entries: []}), do: {:ok, nil, store}
  defp write_node_to_nullable(store, node), do: write_node(store, node)

  # Wraps a CID in `n` empty intermediate nodes (left-pointer only).
  @spec wrap_with_empty_layers(store(), CID.t(), non_neg_integer()) ::
          {:ok, CID.t(), store()} | tree_error()
  defp wrap_with_empty_layers(store, cid, 0), do: {:ok, cid, store}

  defp wrap_with_empty_layers(store, cid, n) when n > 0 do
    wrapper = %Node{left: cid, entries: []}

    with {:ok, wrapper_cid, store} <- write_node(store, wrapper) do
      wrap_with_empty_layers(store, wrapper_cid, n - 1)
    end
  end

  # ---------------------------------------------------------------------------
  # Private — key position helpers
  # ---------------------------------------------------------------------------

  # Returns the position of `key` in the sorted `full_keys` list:
  #   {:found, idx}   — key is at index idx
  #   {:left}         — key < all keys (belongs in left subtree)
  #   {:right, idx}   — key > keys[idx] (belongs in right subtree of entry idx)
  @spec locate([binary()], binary()) ::
          {:found, non_neg_integer()} | {:left} | {:right, non_neg_integer()}
  defp locate([], _key), do: {:left}

  defp locate(keys, key) do
    n = Kernel.length(keys)
    bin_locate(keys, key, 0, n - 1, n)
  end

  @spec bin_locate([binary()], binary(), integer(), integer(), non_neg_integer()) ::
          {:found, non_neg_integer()} | {:left} | {:right, non_neg_integer()}
  defp bin_locate(_keys, _key, lo, hi, _n) when lo > hi do
    if lo == 0, do: {:left}, else: {:right, lo - 1}
  end

  defp bin_locate(keys, key, lo, hi, n) do
    mid = div(lo + hi, 2)
    mid_key = Enum.at(keys, mid)

    cond do
      mid_key == key -> {:found, mid}
      mid_key < key -> bin_locate(keys, key, mid + 1, hi, n)
      true -> bin_locate(keys, key, lo, mid - 1, n)
    end
  end

  # Returns the index of the first key >= `target`, or `length(keys)` if none.
  @spec lower_bound([binary()], binary()) :: non_neg_integer()
  defp lower_bound(keys, target) do
    Enum.find_index(keys, fn k -> k >= target end) || Kernel.length(keys)
  end

  # ---------------------------------------------------------------------------
  # Private — layer inference
  # ---------------------------------------------------------------------------

  # Infer the MST layer of a non-empty node from its first entry's key.
  @spec node_layer(Node.t()) :: non_neg_integer() | nil
  defp node_layer(%Node{entries: []}), do: nil

  defp node_layer(%Node{entries: [first | _]}) do
    Height.for_key(first.key_suffix)
  end

  # Compute the height of a node, walking into children if the node has no
  # entries (empty intermediate nodes).
  @spec require_height(store(), Node.t()) :: non_neg_integer()
  defp require_height(store, node) do
    case node_layer(node) do
      nil ->
        if node.left do
          {:ok, child} = fetch_node(store, node.left)
          require_height(store, child) + 1
        else
          0
        end

      h ->
        h
    end
  end

  # ---------------------------------------------------------------------------
  # Private — construction helpers
  # ---------------------------------------------------------------------------

  @spec leaf_node(binary(), CID.t()) :: Node.t()
  defp leaf_node(key, value) do
    %Node{
      left: nil,
      entries: [%Entry{prefix_len: 0, key_suffix: key, value: value, right: nil}]
    }
  end

  # ---------------------------------------------------------------------------
  # Private — node array conversions
  # ---------------------------------------------------------------------------

  # Converts a node into a parallel-array representation:
  #   {[{key, value}], [subtree_cid | nil]}
  # where subtrees has length(kv_pairs) + 1.
  # subtrees[0] = node.left, subtrees[i+1] = entries[i].right.
  @spec node_to_arrays(Node.t()) :: {[{binary(), CID.t()}], [CID.t() | nil]}
  defp node_to_arrays(node) do
    full_keys = Node.keys(node)
    kv_pairs = Enum.zip(full_keys, Enum.map(node.entries, & &1.value))
    subtrees = [node.left | Enum.map(node.entries, & &1.right)]
    {kv_pairs, subtrees}
  end

  # Converts the parallel-array representation back to a `Node`.
  @spec arrays_to_node([{binary(), CID.t()}], [CID.t() | nil]) :: Node.t()
  defp arrays_to_node(kv_pairs, subtrees) do
    [left | right_ptrs] = subtrees

    triples =
      Enum.zip(kv_pairs, right_ptrs)
      |> Enum.map(fn {{k, v}, r} -> {k, v, r} end)

    entries = Node.compress_entries(triples)
    %Node{left: left, entries: entries}
  end
end