Something went wrong. Try again.
An Elixir implementation of AT Protocol-flavoured Merkle Search Trees (MST)
Something went wrong. Try again.
9.2 kB · 282 lines
Elixir
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283defmodule MST.Node do @moduledoc """ Wire-format representation of a single MST node, plus encode/decode.
An MST node holds an optional left subtree CID (`left`) and an ordered list of `MST.Node.Entry` values, each carrying a key suffix, a value CID, and an optional right subtree CID. This maps exactly to the AT Protocol node schema:
{ l: CID | null, e: [ { p, k, v, t } ] }
Keys inside a node are prefix-compressed: each entry's `key_suffix` is the portion of the full key that follows the bytes it shares with the previous entry's full key. The first entry always has `prefix_len: 0` and carries its full key in `key_suffix`. Prefix compression is mandatory — the serialised form must be deterministic across implementations.
Spec: https://atproto.com/specs/repository#mst-structure """
use TypedStruct
alias DASL.{CID, DRISL} alias MST.Node.Entry
@type encode_error() :: {:error, :encode, atom()} @type decode_error() :: {:error, :decode, atom()}
typedstruct enforce: true do field :left, CID.t() | nil field :entries, [Entry.t()], default: [] end
# --------------------------------------------------------------------------- # Construction helpers # ---------------------------------------------------------------------------
@doc """ Returns an empty MST node — the only valid representation of an empty tree.
## Examples
iex> MST.Node.empty() %MST.Node{left: nil, entries: []}
""" @spec empty() :: t() def empty, do: %__MODULE__{left: nil, entries: []}
# --------------------------------------------------------------------------- # Key expansion # ---------------------------------------------------------------------------
@doc """ Reconstructs the full keys for all entries in the node.
Each entry stores only the suffix of its key relative to the previous entry. This function walks the entry list and accumulates the full key for each.
## Examples
iex> cid = DASL.CID.compute("a") iex> entries = [ ...> %MST.Node.Entry{prefix_len: 0, key_suffix: "foo/bar", value: cid, right: nil}, ...> %MST.Node.Entry{prefix_len: 4, key_suffix: "baz", value: cid, right: nil}, ...> ] iex> MST.Node.keys(%MST.Node{left: nil, entries: entries}) ["foo/bar", "foo/baz"]
""" @spec keys(t()) :: [binary()] def keys(%__MODULE__{entries: entries}), do: expand_keys(entries, "", [])
# --------------------------------------------------------------------------- # CID computation # ---------------------------------------------------------------------------
@doc """ Computes the `:drisl`-codec CID for this node.
Encodes the node to DRISL CBOR bytes and hashes them. Returns an error tuple if encoding fails.
## Examples
iex> {:ok, cid} = MST.Node.cid(MST.Node.empty()) iex> cid.codec :drisl
""" @spec cid(t()) :: {:ok, CID.t()} | encode_error() def cid(node) do with {:ok, bytes} <- encode(node) do {:ok, CID.compute(bytes, :drisl)} end end
# --------------------------------------------------------------------------- # Encoding # ---------------------------------------------------------------------------
@doc """ Encodes an `MST.Node` to DRISL CBOR bytes.
`nil` subtree links are serialised as explicit CBOR `null` — this is mandatory for cross-implementation CID compatibility: skipping a key vs. writing `null` produces different bytes and therefore a different CID.
## Examples
iex> {:ok, bytes} = MST.Node.encode(MST.Node.empty()) iex> is_binary(bytes) true
""" @spec encode(t()) :: {:ok, binary()} | encode_error() def encode(%__MODULE__{left: left, entries: entries}) do with {:ok, entry_maps} <- encode_entries(entries), {:ok, bytes} <- DRISL.encode(%{"e" => entry_maps, "l" => left}) do {:ok, bytes} else {:error, reason} when is_atom(reason) -> {:error, :encode, reason} {:error, :encode, _} = err -> err end end
# --------------------------------------------------------------------------- # Decoding # ---------------------------------------------------------------------------
@doc """ Decodes DRISL CBOR bytes into an `MST.Node`.
## Examples
iex> {:ok, bytes} = MST.Node.encode(MST.Node.empty()) iex> {:ok, node} = MST.Node.decode(bytes) iex> node.entries [] iex> node.left nil
""" @spec decode(binary()) :: {:ok, t()} | decode_error() def decode(bytes) when is_binary(bytes) do with {:ok, term, <<>>} <- DRISL.decode(bytes), {:ok, node} <- decode_term(term) do {:ok, node} else {:ok, _, _leftover} -> {:error, :decode, :trailing_bytes} {:error, reason} when is_atom(reason) -> {:error, :decode, reason} {:error, :decode, _} = err -> err end end
# --------------------------------------------------------------------------- # Compression helpers (used by MST.Tree) # ---------------------------------------------------------------------------
@doc """ Compresses a list of `{full_key, value_cid, right_cid | nil}` tuples into a list of `MST.Node.Entry` structs using the key prefix-compression scheme.
The first entry always has `prefix_len: 0`. Each subsequent entry computes how many leading bytes it shares with the previous full key.
## Examples
iex> cid = DASL.CID.compute("x") iex> entries = MST.Node.compress_entries([{"abc/def", cid, nil}, {"abc/ghi", cid, nil}]) iex> hd(tl(entries)).prefix_len 4
""" @spec compress_entries([{binary(), CID.t(), CID.t() | nil}]) :: [Entry.t()] def compress_entries(triples), do: do_compress(triples, "", [])
# --------------------------------------------------------------------------- # Private helpers # ---------------------------------------------------------------------------
@spec expand_keys([Entry.t()], binary(), [binary()]) :: [binary()] defp expand_keys([], _prev, acc), do: Enum.reverse(acc)
defp expand_keys([entry | rest], prev, acc) do full_key = binary_part(prev, 0, entry.prefix_len) <> entry.key_suffix expand_keys(rest, full_key, [full_key | acc]) end
@spec do_compress([{binary(), CID.t(), CID.t() | nil}], binary(), [Entry.t()]) :: [Entry.t()] defp do_compress([], _prev, acc), do: Enum.reverse(acc)
defp do_compress([{key, value, right} | rest], prev, acc) do plen = common_prefix_length(prev, key) suffix = binary_part(key, plen, byte_size(key) - plen)
entry = %Entry{ prefix_len: plen, key_suffix: suffix, value: value, right: right }
do_compress(rest, key, [entry | acc]) end
@spec common_prefix_length(binary(), binary()) :: non_neg_integer() defp common_prefix_length(a, b), do: cpl(a, b, 0)
defp cpl(<<c, ra::binary>>, <<c, rb::binary>>, n), do: cpl(ra, rb, n + 1) defp cpl(_, _, n), do: n
@spec encode_entries([Entry.t()]) :: {:ok, [map()]} | encode_error() defp encode_entries(entries) do result = Enum.reduce_while(entries, {:ok, []}, fn entry, {:ok, acc} -> {:ok, map} = encode_entry(entry) {:cont, {:ok, [map | acc]}} end)
case result do {:ok, reversed} -> {:ok, Enum.reverse(reversed)} err -> err end end
@spec encode_entry(Entry.t()) :: {:ok, map()} | encode_error() defp encode_entry(%Entry{prefix_len: p, key_suffix: k, value: v, right: t}) do {:ok, %{ "k" => %CBOR.Tag{tag: :bytes, value: k}, "p" => p, "t" => t, "v" => v }} end
@spec decode_term(any()) :: {:ok, t()} | decode_error() defp decode_term(%{"e" => entries_raw, "l" => left_raw}) when is_list(entries_raw) do with {:ok, left} <- decode_cid_or_null(left_raw), {:ok, entries} <- decode_entries(entries_raw) do {:ok, %__MODULE__{left: left, entries: entries}} end end
defp decode_term(_), do: {:error, :decode, :invalid_structure}
@spec decode_entries(list()) :: {:ok, [Entry.t()]} | decode_error() defp decode_entries(entries_raw) do result = Enum.reduce_while(entries_raw, {:ok, []}, fn raw, {:ok, acc} -> case decode_entry(raw) do {:ok, entry} -> {:cont, {:ok, [entry | acc]}} {:error, :decode, _} = err -> {:halt, err} end end)
case result do {:ok, reversed} -> {:ok, Enum.reverse(reversed)} err -> err end end
@spec decode_entry(any()) :: {:ok, Entry.t()} | decode_error() defp decode_entry(%{ "k" => %CBOR.Tag{tag: :bytes, value: k}, "p" => p, "t" => t_raw, "v" => %CID{} = v }) when is_integer(p) and p >= 0 and is_binary(k) do with {:ok, right} <- decode_cid_or_null(t_raw) do {:ok, %Entry{prefix_len: p, key_suffix: k, value: v, right: right}} end end
defp decode_entry(_), do: {:error, :decode, :invalid_entry}
@spec decode_cid_or_null(any()) :: {:ok, CID.t() | nil} | decode_error() defp decode_cid_or_null(nil), do: {:ok, nil} defp decode_cid_or_null(%CID{} = cid), do: {:ok, cid} defp decode_cid_or_null(_), do: {:error, :decode, :invalid_cid_link}end