From bbbbdc455d8cebd868d7c1f090a183602c7382c8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C3=86de?= <ae.s.hoekstra@student.utwente.nl>
Date: Fri, 9 Mar 2018 19:59:15 +0100
Subject: [PATCH] Use header hash for by height index (#761)

* Use header hash for by height index
* Update semantics of save_header_height
---
 chain/src/store.rs           | 11 ++++-------
 chain/tests/store_indices.rs | 23 +++++++++++++++++++++++
 2 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/chain/src/store.rs b/chain/src/store.rs
index 4e88c916..d45f34c4 100644
--- a/chain/src/store.rs
+++ b/chain/src/store.rs
@@ -146,11 +146,12 @@ impl ChainStore for ChainKVStore {
 
 	fn get_header_by_height(&self, height: u64) -> Result<BlockHeader, Error> {
 		option_to_not_found(self.db.get_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, height)))
+			.and_then(|hash| self.get_block_header(&hash))
 	}
 
 	fn save_header_height(&self, bh: &BlockHeader) -> Result<(), Error> {
 		self.db
-			.put_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, bh.height), bh)
+			.put_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, bh.height), &bh.hash())
 	}
 
 	fn delete_header_by_height(&self, height: u64) -> Result<(), Error> {
@@ -229,8 +230,7 @@ impl ChainStore for ChainKVStore {
 			self.delete_header_by_height(n)?;
 		}
 
-		self.db
-			.put_ser(&u64_to_key(HEADER_HEIGHT_PREFIX, header.height), header)?;
+		self.save_header_height(&header)?;
 
 		if header.height > 0 {
 			let mut prev_header = self.get_block_header(&header.previous)?;
@@ -238,10 +238,7 @@ impl ChainStore for ChainKVStore {
 				if let Ok(_) = self.is_on_current_chain(&prev_header) {
 					break;
 				}
-				self.db.put_ser(
-					&u64_to_key(HEADER_HEIGHT_PREFIX, prev_header.height),
-					&prev_header,
-				)?;
+				self.save_header_height(&prev_header)?;
 
 				prev_header = self.get_block_header(&prev_header.previous)?;
 			}
diff --git a/chain/tests/store_indices.rs b/chain/tests/store_indices.rs
index e7b0a981..0a119830 100644
--- a/chain/tests/store_indices.rs
+++ b/chain/tests/store_indices.rs
@@ -24,6 +24,7 @@ use std::fs;
 use chain::{ChainStore, Tip};
 use core::core::hash::Hashed;
 use core::core::Block;
+use core::core::BlockHeader;
 use core::core::target::Difficulty;
 use keychain::Keychain;
 use core::global;
@@ -70,3 +71,25 @@ fn test_various_store_indices() {
 	let block_header = chain_store.get_header_by_height(1).unwrap();
 	assert_eq!(block_header.hash(), block_hash);
 }
+
+#[test]
+fn test_store_header_height() {
+	let _ = env_logger::init();
+	clean_output_dir(".grin");
+
+	let chain_store = &chain::store::ChainKVStore::new(".grin".to_string()).unwrap() as &ChainStore;
+
+	let mut block_header = BlockHeader::default();
+	block_header.height = 1;
+
+	chain_store.save_block_header(&block_header).unwrap();
+	chain_store.save_header_height(&block_header).unwrap();
+
+	let stored_block_header = chain_store.get_header_by_height(1).unwrap();
+	assert_eq!(block_header.hash(), stored_block_header.hash());
+
+	chain_store.delete_header_by_height(1).unwrap();
+
+	let result = chain_store.get_header_by_height(1);
+	assert_eq!(result.is_err(), true);
+}