Merkle search tree

2 files changed, 216 insertions, 0 deletions

diff --git a/lib/data/merklesearchtree.ex b/lib/data/merklesearchtree.ex
new file mode 100644
index 0000000..d23e8a7
--- /dev/null
+++ b/lib/data/merklesearchtree.ex
@@ -0,0 +1,181 @@
+defmodule SData.MerkleSearchTree do
+  @moduledoc"""
+  A Merkle search tree.
+
+  A node of the tree is
+    {
+      level,
+      hash_of_node | nil,
+      [ 
+        { k_low_bound, v, hash_of_node | nil },
+        { k_low_bound, v, hash_of_node | nil },
+        ...
+      }
+    }
+  """
+
+  defstruct [:root, :store]
+
+  def new() do
+    %SData.MerkleSearchTree{
+      root:  nil,
+      store: %{}
+    }
+  end
+
+  def insert(state, key, value) do
+    level = calc_level(key)
+    {hash, store} = insert_at(state.store, state.root, key, level, value)
+    %{ state | root: hash, store: store }
+  end
+
+  def remove(_state, _key) do
+    # TODO
+  end
+
+  defp insert_at(store, root, key, level, value) do
+    {new_page, store} = if root == nil do
+      { { level, nil, [ { key, value, nil } ] }, store }
+    else
+      { plevel, low, lst } = store[root]
+      [ { k0, _, _} | _ ] = lst
+      cond do
+        plevel < level ->
+          {low, high, store} = split(store, root, key)
+          { { level, low, [ { key, value, high } ] }, store }
+        plevel == level ->
+          store = Map.delete(store, root)
+          if key < k0 do
+            {low2a, low2b, store} = split(store, low, key)
+            { { level, low2a, [ { key, value, low2b } | lst] }, store }
+          else
+            {new_lst, store} = aux_insert_after_first(store, lst, key, value)
+            { { plevel, low, new_lst }, store }
+          end
+        plevel > level ->
+          store = Map.delete(store, root)
+          if key < k0 do
+            {new_low, store} = insert_at(store, low, key, level, value)
+            { { plevel, new_low, lst }, store }
+          else 
+            {new_lst, store} = aux_insert_sub_after_first(store, lst, key, level, value)
+            { { plevel, low, new_lst }, store }
+          end
+      end
+    end
+    hash = SData.term_hash new_page
+    store = Map.put(store, hash, new_page)
+    { hash, store }
+  end
+
+  defp split(store, hash, key) do
+    if hash == nil do
+      {nil, nil, store}
+    else
+      { level, low, lst } = store[hash]
+      store = Map.delete(store, hash)
+      [ { k0, _, _} | _ ] = lst
+      cond do
+        key < k0 ->
+          {lowlow, lowhi, store} = split(store, low, key)
+          newp2 = { level, lowhi, lst }
+          newp2h = SData.term_hash newp2
+          store = Map.put(store, newp2h, newp2)
+          {lowlow, newp2h, store}
+        key > k0 ->
+          {lst1, p2, store} = split_aux(store, lst, key, level)
+          newp1 = {level, low, lst1}
+          newp1h = SData.term_hash newp1
+          store = Map.put(store, newp1h, newp1)
+          {newp1h, p2, store}
+      end 
+    end
+  end
+
+  defp split_aux(store, lst, key, level) do
+    case lst do
+      [ {k1, _, _} | _ ] when k1 == key ->
+        raise "Bad logic"
+      [ {k1, v1, r1} ] ->
+        {r1l, r1h, store} = split(store, r1, key)
+        { [{k1, v1, r1l}], r1h, store }
+      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
+        {r1l, r1h, store} = split(store, r1, key)
+        p2 = { level, r1h, [ {k2, v2, r2} | rst ] }
+        p2h = SData.term_hash p2
+        store = store
+          |> Map.put(p2h, p2)
+        { [{k1, v1, r1l}], p2h, store }
+      [ fst | rst ] ->
+        {rst2, hi, store} = split_aux(store, rst, key, level)
+        { [ fst | rst2 ], hi, store }
+    end
+  end
+
+  defp aux_insert_after_first(store, lst, key, value) do
+    case lst do
+      [ {k1, _, r1} | rst ] when k1 == key ->
+        { [ {k1, value, r1} | rst ], store }
+      [ {k1, v1, r1} ] ->
+        {r1a, r1b, new_store} = split(store, r1, key)
+        { [ {k1, v1, r1a}, {key, value, r1b} ], new_store }
+      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
+        {r1a, r1b, new_store} = split(store, r1, key)
+        { [ {k1, v1, r1a}, {key, value, r1b}, {k2, v2, r2} | rst ], new_store }
+      [ fst | rst ] ->
+        {rst2, new_store} = aux_insert_after_first(store, rst, key, value)
+        { [ fst | rst2 ], new_store }
+    end
+  end
+
+  defp aux_insert_sub_after_first(store, lst, key, level, value) do
+    case lst do
+      [ {k1, _, _} | _ ] when k1 == key ->
+        raise "Bad logic"
+      [ {k1, v1, r1} ] ->
+        {r1new, store_new} = insert_at(store, r1, key, level, value)
+        { [{k1, v1, r1new}], store_new }
+      [ {k1, v1, r1}, {k2, v2, r2} | rst ] when key < k2->
+        {r1new, store_new} = insert_at(store, r1, key, level, value)
+        { [{k1, v1, r1new}, {k2, v2, r2} | rst], store_new }
+      [ fst | rst ] ->
+        {rst2, new_store} = aux_insert_sub_after_first(store, rst, key, level, value)
+        { [ fst | rst2 ], new_store }
+    end
+  end
+
+  def dump(state) do
+    dump(state.store, state.root, "")
+  end
+
+  defp dump(store, root, lvl) do
+    case root do
+      nil ->
+        IO.puts(lvl <> "nil")
+      _ ->
+        { level, low, lst} = store[root]
+        IO.puts(lvl <> "#{root|>Base.encode16} (#{level})")
+        dump(store, low, lvl <> "  ")
+        for {k, v, r} <- lst do
+          IO.puts(lvl<>"- #{inspect k} => #{inspect v}")
+          dump(store, r, lvl <> "  ")
+        end
+    end
+  end
+
+  defp calc_level(key) do
+    key
+      |> SData.term_hash
+      |> Base.encode16
+      |> String.to_charlist
+      |> count_leading_zeroes
+  end
+
+  defp count_leading_zeroes('0' ++ rest) do
+    1 + count_leading_zeroes(rest)
+  end
+  defp count_leading_zeroes(_) do
+    0
+  end
+
+end
diff --git a/test/mst_test.exs b/test/mst_test.exs
new file mode 100644
index 0000000..68886c7
--- /dev/null
+++ b/test/mst_test.exs
@@ -0,0 +1,35 @@
+defmodule ShardTest.MST do
+  use ExUnit.Case
+  alias SData.MerkleSearchTree, as: MST
+  doctest Shard.Application
+
+  test "merkle search tree 1" do
+    y = Enum.reduce(0..1000, MST.new(),
+                    fn i, acc -> MST.insert(acc, i, i) end)
+
+
+    z = Enum.reduce(Enum.shuffle(0..1000), MST.new(),
+                    fn i, acc -> MST.insert(acc, i, i) end)
+
+    IO.puts "y.root: #{y.root|>Base.encode16}"
+    IO.puts "z.root: #{z.root|>Base.encode16}"
+    assert y.root == z.root
+  end
+
+  test "merkle search tree 2" do
+    items = for i <- 0..1000 do
+              h = SData.term_hash i
+              {h, SData.term_hash h}
+            end
+
+    y = Enum.reduce(items, MST.new(),
+                    fn {k, v}, acc -> MST.insert(acc, k, v) end)
+
+    z = Enum.reduce(Enum.shuffle(items), MST.new(),
+                    fn {k, v}, acc -> MST.insert(acc, k, v) end)
+
+    IO.puts "y.root: #{y.root|>Base.encode16}"
+    IO.puts "z.root: #{z.root|>Base.encode16}"
+    assert y.root == z.root
+  end
+end