defmodule MPTree do
@external_resource "README.md"
[_ignored_intro, usage_section, _ignored] =
"README.md" |> File.read!() |> String.split(~r/<!---.*moduledoc.*-->/, parts: 3)
@moduledoc usage_section
# Here are some words about Modifield Preorder Tree Traversal.
# - ancestors: the `t:path/0` list minus the last node
# - parent: the parent of the node in question
# - children: ancestors, but only one level deep
# See https://gist.github.com/tmilos/f2f999b5839e2d42d751
#
# DESCENDENTS
#
# All the nodes that have a given node as an ancestor.
# In Modified Preorder Tree lingo,
# a descendent has an `lft` greater that its ancestor
# and a `rgt` less that its ancestor.
use TypedStruct
alias MPTree.Node
alias MPTree.NodeMeta
alias MPTree.Seeker
defmacrop extract_value!(ok_tuple_or_error, error_msg \\ "failed match") do
quote do
case unquote(ok_tuple_or_error) do
{:ok, val} -> val
:error -> raise "MPTree: #{unquote(error_msg)}"
end
end
end
#####################################
# TYPES
@typedoc """
Used by `insert`, `update`, and `fetch` functions to find nodes.
Beware, this will match only the first node it finds.
If there any many nodes that might possibly pass the test,
the results could be unpredictable.
"""
@type match_fn :: (Node.t() -> as_boolean(any))
@typedoc """
Updates a node
"""
@type update_fn :: (Node.t() -> Node.t())
@type nodes :: [Node.t()]
@type transition_path :: [{:exit | :enter, Node.t()}]
typedstruct enforce: true do
field :nodes, [Node.t(), ...]
end
#####################################
# CONSTRUCTORS
@doc """
Create a tree from a [`root_node`](`MPTree.Node`).
"""
@spec from_node(Node.t()) :: t()
def from_node(%{__mptree_node__: %NodeMeta{}} = root_node) do
node_with_default_meta = Map.replace!(root_node, :__mptree_node__, NodeMeta.new())
%__MODULE__{nodes: [node_with_default_meta]}
end
#####################################
# REDUCERS
@doc """
Insert a node or tree as a new child of the node matching `parent_fn`.
In this example, the green node is about to be inserted.

Here is after the insertion.

Returns `:error` if no existing node matches [`parent_fn`](`t:match_fn/0`)
"""
@spec insert(t(), Node.t() | t(), match_fn()) :: {:ok, t()} | :error
def insert(tree, node_or_tree, parent_fn)
def insert(%__MODULE__{} = tree, %{__mptree_node__: %NodeMeta{}} = node, parent_fn)
when is_function(parent_fn, 1) do
subtree = from_node(node)
insert(tree, subtree, parent_fn)
end
def insert(
%__MODULE__{nodes: old_nodes} = tree,
%__MODULE__{nodes: new_nodes},
parent_fn
)
when is_function(parent_fn, 1) do
subtree_starting_id = 1 + __max_node_id__(tree)
old_nodes_addend = 2 * length(new_nodes)
with {:ok, [parent]} <- Seeker.filter(old_nodes, parent_fn, keep: [:matched]) do
parent_rgt = Node.__rgt__(parent)
new_nodes_addend = parent_rgt
incremented_old_nodes =
Stream.map(old_nodes, &Node.__incr_lft_rgt__(&1, old_nodes_addend, parent_rgt))
incremented_new_nodes =
new_nodes
|> Stream.map(&Node.__incr_lft_rgt__(&1, new_nodes_addend))
|> Stream.map(&Node.__incr_id__(&1, subtree_starting_id - 1))
nodes =
[incremented_old_nodes, incremented_new_nodes]
|> Stream.concat()
|> Enum.sort_by(&Node.__lft__/1)
{:ok, %__MODULE__{tree | nodes: nodes}}
end
end
@doc """
Insert a node or tree as a new child of the node matching `parent_fn`.
Throws if no existing node matches [`parent_fn`](`t:match_fn/0`)
"""
@spec insert!(t(), Node.t() | t(), match_fn()) :: t()
def insert!(tree, node_or_tree, parent_fn) when is_function(parent_fn, 1) do
insert(tree, node_or_tree, parent_fn) |> extract_value!
end
@doc """
Apply `t:update_fn/0` to all nodes.
"""
@spec update_nodes(t(), update_fn()) :: t()
def update_nodes(%__MODULE__{nodes: nodes} = tree, update_fn) when is_function(update_fn, 1) do
%__MODULE__{tree | nodes: Enum.map(nodes, update_fn)}
end
@doc """
Apply `t:update_fn/0` to all nodes matching the `t:match_fn/0`.
"""
@spec update_nodes(t(), update_fn(), match_fn()) :: t()
def update_nodes(tree, update_fn, match_fn)
when is_function(update_fn, 1) and is_function(match_fn, 1) do
update_nodes(tree, fn node ->
if match_fn.(node), do: update_fn.(node), else: node
end)
end
#####################################
# CONVERTERS
@doc """
Return ancestors of and the node itself matching `descendent_fn`.
In this example, the blue (ancestor) and red (matched) nodes are returned.

Returns `:error` if no node is found.
"""
@spec fetch_ancestors_and_self(t, MPTree.match_fn()) :: {:ok, nodes()} | :error
def fetch_ancestors_and_self(tree, descendent_fn) do
with {:ok, node} <- fetch_node(tree, descendent_fn) do
ancestors =
tree
|> MPTree.nodes()
|> Enum.filter(&Node.__ancestor_and_descendent__?(&1, node))
{:ok, ancestors ++ [node]}
end
end
@doc """
Return ancestors of and the node itself matching `descendent_fn`.
Throws if no node is found.
"""
@spec fetch_ancestors_and_self!(t, MPTree.match_fn()) :: nodes()
def fetch_ancestors_and_self!(tree, descendent_fn) do
fetch_ancestors_and_self(tree, descendent_fn) |> extract_value!
end
@doc """
Find children of the node matching `parent_fn`
In this example, the red node matches, and the blue nodes are returned
```elixir
[
b,
e
]
```

Returns `:error` if no parent is found.
"""
@spec fetch_children(t(), match_fn) :: {:ok, [Node.t()]} | :error
def fetch_children(%__MODULE__{nodes: nodes} = _tree, parent_fn)
when is_function(parent_fn, 1) do
with {:ok, [parent | tail]} <- Seeker.filter(nodes, parent_fn, keep: [:matched, :descendents]) do
if Node.__leaf__?(parent) do
{:ok, []}
else
{:ok, _children_from_tail(tail, parent)}
end
end
end
@doc """
Find children of the node matching `parent_fn`.
Throws if no parent is found.
"""
@spec fetch_children!(t(), match_fn()) :: [Node.t()]
def fetch_children!(%__MODULE__{} = tree, parent_fn) when is_function(parent_fn, 1) do
fetch_children(tree, parent_fn) |> extract_value!("no node matching #{inspect(parent_fn)}")
end
@doc """
Find descendents of the node matching `ancestor_fn`.
In this example, diagram, the blue (descendent) nodes are returned.
The red (matched) node is not returned.

Returns `:error` if no ancestor is found.
"""
@spec fetch_descendents(t(), match_fn()) :: {:ok, [Node.t()]} | :error
def fetch_descendents(%__MODULE__{nodes: nodes} = _tree, ancestor_fn)
when is_function(ancestor_fn, 1) do
Seeker.filter(nodes, ancestor_fn, keep: [:descendents])
end
@doc """
Find descendents of the node matching `ancestor_fn`.
Throws if no ancestor is found.
"""
@spec fetch_descendents!(t(), match_fn()) :: [Node.t()]
def fetch_descendents!(%__MODULE__{} = tree, ancestor_fn) when is_function(ancestor_fn, 1) do
fetch_descendents(tree, ancestor_fn) |> extract_value!
end
@doc """
Find all ancestors and descendents of the node matching `match_fn`.
In the below diagram, the red (matched node) and blue (ancestor and descendent) nodes are returned.

Returns `:error` if node is not found.
"""
@spec fetch_family_tree(t(), MPTree.match_fn()) :: {:ok, nodes()} | :error
def fetch_family_tree(%__MODULE__{nodes: nodes}, match_fn) when is_function(match_fn, 1) do
Seeker.filter(nodes, match_fn)
end
@doc """
Find all ancestors and descendents of the node matching `match_fn`.
Throws if node is not found.
"""
@spec fetch_family_tree!(t(), MPTree.match_fn()) :: nodes()
def fetch_family_tree!(tree, match_fn) do
fetch_family_tree(tree, match_fn) |> extract_value!()
end
@doc """
Return the first node matching `match_fn`.
In this example, the red and blue nodes would all match, but only the red is returned

Returns `:error` if no node is found.
"""
@spec fetch_node(t(), MPTree.match_fn()) :: {:ok, Node.t()} | :error
def fetch_node(tree, match_fn) do
with node when not is_nil(node) <- tree |> MPTree.nodes() |> Enum.find(match_fn) do
{:ok, node}
else
_ -> :error
end
end
@doc """
Return the first node matching `match_fn`.
Throws if no node is found.
"""
@spec fetch_node!(t(), MPTree.match_fn()) :: Node.t()
def fetch_node!(tree, match_fn) do
fetch_node(tree, match_fn) |> extract_value!
end
@doc """
Find parent of the node matching `child_fn`.
In this example, the blue (parent) node is returned.

Returns `:error` if child is not found or if child is the root node.
"""
@spec fetch_parent(t(), match_fn()) :: {:ok, Node.t()} | :error
def fetch_parent(%__MODULE__{nodes: nodes} = _tree, child_fn) when is_function(child_fn, 1) do
with {:ok, nodes} <- Seeker.filter(nodes, child_fn, keep: [:ancestors, :matched]) do
cond do
length(nodes) >= 2 -> {:ok, Enum.at(nodes, -2)}
true -> :error
end
end
end
@doc """
Find parent of the node matching `child_fn`.
Throws if child is not found or if child is the root node.
"""
@spec fetch_parent!(t(), match_fn()) :: Node.t()
def fetch_parent!(%__MODULE__{} = tree, child_fn) when is_function(child_fn, 1) do
fetch_parent(tree, child_fn) |> extract_value!
end
@doc """
Get the path from a starting node to an ending node.
In this example, the green (starting matched), red (ending matched), and blue (path) nodes are returned
in `:exit` or `:enter` tuples

The above returns
```elixir
[
exit: f,
enter: a,
enter: b
]
```
Returns `:error` if no node is found.
"""
@spec fetch_transition_path(t(), MPTree.match_fn(), MPTree.match_fn()) ::
{:ok, transition_path()} | :error
def fetch_transition_path(tree, start_match_fn, end_match_fn) do
with {:ok, up_path} <- fetch_ancestors_and_self(tree, start_match_fn),
{:ok, down_path} <- fetch_ancestors_and_self(tree, end_match_fn) do
path = do_transition_path(up_path, down_path)
{:ok, path}
end
end
@doc """
Get the path from a starting node to an ending node.
Throws if no node is found.
"""
@spec fetch_transition_path!(t(), MPTree.match_fn(), MPTree.match_fn()) :: transition_path()
def fetch_transition_path!(tree, start_match_fn, end_match_fn) do
fetch_transition_path(tree, start_match_fn, end_match_fn) |> extract_value!
end
@doc """
Get all nodes.
"""
@spec nodes(t()) :: [Node.t(), ...]
def nodes(%__MODULE__{nodes: val} = _tree), do: val
@spec __max_node_id__(t()) :: NodeMeta.id()
def __max_node_id__(%__MODULE__{nodes: [root | _]}) do
Node.__count_self_and_descendents__(root)
end
@spec __node_count__(t()) :: pos_integer
def __node_count__(%__MODULE__{nodes: [root | _]}) do
Node.__count_self_and_descendents__(root)
end
@spec __node_ids__(t()) :: [NodeMeta.id()]
def __node_ids__(%__MODULE__{} = tree) do
Enum.to_list(NodeMeta.__starting_node_id__()..__max_node_id__(tree))
end
@spec __root__(t()) :: Node.t()
def __root__(tree)
def __root__(%__MODULE__{nodes: [root | _]}), do: root
#####################################
# HELPERS
@spec do_transition_path(nodes(), nodes()) :: transition_path()
defp do_transition_path([head1, head2 | state_tail], [head1, head2 | destination_tail]) do
do_transition_path([head2 | state_tail], [head2 | destination_tail])
end
defp do_transition_path([head1 | state_tail], [head1 | destination_tail]) do
state_path_items = Stream.map(state_tail, &{:exit, &1})
destination_path_items = Enum.map(destination_tail, &{:enter, &1})
Enum.reduce(state_path_items, destination_path_items, &[&1 | &2])
end
defp _children_from_tail(tail, parent, children \\ [])
defp _children_from_tail([], _, children) do
children
end
defp _children_from_tail([child | rest], parent, children) do
if Node.__parent_and_last_child__?(parent, child) do
[child | children]
else
child_descendent_count = Node.__count_descendents__(child)
tail = Enum.drop(rest, child_descendent_count)
_children_from_tail(tail, parent, [child | children])
end
end
end