From fc024e95b6879122484517b17d640afe424aaaef Mon Sep 17 00:00:00 2001
From: Kresten Krab Thorup <krab@trifork.com>
Date: Fri, 20 Jan 2012 14:08:07 +0100
Subject: [PATCH] Make proper range fold in nursery

---
 src/lsm_btree.erl         | 10 ++++++----
 src/lsm_btree.hrl         |  4 ++++
 src/lsm_btree_nursery.erl | 10 ++++++----
 src/lsm_btree_reader.erl  |  8 +++++---
 4 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/src/lsm_btree.erl b/src/lsm_btree.erl
index 03323d3..80929cc 100644
--- a/src/lsm_btree.erl
+++ b/src/lsm_btree.erl
@@ -36,14 +36,16 @@ delete(Ref,Key) when is_binary(Key) ->
 put(Ref,Key,Value) when is_binary(Key), is_binary(Value) ->
     gen_server:call(Ref, {put, Key, Value}).
 
-sync_range(Ref,FromKey,ToKey) when is_binary(FromKey), is_binary(ToKey) ->
+sync_range(Ref,FromKey,ToKey) when FromKey == undefined orelse is_binary(FromKey),
+                                    ToKey == undefined orelse is_binary(ToKey) ->
     gen_server:call(Ref, {sync_range, self(), FromKey, ToKey}).
 
 sync_fold_range(Ref,Fun,Acc0,FromKey,ToKey) ->
     {ok, PID} = sync_range(Ref,FromKey,ToKey),
     receive_fold_range(PID,Fun,Acc0).
 
-async_range(Ref,FromKey,ToKey) when is_binary(FromKey), is_binary(ToKey) ->
+async_range(Ref,FromKey,ToKey) when FromKey == undefined orelse is_binary(FromKey),
+                                    ToKey == undefined orelse is_binary(ToKey) ->
     gen_server:call(Ref, {async_range, self(), FromKey, ToKey}).
 
 async_fold_range(Ref,Fun,Acc0,FromKey,ToKey) ->
@@ -139,13 +141,13 @@ code_change(_OldVsn, State, _Extra) ->
 
 handle_call({async_range, Sender, FromKey, ToKey}, _From, State=#state{ top=TopLevel, nursery=Nursery }) ->
     {ok, FoldWorkerPID} = lsm_btree_fold_worker:start(Sender),
-    lsm_btree_nursery:do_level_fold(Nursery, FoldWorkerPID),
+    lsm_btree_nursery:do_level_fold(Nursery, FoldWorkerPID, FromKey, ToKey),
     Result = lsm_btree_level:async_range(TopLevel, FoldWorkerPID, FromKey, ToKey),
     {reply, Result, State};
 
 handle_call({sync_range, Sender, FromKey, ToKey}, _From, State=#state{ top=TopLevel, nursery=Nursery }) ->
     {ok, FoldWorkerPID} = lsm_btree_fold_worker:start(Sender),
-    lsm_btree_nursery:do_level_fold(Nursery, FoldWorkerPID),
+    lsm_btree_nursery:do_level_fold(Nursery, FoldWorkerPID, FromKey, ToKey),
     Result = lsm_btree_level:sync_range(TopLevel, FoldWorkerPID, FromKey, ToKey),
     {reply, Result, State};
 
diff --git a/src/lsm_btree.hrl b/src/lsm_btree.hrl
index 5360e4f..b889e22 100644
--- a/src/lsm_btree.hrl
+++ b/src/lsm_btree.hrl
@@ -6,3 +6,7 @@
 
 -define(TOMBSTONE, 'deleted').
 
+-define(KEY_IN_RANGE(Key,FromKey,ToKey),
+        (((FromKey == undefined) orelse (FromKey =< Key))
+         and
+         ((ToKey == undefined) orelse (Key < ToKey)))).
diff --git a/src/lsm_btree_nursery.erl b/src/lsm_btree_nursery.erl
index 25d0d04..4def7d9 100644
--- a/src/lsm_btree_nursery.erl
+++ b/src/lsm_btree_nursery.erl
@@ -1,7 +1,7 @@
 -module(lsm_btree_nursery).
 
 -export([new/1, recover/2, add/3, finish/2, lookup/2, add_maybe_flush/4]).
--export([do_level_fold/2]).
+-export([do_level_fold/4]).
 
 -include("lsm_btree.hrl").
 -include_lib("kernel/include/file.hrl").
@@ -156,11 +156,13 @@ add_maybe_flush(Key, Value, Nursery=#nursery{ dir=Dir }, Top) ->
             lsm_btree_nursery:new(Dir)
     end.
 
-do_level_fold(#nursery{ cache=Cache }, FoldWorkerPID) ->
+do_level_fold(#nursery{ cache=Cache }, FoldWorkerPID, FromKey, ToKey) ->
     Ref = erlang:make_ref(),
     FoldWorkerPID ! {prefix, [Ref]},
-    lists:foreach(fun({Key,Value}) ->
-                          FoldWorkerPID ! {level_result, Ref, Key, Value}
+    lists:foreach(fun({Key,Value}) when ?KEY_IN_RANGE(Key,FromKey,ToKey) ->
+                          FoldWorkerPID ! {level_result, Ref, Key, Value};
+                     (_) ->
+                          ok
                   end,
                   gb_trees:to_list(Cache)),
     FoldWorkerPID ! {level_done, Ref},
diff --git a/src/lsm_btree_reader.erl b/src/lsm_btree_reader.erl
index 42889b4..b00bb9e 100644
--- a/src/lsm_btree_reader.erl
+++ b/src/lsm_btree_reader.erl
@@ -1,6 +1,7 @@
 -module(lsm_btree_reader).
 
 -include_lib("kernel/include/file.hrl").
+-include("lsm_btree.hrl").
 
 -export([open/1,close/1,lookup/2,fold/3,range_fold/5]).
 -export([first_node/1,next_node/1]).
@@ -44,7 +45,8 @@ fold1(File,Fun,Acc0) ->
             fold0(File,Fun,Node,Acc0)
     end.
 
-range_fold(Fun, Acc0, #index{file=File,root=Root}, FromKey, ToKey) ->
+range_fold(Fun, Acc0, #index{file=File,root=Root}, FromKey0, ToKey) ->
+    FromKey = if FromKey0 == undefined -> <<>>; true -> FromKey0 end,
     case lookup_node(File,FromKey,Root,0) of
         {ok, {Pos,_}} ->
             file:position(File, Pos),
@@ -63,7 +65,7 @@ do_range_fold(Fun, Acc0, File, FromKey, ToKey) ->
 
         {ok, #node{members=Members}} ->
             Acc1 =
-                lists:foldl(fun({Key,Value}, Acc) when Key >= FromKey, Key < ToKey ->
+                lists:foldl(fun({Key,Value}, Acc) when ?KEY_IN_RANGE(Key, FromKey, ToKey) ->
                                     Fun(Key, Value, Acc);
                                (_,Acc) ->
                                     Acc
@@ -72,7 +74,7 @@ do_range_fold(Fun, Acc0, File, FromKey, ToKey) ->
                             Members),
 
             case lists:last(Members) of
-                {LastKey,_} when LastKey < ToKey ->
+                {LastKey,_} when LastKey < ToKey; ToKey == undefined ->
                     do_range_fold(Fun, Acc1, File, FromKey, ToKey);
                 _ ->
                     Acc1