From 8c3946309588ede9c7efd9b37eee5425101bd865 Mon Sep 17 00:00:00 2001
From: Ramon Snir <ramon@dynamicyield.com>
Date: Sat, 11 Jul 2015 09:55:51 +0300
Subject: [PATCH] fix Geef.Tree enumerable

---
 lib/geef/tree.ex | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/lib/geef/tree.ex b/lib/geef/tree.ex
index 4d190f9..854e695 100644
--- a/lib/geef/tree.ex
+++ b/lib/geef/tree.ex
@@ -66,30 +66,31 @@ defmodule Geef.TreeError do
   end
 end
 
-defimpl Enumerable, for: Geef.Tree do
+defimpl Enumerable, for: Geef.Object do
+  alias Geef.Object
   alias Geef.Tree
   alias Geef.TreeError
 
-  def count(tree) do
+  def count(tree = %Object{type: :tree}) do
     {:ok, Tree.count(tree)}
   end
 
-  def member?(tree, key) do
+  def member?(tree = %Object{type: :tree}, key) do
     case Tree.get(tree, key) do
       {:ok, _} -> {:ok, true}
       _ -> {:ok, false}
     end
   end
 
-  def reduce(tree, acc, fun) do
+  def reduce(tree = %Object{type: :tree}, acc, fun) do
     reduce(tree, 0, Tree.count(tree), acc, fun)
   end
 
   # We're done when the index is equal to the number of entries
-  defp reduce(_, idx, idx, acc, _), do: acc
+  defp reduce(_, idx, idx, {:cont, acc}, _), do: {:done, acc}
 
   # Call the user-passed function and recurse with the next index
-  defp reduce(tree, idx, count, acc, fun) do
+  defp reduce(tree, idx, count, {:cont, acc}, fun) do
     case Tree.nth(tree, idx) do
       {:ok, entry} ->
         reduce(tree, idx + 1, count, fun.(entry, acc), fun)
@@ -98,4 +99,9 @@ defimpl Enumerable, for: Geef.Tree do
     end
   end
 
+  # Almost default Enumerable implementation
+  defp reduce(_, _, _, {:halt, acc}, _), do: {:halted, acc}
+  defp reduce(tree, idx, count, {:suspend, acc}, fun) do
+    {:suspended, acc, &reduce(tree, idx, count, &1, fun)}
+  end
 end