From a197550bc0397f4abcd6e9bbc28f6b061937d350 Mon Sep 17 00:00:00 2001
From: Oleg Kovalov <oleg@hey.com>
Date: Wed, 9 Oct 2024 10:37:25 +0200
Subject: [PATCH] introduct getByHeight

---
 store/store.go      | 12 +++++++++---
 store/store_test.go |  3 +++
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/store/store.go b/store/store.go
index 9afbfe7..d7e700c 100644
--- a/store/store.go
+++ b/store/store.go
@@ -243,12 +243,16 @@ func (s *Store[H]) GetByHeight(ctx context.Context, height uint64) (H, error) {
 		return h, nil
 	}
 
+	return s.getByHeight(ctx, height)
+}
+
+func (s *Store[H]) getByHeight(ctx context.Context, height uint64) (H, error) {
+	var zero H
 	hash, err := s.heightIndex.HashByHeight(ctx, height)
 	if err != nil {
 		if errors.Is(err, datastore.ErrNotFound) {
 			return zero, header.ErrNotFound
 		}
-
 		return zero, err
 	}
 
@@ -519,7 +523,8 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 		if headers[i].Height() != currHeight+1 {
 			break
 		}
-		s.writeHead.Store(&headers[i])
+		newHead := headers[i]
+		s.writeHead.Store(&newHead)
 		currHeight++
 	}
 
@@ -528,10 +533,11 @@ func (s *Store[H]) tryAdvanceHead(ctx context.Context, headers ...H) {
 
 	// advance based on already written headers.
 	for {
-		newHead, err := s.GetByHeight(ctx, currHeight+1)
+		h, err := s.getByHeight(ctx, currHeight+1)
 		if err != nil {
 			break
 		}
+		newHead := h
 		s.writeHead.Store(&newHead)
 		currHeight++
 	}
diff --git a/store/store_test.go b/store/store_test.go
index ef9d017..7d123cf 100644
--- a/store/store_test.go
+++ b/store/store_test.go
@@ -190,6 +190,7 @@ func TestStore_Append(t *testing.T) {
 
 	head, err = store.Head(ctx)
 	assert.NoError(t, err)
+	assert.Equal(t, head.Height(), headers[len(headers)-1].Height())
 	assert.Equal(t, head.Hash(), headers[len(headers)-1].Hash())
 }
 
@@ -222,6 +223,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// head is advanced to the last known header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), wantHead.Height())
 		assert.Equal(t, head.Hash(), wantHead.Hash())
 
 		// check that store height is aligned with the head.
@@ -253,6 +255,7 @@ func TestStore_Append_stableHeadWhenGaps(t *testing.T) {
 		// after appending missing headers we're on the latest header.
 		head, err := store.Head(ctx)
 		require.NoError(t, err)
+		assert.Equal(t, head.Height(), latestHead.Height())
 		assert.Equal(t, head.Hash(), latestHead.Hash())
 
 		// check that store height is aligned with the head.