From 102e250fd92e5ccaad5bdaa0f5a378cd84e7beab Mon Sep 17 00:00:00 2001
From: Robert Newson <rnewson@apache.org>
Date: Tue, 28 Jul 2020 21:38:10 +0100
Subject: Support inclusive_start/end for group_reduce

---
 src/ebtree/src/ebtree.erl | 28 ++++++++++++++++++++++------
 1 file changed, 22 insertions(+), 6 deletions(-)

diff --git a/src/ebtree/src/ebtree.erl b/src/ebtree/src/ebtree.erl
index b227be035..edbc9c607 100644
--- a/src/ebtree/src/ebtree.erl
+++ b/src/ebtree/src/ebtree.erl
@@ -302,6 +302,8 @@ group_reduce(Db, #tree{} = Tree, StartKey, EndKey, GroupKeyFun, UserAccFun, User
 %% @returns the final accumulator.
 -type group_key() :: term().
 
+-type group_option() :: [{inclusive_start, boolean()} | {inclusive_end, boolean()}].
+
 -spec group_reduce(
     Db :: term(),
     Tree :: #tree{},
@@ -310,7 +312,7 @@ group_reduce(Db, #tree{} = Tree, StartKey, EndKey, GroupKeyFun, UserAccFun, User
     GroupKeyFun :: fun((term()) -> group_key()),
     UserAccFun :: fun(({group_key(), GroupValue :: term()}, Acc0 :: term()) -> Acc1 :: term()),
     UserAcc0 :: term(),
-    Options :: [fold_option()]) -> Acc1 :: term().
+    Options :: [fold_option() | group_option()]) -> Acc1 :: term().
 group_reduce(Db, #tree{} = Tree, StartKey, EndKey, GroupKeyFun, UserAccFun, UserAcc0, Options) ->
     Dir = proplists:get_value(dir, Options, fwd),
     NoGroupYet = ?MIN,
@@ -318,7 +320,7 @@ group_reduce(Db, #tree{} = Tree, StartKey, EndKey, GroupKeyFun, UserAccFun, User
         ({visit, Key, Value}, {CurrentGroup, UserAcc, MapAcc, ReduceAcc}) ->
             BeforeStart = less_than(Tree, Key, StartKey),
             AfterEnd = greater_than(Tree, Key, EndKey),
-            InRange = in_range(Tree, StartKey, Key, EndKey),
+            InRange = in_range(Tree, StartKey, Key, EndKey, Options),
             KeyGroup = GroupKeyFun(Key),
             SameGroup = CurrentGroup =:= KeyGroup,
             if
@@ -341,8 +343,8 @@ group_reduce(Db, #tree{} = Tree, StartKey, EndKey, GroupKeyFun, UserAccFun, User
             BeforeStart = less_than(Tree, LastKey, StartKey),
             AfterEnd = greater_than(Tree, FirstKey, EndKey),
             Whole = CurrentGroup =:= GroupKeyFun(FirstKey) andalso CurrentGroup =:= GroupKeyFun(LastKey),
-            FirstInRange = in_range(Tree, StartKey, FirstKey, EndKey),
-            LastInRange = in_range(Tree, StartKey, LastKey, EndKey),
+            FirstInRange = in_range(Tree, StartKey, FirstKey, EndKey, Options),
+            LastInRange = in_range(Tree, StartKey, LastKey, EndKey, Options),
             if
                 Dir == fwd andalso BeforeStart ->
                     {skip, {CurrentGroup, UserAcc, MapAcc, ReduceAcc}};
@@ -876,8 +878,22 @@ reduce_values(#tree{} = Tree, Values, Rereduce) when is_list(Values) ->
 
 %% collation functions
 
-in_range(#tree{} = Tree, StartOfRange, Key, EndOfRange) ->
-    greater_than_or_equal(Tree, Key, StartOfRange) andalso less_than_or_equal(Tree, Key, EndOfRange).
+in_range(#tree{} = Tree, StartOfRange, Key, EndOfRange, Options) ->
+    InclusiveStart = proplists:get_value(inclusive_start, Options, true),
+    InclusiveEnd = proplists:get_value(inclusive_end, Options, true),
+    Left = if
+        InclusiveStart ->
+            greater_than_or_equal(Tree, Key, StartOfRange);
+        true ->
+            greater_than(Tree, Key, StartOfRange)
+    end,
+    Right = if
+        InclusiveEnd ->
+            less_than_or_equal(Tree, Key, EndOfRange);
+        true ->
+            less_than(Tree, Key, EndOfRange)
+    end,
+    Left andalso Right.
 
 
 greater_than(#tree{} = Tree, A, B) ->
-- 
cgit v1.2.1