From 8d2f54f913f597ac1868fe1d99658039049ca89b Mon Sep 17 00:00:00 2001
From: "tung.tq" <tung.tq@teko.vn>
Date: Wed, 4 Oct 2023 18:02:20 +0700
Subject: [PATCH] Fix MMap Filler Hash Range

---
 mmap/filler.go      | 86 +++++++++++++++++++++++++++++++++++++--------
 mmap/filler_test.go | 44 +++++++++++++++++++++++
 2 files changed, 116 insertions(+), 14 deletions(-)

diff --git a/mmap/filler.go b/mmap/filler.go
index efa4d3a..76b53dc 100644
--- a/mmap/filler.go
+++ b/mmap/filler.go
@@ -2,6 +2,7 @@ package mmap
 
 import (
 	"context"
+	"sort"
 )
 
 // FillKey ...
@@ -10,17 +11,12 @@ type FillKey[R comparable] struct {
 	Range   HashRange
 }
 
-type multiGetState[T any, R comparable] struct {
-	keys   []FillKey[R]
-	result map[R][]T
-	err    error
-}
-
 // NewMultiGetFiller converts from function often using SELECT WHERE IN
 // into a Filler[T, R] that allow to be passed to New
-func NewMultiGetFiller[T any, R comparable](
+func NewMultiGetFiller[T any, R comparable, K Key](
 	multiGetFunc func(ctx context.Context, keys []FillKey[R]) ([]T, error),
 	getRootKey func(v T) R,
+	getKey func(v T) K,
 ) Filler[T, R] {
 	var state *multiGetState[T, R]
 
@@ -44,11 +40,7 @@ func NewMultiGetFiller[T any, R comparable](
 				if err != nil {
 					s.err = err
 				} else {
-					for _, v := range values {
-						rootKey := getRootKey(v)
-						prev := s.result[rootKey]
-						s.result[rootKey] = append(prev, v)
-					}
+					collectStateValues(s, values, getRootKey, getKey)
 				}
 			}
 
@@ -56,8 +48,74 @@ func NewMultiGetFiller[T any, R comparable](
 				return nil, s.err
 			}
 
-			result := s.result[rootKey]
-			return result, nil
+			valuesByRootKey := s.result[rootKey]
+			lowerBound := findLowerBound(valuesByRootKey, getKey, hashRange.Begin)
+
+			return computeValuesInHashRange(valuesByRootKey, lowerBound, hashRange, getKey), nil
+		}
+	}
+}
+
+type multiGetState[T any, R comparable] struct {
+	keys   []FillKey[R]
+	result map[R][]T
+	err    error
+}
+
+func findLowerBound[T any, K Key](
+	values []T,
+	getKey func(v T) K,
+	lowerBound uint64,
+) int {
+	// similar to std::lower_bound of C++
+	first := 0
+	last := len(values)
+	for first != last {
+		mid := (first + last) / 2
+
+		val := values[mid]
+		if getKey(val).Hash() >= lowerBound {
+			last = mid
+		} else {
+			first = mid + 1
+		}
+	}
+	return first
+}
+
+func computeValuesInHashRange[T any, K Key](
+	values []T,
+	lowerBound int,
+	hashRange HashRange,
+	getKey func(T) K,
+) []T {
+	var result []T
+	for i := lowerBound; i < len(values); i++ {
+		v := values[i]
+		if getKey(v).Hash() > hashRange.End {
+			break
 		}
+		result = append(result, v)
+	}
+	return result
+}
+
+func collectStateValues[T any, R comparable, K Key](
+	s *multiGetState[T, R],
+	values []T,
+	getRootKey func(T) R,
+	getKey func(T) K,
+) {
+	for _, v := range values {
+		rootKey := getRootKey(v)
+		prev := s.result[rootKey]
+		s.result[rootKey] = append(prev, v)
+	}
+
+	// sort by hash
+	for _, v := range s.result {
+		sort.Slice(v, func(i, j int) bool {
+			return getKey(v[i]).Hash() < getKey(v[j]).Hash()
+		})
 	}
 }
diff --git a/mmap/filler_test.go b/mmap/filler_test.go
index 4fb088f..92e3b32 100644
--- a/mmap/filler_test.go
+++ b/mmap/filler_test.go
@@ -25,6 +25,7 @@ func newMultiGetFillerTest() *multiGetFillerTest {
 			return f.fillFunc(ctx, keys)
 		},
 		stockLocation.getRootKey,
+		stockLocation.getKey,
 	)
 
 	return f
@@ -180,4 +181,47 @@ func TestNewMultiGetFiller(t *testing.T) {
 			},
 		}, f.fillKeys)
 	})
+
+	t.Run("multiple same root key", func(t *testing.T) {
+		f := newMultiGetFillerTest()
+
+		stock1 := stockLocation{
+			Sku:      sku1,
+			Location: loc1,
+			Hash:     hash1.Begin + 100,
+			Quantity: 41,
+		}
+		stock2 := stockLocation{
+			Sku:      sku1,
+			Location: loc2,
+			Hash:     hash2.Begin + 100,
+			Quantity: 42,
+		}
+
+		f.fillFunc = func(ctx context.Context, keys []FillKey[stockLocationRootKey]) ([]stockLocation, error) {
+			return []stockLocation{stock2, stock1}, nil
+		}
+
+		fn1 := f.filler(context.Background(), stock1.getRootKey(), hash1)
+		fn2 := f.filler(context.Background(), stock2.getRootKey(), hash2)
+
+		resp, err := fn1()
+		assert.Equal(t, nil, err)
+		assert.Equal(t, []stockLocation{
+			stock1,
+		}, resp)
+
+		resp, err = fn2()
+		assert.Equal(t, nil, err)
+		assert.Equal(t, []stockLocation{
+			stock2,
+		}, resp)
+
+		assert.Equal(t, [][]FillKey[stockLocationRootKey]{
+			{
+				{RootKey: stock1.getRootKey(), Range: hash1},
+				{RootKey: stock2.getRootKey(), Range: hash2},
+			},
+		}, f.fillKeys)
+	})
 }