1 files changed, 112 insertions, 77 deletions

diff --git a/lib/data/merklesearchtree.ex b/lib/data/merklesearchtree.ex
index 963653e..0554f82 100644
--- a/lib/data/merklesearchtree.ex
+++ b/lib/data/merklesearchtree.ex
@@ -7,8 +7,8 @@ defmodule SData.MerkleSearchTree do
       level,
       hash_of_node | nil,
       [ 
-        { k_low_bound, v, hash_of_node | nil },
-        { k_low_bound, v, hash_of_node | nil },
+        { item_low_bound, hash_of_node | nil },
+        { item_low_bound, hash_of_node | nil },
         ...
       }
     }
@@ -16,8 +16,22 @@ defmodule SData.MerkleSearchTree do
 
   alias SData.PageStore, as: Store
 
-  defstruct [:root, :store]
+  @doc"""
+  Create a new Merkle search tree.
 
+  This structure can be used as a set with only true keys,
+  or as a map if a merge function is given.
+
+  `cmp` is a compare function for keys as defined in module `SData`.
+
+  `merge` is a function for merging two items that have the same key.
+  """
+  defstruct root: nil,
+            store: SData.LocalStore.new,
+            cmp: &SData.cmp_term/2,
+            merge: &SData.merge_true/2
+
+          
   defmodule Page do
     defstruct [:level, :low, :list]
   end
@@ -34,17 +48,10 @@ defmodule SData.MerkleSearchTree do
   end
 
 
-  def new(store \\ nil) do
-    store = store || SData.LocalStore.new
-    %SData.MerkleSearchTree{
-      root:  nil,
-      store: store
-    }
-  end
 
-  def insert(state, key, value) do
+  def insert(state, key, value \\ true) do
     level = calc_level(key)
-    {hash, store} = insert_at(state.store, state.root, key, level, value)
+    {hash, store} = insert_at(state, state.store, state.root, key, level, value)
     %{ state | root: hash, store: store }
   end
 
@@ -52,7 +59,7 @@ defmodule SData.MerkleSearchTree do
     # TODO
   end
 
-  defp insert_at(store, root, key, level, value) do
+  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 }
     else
@@ -60,46 +67,48 @@ defmodule SData.MerkleSearchTree do
       [ { k0, _, _} | _ ] = lst
       cond do
         plevel < level ->
-          {low, high, store} = split(store, root, key)
+          {low, high, store} = split(s, store, root, key)
           { %Page{ level: level, low: low, list: [ { key, value, high } ] }, store }
         plevel == level ->
           store = Store.free(store, root)
-          if key < k0 do
-            {low2a, low2b, store} = split(store, low, key)
-            { %Page{ level: level, low: low2a, list: [ { key, value, low2b } | lst] }, store }
-          else
-            {new_lst, store} = aux_insert_after_first(store, lst, key, value)
-            { %Page{ level: plevel, low: low, list: new_lst }, store }
+          case s.cmp.(key, k0) do
+            :before ->
+              {low2a, low2b, store} = split(s, store, low, key)
+              { %Page{ level: level, low: low2a, list: [ { key, value, low2b } | lst] }, store }
+            _ ->
+              {new_lst, store} = aux_insert_after_first(s, store, lst, key, value)
+              { %Page{ level: plevel, low: low, list: new_lst }, store }
           end
         plevel > level ->
           store = Store.free(store, root)
-          if key < k0 do
-            {new_low, store} = insert_at(store, low, key, level, value)
-            { %Page{ level: plevel, low: new_low, list: lst }, store }
-          else 
-            {new_lst, store} = aux_insert_sub_after_first(store, lst, key, level, value)
-            { %Page{ level: plevel, low: low, list: new_lst }, store }
+          case s.cmp.(key, k0) do
+            :before ->
+              {new_low, store} = insert_at(s, store, low, key, level, value)
+              { %Page{ level: plevel, low: new_low, list: lst }, store }
+            :after ->
+              {new_lst, store} = aux_insert_sub_after_first(s, store, lst, key, level, value)
+              { %Page{ level: plevel, low: low, list: new_lst }, store }
           end
       end
     end
     Store.put(store, new_page)    # returns {hash, store}
   end
 
-  defp split(store, hash, key) do
+  defp split(s, store, hash, key) do
     if hash == nil do
       {nil, nil, store}
     else
       %Page{ level: level, low: low, list: lst } = Store.get(store, hash)
       store = Store.free(store, hash)
       [ { k0, _, _} | _ ] = lst
-      cond do
-        key < k0 ->
-          {lowlow, lowhi, store} = split(store, low, key)
+      case s.cmp.(key, k0) do
+        :before ->
+          {lowlow, lowhi, store} = split(s, store, low, key)
           newp2 = %Page{ level: level, low: lowhi, list: lst }
           {newp2h, store} = Store.put(store, newp2)
           {lowlow, newp2h, store}
