feat(indexes): add get_n_height_tips method to height index

hathor/indexes/height_index.py

@@ -19,6 +19,7 @@
 from hathor.indexes.scope import Scope
 from hathor.transaction import BaseTransaction, Block
 from hathor.transaction.genesis import BLOCK_GENESIS
+from hathor.types import VertexId
 from hathor.util import not_none
 
 SCOPE = Scope(
@@ -34,6 +35,12 @@ class IndexEntry(NamedTuple):
     timestamp: int
 
 
+class HeightInfo(NamedTuple):
+    """Used by a few methods to represent a (height, hash) tuple."""
+    height: int
+    id: VertexId
+
+
 BLOCK_GENESIS_ENTRY: IndexEntry = IndexEntry(not_none(BLOCK_GENESIS.hash), BLOCK_GENESIS.timestamp)
 
 
@@ -84,11 +91,19 @@ def get_tip(self) -> bytes:
         raise NotImplementedError
 
     @abstractmethod
-    def get_height_tip(self) -> tuple[int, bytes]:
+    def get_height_tip(self) -> HeightInfo:
         """ Return the best block height and hash, it returns the genesis when there is no other block
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def get_n_height_tips(self, n_blocks: int) -> list[HeightInfo]:
+        """ Return the n best block height and hash list, it returns the genesis when there is no other block
+
+        The returned list starts at the highest block and goes down in reverse height order.
+        """
+        raise NotImplementedError
+
     def update_new_chain(self, height: int, block: Block) -> None:
         """ When we have a new winner chain we must update all the height index
             until the first height with a common block

hathor/indexes/memory_height_index.py

@@ -14,7 +14,7 @@
 
 from typing import Optional
 
-from hathor.indexes.height_index import BLOCK_GENESIS_ENTRY, HeightIndex, IndexEntry
+from hathor.indexes.height_index import BLOCK_GENESIS_ENTRY, HeightIndex, HeightInfo, IndexEntry
 
 
 class MemoryHeightIndex(HeightIndex):
@@ -68,6 +68,15 @@ def get(self, height: int) -> Optional[bytes]:
     def get_tip(self) -> bytes:
         return self._index[-1].hash
 
-    def get_height_tip(self) -> tuple[int, bytes]:
+    def get_height_tip(self) -> HeightInfo:
         height = len(self._index) - 1
-        return height, self._index[height].hash
+        return HeightInfo(height, self._index[height].hash)
+
+    def get_n_height_tips(self, n_blocks: int) -> list[HeightInfo]:
+        if n_blocks < 1:
+            raise ValueError('n_blocks must be a positive, non-zero, integer')
+        # highest height that is included, will be the first element
+        h_high = len(self._index) - 1
+        # lowest height that is not included, -1 if it reaches the genesis
+        h_low = max(h_high - n_blocks, -1)
+        return [HeightInfo(h, self._index[h].hash) for h in range(h_high, h_low, -1)]

hathor/indexes/rocksdb_height_index.py

@@ -12,12 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 from structlog import get_logger
 
 from hathor.conf import HathorSettings
-from hathor.indexes.height_index import BLOCK_GENESIS_ENTRY, HeightIndex, IndexEntry
+from hathor.indexes.height_index import BLOCK_GENESIS_ENTRY, HeightIndex, HeightInfo, IndexEntry
 from hathor.indexes.rocksdb_utils import RocksDBIndexUtils
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -141,11 +141,28 @@ def get_tip(self) -> bytes:
         assert value is not None  # must never be empty, at least genesis has been added
         return self._from_value(value).hash
 
-    def get_height_tip(self) -> tuple[int, bytes]:
+    def get_height_tip(self) -> HeightInfo:
         it = self._db.iteritems(self._cf)
         it.seek_to_last()
         (_, key), value = it.get()
         assert key is not None and value is not None  # must never be empty, at least genesis has been added
         height = self._from_key(key)
         entry = self._from_value(value)
-        return height, entry.hash
+        return HeightInfo(height, entry.hash)
+
+    def get_n_height_tips(self, n_blocks: int) -> list[HeightInfo]:
+        if n_blocks < 1:
+            raise ValueError('n_blocks must be a positive, non-zero, integer')
+        info_list: list[HeightInfo] = []
+        # we need to iterate in reverse order
+        it: Any = reversed(self._db.iteritems(self._cf))  # XXX: mypy doesn't know what reversed does to this iterator
+        it.seek_to_last()
+        for (_, key), value in it:
+            # stop when we have enough elements, otherwise the iterator will stop naturally when it reaches the genesis
+            if len(info_list) == n_blocks:
+                break
+            assert key is not None and value is not None  # must never be empty, at least genesis has been added
+            height = self._from_key(key)
+            entry = self._from_value(value)
+            info_list.append(HeightInfo(height, entry.hash))
+        return info_list

tests/tx/test_indexes.py

@@ -842,6 +842,26 @@ def test_addresses_index_last(self):
         self.assertTrue(addresses_indexes.is_address_empty(address))
         self.assertEqual(addresses_indexes.get_sorted_from_address(address), [])
 
+    def test_height_index(self):
+        from hathor.indexes.height_index import HeightInfo
+
+        # make height 100
+        H = 100
+        blocks = add_new_blocks(self.manager, H - settings.REWARD_SPEND_MIN_BLOCKS, advance_clock=15)
+        height_index = self.manager.tx_storage.indexes.height
+        self.assertEqual(height_index.get_height_tip(), HeightInfo(100, blocks[-1].hash))
+        self.assertEqual(height_index.get_n_height_tips(1), [HeightInfo(100, blocks[-1].hash)])
+        self.assertEqual(height_index.get_n_height_tips(2),
+                         [HeightInfo(100, blocks[-1].hash), HeightInfo(99, blocks[-2].hash)])
+        self.assertEqual(height_index.get_n_height_tips(3),
+                         [HeightInfo(100, blocks[-1].hash),
+                          HeightInfo(99, blocks[-2].hash),
+                          HeightInfo(98, blocks[-3].hash)])
+        self.assertEqual(len(height_index.get_n_height_tips(100)), 100)
+        self.assertEqual(len(height_index.get_n_height_tips(101)), 101)
+        self.assertEqual(len(height_index.get_n_height_tips(102)), 101)
+        self.assertEqual(height_index.get_n_height_tips(103), height_index.get_n_height_tips(104))
+
 
 class BaseMemoryIndexesTest(BaseIndexesTest):
     def setUp(self):
@@ -860,7 +880,7 @@ def setUp(self):
 
         # this makes sure we can spend the genesis outputs
         self.manager = self.create_peer('testnet', tx_storage=self.tx_storage, unlock_wallet=True, wallet_index=True,
-                                        utxo_index=True)
+                                        use_memory_index=True, utxo_index=True)
         self.blocks = add_blocks_unlock_reward(self.manager)
         self.last_block = self.blocks[-1]