1 files changed, 95 insertions, 14 deletions

diff --git a/lib/data/merklesearchtree.ex b/lib/data/merklesearchtree.ex
index 0554f82..b751a08 100644
--- a/lib/data/merklesearchtree.ex
+++ b/lib/data/merklesearchtree.ex
@@ -48,17 +48,15 @@ defmodule SData.MerkleSearchTree do
   end
 
 
-
+  @doc"""
+    Insert an item into the search tree.
+  """
   def insert(state, key, value \\ true) do
     level = calc_level(key)
     {hash, store} = insert_at(state, state.store, state.root, key, level, value)
     %{ state | root: hash, store: store }
   end
 
-  def remove(_state, _key) do
-    # TODO
-  end
-
   defp insert_at(s, store, root, key, level, value) do
     {new_page, store} = if root == nil do
       { %Page{ level: level, low: nil, list: [ { key, value, nil } ] }, store }
@@ -190,34 +188,118 @@ defmodule SData.MerkleSearchTree do
     end
   end
 
-  def get(state, key) do
-    get(state, state.store, state.root, key)
+  @doc"""
+    Merge values from another MST in this MST
+  """
+  def merge(to, from, callback \\ fn _, _ -> nil end) do
+    { store, root } = merge_aux(to, from, to.store, to.root, from.root, callback)
+    %{ to | store: store, root: root }
+  end
+
+  defp merge_aux(s1, s2, store, r1, r2, callback) do
+    case {r1, r2} do
+      {_, nil} ->
+        { store, r1 }
+      {nil, _} ->
+        store = Store.copy(store, s2.store, r2)
+        rec_callback(store, r2, callback)
+        { store, r2 }
+      _ ->
+        IO.puts("not implemented: complex merge step")
+        #TODO
+        { store, r1 }
+    end
   end
 
-  defp get(s, store, root, key) do
+  defp rec_callback(store, root, callback) do
     case root do
       nil -> nil
       _ ->
         %Page{ level: _, low: low, list: lst } = Store.get(store, root)
-        get_aux(s, store, low, lst, key)
+        rec_callback(store, low, callback)
+        for {k, v, rst} <- lst do
+          callback.(k, v)
+          rec_callback(store, rst, callback)
+        end
+    end
+  end
+
+  @doc"""
+    Get value for a specific key in search tree.
+  """
+  def get(state, key) do
+    get(state, state.root, key)
+  end
+
+  defp get(s, root, key) do
+    case root do
+      nil -> nil
+      _ ->
+        %Page{ level: _, low: low, list: lst } = Store.get(s.store, root)
+        get_aux(s, low, lst, key)
     end
   end
 
-  defp get_aux(s, store, low, lst, key) do
+  defp get_aux(s, low, lst, key) do
     case lst do
       [] ->
-        get(s, store, low, key)
+        get(s, low, key)
       [ {k, v, low2} | rst ] ->
         case s.cmp.(key, k) do
           :duplicate -> v
           :before ->
-            get(s, store, low, key)
+            get(s, low, key)
           :after -> 
-            get_aux(s, store, low2, rst, key)
+            get_aux(s, low2, rst, key)
         end
     end
   end
 
+  @doc"""
+    Get the last n items of the tree, or the last n items
+    strictly before given upper bound if non nil
+  """
+  def last(state, top_bound, num) do
+    last(state, state.root, top_bound, num)
+  end
+
+  defp last(s, root, top_bound, num) do
+    case root do
+      nil -> []
+      _ ->
+        %Page{ level: _, low: low, list: lst } = Store.get(s.store, root)
+        last_aux(s, low, lst, top_bound, num)
+    end
+  end
+
+  defp last_aux(s, low, lst, top_bound, num) do
+    case lst do
+      [] ->
+        last(s, low, top_bound, num)
+      [ {k, v, low2} | rst ] ->
+        if top_bound == nil or s.cmp.(top_bound, k) == :after do
+          items = last_aux(s, low2, rst, top_bound, num)
+          items = if Enum.count(items) < num do
+            [ {k, v} | items ]
+          else
+            items
+          end
+          cnt = Enum.count items
+          if cnt < num do
+            last(s, low, top_bound, num - cnt) ++ items
+          else
+            items
+          end
+        else
+          last(s, low, top_bound, num)
+        end
+    end
+  end
+
+
+  @doc"""
+    Dump Merkle search tree structure.
+  """
   def dump(state) do
     dump(state.store, state.root, "")
   end
@@ -251,5 +333,4 @@ defmodule SData.MerkleSearchTree do
   defp count_leading_zeroes(_) do
     0
   end
-
 end