-        key > k0 ->
-          {lst1, p2, store} = split_aux(store, lst, key, level)
+        :after ->
+          {lst1, p2, store} = split_aux(s, store, lst, key, level)
           newp1 = %Page{ level: level, low: low, list: lst1}
           {newp1h, store} = Store.put(store, newp1)
           {newp1h, p2, store}
@@ -107,79 +116,105 @@ defmodule SData.MerkleSearchTree do
     end
   end
 
-  defp split_aux(store, lst, key, level) do
+  defp split_aux(s, 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)
+        if s.cmp.(k1, key) == :duplicate do
+          raise "Bad logic"
+        end
+        {r1l, r1h, store} = split(s, 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 = %Page{ level: level, low: r1h, list: [ {k2, v2, r2} | rst ] }
-        {p2h, store} = Store.put(store, p2)
-        { [{k1, v1, r1l}], p2h, store }
-      [ fst | rst ] ->
-        {rst2, hi, store} = split_aux(store, rst, key, level)
-        { [ fst | rst2 ], hi, store }
+      [ {k1, v1, r1} = fst, {k2, v2, r2} = rst1 | rst ] ->
+        case s.cmp.(key, k2) do
+          :before ->
+            {r1l, r1h, store} = split(s, store, r1, key)
+            p2 = %Page{ level: level, low: r1h, list: [ {k2, v2, r2} | rst ] }
+            {p2h, store} = Store.put(store, p2)
+            { [{k1, v1, r1l}], p2h, store }
+          :after ->
+            {rst2, hi, store} = split_aux(s, store, [rst1 | rst], key, level)
+            { [ fst | rst2 ], hi, store }
+          :duplicate ->
+            raise "Bad logic"
+        end
     end
   end
 
-  defp aux_insert_after_first(store, lst, key, value) do
+  defp aux_insert_after_first(s, store, lst, key, value) do
     case lst do
-      [ {k1, _old_value, 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 }
+        case s.cmp.(k1, key) do
+          :duplicate ->
+            { [ {k1, s.merge.(v1, value), r1} ], store }
+          :before ->
+            {r1a, r1b, new_store} = split(s, store, r1, key)
+            { [ {k1, v1, r1a}, {key, value, r1b} ], new_store }
+        end
+      [ {k1, v1, r1} = fst, {k2, v2, r2} = rst1 | rst ] ->
+        case s.cmp.(k1, key) do
+          :duplicate ->
+            { [ {k1, s.merge.(v1, value), r1}, rst1 | rst ], store }
+          :before ->
+            case s.cmp.(k2, key) do
+              :after ->
+                {r1a, r1b, new_store} = split(s, store, r1, key)
+                { [ {k1, v1, r1a}, {key, value, r1b}, {k2, v2, r2} | rst ], new_store }
+              _ ->
+                {rst2, new_store} = aux_insert_after_first(s, store, [rst1 | rst], key, value)
+                { [ fst | rst2 ], new_store }
+            end
+        end
+      end
     end
-  end
 
-  defp aux_insert_sub_after_first(store, lst, key, level, value) do
+  defp aux_insert_sub_after_first(s, 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)
+        if s.cmp.(k1, key) == :duplicate do
+          raise "Bad logic"
+        end
+        {r1new, store_new} = insert_at(s, 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 }
+      [ {k1, v1, r1} = fst, {k2, _, _} = rst1 | rst ] ->
+        if s.cmp.(k1, key) == :duplicate do
+          raise "Bad logic"
+        end
+        case s.cmp.(key, k2) do
+          :before ->
+            {r1new, store_new} = insert_at(s, store, r1, key, level, value)
+            { [{k1, v1, r1new}, rst1 | rst], store_new }
+          _ ->
+            {rst2, new_store} = aux_insert_sub_after_first(s, store, [rst1 |rst], key, level, value)
+            { [ fst | rst2 ], new_store }
+        end
     end
   end
 
   def get(state, key) do
-    get(state.store, state.root, key)
+    get(state, state.store, state.root, key)
   end
 
-  defp get(store, root, key) do
+  defp get(s, store, root, key) do
     case root do
       nil -> nil
       _ ->
         %Page{ level: _, low: low, list: lst } = Store.get(store, root)
-        get_aux(store, low, lst, key)
+        get_aux(s, store, low, lst, key)
     end
   end
 
-  defp get_aux(store, low, lst, key) do
+  defp get_aux(s, store, low, lst, key) do
     case lst do
       [] ->
-        get(store, low, key)
-      [ {k, v, _} | _ ] when key == k ->
-        v
-      [ {k, _, _} | _ ] when key < k ->
-        get(store, low, key)
-      [ {k, _, low2} | rst ] when key > k ->
-        get_aux(store, low2, rst, key)
+        get(s, store, low, key)
+      [ {k, v, low2} | rst ] ->
+        case s.cmp.(key, k) do
+          :duplicate -> v
+          :before ->
+            get(s, store, low, key)
+          :after -> 
+            get_aux(s, store, low2, rst, key)
+        end
     end
   end