flat_tree

defmodule FlatTree do
  @moduledoc """
  You can represent a binary tree in a simple flat list using the following structure:

            3
        1       5
      0   2   4   6  ...

  `FlatTree` exposes a series of functions to help you build and maintain this data structure.

      import FlatTree

      list = []
      i = FlatTree.index(0, 0) # get array index for depth: 0, offset: 0
      j = FlatTree.index(1, 0) # get array index for depth: 1, offset: 0

      # use these indexes to store some data
      list = List.insert_at(list, i, "a")
      list = List.insert_at(list, j, "b")
      list = List.insert_at(list, FlatTree.parent(i), "parent of a and b")
  """

  use Bitwise, only_operators: true

  import Integer, only: [floor_div: 2]

  @doc """
  Returns a tuple {left_child, right_child} with the indexes of this elements
  children. If this element does not have any children is returns `nil`.

  ## Examples

      iex> FlatTree.children(0)
      nil

      iex> FlatTree.children(3)
      {1, 5}
  """
  @spec children(index :: integer, depth :: integer | nil) :: {integer, integer}
  def children(index, depth \\ nil)
  def children(index, _) when is_integer(index) and (index &&& 1) == 0, do: nil
  def children(index, depth) when is_nil(depth), do: children(index, depth(index))

  def children(index, depth) when is_integer(index) and is_integer(depth) do
    offset = offset(index, depth) * 2
    {index(depth - 1, offset), index(depth - 1, offset + 1)}
  end

  @doc """
  Returns how many nodes (including parent nodes) a tree contains.

  ## Examples

      iex> FlatTree.count(1)
      3
  """
  @spec count(index :: integer, depth :: integer | nil) :: integer
  def count(index, depth \\ nil)
  def count(index, _) when is_integer(index) and (index &&& 1) == 0, do: 1
  def count(index, depth) when is_nil(depth), do: count(index, depth(index))

  def count(index, depth) when is_integer(index) and is_integer(depth) do
    two_pow(depth + 1) - 1
  end

  @doc """
  Returns the depth of an element.

  ## Examples

      iex> FlatTree.depth(0)
      0

      iex> FlatTree.depth(3)
      2
  """
  @spec depth(index :: integer, depth :: integer) :: integer
  def depth(index, depth \\ 0), do: do_depth(index + 1, depth)

  @doc """
  Returns a list of all the full roots (subtrees where all nodes have either 2
  or 0 children) < index. For example `full_roots(8)` returns `[3]` since the
  subtree rooted at `3` spans `0 -> 6` and the tree rooted at `7` has a child
  located at `9` which is `>= 8`.

  ## Examples

      iex> FlatTree.full_roots(0)
      []

      iex> FlatTree.full_roots(18)
      [7, 16]
  """
  @spec full_roots(index :: integer, result :: [integer]) :: [integer]
  def full_roots(index, result \\ [])

  def full_roots(index, _) when is_integer(index) and (index &&& 1) == 1 do
    raise("You can only look up roots for depth(0) blocks")
  end

  def full_roots(index, result) do
    index = floor_div(index, 2)
    do_full_roots(result, index)
  end

  @doc """
  Returns an array index for the tree element at the given depth and offset.

  ## Examples

      iex> FlatTree.index(0, 0)
      0

      iex> FlatTree.index(0, 2)
      4
  """
  @spec index(depth :: integer, offset :: integer) :: integer
  def index(depth, offset) when is_integer(depth) and is_integer(offset) do
    (1 + 2 * offset) * two_pow(depth) - 1
  end

  @doc """
  Returns the index of the left child of this element. If this element does not
  have a left child it returns `-1`.

  ## Examples

      iex> FlatTree.left_child(0)
      -1

      iex> FlatTree.left_child(3)
      1
  """
  @spec left_child(index :: integer, depth :: integer | nil) :: integer
  def left_child(index, depth \\ nil)
  def left_child(index, _) when (index &&& 1) == 0, do: -1
  def left_child(index, depth) when is_nil(depth), do: left_child(index, depth(index))

  def left_child(index, depth) when is_integer(index) and is_integer(depth) do
    index(depth - 1, offset(index, depth) * 2)
  end

  @doc """
  Returns the left spanning in index in the tree `index` spans.

  ## Examples

      iex> FlatTree.left_span(0)
      0

      iex> FlatTree.left_span(23)
      16
  """
  @spec left_span(index :: integer, depth :: integer | nil) :: integer
  def left_span(index, depth \\ nil)
  def left_span(index, _) when (index &&& 1) == 0, do: index
  def left_span(index, depth) when is_nil(depth), do: left_span(index, depth(index))

  def left_span(index, depth) when is_integer(index) and is_integer(depth) do
    offset(index, depth) * two_pow(depth + 1)
  end

  @doc """
  Returns the relative offset of an element.

  ## Examples

      iex> FlatTree.offset(1)
      0

      iex> FlatTree.offset(4)
      2
  """
  @spec offset(index :: integer, depth :: integer | nil) :: integer
  def offset(index, depth \\ nil)
  def offset(index, _) when (index &&& 1) == 0, do: floor_div(index, 2)
  def offset(index, depth) when is_nil(depth), do: offset(index, depth(index))

  def offset(index, depth) when is_integer(index) and is_integer(depth) do
    (index + 1)
    |> floor_div(two_pow(depth))
    |> Kernel.-(1)
    |> floor_div(2)
  end

  @doc """
  Returns the index of the parent element in tree.

  ## Examples

      iex> FlatTree.parent(0)
      1

      iex> FlatTree.parent(1)
      3
  """
  @spec parent(index :: integer, depth :: integer | nil) :: integer
  def parent(index, depth \\ nil)
  def parent(index, depth) when is_nil(depth), do: parent(index, depth(index))

  def parent(index, depth) when is_integer(index) and is_integer(depth) do
    offset = offset(index, depth)
    index(depth + 1, right_shift(offset))
  end

  @doc """
  Returns the index of the right child of this element. If this element does not
  have a right child it returns `-1`.

  ## Examples

      iex> FlatTree.right_child(0)
      -1

      iex> FlatTree.right_child(3)
      5 
  """
  @spec right_child(index :: integer, depth :: integer | nil) :: integer
  def right_child(index, depth \\ nil)
  def right_child(index, _) when (index &&& 1) == 0, do: -1
  def right_child(index, depth) when is_nil(depth), do: right_child(index, depth(index))

  def right_child(index, depth) when is_integer(index) and is_integer(depth) do
    index(depth - 1, 1 + offset(index, depth) * 2)
  end

  @doc """
  Returns the right spanning in index in the tree `index` spans.

  ## Examples

      iex> FlatTree.right_span(0)
      0

      iex> FlatTree.right_span(23)
      30
  """
  @spec right_span(index :: integer, depth :: integer | nil) :: integer
  def right_span(index, depth \\ nil)
  def right_span(index, _) when (index &&& 1) == 0, do: index
  def right_span(index, depth) when is_nil(depth), do: right_span(index, depth(index))

  def right_span(index, depth) when is_integer(index) and is_integer(depth) do
    (offset(index, depth) + 1) * two_pow(depth + 1) - 2
  end

  @doc """
  Return the index of this elements sibling.

  ## Examples

      iex> FlatTree.sibling(0)
      2

      iex> FlatTree.sibling(2)
      0
  """
  @spec sibling(index :: integer, depth :: integer | nil) :: integer
  def sibling(index, depth \\ nil)
  def sibling(index, depth) when is_nil(depth), do: sibling(index, depth(index))

  def sibling(index, depth) when is_integer(index) and is_integer(depth) do
    offset = offset(index, depth)
    offset = if (offset &&& 1) == 1, do: offset - 1, else: offset + 1
    index(depth, offset)
  end

  @doc """
  Returns the range (inclusive) the tree root at `index` spans.

  ## Examples

      iex> FlatTree.spans(0)
      {0, 0}

      iex> FlatTree.spans(3)
      {0, 6}
  """
  @spec spans(index :: integer, depth :: integer | nil) :: {integer, integer}
  def spans(index, depth \\ nil)
  def spans(index, _) when is_integer(index) and (index &&& 1) == 0, do: {index, index}
  def spans(index, depth) when is_nil(depth), do: spans(index, depth(index))

  def spans(index, depth) when is_integer(index) and is_integer(depth) do
    offset = offset(index, depth)
    width = two_pow(depth + 1)

    {offset * width, (offset + 1) * width - 2}
  end

  @spec do_depth(index :: integer, depth :: integer) :: integer
  defp do_depth(index, depth) when (index &&& 1) == 1, do: depth
  defp do_depth(index, depth), do: index |> right_shift() |> do_depth(depth + 1)

  @spec do_full_roots(result :: [integer], index :: integer, offset :: integer, factor :: integer) ::
          [integer]
  defp do_full_roots(result, index, offset \\ 0, factor \\ 1)

  defp do_full_roots(result, index, _, _) when is_list(result) and index == 0 do
    Enum.reverse(result)
  end

  defp do_full_roots(result, index, offset, factor)
       when is_list(result) and is_integer(index) and is_integer(offset) and is_integer(factor) do
    factor = get_factor(index, factor)
    result = [offset + factor - 1 | result]
    offset = offset + 2 * factor
    index = index - factor
    do_full_roots(result, index, offset)
  end

  @spec get_factor(index :: integer, factor :: integer) :: integer
  defp get_factor(index, factor) when factor * 2 > index, do: factor

  defp get_factor(index, factor) when is_integer(index) and is_integer(factor) do
    get_factor(index, factor * 2)
  end

  @spec right_shift(n :: integer) :: integer
  defp right_shift(n), do: floor_div(n - (n &&& 1), 2)

  @spec two_pow(n :: integer) :: integer
  defp two_pow(n) when n < 31, do: 1 <<< n
  defp two_pow(n), do: (1 <<< 30) * (1 <<< (n - 30))
end