diff --git a/lib/merkle_map.ex b/lib/merkle_map.ex
index 9f3a9b6..eb6e29d 100644
--- a/lib/merkle_map.ex
+++ b/lib/merkle_map.ex
@@ -4,6 +4,7 @@ defmodule MerkleMap do
   """
 
   alias MerkleMap.MerkleTree
+  alias MerkleMap.MerkleTree.Diff
 
   defstruct map: %{},
             merkle_tree: MerkleTree.new()
@@ -94,55 +95,19 @@ defmodule MerkleMap do
     {:ok, MerkleTree.diff_keys(mm1.merkle_tree, mm2.merkle_tree)}
   end
 
-  def prepare_partial_diff(mm, depth) do
-    {:continue,
-     %MerkleTree.Diff{trees: [{<<>>, MerkleTree.subtree(mm.merkle_tree, <<>>, depth)}]}}
+  def prepare_partial_diff(%__MODULE__{} = mm, depth) do
+    MerkleTree.prepare_partial_diff(mm.merkle_tree, depth)
   end
 
-  def continue_partial_diff(mm1, %__MODULE__.MerkleTree.Diff{} = partial, depth) do
-    continue_partial_diff(partial, mm1, depth)
+  def continue_partial_diff(%__MODULE__{} = mm, %Diff{} = partial, depth) do
+    MerkleTree.continue_partial_diff(mm.merkle_tree, partial, depth)
   end
 
-  def continue_partial_diff(%__MODULE__.MerkleTree.Diff{} = partial, %__MODULE__{} = mm, depth)
-      when is_integer(depth) and depth > 0 do
-    diff_keys =
-      Enum.reduce(partial.trees, [], fn {loc, tree}, acc_keys ->
-        sub_tree = MerkleTree.subtree(mm.merkle_tree, loc, depth)
-
-        diff_keys = MerkleTree.diff_keys(sub_tree, tree, bit_size(loc))
-
-        Enum.map(diff_keys, fn
-          {:partial, partial_loc} -> {:partial, <<loc::bitstring, partial_loc::bitstring>>}
-          other -> other
-        end)
-
-        [diff_keys | acc_keys]
-      end)
-
-    {partials, keys} =
-      List.flatten(diff_keys)
-      |> Enum.split_with(fn
-        {:partial, _loc} -> true
-        _ -> false
-      end)
-
-    trees =
-      Enum.map(partials, fn {:partial, loc} ->
-        sub_tree = MerkleTree.subtree(mm.merkle_tree, loc, depth)
-        {loc, sub_tree}
-      end)
-
-    case trees do
-      [] -> {:ok, partial.keys ++ keys}
-      trees -> {:continue, %__MODULE__.MerkleTree.Diff{keys: partial.keys ++ keys, trees: trees}}
-    end
+  def continue_partial_diff(%Diff{} = partial, %__MODULE__{} = mm, depth) do
+    MerkleTree.continue_partial_diff(mm.merkle_tree, partial, depth)
   end
 
-  def truncate_diff(%__MODULE__.MerkleTree.Diff{} = diff, amount) do
-    keys = Enum.take(diff.keys, amount)
-    trees = Enum.take(diff.trees, amount - length(keys))
-    %{diff | keys: keys, trees: trees}
-  end
+  defdelegate truncate_diff(diff, amount), to: Diff
 
   def merge(mm1, mm2) do
     {:ok, diff_keys} = diff_keys(mm1, mm2)
diff --git a/lib/merkle_map/merkle_tree.ex b/lib/merkle_map/merkle_tree.ex
index 8e42300..2cd90c9 100644
--- a/lib/merkle_map/merkle_tree.ex
+++ b/lib/merkle_map/merkle_tree.ex
@@ -8,6 +8,7 @@ defmodule MerkleMap.MerkleTree do
   defstruct [:tree]
 
   alias MerkleMap.MerkleTreeImpl
+  alias MerkleMap.MerkleTree.Diff
 
   @spec new(Enumerable.t()) :: t()
   def new(enum) do
@@ -59,4 +60,39 @@ defmodule MerkleMap.MerkleTree do
   def update_hashes(%__MODULE__{tree: tree}) do
     %__MODULE__{tree: MerkleTreeImpl.calculate_hashes(tree)}
   end
+
+  def prepare_partial_diff(merkle_tree, depth) do
+    {:continue, %Diff{trees: [{<<>>, subtree(merkle_tree, <<>>, depth)}]}}
+  end
+
+  def continue_partial_diff(merkle_tree, %Diff{} = partial, depth)
+      when is_integer(depth) and depth > 0 do
+    {partials, keys} =
+      partial.trees
+      |> Enum.flat_map(fn {loc, tree} ->
+        merkle_tree
+        |> subtree(loc, depth)
+        |> diff_keys(tree, bit_size(loc))
+      end)
+      |> Enum.split_with(fn
+        {:partial, _loc} -> true
+        _ -> false
+      end)
+
+    trees =
+      Enum.map(partials, fn {:partial, loc} ->
+        {loc, subtree(merkle_tree, loc, depth)}
+      end)
+
+    case trees do
+      [] -> {:ok, partial.keys ++ keys}
+      trees -> {:continue, %Diff{keys: partial.keys ++ keys, trees: trees}}
+    end
+  end
+
+  def truncate_diff(%Diff{} = diff, amount) do
+    keys = Enum.take(diff.keys, amount)
+    trees = Enum.take(diff.trees, amount - length(keys))
+    %{diff | keys: keys, trees: trees}
+  end
 end
diff --git a/lib/merkle_map/merkle_tree_diff.ex b/lib/merkle_map/merkle_tree_diff.ex
index 6c4897c..4165b1b 100644
--- a/lib/merkle_map/merkle_tree_diff.ex
+++ b/lib/merkle_map/merkle_tree_diff.ex
@@ -2,4 +2,10 @@ defmodule MerkleMap.MerkleTree.Diff do
   @moduledoc false
 
   defstruct trees: [], keys: []
+
+  def truncate_diff(%__MODULE__{} = diff, amount) do
+    keys = Enum.take(diff.keys, amount)
+    trees = Enum.take(diff.trees, amount - length(keys))
+    %{diff | keys: keys, trees: trees}
+  end
 end
diff --git a/mix.exs b/mix.exs
index b0bdfbe..8b6b6c0 100644
--- a/mix.exs
+++ b/mix.exs
@@ -32,7 +32,7 @@ defmodule MerkleMap.MixProject do
   # Run "mix help deps" to learn about dependencies.
   defp deps do
     [
-      {:stream_data, "~> 0.4", only: :test},
+      {:stream_data, "~> 0.5", only: :test},
       {:benchee, "> 0.0.0", only: :dev},
       {:ex_doc, "> 0.0.0", only: :dev}
     ]
diff --git a/mix.lock b/mix.lock
index 3c09c04..2831131 100644
--- a/mix.lock
+++ b/mix.lock
@@ -1,10 +1,10 @@
 %{
-  "benchee": {:hex, :benchee, "1.0.1", "66b211f9bfd84bd97e6d1beaddf8fc2312aaabe192f776e8931cb0c16f53a521", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}], "hexpm"},
-  "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm"},
-  "earmark": {:hex, :earmark, "1.3.2", "b840562ea3d67795ffbb5bd88940b1bed0ed9fa32834915125ea7d02e35888a5", [:mix], [], "hexpm"},
-  "ex_doc": {:hex, :ex_doc, "0.20.1", "88eaa16e67c505664fd6a66f42ddb962d424ad68df586b214b71443c69887123", [:mix], [{:earmark, "~> 1.3", [hex: :earmark, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.10", [hex: :makeup_elixir, repo: "hexpm", optional: false]}], "hexpm"},
-  "makeup": {:hex, :makeup, "0.8.0", "9cf32aea71c7fe0a4b2e9246c2c4978f9070257e5c9ce6d4a28ec450a839b55f", [:mix], [{:nimble_parsec, "~> 0.5.0", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm"},
-  "makeup_elixir": {:hex, :makeup_elixir, "0.13.0", "be7a477997dcac2e48a9d695ec730b2d22418292675c75aa2d34ba0909dcdeda", [:mix], [{:makeup, "~> 0.8", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm"},
-  "nimble_parsec": {:hex, :nimble_parsec, "0.5.0", "90e2eca3d0266e5c53f8fbe0079694740b9c91b6747f2b7e3c5d21966bba8300", [:mix], [], "hexpm"},
-  "stream_data": {:hex, :stream_data, "0.4.3", "62aafd870caff0849a5057a7ec270fad0eb86889f4d433b937d996de99e3db25", [:mix], [], "hexpm"},
+  "benchee": {:hex, :benchee, "1.0.1", "66b211f9bfd84bd97e6d1beaddf8fc2312aaabe192f776e8931cb0c16f53a521", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}], "hexpm", "3ad58ae787e9c7c94dd7ceda3b587ec2c64604563e049b2a0e8baafae832addb"},
+  "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"},
+  "earmark": {:hex, :earmark, "1.3.2", "b840562ea3d67795ffbb5bd88940b1bed0ed9fa32834915125ea7d02e35888a5", [:mix], [], "hexpm", "e3be2bc3ae67781db529b80aa7e7c49904a988596e2dbff897425b48b3581161"},
+  "ex_doc": {:hex, :ex_doc, "0.20.1", "88eaa16e67c505664fd6a66f42ddb962d424ad68df586b214b71443c69887123", [:mix], [{:earmark, "~> 1.3", [hex: :earmark, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.10", [hex: :makeup_elixir, repo: "hexpm", optional: false]}], "hexpm", "704b758b659a143f2f1b6f20cc3db2107a899c86c8485458ae2fcb0773651e59"},
+  "makeup": {:hex, :makeup, "0.8.0", "9cf32aea71c7fe0a4b2e9246c2c4978f9070257e5c9ce6d4a28ec450a839b55f", [:mix], [{:nimble_parsec, "~> 0.5.0", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "5fbc8e549aa9afeea2847c0769e3970537ed302f93a23ac612602e805d9d1e7f"},
+  "makeup_elixir": {:hex, :makeup_elixir, "0.13.0", "be7a477997dcac2e48a9d695ec730b2d22418292675c75aa2d34ba0909dcdeda", [:mix], [{:makeup, "~> 0.8", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "adf0218695e22caeda2820eaba703fa46c91820d53813a2223413da3ef4ba515"},
+  "nimble_parsec": {:hex, :nimble_parsec, "0.5.0", "90e2eca3d0266e5c53f8fbe0079694740b9c91b6747f2b7e3c5d21966bba8300", [:mix], [], "hexpm", "5c040b8469c1ff1b10093d3186e2e10dbe483cd73d79ec017993fb3985b8a9b3"},
+  "stream_data": {:hex, :stream_data, "0.5.0", "b27641e58941685c75b353577dc602c9d2c12292dd84babf506c2033cd97893e", [:mix], [], "hexpm", "012bd2eec069ada4db3411f9115ccafa38540a3c78c4c0349f151fc761b9e271"},
 }