From e79e9f43a31e3b833d54411b6703f7326d0ecdaa Mon Sep 17 00:00:00 2001
From: Raul Jordan <rauljordan@college.harvard.edu>
Date: Fri, 4 May 2018 14:18:48 -0500
Subject: [PATCH] sharding: fix shard methods and lookup keys

Former-commit-id: 5cba9cefd8d480ec9856ce3059b089f32ca4efce [formerly 24ec21a5989261e155d88b19c4df07b501b97942]
Former-commit-id: 97a654f56e99a83eb4aeb11349e5c0ac9ad0caec
---
 sharding/db.go    | 12 ++++++++++++
 sharding/shard.go | 44 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/sharding/db.go b/sharding/db.go
index b016f7f8e..04bc96ec1 100644
--- a/sharding/db.go
+++ b/sharding/db.go
@@ -10,6 +10,10 @@ type shardKV struct {
 	kv map[common.Hash][]byte
 }
 
+func makeShardKV() *shardKV {
+	return &shardKV{kv: make(map[common.Hash][]byte)}
+}
+
 func (sb *shardKV) Get(k common.Hash) ([]byte, error) {
 	v := sb.kv[k]
 	if v == nil {
@@ -18,6 +22,14 @@ func (sb *shardKV) Get(k common.Hash) ([]byte, error) {
 	return v, nil
 }
 
+func (sb *shardKV) Has(k common.Hash) bool {
+	v := sb.kv[k]
+	if v == nil {
+		return false
+	}
+	return true
+}
+
 func (sb *shardKV) Put(k common.Hash, v []byte) {
 	sb.kv[k] = v
 	return
diff --git a/sharding/shard.go b/sharding/shard.go
index 38bbac376..2127f979e 100644
--- a/sharding/shard.go
+++ b/sharding/shard.go
@@ -16,8 +16,8 @@ type Shard struct {
 
 // MakeShard creates an instance of a Shard struct given a shardID.
 func MakeShard(shardID *big.Int) *Shard {
-	// Swappable.
-	shardDB := &shardKV{kv: make(map[common.Hash][]byte)}
+	// Swappable - can be makeShardLevelDB, makeShardSparseTrie, etc.
+	shardDB := makeShardKV()
 
 	return &Shard{
 		shardID: shardID,
@@ -81,6 +81,31 @@ func (s *Shard) GetCollationByHash(hash common.Hash) (*Collation, error) {
 	return &Collation{header: header, body: body}, nil
 }
 
+// GetCanonicalCollationHash gets a collation header hash that has been set as canonical for
+// shardID/period pair
+func (s *Shard) GetCanonicalCollationHash(shardID *big.Int, period *big.Int) (*common.Hash, error) {
+	key := canonicalCollationLookupKey(shardID, period)
+	collationHashBytes, err := s.shardDB.Get(common.BytesToHash(key))
+	if err != nil {
+		return nil, fmt.Errorf("Error: No Canonical Collation Set for Period/ShardID")
+	}
+	collationHash := common.BytesToHash(collationHashBytes)
+	return &collationHash, nil
+}
+
+// GetCanonicalCollation fetches the collation set as canonical in the shardDB.
+func (s *Shard) GetCanonicalCollation(shardID *big.Int, period *big.Int) (*Collation, error) {
+	h, err := s.GetCanonicalCollationHash(shardID, period)
+	if err != nil {
+		return nil, fmt.Errorf("Error: No Hash Found")
+	}
+	collation, err := s.GetCollationByHash(*h)
+	if err != nil {
+		return nil, fmt.Errorf("Error: No Canonical Collation Found for Hash")
+	}
+	return collation, nil
+}
+
 // GetBodyByChunkRoot fetches a collation body.
 func (s *Shard) GetBodyByChunkRoot(chunkRoot common.Hash) ([]byte, error) {
 	body, err := s.shardDB.Get(chunkRoot)
@@ -99,3 +124,18 @@ func (s *Shard) CheckAvailability(header *CollationHeader) bool {
 func (s *Shard) SetAvailability(chunkRoot *common.Hash) error {
 	return nil
 }
+
+// dataAvailabilityLookupKey formats a string that will become a lookup
+// key in the shardDB.
+func dataAvailabilityLookupKey(chunkRoot *common.Hash) []byte {
+	return []byte(fmt.Sprintf("availability-lookup:%s", chunkRoot.Str()))
+}
+
+// dataAvailabilityLookupKey formats a string that will become a lookup key
+// in the shardDB that takes into account the shardID and the period
+// of the shard.
+func canonicalCollationLookupKey(shardID *big.Int, period *big.Int) []byte {
+	str := "canonical-collation-lookup:shardID=%s,period=%s"
+	key := fmt.Sprintf(str, shardID.String(), period.String())
+	return []byte(key)
+}